Model Tells You Where to Merge: Adaptive KV Cache Merging for LLMs on Long-Context Tasks

如何高效地为大型语言模型（LLMs）提供服务已成为一个紧迫的问题，因为它们在自回归生成过程中的计算成本非常高。为了减少计算成本，LLMs通常采用KV缓存技术来提高生成速度。虽然提高了计算效率，但KV缓存的存储需求很大，特别是在长上下文场景下，导致显著的内存消耗。现有的KV缓存清理方法往往会因为清理引入的信息丢失而降低LLMs在长上下文场景下的性能。在本文中，我们提出了一种新的KV缓存合并方法，称为KVMerger，以实现自适应KV缓存压缩，以适应内存预算限制下的长上下文任务，而不会显著降低性能。我们的方法受到一个有趣的观察的启发，即在单个序列内，关键状态在令牌级别上表现出高度相似性。为了促进合并，我们开发了一种有效而直接的合并集合识别算法，以识别适合合并的KV状态。我们的合并集合识别算法刺激了第二个观察，即从相似性的角度看，KV缓存的稀疏性独立于数据集，并在模型级别上保持持久。随后，我们提出了一种高斯核加权合并算法，以选择性地合并每个合并集合中的所有状态。我们进行了大量实验，以证明KVMerger在内存预算受限的长上下文任务中的有效性，将其应用于包括Llama2-7B-chat和Llama2-13B-chat在内的模型。使用LongBench和ZeroScroll基准测试，我们将我们的方法与其他KV缓存压缩技术进行比较，包括H2O和CaM，显示出我们的方法在50％和35％的KV缓存预算的任务中都实现了优越的性能。

本文旨在解决大型语言模型在自回归生成过程中的高计算成本和高存储要求问题，提出了一种名为KVMerger的KV缓存合并方法，实现了适应性KV缓存压缩，以满足长序列任务下的内存限制。

KVMerger方法的关键思路是利用同一序列中键状态在标记级别上的高相似度，开发了一种有效而简单的合并集合识别算法来识别适合合并的KV状态，并提出了高斯核加权合并算法来选择性地合并每个合并集合中的所有状态。

本文通过在Llama2-7B-chat和Llama2-13B-chat等模型上进行实验，使用LongBench和ZeroScroll基准测试比较了KVMerger方法与其他KV缓存压缩技术，包括H2O和CaM，证明了KVMerger方法在50％和35％ KV缓存预算下的长序列任务中具有优越的性能。本文的方法可以为解决LLM计算和存储问题提供新思路。

相关研究包括KV缓存的其他压缩方法，例如H2O和CaM。