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

AI summary

Overview

• Research area: Natural Language Processing / efficient LLM inference, specifically KV cache compression for long-context reasoning models.
• Technical level: Intermediate. The paper builds on standard Transformer attention and KV cache concepts, plus supervised learning and reinforcement learning (GRPO), but explains each component in sequence.
• Scope: The paper proposes a two-stage trained scoring model that decides which key-value pairs to permanently evict during long reasoning generation, and evaluates it on three reasoning LLMs and five benchmarks.

What This Paper Is About

As reasoning models generate long chains of thought, their KV cache grows linearly, increasing memory and latency costs. Existing eviction methods rely on fixed rules (attention scores, positions, KV features) or make a single one-time judgment of importance, and they struggle when the importance of a KV pair changes as the reasoning trace evolves. ForesightKV instead trains a lightweight scorer to predict the long-term contribution of each KV pair and to make adaptive eviction decisions under a fixed cache budget.

Key Contributions

1. A dynamic, learned eviction framework. ForesightKV defines long-context generation as a periodic eviction process with a cache budget B and eviction length L, and uses a small MLP scorer over concatenated key, value, and attention features to score each KV pair.
2. Golden Eviction plus supervised training. A Golden Eviction algorithm partitions the attention matrix into blocks along the query dimension, pools across queries and attention groups, and uses the maximum future block score as a label. The scorer is trained to match the ranking of these labels with a Pairwise Ranking Loss.
3. Reinforcement learning formulation. Eviction is cast as a Markov Decision Process with state (current KV cache), action (which KV pairs to retain), policy (the scoring models) and a sequence-level reward; GRPO is used to optimize it, with the reward built from loss increases on low-entropy tokens.
4. A Top-K multinomial action parameterization. Eviction samples L pairs from the 2L lowest-scoring candidates, which the authors describe as balancing the stability of greedy top-K eviction against the exploration of pure multinomial sampling.

Main Findings

• Motivating measurement on KV cache size: For Qwen3-4B, at a sequence length of 32K, a single instance's KV cache consumes 4.5 GB at BFloat16 precision.
• Attention patterns fall into three types: Using Qwen3-4B on questions and self-generated reasoning traces (with STILL dataset questions), the authors categorize KV pairs as global (vertical lines in the attention map), position-dependent (locality; attention to earlier pairs decreases as decoding lengthens), and semantic-dependent (block-wise patterns with shifting high-attention regions, and pairs that become permanently irrelevant after a certain query). The patterns can co-occur within a single head.
• Low-entropy tokens suffer the most after eviction: With R-KV at a 1024-token budget, loss increases were +147% (math), +75% (code) and +187% (summarization) for low-entropy tokens, versus +52%, +1% and +142% for high-entropy tokens.
• Golden Eviction beats rule-based methods on loss: With Qwen3-4B, loss ratios relative to the original model were 1.0711 (1024,256), 1.0166 (2048,256), 1.0715 (1024,128) and 1.0185 (2048,128) for Golden, compared with 1.4101, 1.1606, 1.4750 and 1.1814 for R-KV; 1.4091, 1.1281, 1.4214 and 1.1343 for SnapKV; and 1.2730, 1.0948, 1.4106 and 1.1578 for H2O.
• Better accuracy at half the budget: On AIME2024 with Qwen3-4B, ForesightKV with a 1K budget reaches 54.5 versus 44.8 for R-KV with a 2K budget. The abstract and introduction report that ForesightKV preserves 92% and 99% of the original model's performance under 2K and 4K budget constraints respectively.
• Throughput and batch size gains (Qwen3-4B, one A800 GPU): The full model supports 48 concurrent batches and 139.39 throughput at 8K, 24 and 73.40 at 16K, and 11 and 37.73 at 32K. ForesightKV-1K supports 96 concurrent batches at every generation length, with throughput 375.10 (2.69x), 372.35 (5.07x) and 369.43 (9.79x) at 8K, 16K and 32K. ForesightKV-2K gives 70 concurrent batches with 272.48 (1.95x), 270.65 (3.69x) and 268.36 (7.11x). ForesightKV-4K gives 48 with 198.58 (1.42x), 196.32 (2.67x) and 193.95 (5.14x). ForesightKV-8K gives 36 with 114.32 (1.56x) and 112.99 (3.03x) at 16K and 32K.
• Reward design matters: With Qwen3-4B at a 1K budget, AIME24/AIME25 scores were 51.7/40.9 with no reward term, 50.6/40.0 for L_all, 53.5/40.4 for L_low, 49.6/35.4 for L_high, 53.8/42.3 for L_low,large, and 54.5/42.3 for the paper's MSE-based L_ours. Optimizing all-token loss did not help, and targeting high-entropy tokens degraded performance.
• Both inputs and sampling matter: At a 1K budget, attention plus KV features with Top-K+multinomial sampling scores 51.7/40.9 (AIME24/AIME25); attention features only scores 37.5/22.9; pure multinomial sampling scores 16.5/13.8; pure Top-K scores 46.0/37.7.
• Generalization to other reasoning domains: On GPQA and LiveCodeBench with Qwen3-4B, the full model scores 54.6 and 63.4. ForesightKV scores 45.2, 51.3, 53.7 on GPQA at 1K/2K/4K and 55.7, 61.1 on LiveCodeBench at 1K/2K. ForesightKV without RL scores 44.2, 51.2, 52.4 and 55.7, 61.5.
• Generalization to long-input tasks: On LongBench with Qwen3-4B at a 1K budget in non-thinking mode (compression applied once after prefilling), the overall averages are 39.41 (full), 38.95 (ForesightKV), 37.50 (SnapKV), 37.11 (R-KV) and 35.74 (H2O). ForesightKV leads on code completion (3.41 versus 2.29 for the full model) and is closest to the full model on single-document QA (41.07 versus 41.88) and summarization (27.39 versus 27.48).
• Budget generalization: The scoring models were trained with budgets B <= 2K in both stages, yet generalize to settings such as a 4K budget.

