Learning to Evict from Key-Value Cache

AI summary

Overview

Research area: Efficient large language model inference, specifically Key-Value (KV) cache management, framed as a reinforcement learning and learning-to-rank problem.

Technical level: Advanced. The paper assumes familiarity with transformer attention, autoregressive decoding, RL policy gradients, and learning-to-rank.

Scope: The paper introduces KVP, a framework that trains lightweight per-head RL agents to rank cached KV entries by their predicted future usefulness, enabling budget-agnostic eviction that outperforms heuristic and attention-based baselines across short and long-context benchmarks.

What This Paper Is About

Autoregressive LLM inference stores the keys and values of every previously seen token in a growing KV cache, which can consume tens or hundreds of gigabytes for long contexts. Prior eviction methods decide what to discard using heuristics like "keep recent tokens" or "keep tokens that received high past attention," but these are indirect proxies for what a token will actually contribute to future decoding. This paper reframes eviction as a reinforcement learning problem, training small per-head agents to directly predict a token's future utility and rank all cached entries accordingly.

Key Contributions

1. Reframing eviction as ranking. The authors prove that under a uniqueness and nestedness condition (which their additive reward guarantees), the optimal subset for every cache budget is realized by a single budget-agnostic total ordering of tokens. This reduces a combinatorial subset-selection problem to learning one scoring function.

2. The KVP framework. A system of lightweight per-head RL agents, each an MLP scoring function over (key, value, position) only, trained with a Plackett-Luce ranking policy and the REINFORCE-with-RLOO estimator. No queries, no attention scores, no architectural changes to the LLM are required.

3. A budget-agnostic reward. A holistic reward equal to the negated normalized area under the eviction-cost-vs-budget curve, computed entirely from pre-collected generation traces and future attention statistics, so training requires no live LLM inference.

4. Empirical validation at scale. KVP significantly outperforms strong attention-based (TOVA, SnapKV) and attention-free (StreamingLLM, KeyDiff, K-Norm, LagKV, Random) baselines on RULER-4k, OASST2-4k, and generalizes zero-shot to RULER-128K (32× the training length), LongBench passage retrieval (EN/ZH), GovReport, BoolQ, ARC-Challenge, MMLU, and HellaSwag.

Main Findings

• Beats baselines at every budget. On RULER-4k accuracy and OASST2-4k perplexity, KVP retains the highest accuracy and lowest perplexity across nearly all cache sizes, including beating attention-based methods that use query-specific information KVP is not allowed to see.

• Distills attention behavior without attention. In per-budget eviction-cost analysis on OASST2, KVP lands in the same low-cost cluster as TOVA and SnapKV despite using only static keys, values, and positions, showing the learned policy recovers much of the benefit of query-aware attention ranking.

• Length extrapolation works. Trained on roughly 4K-token sequences, KVP retains its lead on RULER-128K across 13 subtasks and every tested cache budget, indicating the ranking captures length-agnostic token-utility structure rather than training-context artifacts.

• Zero-shot transfer is strong and specialized. KVP_S (trained on OASST2 conversational data) leads on BoolQ and GovReport; KVP_R (trained on RULER) leads on LongBench passage retrieval in English and Chinese. Both rank at or near the top across benchmarks, showing the policy transfers out of distribution while still specializing to the training workload.

• Per-head policies are necessary. Some fixed heuristics work well for certain heads (e.g., layer 22, head 0 with StreamingLLM) but fail badly on others (layer 19, head 0; layer 2, head 0 with KeyDiff), confirming that a single global heuristic cannot capture the diversity of head-level attention patterns.

• RL beats supervised learning-to-rank. Against pointwise regression, RankNet, ListNet, Soft-Sort, and Weighted Soft-Sort, RL achieves the lowest eviction cost across nearly the entire budget range. Magnitude-sensitive losses waste capacity on a few dominant attention sinks and never learn the ordering of non-sink tokens; the magnitude-invariant Soft-Sort surrogate tracks RL in the bulk of the curve but spikes at tight budgets where top-ranked swaps matter most.

Methodology in Plain English

The authors start by showing that if each budget's ideal token subset nests inside the next-larger budget's ideal subset, then finding the best subset for every possible cache size is the same as producing one good sorted list of tokens. This lets them replace a hard subset-selection problem with a ranking problem.

They then train a small neural network for each attention head. The network takes only the key vector, the value vector, and the token's position, and outputs a score. Higher-scored tokens are kept; lower-scored tokens are evicted when the budget runs short. To turn scores into a distribution over orderings, they use the Plackett-Luce model, and they sample entire permutations in one shot via Gumbel-Sort, which makes training parallelizable.

