arXivDaily arXiv每日学术速递 周一至周五更新
arXiv周末暂无论文更新,休息一下吧,周末愉快~~
arXiv 2602.03203cs.CLcs.LG
收藏

ForesightKV: 通过学习长期贡献优化推理模型的KV缓存驱逐

ForesightKV: Optimizing KV Cache Eviction for Reasoning Models by Learning Long-Term Contribution

Zhejiang University(浙江大学)

Zican Dong, Peiyu Liu, Junyi Li, Zhipeng Chen, Han Peng, Shuo Wang, Wayne Xin Zhao

AI总结:

提出ForesightKV框架,通过监督学习和强化学习预测KV对在长文本生成中的重要性,在仅用一半缓存预算下优于现有方法。

摘要:

近期,大型语言模型通过生成长推理轨迹展现出卓越的推理能力。然而,随着序列长度增长,键值缓存线性扩展,导致显著的内存和计算成本。现有的KV缓存驱逐方法通过丢弃不重要的KV对来缓解此问题,但往往难以捕捉复杂的KV依赖关系,导致性能下降。为了更好地平衡效率与性能,我们引入了ForesightKV,一种基于训练的KV缓存驱逐框架,学习在长文本生成过程中预测哪些KV对应该被驱逐。我们首先设计了Golden Eviction算法,该算法利用未来注意力分数在每一步识别最优的驱逐KV对。然后通过监督训练,使用成对排序损失将这些轨迹和每一步的分数进行蒸馏。此外,我们将缓存驱逐建模为马尔可夫决策过程,并应用GRPO算法来缓解低熵令牌上语言模型损失显著增加的问题。在三个推理模型的AIME2024和AIME2025基准测试上的实验表明,ForesightKV在仅用一半缓存预算的情况下始终优于先前方法,同时从监督学习和强化学习方法中协同受益。代码可在https://github.com/RUCAIBox/ForesightKV获取。

英文摘要:

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. Code is available at https://github.com/RUCAIBox/ForesightKV.

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

↑