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

📄 arXiv: 2602.03203v1 📥 PDF

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

分类: cs.CL, cs.LG

发布日期: 2026-02-03

💡 一句话要点

ForesightKV：通过学习长期贡献优化推理模型KV缓存淘汰

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: KV缓存淘汰 长文本推理 大型语言模型 注意力机制 强化学习

📋 核心要点

  1. 现有KV缓存淘汰方法难以捕捉复杂KV依赖关系，导致长文本推理性能下降。
  2. ForesightKV通过学习预测KV对的长期贡献，从而优化缓存淘汰策略。
  3. 实验表明，ForesightKV在更小的缓存预算下，显著优于现有方法，提升推理性能。

📝 摘要（中文）

大型语言模型（LLMs）通过生成长推理轨迹展现了卓越的推理能力。然而，随着序列长度的增长，键值（KV）缓存线性扩展，导致显著的内存和计算成本。现有的KV缓存淘汰方法通过丢弃不太重要的KV对来缓解这个问题，但通常无法捕捉复杂的KV依赖关系，从而导致性能下降。为了更好地平衡效率和性能，我们引入了ForesightKV，这是一个基于训练的KV缓存淘汰框架，它学习预测在长文本生成过程中要淘汰哪些KV对。我们首先设计了Golden Eviction算法，该算法使用未来的注意力分数来识别每个步骤的最佳淘汰KV对。然后，通过监督训练和Pairwise Ranking Loss来提炼这些轨迹和每个步骤的分数。此外，我们将缓存淘汰形式化为马尔可夫决策过程，并应用GRPO算法来减轻低熵token上显著的语言建模损失增加。在三个推理模型的AIME2024和AIME2025基准测试中进行的实验表明，ForesightKV在只有一半缓存预算的情况下始终优于先前的方法，同时受益于监督学习和强化学习方法的协同作用。

🔬 方法详解

问题定义：论文旨在解决大型语言模型在长文本推理时，KV缓存线性增长带来的内存和计算成本问题。现有KV缓存淘汰方法无法有效捕捉KV对之间的复杂依赖关系，导致推理性能下降，尤其是在需要长期记忆的推理任务中。

核心思路：论文的核心思路是通过学习KV对的长期贡献度来指导缓存淘汰。不同于以往只关注当前时刻KV对重要性的方法，ForesightKV预测未来注意力分数，从而评估KV对在整个推理过程中的价值，并优先淘汰贡献度低的KV对。

技术框架：ForesightKV框架包含以下几个主要阶段：1) Golden Eviction算法：利用未来注意力分数，离线确定每个步骤的最佳淘汰KV对，生成训练数据。2) 监督训练：使用Pairwise Ranking Loss，训练模型预测KV对的重要性排序，模仿Golden Eviction算法的行为。3) 强化学习优化：将缓存淘汰过程建模为马尔可夫决策过程，利用GRPO算法优化低熵token上的语言建模损失，进一步提升性能。

关键创新：ForesightKV的关键创新在于其前瞻性（foresight）的KV对重要性评估机制。它不是简单地基于当前注意力分数进行淘汰，而是预测KV对在未来推理步骤中的贡献，从而做出更明智的淘汰决策。这种方法能够更好地保留对长期推理有用的信息，提高整体性能。

关键设计：在监督训练阶段，使用了Pairwise Ranking Loss来学习KV对的重要性排序。在强化学习阶段，采用了GRPO算法来缓解低熵token上的损失增加问题。Golden Eviction算法通过计算未来多个步骤的注意力分数来确定最佳淘汰KV对，具体实现细节（例如，未来步数、注意力分数计算方式）需要在实验中进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，在AIME2024和AIME2025基准测试中，ForesightKV在仅使用一半缓存预算的情况下，始终优于现有的KV缓存淘汰方法。同时，监督学习和强化学习的结合进一步提升了ForesightKV的性能，验证了该方法的有效性。

🎯 应用场景

ForesightKV可应用于各种需要长文本推理的大型语言模型，例如问答系统、文档摘要、代码生成等。通过降低内存和计算成本，该方法可以使这些模型在资源受限的环境中运行，并提高推理效率。此外，该研究为KV缓存管理提供了新的思路，可以促进相关领域的发展。

📄 摘要（原文）

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.