arXiv:2412.02252cs.CL2024-12被引 12

通过共享远距离词元的键状态,35%压缩注意力缓存且不丢重要信息。

Compressing KV Cache for Long-Context LLM Inference with Inter-Layer Attention Similarity

  • 按词元重要性分配内存,近端词元保留完整缓存,远端共享键状态
  • 在真实与合成数据集上实现最高35%的缓存压缩,性能零损失
  • 适合长文本推理场景,可与现有删减方法结合使用

大型语言模型上下文窗口的快速扩展使其能处理更复杂的长文档任务,但推理时内存消耗随之激增,主要源于键值(KV)缓存的线性增长。现有压缩方法常舍弃较不相关的词元,可能导致关键信息丢失、性能下降。本文提出 extsc{PoD}(近端优于远端),一种基于词元重要性的新压缩框架:近端词元(首尾部分)保留完整KV缓存,远端词元则跨层共享键状态。该方法基于两个观察:(1)近端词元对下一个词预测更为重要;(2)远端词元在连续层间的注意力分数高度冗余。通过共享键状态,多个层可复用同一组远端键,大幅降低内存占用而无需丢弃上下文。我们引入轻量级微调适配模型以适应新结构。在合成(针在草堆中)与真实长文本基准测试中, extsc{PoD} 实现最高35%的缓存压缩,性能无损。该方法与已有词元选择技术正交,可进一步联合使用。

原文摘要 · Abstract (English)

The rapid expansion of context window sizes in Large Language Models~(LLMs) has enabled them to tackle increasingly complex tasks involving lengthy documents. However, this progress comes at the cost of a substantial increase in memory usage during inference, primarily due to the linear growth of the key-value~(KV) cache. Existing KV cache compression methods often discard less relevant tokens, which can lead to significant performance degradation when critical information is lost. In this paper, we propose \textsc{PoD}~(Proximal tokens over Distant tokens), a novel KV cache compression framework that allocates memory according to token importance, retaining less important tokens in a more compact, shared form rather than discarding them entirely. Our approach is motivated by two key observations: (1) proximal tokens -- those at the beginning and end of the context -- are significantly more important for next-token prediction, and (2) attention scores for distant tokens are highly redundant across consecutive layers. Leveraging these insights, \textsc{PoD} preserves the full KV cache for proximal tokens, while for distant tokens, it shares key states across layers. Since attention scores are determined by both queries and keys, sharing key states enables multiple layers to reuse a single set of keys for distant tokens, substantially reducing KV cache memory without discarding essential context. We further introduce a lightweight post-training adaptation to enable the model to adjust to this new attention-sharing structure. Extensive experiments on both synthetic~(Needle in a Haystack) and real-world long-context benchmarks demonstrate that \textsc{PoD} reduces KV cache memory usage by up to 35\% without compromising performance. Our method is orthogonal to existing token-selection-based techniques and can be combined with them for further KV cache compression.

KV缓存压缩长文本推理注意力机制

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。