Methodology in Plain English

The framework runs generation with a fixed KV cache budget B and re-checks the cache every L new tokens. When the cache reaches B + L, the most recent L pairs are kept and a scorer decides which of the remaining pairs to discard, bringing the cache back to B - L retained pairs.

The scorer is a small MLP (intermediate size 16) that reads each KV pair's key, value, and attention features and outputs an importance score. Eviction is stochastic: the 2L lowest-scoring candidates are found, and L of them are sampled with probability proportional to their negative scores.

Training happens in two stages, with the language model frozen throughout:

1. Supervised stage. Golden Eviction runs on full reasoning traces to build labels. The attention matrix is split into blocks along the query dimension with stride L, starting at the first eviction position B + L. Block scores are pooled across queries and across heads within a group. For each eviction step, each KV pair gets a "future score" equal to its maximum block score over all future blocks, and the pairs with the largest future scores are treated as the ones that should be retained. The scorer is then trained with a pairwise ranking loss with margin m, so that its predicted ordering matches the ordering of the future scores.
2. Reinforcement stage. The authors argue the supervised labels do not reflect the distribution shift caused by eviction at inference time, so they refine the scorer with GRPO. For each sequence, G eviction traces are sampled with the old policy across all layers and heads, group-relative normalization converts rewards into advantages, and the standard clipped objective with KL regularization against the reference (supervised-initialized) policy is optimized. The reward is the negative mean squared loss increase over a subset of tokens that are both in the bottom 80% of original entropy and whose loss increase exceeds a threshold; this targets low-entropy tokens where eviction causes the largest damage.

Experiments use DeepSeek-R1-Distill-Qwen-7B, Qwen3-4B and Qwen3-1.7B on AIME2024 and AIME2025, with temperature 0.6, top-k 20, top-p 0.95, averaged pass@1 over 32 independent evaluations per benchmark. Budgets tested are 1024, 2048 and 4096 with eviction length 256. Baselines are SnapKV, H2O and R-KV, compressed every L steps as in R-KV.

Why This Matters

The paper reframes KV cache eviction as a learning problem coupled to the downstream outcome (reasoning accuracy and loss on fragile tokens), rather than a fixed heuristic. It reports the first joint use of supervised and reinforcement learning for the eviction process, and it explicitly ties the reward design to the finding that low-entropy tokens absorb most of the damage from eviction. Practically, it reports that a small learned scorer can hold batch size constant across generation lengths while delivering up to 9.79x throughput at 32K tokens.

Real-world applications:

• Long chain-of-thought reasoning services: Serving math, science and code reasoning assistants where output traces are long and batched concurrency is limited by KV memory.
• Cost-controlled inference platforms: Operators can fix a memory budget per request and predict the accuracy tradeoff, since the method generalizes from 2K-trained budgets to 4K.
• Long-document question answering and summarization: Evaluated on LongBench categories with a single post-prefill compression pass.
• Code assistants: LiveCodeBench-V3 and LongBench code completion are both reported, with code completion being the category where ForesightKV exceeds the full model (3.41 versus 2.29).

Industry relevance: the paper offers a deployable efficiency lever that does not require retraining or fine-tuning the base LLM, since only the lightweight scorers are updated; code is released at https://github.com/RUCAIBox/ForesightKV.

Future Directions

• Reducing the cost of the Golden Eviction label construction, which requires computing full attention over complete reasoning traces and pooling over blocks.
• Broadening training domains. The scoring models were trained on math tasks; LongBench and GPQA/LiveCodeBench results are presented as generalization evidence, so training directly on other domains is an open question.
• Extending the reward and eviction formulation. The authors frame the reinforcement learning perspective as opening a new direction for adaptive KV cache management; alternative reward signals, entropy thresholds and eviction frequencies remain to be explored.
• Scaling and deployment questions not reported here. The paper does not report results for budgets below 1K, for generation lengths beyond 32K, or on GPUs other than the A800 used in the efficiency study.

Target Audience

Researchers and engineers working on LLM inference efficiency, KV cache compression, and long-context reasoning. It is most useful to readers who already understand Transformer attention and KV caching and want to see how supervised ranking plus GRPO can be applied to a systems-level decision problem. Readers focused purely on model architecture or pretraining will find it less directly relevant, though the attention-pattern taxonomy in the empirical study is broadly useful.

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