💬 NLP

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

ForesightKV is a training-based framework that optimizes KV cache eviction for reasoning models by combining a Golden Eviction algorithm with supervised learning and reinforcement learning (GRPO) to predict and retain the most critical KV pairs, thereby significantly reducing memory costs while maintaining high performance on long reasoning tasks.

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

Published 2026-02-04

📖 4 min read☕ Coffee break read

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a brilliant detective (the AI) trying to solve a very long, complex mystery. To do this, you need to keep a massive notebook of clues (the KV Cache) on your desk. As the story gets longer, your notebook grows thicker and thicker. Eventually, your desk becomes so cluttered with papers that you can't move your hands, and it takes forever to find the specific clue you need to make your next deduction. This is the problem of "memory overload" in large language models.

Traditionally, to clear the desk, people use simple rules: "Throw away the oldest papers" or "Throw away the papers with the least ink on them." But the paper ForesightKV argues that these rules are too dumb. Sometimes, an old paper or a paper with faint ink is actually the most critical clue for solving the mystery later in the story. Throwing it away causes the detective to make a mistake that ruins the rest of the investigation.

Here is how ForesightKV solves this problem, explained in three simple steps:

1. The Problem: The "One-Pass" Mistake

Imagine you are packing for a trip. A standard method might say, "Keep the last 10 items you packed, and throw away the rest." But what if the first item you packed (a map) is essential for the whole trip, even though it's far back in the suitcase?
Existing methods try to guess which clues to keep based on simple patterns (like "keep the most recent" or "keep the most popular"). The paper found that in complex reasoning (like math problems), clues have three types of personalities:

The Global Star: Always important, no matter what.
The Local Neighbor: Only important for a short time.
The Semantic Chameleon: This is the tricky one. A clue might seem useless now, but in 50 steps, it becomes the key to the whole puzzle. Simple rules miss these "chameleons."

2. The Solution: A "Time-Traveling" Coach

The authors created a new system called ForesightKV. Instead of using a rigid rulebook, they trained a tiny, smart assistant (a Scoring Model) to act like a coach who can see the future.

This coach doesn't just look at the clues right now; it learns to predict which clues will matter most in the future. It does this through a two-stage training process:

Stage 1: The "Golden Standard" (Supervised Learning)

First, the researchers ran the AI through a math problem without throwing anything away. They watched the AI's attention and asked: "If we had to throw away some clues right now, which ones would cause the least damage to the final answer?"
They used a method called Golden Eviction to find the perfect set of clues to keep. They then taught the tiny assistant to mimic this "perfect" decision-making process. It's like showing a student the answer key and saying, "Learn how to pick the right answers."

Stage 2: The "Real-World Simulation" (Reinforcement Learning)

However, the "perfect" answer key isn't always perfect when the AI is actually solving a problem live, because the AI's mind changes as it thinks.
So, the second stage is like a video game simulation. The assistant is told: "Go ahead and throw away some clues. If the AI makes a mistake on a specific type of word (like a number or a symbol that is easy to get wrong), you get a penalty. If the AI keeps solving correctly, you get a reward."
The assistant learns to be even smarter, realizing that keeping certain "boring" words (low-entropy tokens) is actually vital to prevent the AI from hallucinating numbers later on.

3. The Result: A Smarter, Lighter Desk

The paper tested this on three different AI models using hard math benchmarks (AIME 2024 and 2025).

The Claim: ForesightKV can solve these math problems just as well as the full, heavy notebook, but using only half the memory space.
The Analogy: It's like being able to carry a backpack that is 50% lighter, yet you still remember every single clue you need to win the game.
Speed: Because the backpack is lighter, the AI can think faster and handle more people at once (higher throughput).

Summary

ForesightKV is a new way to manage an AI's memory. Instead of blindly throwing away old notes, it uses a smart, trained assistant that learns to predict which notes will be crucial for the long-term solution. By combining a "perfect example" training phase with a "learn by doing" game phase, it keeps the AI fast and efficient without making it forget the important details needed to solve complex reasoning puzzles.

Technical Summary: ForesightKV

Problem Statement

Large Language Models (LLMs), particularly reasoning models, have demonstrated exceptional capabilities in generating long reasoning traces. However, as the sequence length increases, the Key-Value (KV) cache expands linearly, leading to significant memory overhead and increased latency due to memory-bound decoding. While existing KV cache eviction methods attempt to mitigate this by discarding less important KV pairs based on heuristic rules (e.g., attention scores, positions, or features), they often fail to capture complex, dynamic dependencies across different attention heads. Specifically, these training-free methods struggle with "semantic-dependent" patterns where the importance of a KV pair changes dynamically based on the evolving context. Furthermore, the eviction of critical KV pairs can cause sharp increases in loss for low-entropy tokens (deterministic parts of reasoning), leading to factual errors that distort subsequent reasoning steps.

