arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2510.08525cs.CL

哪些注意力头对推理重要?RL引导的KV缓存压缩

Which Heads Matter for Reasoning? RL-Guided KV Cache Compression

Westlake University(西华大学);McGill University(麦吉尔大学);Mila - Quebec AI Institute(魁北克AI研究院);Zhejiang University(浙江大学);Mohamed bin Zayed University of Artificial Intelligence(穆罕默德·本·扎耶德智能大学)

Wenjie Du, Li Jiang, Keda Tao, Xue Liu, Huan Wang

AI总结:

提出RLKV方法,利用强化学习识别对推理质量关键的注意力头,并对其保留完整KV缓存而对其他头进行激进压缩,实现20-60%缓存减少且性能近乎无损。

摘要:

推理型大语言模型通过扩展的思维链生成展现出复杂的推理行为,这些行为在解码过程中对信息损失高度敏感,给KV缓存压缩带来了关键挑战。现有的token丢弃方法通过移除中间步骤直接破坏推理链,而为检索任务设计的头重分配方法无法保留对生成推理至关重要的注意力头。然而,现有方法均无法识别哪些注意力头真正维持推理一致性并控制生成终止。为解决此问题,我们提出RLKV,它使用强化学习作为探针,通过直接优化注意力头缓存使用与实际生成结果的关系,发现哪些头对推理质量有贡献。这一发现自然引出了高效的压缩策略:我们对推理关键的头分配完整KV缓存,同时对其他头使用固定大小的KV缓存进行激进压缩。实验表明,少数头对推理至关重要,使得在多种任务和模型上实现20-60%的缓存减少且性能近乎无损,在60%压缩率下实现高达2.06倍的端到端加速。

英文摘要:

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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