Research map / Large language models / Reasoning and chain of thought

Learning to Evict from Key-Value Cache

ICML 2026 · Large language models › Reasoning and chain of thought · approach: KV cache compression

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes KVP (KV Policy), per-head lightweight RL agents that learn to rank KV cache entries by predicted future attention using only keys, values and positions, trained offline on pre-computed traces with a budget-agnostic AUC-style reward.

Problem

KV cache eviction in LLMs relies on heuristics (recency, past attention scores) that are only indirect proxies for a token's future utility, and attention-based methods add overhead and are incompatible with FlashAttention.

Approach

Reframe eviction as learning a budget-agnostic ranking of cache entries via RL (REINFORCE with RLOO on a Plackett-Luce/Gumbel-Sort policy), with a reward equal to the negative normalized area under the eviction-cost-vs-budget curve, computed offline from future attention in training traces; one small MLP per KV head uses only K, V and position.

Key results

• On RULER-4k, KVP consistently outperforms all attention-free baselines and retains higher accuracy as cache budget shrinks (in-distribution).
• Trained at ~4K context, KVP still dominates every baseline at every budget on RULER-128K (32x extrapolation).
• Per-budget eviction cost of KVP matches attention-based TOVA and SnapKV on OASST2 test set while using no query information; at 10k context, a single-layer compression takes 0.71 ms vs 404 ms for full-model prefill.

Surprising findings

• Supervised magnitude-sensitive learning-to-rank losses (regression, pairwise, listwise) underperform RL because they overfit attention sinks, while a rank-only Soft-Sort surrogate spikes at tight budgets.
• Per-head agents learn markedly different rankings (mean rank variance per position 294.8), indicating heads need distinct policies.

Limitations

• Uniform compression budget across heads and layers; no adaptive budget allocation.
• Separate agent per head ignores cross-head dependencies.
• Reward uses future attention as proxy for task utility rather than direct downstream signal; implementation lacks optimized end-to-end inference pipeline.

Experimental setup

Datasets: RULER-4k OASST2-4k RULER-128K LongBench BoolQ ARC-Challenge MMLU HellaSwag

Benchmarks: RULER-4k RULER-128K OASST2-4k LongBench GovReport LongBench passage retrieval (EN/ZH) BoolQ ARC-Challenge MMLU HellaSwag

Metrics: accuracy perplexity L-Rouge per-budget eviction cost (-R^b) wall-clock latency FLOPs

Baselines: TOVA SnapKV StreamingLLM LagKV KeyDiff K-Norm Random JudgeQ RankNet ListNet Soft-Sort

Builds on: REINFORCE RLOO Gumbel-Sort Plackett-Luce model StreamingLLM H2O

Closest papers by meaning

• ForesightKV: Optimizing KV Cache Eviction for Reasoning Models by Learning Long-Term Contribution (ICML 2026 · similarity 0.92)
• Taming the Fragility of KV Cache Eviction in LLM Inference (ICLR 2026 · similarity 0.90)
• CAKE: Cascading and Adaptive KV Cache Eviction with Layer Preferences (ICLR 2025 · similarity 0.90)
• Ada-KV: Optimizing KV Cache Eviction by Adaptive Budget Allocation for Efficient LLM Inference (NeurIPS 2025 · similarity 0.90)
• Get More with LESS: Synthesizing Recurrence with KV Cache Compression for Efficient LLM Inference (ICML 2024 · similarity 0.89)
• KVzip: Query-Agnostic KV Cache Compression with Context Reconstruction (NeurIPS 2025 · similarity 0.89)