Methodology

The authors propose ForesightKV, a training-based framework designed to learn the long-term contribution of KV pairs and adaptively evict less important ones during generation. The framework operates on a fixed memory budget and employs a two-stage training paradigm to optimize a lightweight scoring model ( $\pi_\theta$ ) that predicts the importance of each KV pair. Crucially, the LLM parameters remain frozen throughout; only the scoring model is updated.

1. Dynamic Eviction Process

The eviction is triggered periodically (every $L$ tokens). When the cache size reaches $B + L$ (where $B$ is the target budget), the system retains the most recent $L$ tokens and uses the scoring model to prune the remaining cache back to $B - L$ . The scoring model takes as input the concatenation of the key ( $k_n$ ), value ( $v_n$ ), and attention features ( $a_n$ ) to predict an importance score ( $\phi_n$ ). Eviction decisions are made via a Top-K followed by multinomial sampling strategy to introduce stochasticity.

2. Stage One: Supervised Learning with "Golden Eviction"

To train the scoring model, the authors introduce the Golden Eviction algorithm to construct supervision labels.

Mechanism: For a full reasoning trace, the attention matrix is partitioned into blocks. The algorithm computes "block scores" by aggregating attention scores within blocks and across attention heads.
Future Awareness: For each eviction step, the algorithm identifies the maximum block score across all future steps for each KV pair. This "future score" represents the optimal retention strategy that minimizes the impact on future attention computations.
Training: The scoring model is trained using a Pairwise Ranking Loss to ensure that the predicted scores for any two KV pairs align with the ranking of their future scores derived from the Golden Eviction algorithm.

3. Stage Two: Reinforcement Learning (RL)

Supervised training alone may not account for the distribution shift caused by eviction during inference. To bridge this gap, the authors formulate KV cache eviction as a Markov Decision Process (MDP) and apply the GRPO (Group Relative Policy Optimization) algorithm.

State: The current remaining KV cache.
Action: Selecting a subset of KV pairs to retain.
Reward: The reward function is designed to mitigate performance degradation on low-entropy tokens. Specifically, it focuses on tokens that originally had low entropy but experienced a significant increase in loss after eviction. The reward is defined as the negative Mean Squared Error (MSE) of the loss increases for this specific subset of tokens.
Optimization: The policy model is optimized to maximize the cumulative reward across the sequence, effectively learning to preserve the critical context needed for deterministic reasoning steps.

Key Contributions

ForesightKV Framework: A novel two-stage training method (Supervised Learning + RL) that learns to predict the long-term contribution of KV pairs, moving beyond static heuristic rules.
Golden Eviction Algorithm: A method to generate optimal eviction traces using future attention information, serving as high-quality supervision for the scoring model.
MDP Formulation for Eviction: Modeling the eviction process as an MDP and utilizing GRPO to specifically target the mitigation of loss spikes in low-entropy tokens, which are critical for reasoning accuracy.
Efficiency: The framework trains only a lightweight scoring model without requiring backward operations or parameter updates on the base LLM.

Experimental Results

The authors evaluated ForesightKV on three reasoning models (Qwen3-1.7B, Qwen3-4B, and DeepSeek-R1-Distill-Qwen-7B) using the AIME2024 and AIME2025 math benchmarks.

Performance: ForesightKV consistently outperformed existing baselines (SnapKV, H2O, R-KV). Notably, it achieved comparable or superior performance to full-cache baselines while using only half the KV cache budget. For instance, on AIME2024 with Qwen3-4B, ForesightKV with a 1K budget outperformed R-KV with a 2K budget (54.5 vs. 44.8 Pass@1).
Throughput: The method significantly improved throughput. With a 1K budget on a 32K sequence, ForesightKV achieved up to a 9.79× throughput gain compared to the full KV cache.
Generalization: Despite being trained primarily on math tasks, ForesightKV demonstrated strong generalization on GPQA (science) and LiveCodeBench (coding) benchmarks, outperforming baselines with half the cache budget.
Ablation Studies: The results confirmed that both the supervised training stage (providing initial predictive capability) and the reinforcement learning stage (refining the policy for low-entropy token preservation) are crucial for success. The specific reward design focusing on low-entropy tokens with large loss increases was found to be the most effective.

Significance and Claims

The paper claims that ForesightKV offers a promising new direction for adaptive KV cache management in long-context reasoning. By shifting from rule-based heuristics to a learning-based approach that explicitly models the long-term contribution of KV pairs and targets the preservation of low-entropy tokens, the method achieves a superior trade-off between memory efficiency and generation quality. The authors emphasize that their approach captures the intrinsic importance of KV cache entries rather than overfitting to specific budget constraints, allowing it to generalize well across different budget settings and task domains.