Recently, large language models (LLMs) have shown remarkable reasoning abilities by producing long reasoning traces. However, as the sequence length grows, the key-value (KV) cache expands linearly, incurring significant memory and computation costs. Existing KV cache eviction methods mitigate this issue by discarding less important KV pairs, but often fail to capture complex KV dependencies, resulting in performance degradation. To better balance efficiency and performance, we introduce ForesightKV, a training-based KV cache eviction framework that learns to predict which KV pairs to evict during long-text generations. We first design the Golden Eviction algorithm, which identifies the optimal eviction KV pairs at each step using future attention scores. These traces and the scores at each step are then distilled via supervised training with a Pairwise Ranking Loss. Furthermore, we formulate cache eviction as a Markov Decision Process and apply the GRPO algorithm to mitigate the significant language modeling loss increase on low-entropy tokens. Experiments on AIME2024 and AIME2025 benchmarks of three reasoning models demonstrate that ForesightKV consistently outperforms prior methods under only half the cache budget, while benefiting synergistically from both supervised and reinforcement learning approaches.


翻译:近年来,大型语言模型通过生成长推理轨迹展现出卓越的推理能力。然而,随着序列长度增加,键值缓存呈线性扩张,导致显著的内存与计算开销。现有的键值缓存淘汰方法通过丢弃重要性较低的键值对来缓解此问题,但往往难以捕捉复杂的键值依赖关系,导致性能下降。为更好地平衡效率与性能,本文提出ForesightKV——一种基于训练的键值缓存淘汰框架,通过学习预测长文本生成过程中应淘汰的键值对。我们首先设计Golden Eviction算法,该算法利用未来注意力分数在每一步识别最优淘汰键值对。随后通过监督训练与成对排序损失函数,对这些轨迹及每一步的分数进行知识蒸馏。此外,我们将缓存淘汰建模为马尔可夫决策过程,并应用GRPO算法以缓解低熵词元上语言建模损失的大幅上升。在AIME2024与AIME2025基准测试中对三种推理模型的实验表明,ForesightKV在仅使用半数缓存预算的情况下持续优于现有方法,同时从监督学习与强化学习方法中获得了协同增益。

0
下载
关闭预览

相关内容

在数学和计算机科学之中,算法(Algorithm)为一个计算的具体步骤,常用于计算、数据处理和自动推理。精确而言,算法是一个表示为有限长列表的有效方法。算法应包含清晰定义的指令用于计算函数。 来自维基百科: 算法
小型推理模型简要综述:训练、推理、应用与研究方向
专知会员服务
42+阅读 · 2025年4月16日
大模型的模型压缩与有效推理综述
专知会员服务
43+阅读 · 2024年7月8日
专知会员服务
16+阅读 · 2021年6月4日
一种关键字提取新方法
1号机器人网
21+阅读 · 2018年11月15日
关系推理:基于表示学习和语义要素
计算机研究与发展
19+阅读 · 2017年8月22日
国家自然科学基金
18+阅读 · 2017年12月31日
国家自然科学基金
2+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
18+阅读 · 2012年12月31日
VIP会员
相关VIP内容
小型推理模型简要综述:训练、推理、应用与研究方向
专知会员服务
42+阅读 · 2025年4月16日
大模型的模型压缩与有效推理综述
专知会员服务
43+阅读 · 2024年7月8日
专知会员服务
16+阅读 · 2021年6月4日
相关基金
国家自然科学基金
18+阅读 · 2017年12月31日
国家自然科学基金
2+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
18+阅读 · 2012年12月31日
Top
微信扫码咨询专知VIP会员