并非所有注意力头都重要：一种融合检索与推理的头部级KV缓存压缩方法
Not All Heads Matter: A Head-Level KV Cache Compression Method with Integrated Retrieval and Reasoning
AI总结：
提出头部级KV缓存压缩方法HeadKV及HeadKV-R2，通过估计注意力头在检索与推理任务中的重要性，仅保留1.5%缓存即可达到完整缓存97%的性能。
摘要：
键值（KV）缓存是一种常见技术，用于提升大型语言模型（LLMs）的计算效率，但其内存开销随输入长度迅速增长。先前工作表明，并非所有令牌对文本生成同等重要，提出了层级KV缓存压缩以选择性保留关键信息。鉴于注意力头在生成中的不同作用，我们提出HeadKV，一种头部级KV缓存压缩方法，以及HeadKV-R2，其利用一种新颖的上下文推理能力估计进行压缩。我们的方法在单个注意力头层面操作，估计它们对需要检索和推理能力的上下文问答任务的重要性。跨多样基准（LongBench、LooGLE）、模型架构（如Llama-3-8B-Instruct、Mistral-7B-Instruct）和长上下文能力测试的大量实验表明，我们的头部级KV缓存压缩显著优于强基线，尤其在低资源设置（KV大小=64和128）下。值得注意的是，我们的方法仅保留1.5%的KV缓存，即在上下文问答基准上实现了完整KV缓存性能的97%。代码可在https://github.com/FYYFU/HeadKV获取。
英文摘要：
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
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
