按注意力头粒度压缩键值缓存，大幅降低内存占用同时保持高问答准确率。
Not All Heads Matter: A Head-Level KV Cache Compression Method with Integrated Retrieval and Reasoning
- 按注意力头重要性动态筛选缓存，而非整层压缩
- 仅用1.5%缓存即达全量97%性能，低资源下优势明显
- 适合长文本问答与资源受限场景的模型部署
键值（KV）缓存是提升大语言模型推理效率的常用技术，但其内存开销随输入长度快速增长。已有研究发现，并非所有词元对生成都同样重要，提出在层粒度上进行缓存压缩以保留关键信息。本文进一步观察到不同注意力头在生成任务中角色各异，提出HeadKV方法，在头级别进行缓存压缩，并引入新方法HeadKV-R2，通过上下文推理能力估计实现更优压缩。该方法基于单个注意力头的重要性评估，用于需要检索与推理结合的上下文问答任务。在多个基准测试（LongBench、LooGLE）、模型架构（如Llama-3-8B-Instruct、Mistral-7B-Instruct）及长上下文能力测试中，本方法显著优于强基线，尤其在低资源设置下（缓存大小为64和128时）。值得注意的是，仅保留1.5%的缓存即可达到全量缓存97%的性能表现。代码已开源：https://github.com/FYYFU/HeadKV
原文摘要 · Abstract (English)
Key-Value (KV) caching is a common technique to enhance the computational efficiency of Large Language Models (LLMs), but its memory overhead grows rapidly with input length. Prior work has shown that not all tokens are equally important for text generation, proposing layer-level KV cache compression to selectively retain key information. Recognizing the distinct roles of attention heads in generation, we propose HeadKV, a head-level KV cache compression method, and HeadKV-R2, which leverages a novel contextual reasoning ability estimation for compression. Our approach operates at the level of individual heads, estimating their importance for contextual QA tasks that require both retrieval and reasoning capabilities. Extensive experiments across diverse benchmarks (LongBench, LooGLE), model architectures (e.g., Llama-3-8B-Instruct, Mistral-7B-Instruct), and long-context abilities tests demonstrate that our head-level KV cache compression significantly outperforms strong baselines, particularly in low-resource settings (KV size = 64 & 128). Notably, our method retains just 1.5% of the KV cache while achieving 97% of the performance of the full KV cache on the contextual question answering benchmark. Codes are available at https://github.com/FYYFU/HeadKV
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。