Model Tells You Where to Merge: Adaptive KV Cache Merging for LLMs on Long-Context Tasks
作者: Zheng Wang, Boxiao Jin, Zhongzhi Yu, Minjia Zhang
分类: cs.CL
发布日期: 2024-07-11 (更新: 2024-07-21)
💡 一句话要点
提出KVMerger，自适应合并KV缓存，提升LLM在长文本任务中的性能。
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: KV缓存压缩 长文本处理 大型语言模型 内存优化 自适应合并
📋 核心要点
- 现有KV缓存淘汰策略在长文本任务中导致信息损失，降低LLM性能，是亟待解决的核心问题。
- KVMerger通过识别token级别相似的关键状态，自适应合并KV缓存，在压缩的同时尽量保留有效信息。
- 实验表明，在LongBench和ZeroScroll基准测试中，KVMerger优于H2O和CaM等现有KV缓存压缩技术。
📝 摘要（中文）
大型语言模型（LLM）的自回归生成过程计算成本巨大，如何高效地服务LLM已成为一个紧迫的问题。为了降低计算成本，LLM通常采用KV缓存技术来提高生成速度。在提高计算效率的同时，KV缓存的存储需求也很高，尤其是在长文本场景中，导致显著的内存消耗。现有的KV缓存淘汰方法通常会降低LLM在长文本场景中的性能，这是由于淘汰引入的信息损失。在本文中，我们提出了一种新颖的KV缓存合并方法，称为KVMerger，以实现长文本任务的自适应KV缓存压缩，在有限的内存预算下不会显著降低性能。我们的方法受到一个有趣的观察的启发，即关键状态在单个序列中的token级别表现出高度相似性。为了方便合并，我们开发了一种有效而直接的合并集识别算法，以识别适合合并的KV状态。我们的合并集识别算法激发了第二个观察，即从相似性的角度来看，KV缓存的稀疏性与数据集无关，并且在模型级别保持持久性。随后，我们提出了一种高斯核加权合并算法，以选择性地合并每个合并集中的所有状态。我们进行了广泛的实验，以证明KVMerger在有限内存预算下对长文本任务的有效性，并将其应用于包括Llama2-7B-chat和Llama2-13B-chat在内的模型。使用LongBench和ZeroScroll基准，我们将我们的方法与其他KV缓存压缩技术（包括H2O和CaM）进行了比较，表明我们的方法在50％和35％的KV缓存预算下，在各项任务中均实现了卓越的性能。
🔬 方法详解
问题定义：论文旨在解决大型语言模型在长文本处理中，KV缓存占用大量内存的问题。现有的KV缓存淘汰策略会造成信息丢失，影响模型性能，尤其是在长文本任务中，模型需要访问更早期的信息。因此，如何在限制内存预算的同时，尽可能保留关键信息，是本论文要解决的核心问题。
核心思路：论文的核心思路是，通过合并相似的KV状态来压缩KV缓存，而不是直接淘汰。论文观察到，在单个序列中，不同token对应的key状态具有很高的相似性。因此，可以将相似的KV状态合并，从而减少内存占用，同时尽量保留重要的历史信息。
技术框架：KVMerger主要包含两个阶段：合并集识别和状态合并。首先，合并集识别算法用于识别适合合并的KV状态集合。该算法基于KV缓存的稀疏性，这种稀疏性在模型层面是持久的，与数据集无关。然后，使用高斯核加权合并算法，对每个合并集中的所有状态进行选择性合并。
关键创新：论文的关键创新在于提出了一种自适应的KV缓存合并方法，该方法能够根据KV状态的相似性，动态地调整合并策略。与传统的KV缓存淘汰方法相比，KVMerger能够更好地保留长文本中的关键信息，从而提高模型在长文本任务中的性能。此外，合并集识别算法利用了KV缓存稀疏性的特性，使得合并过程更加高效。
关键设计：合并集识别算法的设计是关键。论文提出了一种简单有效的算法，用于识别相似的KV状态。高斯核加权合并算法使用高斯核函数来计算合并权重，从而更好地保留重要的KV状态。具体的参数设置和损失函数在论文中没有详细描述，属于未知信息。
🖼️ 关键图片
📊 实验亮点
实验结果表明，在LongBench和ZeroScroll基准测试中，KVMerger在50%和35%的KV缓存预算下，均优于H2O和CaM等现有KV缓存压缩技术。具体性能提升数据未在摘要中给出，属于未知信息。该结果证明了KVMerger在长文本任务中，能够在有限内存预算下有效提升LLM的性能。
🎯 应用场景
该研究成果可广泛应用于需要处理长文本的自然语言处理任务中，例如长文档摘要、机器翻译、问答系统等。通过降低LLM的内存需求，KVMerger能够使LLM在资源受限的设备上运行，并提高LLM在长文本任务中的推理效率。未来，该技术有望推动LLM在移动设备和边缘计算等领域的应用。
📄 摘要（原文）
How to efficiently serve Large Language Models (LLMs) has become a pressing issue because of their huge computational cost in their autoregressive generation process. To mitigate computational costs, LLMs often employ the KV Cache technique to improve the generation speed. While improving the computational efficiency, the storage requirements of the KV cache are substantial, particularly in long-context scenarios, leading to significant memory consumption. Existing KV cache eviction methods often degrade the performance of LLMs in long-context scenarios due to the information loss introduced by eviction. In this paper, we propose a novel KV cache merging approach, called KVMerger, to achieve adaptive KV cache compression for long-context tasks without significant performance degradation under constrained memory budgets. Our approach is inspired by the intriguing observation that key states exhibit high similarity at the token level within a single sequence. To facilitate merging, we develop an effective yet straightforward merging set identification algorithm to identify suitable KV states for merging. Our merging set identification algorithm stimulates the second observation that KV cache sparsity, from similarity perspective, is independent of the dataset and remains persistent at the model level. Subsequently, we propose a Gaussian kernel weighted merging algorithm to selectively merge all states within each merging set. We conduct extensive experiments to demonstrate the effectiveness of KVMerger for long-context tasks under constrained memory budgets, applying it to models including Llama2-7B-chat and Llama2-13B-chat. Using the LongBench and ZeroScroll benchmarks, we compare our method with other KV cache compression techniques, including H2O and CaM, showing that our method achieves superior performance across tasks with both 50% and 35% KV cache budgets.