The reward is designed so the agent is penalized for putting tokens that will receive lots of future attention low in the ranking. Concretely, they measure how much future attention mass would be lost if only the top-b tokens were kept, integrate this cost over all budgets b, and normalize by the cost of the optimal ordering. This single number summarizes how good the ranking is for every possible memory limit at once.

Training is fully offline. The authors first run the base LLM over a training corpus and store queries, keys, and values. They then repeatedly sample sequences from this dataset, have each per-head agent generate several candidate rankings, compute the reward from stored future attention, and update the agent with REINFORCE using a leave-one-out baseline for variance reduction. At inference, agents simply score and sort the cached entries, and the lowest-ranked tokens are discarded to meet the target budget.

Why This Matters

Impact on research. This is one of the first frameworks to treat KV cache eviction as a genuine reinforcement learning and ranking problem rather than a heuristic or backward-looking compression scheme. It shows that a query-free, offline-trained policy can match attention-based methods that use strictly more information, which shifts the design space for future work on memory-efficient inference. It also provides a principled signal that could feed into orthogonal techniques such as offloading to CPU memory, dynamic per-head budget allocation, and context gisting.

Real-world applications.

• Long-document summarization and question answering where the context cannot fit comfortably in GPU memory.
• Multi-turn chatbots that must keep a dialogue history while staying within fixed hardware limits.
• Retrieval-augmented generation with large retrieved passages, where the same document may be reused across many different queries.
• On-device or edge LLM deployment, where memory budgets are especially tight and per-token costs are directly tied to feasibility.

Industry relevance. The method is drop-in: it requires no changes to the LLM architecture, no attention recomputation at inference, and no extra forward passes. It is validated on models from two families (Qwen2.5-7B-Chat and Phi-4 14B), a natural fit for production serving stacks that already prefill KV caches and need to trade memory for quality on a per-request basis. The Apple affiliation and released code suggest near-term applicability to on-device inference, where per-token KV memory is a dominant cost.

Future Directions

• Dynamic, non-uniform budgets. Since KVP produces per-head rankings, it could drive head- or layer-specific cache budgets instead of the uniform allocation used in these experiments, potentially unlocking further savings.
• Bridging the RL vs. supervised gap. The authors note that a rank-dependent rescaling of supervised losses could in principle match RL, but its shape depends on per-head importance distributions and would need re-tuning, motivating future work on adaptive supervised surrogates.
• Composition with other techniques. Combining KVP with quantization, low-rank KV approximations, state merging (e.g., MorphKV), or memory-hierarchy offloading is explicitly flagged as a promising orthogonal direction.
• Beyond eviction. The learned future-utility ranking could inform context gisting (choosing which tokens to summarize), page-level offloading decisions, and summary-token compression methods such as Activation Beacon.

Target Audience

Researchers and engineers working on efficient LLM inference, KV cache management, and long-context modeling, especially those interested in reinforcement learning for systems problems or learning-to-rank applied to attention. It will also appeal to practitioners deploying LLMs under strict memory budgets (mobile, edge, or single-GPU serving) who want a drop-in alternative to heuristic eviction, and to readers tracking how RL can be used offline on pre-collected traces to train small, task-specialized auxiliary models.

Authors’ abstract

The growing size of Large Language Models (LLMs) makes efficient inference challenging, primarily due to the memory demands of the autoregressive Key-Value (KV) cache. Existing eviction or compression methods reduce cost but rely on heuristics, such as recency or past attention scores, which serve only as indirect proxies for a token's future utility and introduce computational overhead. We reframe KV cache eviction as a reinforcement learning (RL) problem: learning to rank tokens by their predicted usefulness for future decoding. To this end, we introduce KV Policy (KVP), a framework of lightweight per-head RL agents trained on pre-computed generation traces using only key and value vectors. Each agent learns a specialized eviction policy guided by a holistic reward, derived from future utility, that evaluates the quality of the ranking across all cache budgets, requiring no modifications to the underlying LLM or additional inference. Evaluated across two model families on the long-context benchmark RULER (up to 128K tokens) and the multi-turn dialogue benchmark OASST2-4k, KVP significantly outperforms strong baselines. Zero-shot tests on standard downstream tasks (BoolQ, LongBench passage retrieval, GovReport) further show that KVP generalizes beyond its training distribution and to considerably longer sequence lengths. These results demonstrate that learning to predict future token utility is a powerful and scalable paradigm for adaptive KV cache management.