用强化学习找出推理关键注意力头，实现高效缓存压缩。
Which Heads Matter for Reasoning? RL-Guided KV Cache Compression
- 用强化学习探测哪些注意力头影响推理质量
- 仅保留关键头完整缓存，其他头压缩至固定大小
- 在多个模型任务中实现60%缓存压缩，几乎无性能损失
推理型大语言模型通过长链式思维生成复杂行为，在解码过程中对信息丢失极为敏感，给KV缓存压缩带来挑战。现有逐标记删除方法会破坏推理链，头重分配方法虽用于检索任务，却无法保留生成推理所需的关键注意力头。现有方法均无法识别真正维持推理一致性与生成终止控制的注意力头。为此，我们提出RLKV，利用强化学习作为探测工具，通过优化实际生成结果来发现哪些头对推理质量有贡献。该发现自然导出一种高效压缩策略：将完整KV缓存分配给推理关键头，其余头则采用固定大小缓存进行激进压缩。实验表明，少数头对推理至关重要，可在多种任务和模型上实现20%–60%的缓存减少，性能近乎无损，且在60%压缩下达到2.06倍端到端加速。
原文摘要 · Abstract (English)
Reasoning large language models exhibit complex reasoning behaviors via extended chain-of-thought generation that are highly fragile to information loss during decoding, creating critical challenges for KV cache compression. Existing token-dropping methods directly disrupt reasoning chains by removing intermediate steps, while head-reallocation methods, designed for retrieval tasks, fail to preserve the heads essential for generative reasoning. However, no existing method can identify which attention heads genuinely maintain reasoning consistency and control generation termination. To address this, we propose RLKV, which uses reinforcement learning as a probe to discover which heads contribute to reasoning quality by directly optimizing their cache usage against actual generation outcomes. This discovery naturally leads to an efficient compression strategy: we allocate full KV cache to reasoning-critical heads while aggressively compressing others with constant-size KV cache. Experiments reveal that a fraction of heads proves essential for reasoning, enabling 20--60% cache reduction with near-lossless performance across diverse tasks and models, and up to 2.06x end-to-end speedup at 60% reduction.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。