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. Experiments reveal that a fraction of heads proves essential for reasoning, enabling 20--50% cache reduction with near-lossless performance and up to 1.21x speedup.


翻译:推理型大语言模型通过扩展的思维链生成展现出复杂的推理行为,这些行为对解码过程中的信息丢失高度敏感,这为KV缓存压缩带来了严峻挑战。现有的令牌丢弃方法通过移除中间步骤直接破坏了推理链,而为检索任务设计的注意力头重分配方法则无法保留生成式推理所需的关键注意力头。然而,尚无现有方法能够识别哪些注意力头真正维持了推理一致性并控制生成终止。为解决此问题,我们提出了RLKV,该方法利用强化学习作为探针,通过直接优化注意力头的缓存使用以匹配实际生成结果,从而发现哪些注意力头对推理质量有贡献。这一发现自然引出了一项高效的压缩策略:我们为推理关键的头分配完整的KV缓存,同时对其他头进行激进压缩。实验表明,一小部分注意力头对推理至关重要,这使得在性能近乎无损的情况下实现20-50%的缓存缩减,并带来最高1.21倍的加速。

0
下载
关闭预览

相关内容

强化学习(RL)是机器学习的一个领域,与软件代理应如何在环境中采取行动以最大化累积奖励的概念有关。除了监督学习和非监督学习外,强化学习是三种基本的机器学习范式之一。 强化学习与监督学习的不同之处在于,不需要呈现带标签的输入/输出对,也不需要显式纠正次优动作。相反,重点是在探索(未知领域)和利用(当前知识)之间找到平衡。 该环境通常以马尔可夫决策过程(MDP)的形式陈述,因为针对这种情况的许多强化学习算法都使用动态编程技术。经典动态规划方法和强化学习算法之间的主要区别在于,后者不假设MDP的确切数学模型,并且针对无法采用精确方法的大型MDP。

知识荟萃

精品入门和进阶教程、论文和代码整理等

更多

查看相关VIP内容、论文、资讯等
大模型推理的天花板在哪里?
专知会员服务
15+阅读 · 2025年6月12日
大型语言模型推理引擎的综述:优化与效率的视角
专知会员服务
21+阅读 · 2025年5月13日
TransMLA:多头潜在注意力(MLA)即为所需
专知会员服务
23+阅读 · 2025年2月13日
迈向大型推理模型:基于大型语言模型的强化推理综述
专知会员服务
49+阅读 · 2025年1月17日
LlamaV-o1: 重新思考大语言模型中的逐步视觉推理
专知会员服务
17+阅读 · 2025年1月14日
大模型的模型压缩与有效推理综述
专知会员服务
43+阅读 · 2024年7月8日
大型语言模型的模型压缩与高效推理:综述
专知会员服务
94+阅读 · 2024年2月17日
通过集成 XNNPACK 实现推理速度飞跃
TensorFlow
26+阅读 · 2020年7月30日
注意力机制可解释吗?这篇ACL 2019论文说……
机器之心
11+阅读 · 2019年6月16日
基础 | 基于注意力机制的seq2seq网络
黑龙江大学自然语言处理实验室
16+阅读 · 2018年3月7日
深度学习中的注意力机制
CSDN大数据
24+阅读 · 2017年11月2日
关系推理:基于表示学习和语义要素
计算机研究与发展
19+阅读 · 2017年8月22日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
4+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
4+阅读 · 2015年12月31日
国家自然科学基金
2+阅读 · 2015年12月31日
国家自然科学基金
18+阅读 · 2012年12月31日
国家自然科学基金
18+阅读 · 2012年12月31日
VIP会员
相关VIP内容
大模型推理的天花板在哪里?
专知会员服务
15+阅读 · 2025年6月12日
大型语言模型推理引擎的综述:优化与效率的视角
专知会员服务
21+阅读 · 2025年5月13日
TransMLA:多头潜在注意力(MLA)即为所需
专知会员服务
23+阅读 · 2025年2月13日
迈向大型推理模型:基于大型语言模型的强化推理综述
专知会员服务
49+阅读 · 2025年1月17日
LlamaV-o1: 重新思考大语言模型中的逐步视觉推理
专知会员服务
17+阅读 · 2025年1月14日
大模型的模型压缩与有效推理综述
专知会员服务
43+阅读 · 2024年7月8日
大型语言模型的模型压缩与高效推理:综述
专知会员服务
94+阅读 · 2024年2月17日
相关基金
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
4+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
4+阅读 · 2015年12月31日
国家自然科学基金
2+阅读 · 2015年12月31日
国家自然科学基金
18+阅读 · 2012年12月31日
国家自然科学基金
18+阅读 · 2012年12月31日
Top
微信扫码咨询专知VIP会员