Abstract
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.
Abstract (translated)
近期，大型语言模型（LLM）通过生成长推理轨迹展示了卓越的推理能力。然而，随着序列长度的增长，关键值（KV）缓存会线性扩展，导致显著增加内存和计算成本。现有的KV缓存淘汰方法通过丢弃不重要的KV对来缓解这一问题，但通常无法捕捉复杂的KV依赖关系，从而导致性能下降。为了更好地平衡效率与性能，我们引入了ForesightKV，这是一种基于训练的KV缓存淘汰框架，能够学习在长文本生成过程中预测应被淘汰的KV对。 首先，我们设计了Golden Eviction算法，在每个步骤使用未来注意力分数来识别最佳的淘汰KV对。然后，通过带有成对排名损失（Pairwise Ranking Loss）的监督训练，将这些轨迹和每个步骤中的分数进行提炼。此外，我们将缓存淘汰形式化为马尔可夫决策过程，并应用GRPO算法以减轻低熵令牌上显著的语言建模损耗增加。 在三个推理模型上的AIME2024和AIME2025基准测试中，实验结果表明，ForesightKV在只有半的缓存预算下始终优于先前的方法，并且能够从监督学习和强化学习方法中协同受益。
URL
https://arxiv.org/abs/2602.03203
