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缓存预算的任务中均实现了卓越的性能。

本论文旨在解决大型语言模型（LLMs）在自回归生成过程中的巨大计算成本问题，提出一种新的KV缓存合并方法KVMerger来实现自适应KV缓存压缩，以在受限内存预算下不会显著降低性能的情况下，为长上下文任务提供支持。

论文提出的KVMerger方法通过在单个序列中相似的键状态之间进行合并，实现自适应的KV缓存压缩，从而在长上下文任务中提高计算效率。

论文设计了有效的合并集识别算法来识别合适的KV状态进行合并，并提出了高斯核加权合并算法来选择性地合并每个合并集内的所有状态。在Llama2-7B-chat和Llama2-13B-chat等模型上进行了广泛的实验，使用了LongBench和ZeroScroll基准测试，并将KVMerger与其他KV缓存压缩技术进行了比较，证明了KVMerger在50％和35％的KV缓存预算下都能实现卓越的性能。

与本论文相关的研究包括H2O和CaM等KV缓存压缩技术，以及其他面向长上下文任务的语言模型优化方法，如XLNet和Transformer-XL等。