arXiv:2602.03203cs.CLcs.LG2026-02被引 3

通过学习长期贡献优化推理模型的键值缓存淘汰策略。

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

  • 基于未来注意力分数设计黄金淘汰算法,精准识别应淘汰的键值对。
  • 在半量缓存预算下,三个推理模型在AIME2024/2025上均超越现有方法。
  • 融合监督与强化学习,适合需要高效长文本推理的场景。

大型语言模型(LLMs)在生成长推理链时展现出强大能力,但序列长度增长导致键值(KV)缓存线性膨胀,带来显著内存和计算开销。现有缓存淘汰方法常因忽略复杂的KV依赖关系而性能下降。为此,我们提出ForesightKV,一种基于训练的缓存淘汰框架,可学习预测长文本生成中应淘汰的KV对。首先设计黄金淘汰(Golden Eviction)算法,利用未来注意力分数每步识别最优淘汰对象;再通过成对排序损失进行监督蒸馏。此外,将缓存淘汰建模为马尔可夫决策过程,并采用GRPO算法缓解低熵词上的语言建模损失激增问题。在三种推理模型的AIME2024和AIME2025基准测试中,ForesightKV在仅一半缓存预算下持续优于先前方法,且受益于监督与强化学习的协同作用。代码已开源:https://github.com/RUCAIBox/ForesightKV。

原文摘要 · Abstract (English)

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.

缓存优化推理模型强化学习大模型

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。