arXiv:2508.02401cs.CLcs.AI2025-08被引 10

通过识别不重要的词元,智能压缩注意力缓存,提升长文本处理效率。

CompressKV: Semantic Retrieval Heads Know What Tokens are Not Important Before Generation

  • 利用特定注意力头提前识别无关词元,精准保留关键信息。
  • 在长文本测试中,内存占用减少30%仍保持高准确率。
  • 适合需要高效处理长文档的模型部署场景。

大型语言模型(LLM)在长上下文处理方面取得显著进展,但键值(KV)缓存规模增大导致内存和执行效率问题。现有压缩方法依赖组查询注意力(GQA)中所有注意力头进行启发式词元剔除,忽视了各注意力头的功能差异,常误删关键词元,降低模型性能。为此,本文首先识别每层中能有效检索提示开头、结尾及文本内重要词元,并关注其语义上下文的注意力头;随后利用这些头部判断重要词元,保留对应KV缓存对。此外,针对各层缓存剔除误差进行分析,提出分层自适应缓存分配策略。实验表明,在LongBench与Needle-in-a-Haystack基准上,CompressKV在多种内存预算下持续优于当前最优方法。代码已开源。

原文摘要 · Abstract (English)

Recent advances in large language models (LLMs) have significantly boosted long-context processing. However, the increasing key-value (KV) cache size poses critical challenges to memory and execution efficiency. Most KV cache compression methods rely on heuristic token eviction using all attention heads in Grouped Query Attention (GQA)-based LLMs. This method ignores the different functionalities of attention heads, leading to the eviction of critical tokens and thus degrades the performance of LLMs. To address the issue above, instead of using all the attention heads in GQA-based LLMs to determine important tokens as in the previous work, we first identify the attention heads in each layer that are not only capable of retrieving the initial and final tokens of a prompt, but also capable of retrieving important tokens within the text and attending to their surrounding semantic context. Afterwards, we exploit such heads to determine the important tokens and retain their corresponding KV cache pairs. Furthermore, we analyze the cache eviction error of each layer individually and introduce a layer-adaptive KV cache allocation strategy. Experimental results demonstrate the proposed CompressKV consistently outperforms state-of-the-art approaches under various memory budgets on LongBench and Needle-in-a-Haystack benchmarks. Our code is publicly available at: https://github.com/TUDa-HWAI/CompressKV.git.

KV缓存注意力头长文本压缩

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