PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning to Evict from Key-Value Cache

Luca Moschella, Laura Manduchi, Ozan Sener

arXiv:2602.10238v2cs.CLcs.LG

TL;DR

Long-context inference is constrained by KV-cache memory, and existing eviction heuristics use indirect proxies for future token usefulness. KVP learns lightweight per-head rankings from cached features, and it outperforms strong heuristic baselines while generalizing to longer contexts and unseen tasks.

  • Problem

    KV caches grow with sequence length and can consume tens or hundreds of gigabytes, while heuristic eviction relies on backward-looking proxies rather than directly optimizing tokens’ future utility.

    Source: §1 · pp. 1–2

  • Method

    KVP trains lightweight per-head RL agents offline to rank tokens using cached keys, values, and positions, with a reward evaluating rankings across cache budgets.

    Source: §1 · p. 2, §3.1 · p. 5, §3.2 · p. 6

  • Results

    KVP significantly outperforms strong heuristic baselines across two model families and long-context and dialogue benchmarks, with strong zero-shot performance out of distribution at up to 32× longer contexts.

    Source: §1 · pp. 2–3, §5 · p. 11

  • Takeaways & Limitations

    The evaluations support learned, head-specific token ranking as an offline alternative to inference-time attention access and hand-tuned eviction heuristics.

    Source: §5 · p. 11

    KVP uses future attention as a proxy for task-level utility rather than direct downstream task feedback.

    Source: §5 · p. 11

Abstract

from arXiv · show

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.

1 Introduction

KV cache growth makes long-context inference memory-intensive, while heuristic eviction policies use indirect, backward-looking signals for future token utility. KVP recasts eviction as learning to rank tokens and reports stronger long-context performance than heuristic baselines, including generalization beyond its training distribution.

1 Introduction

  • KV caches grow with sequence length and can consume tens or hundreds of gigabytes, creating a memory bottleneck for long-context inference.Existing management methods include eviction, quantization, and low-rank approximations.Source: §1 · p. 1
  • Hand-crafted eviction policies use indirect proxies for future importance, which can miss dynamic, content-dependent utility and require costly attention-statistics computation.The paper also notes that post-hoc attention-based compression is incompatible with efficient implementations such as FlashAttention.Source: §1 · pp. 1–2
  • The learned ranking targets future token utility, which the paper illustrates as a complex, non-local structure that a fixed sink-and-recency heuristic does not closely recover.The figure compares true cumulative future attention with the heuristic and learned rankings.Source: §1 · p. 2, Figure 1 · p. 2
  • KVP significantly outperforms strong heuristic baselines across RULER and OASST2-4k evaluations spanning two model families, and retains strong out-of-distribution performance at up to 32× longer contexts.The introduction also reports scaling to contexts up to 128K tokens and generalization to unseen domains.Source: §1 · pp. 2–3
  • The method trains lightweight, specialized RL agents per KV head using only keys and values, with no attention information or architectural changes to the LLM.Agents train efficiently on pre-computed generation traces without additional inference.Source: §1 · p. 2
  • Its reward evaluates a token ranking across all cache budgets by measuring relevant future information that would be erroneously evicted.Eviction then discards the lowest-ranked entries to meet a chosen memory budget.Source: §1 · p. 2

2 Related Work

Prior KV-cache management includes eviction, memory offloading, and representation compression, while learned approaches are less common. KVP differs by learning a fine-grained, per-head token ranking offline and can complement other cache-management techniques.

2 Related Work

  • Existing attention-based eviction uses sinks, recency, or attention structure, while KVP learns a forward-looking, query-independent policy instead of inferring utility from past attention or the current query.The discussion also situates machine-learning approaches to dynamic Transformer KV-cache management as an emerging area.Source: §2 · p. 3
  • Memory-hierarchy methods offload KV entries to slower storage and incur reload latency, while representation compression reduces per-entry cost; the paper describes these approaches as complementary to eviction.It suggests filtering less useful tokens before applying state merging to the remaining entries.Source: §2 · p. 3
  • Learned context compression may summarize tokens, whereas KVP ranks tokens individually; unlike query-dependent methods, KVP uses cached keys, values, and positions and trains offline on pre-collected traces.The related-work discussion contrasts KVP with methods that use query vectors or require full LLM forward passes.Source: §2 · p. 3
  • KVP frames eviction as an RL ranking problem, using lightweight per-head policies and a budget-agnostic reward to predict future token utility.The paper presents the ranking as complementary to offloading, context gisting, and dynamic budget allocation.Source: §2 · p. 4

3 KV Policy (KVP)

KVP turns KV-cache eviction into learning a single ranking that can serve arbitrary cache budgets. Per-head agents score cache entries, and an offline reward evaluates the ranking across budgets using future attention importance.

3 KV Policy (KVP)

  • KVP learns a budget-agnostic ranking whose top-b tokens form the cache for any budget, relying on a reward designed for unique, nested optimal subsets.Under these conditions, eviction is equivalent to ranking.Source: §3 · p. 4

3.1 KV Cache Eviction Agent

  • Each per-head agent scores tokens from their key, value, and original position vectors, then samples rankings with Gumbel-Sort for parallel, single-pass sampling.The scoring policy needs no future information, previous attention scores, or query embedding.Source: §3.1 · pp. 4–5, §3 · p. 4

3.1.1 Global Reward for Offline RL

  • The global reward measures the future-attention importance of evicted tokens across cache budgets and evaluates a ranking by its cost-versus-budget area under the curve.The reward is normalized by the cost of the optimal ranking and requires no additional LLM inference.Source: §3.1.1 · p. 5

3.2 Per-Head RL Agent and Efficient Training

  • Agents train offline on precomputed traces with policy gradients and a leave-one-out baseline, then rank each head’s KV entries for budget-based eviction.The training traces store queries, keys, and values, but not attention matrices.Source: §3.2 · p. 6

4 Evaluation

KVP is evaluated against heuristic and attention-based cache-management methods across long-context, dialogue, and zero-shot downstream benchmarks. It performs strongly across cache budgets and contexts, while ablations examine its reward design and use of reinforcement learning.

4 Evaluation

  • The evaluation uses Qwen2.5-7B-Chat, training agents on RULER-4k and OASST2-4k, then testing on longer-context and downstream benchmarks.KVP is not trained or fine-tuned on the generalization tasks.Source: §4 · pp. 6–7

4.1 Downstream Performance

  • KVP outperforms existing heuristics across budgets, with higher RULER-4k accuracy and lower OASST2-4k perplexity across most cache sizes.The comparison includes attention-based methods that use query-specific information at compression time.Source: §4 · p. 6, Figure 2 · p. 7, §4.1 · p. 8

4 Evaluation

  • Short-context tests compress once after prefill, whereas long-context tests repeatedly compress each chunk while maintaining the target cache budget.Budgets are uniform across heads and layers; attention-based baselines use prefill attention scores, including query-specific information unavailable to KVP.Source: §4 · pp. 7–8

4.1 Downstream Performance

  • Trained on ∼4K-token sequences, KVP leads every baseline at every cache budget on RULER-128K, a 32× longer context.Evaluation uses chunked prefill-compress with physical tensor compaction.Source: §4.1 · p. 8, Figure 3 · p. 9
  • On zero-shot tasks, KVPS leads on BoolQ and GovReport, while KVPR leads on English and Chinese LongBench retrieval, aligning with their training distributions.Both agents remain competitive across the diverse benchmarks as cache budgets shrink.Source: Figure 4 · p. 9, §4.1 · p. 9, Figure 5 · p. 10

4.2 Ablations

  • KVP reaches loss comparable to TOVA and SnapKV without their query information, while head-level analyses show fixed heuristics can help some heads and harm others.The reward ablation measures negative per-budget reward on unseen test data.Source: §4.2 · pp. 10–11
  • Supervised surrogates exhibit distinct budget tradeoffs: magnitude-sensitive losses miss non-sink ordering, while rank-only Soft-Sort spikes at tight budgets.Policy gradients optimize eviction-cost AUC directly, weighting swaps by their contribution to budget cost.Source: §4.2 · p. 11

5 Conclusions

KVP frames KV cache eviction as learnable sorting by future token utility, using lightweight per-head policies trained offline. The paper reports that these policies outperform heuristic baselines while preserving model capabilities, and identifies design choices and evaluation costs that remain open.

5 Conclusions

  • KVP trains lightweight, per-head policies to rank cache entries by future utility through an inference-free RL objective, without architectural changes.The conclusion presents this as adaptive memory management learned offline.Source: §5 · p. 11
  • KVP consistently outperforms heuristic baselines across short- and long-context evaluations while preserving the underlying model’s capabilities.Evaluations include RULER at 128K tokens and LongBench passage retrieval.Source: §5 · p. 11
  • The current design uses uniform compression budgets across heads and layers, does not model cross-head dependencies, and uses future attention as a proxy for task-level utility.The authors identify adaptive budget allocation, joint policy optimization, and task feedback as future-work directions.Source: §5 · p. 11
  • The implementation lacks an optimized end-to-end inference pipeline, so the paper reports policy-level wall-clock costs rather than full-system latency.Source: §5 · p. 11

A.1 Missing Proof of Proposition 1

The proof constructs an ordering from a nested sequence of uniquely specified subsets, then defines a scoring function consistent with that ordering.

A.1 Missing Proof of Proposition 1

  • Uniqueness implies that the difference between consecutive specified subsets contains exactly one element.Otherwise, two distinct intermediate subsets would contradict uniqueness.Source: §A.1 · p. 16
  • The construction yields a sequence of distinct elements whose first b elements form the specified subset S⋆.Source: §A.1 · p. 16
  • The resulting bijection π induces a strictly decreasing scoring function whose top-b elements are precisely the specified subset S⋆.Source: §A.1 · p. 16

A.2 Implementation details

The implementation uses small per-head MLP agents, trained with RLOO on precomputed activations, and evaluates eviction while retaining designated boundary tokens.

A.2 Implementation details

  • Each KVP agent is a lightweight two-layer MLP with 256 hidden units, trained using RLOO.Source: §A.2 · p. 16
  • For Qwen2.5-7B-Chat, the implementation trains one approximately 650K-parameter agent for each of four KV heads across 28 layers, totaling 112 agents.Source: §A.2 · p. 16
  • Agents optimize the reward in Equation (3.7) using gradient clipping and normalized advantages, with AdamW and a cosine learning-rate schedule.Training runs for 4,000 steps on precomputed activations; the learning rate starts at 5 × 10^-5 and decays to 1 × 10^-6.Source: §A.2 · p. 16
  • During evaluation, the policy ranks all tokens except the first four and last 16, which are always retained.Eviction is emulated with custom attention masks on benchmarks that fit in memory and with physical tensor compaction on long-context benchmarks.Source: §A.2 · p. 16

A.3 Training Efficiency and Inference Overhead

KVP adds no FLOPs during autoregressive decoding in the prefill-then-compress setting, while its reported prefill cost rises by 1%. The paper also describes modest offline training requirements and reports low policy-level latency overhead.

A.3 Training Efficiency and Inference Overhead

  • Training uses approximately 6,000 RULER-4k samples and 4,500 OASST2 samples, with trace collection requiring a one-time forward pass over the training data.The resulting traces occupy approximately 1.2 TB on cluster storage.Source: §A.3 · pp. 16–17
  • Each approximately 650K-parameter agent has a 2.6MB checkpoint, and agent training completes in under 30 minutes on one node of eight NVIDIA H100 GPUs.Source: §A.3 · p. 17
  • KVP adds zero FLOPs during autoregressive decoding in the prefill-then-compress setting, with throughput identical to other eviction methods at the same cache budget.The paper states that KVP’s contribution is superior accuracy at that throughput level.Source: §A.3 · p. 17
  • 14.15 GFLOPs/token is the reported prefill cost with KVP, a 1% increase over the base model.The stated base-model cost is 14.00 GFLOPs/token.Source: §A.3 · p. 17
  • KVP policy-level wall-clock measurements show negligible latency compared with LLM prefill time, and the per-cache operations can be parallelized.The paper reports these measurements as validation of its theoretical estimates.Source: §A.3 · p. 17

A.4 Comparison with the learned baseline JudgeQ

KVP and JudgeQ both use future-generation attention as supervision, but KVP trains an offline, head-local scorer on pre-collected traces. KVP outperforms JudgeQ across long-context comparisons and is on par with or above it on BoolQ at every cache budget.

A.4 Comparison with the learned baseline JudgeQ

  • KVP trains a head-local scorer offline on pre-collected (K, V) traces, whereas JudgeQ learns soft tokens through supervised MSE and requires full LLM forward passes during training.At inference, JudgeQ couples scoring to prefill while KVP operates on cached (K, V, position) features.Source: §A.4 · p. 17
  • KVP outperforms JudgeQ at every cache budget across RULER-128K and English and Chinese LongBench passage retrieval.The comparison uses JudgeQ re-implemented and trained on the same data as KVP on Qwen2.5-7B-Chat.Source: §A.4 · p. 17
  • On BoolQ, KVP is on par with or above JudgeQ at every cache budget despite training without any LLM forward pass.Source: §A.4 · p. 17

A.5 Why supervised ranking fails: per-head importance heterogeneity

Per-head oracle-importance distributions range from sharp-Zipfian to near-flat, so one magnitude-sensitive supervised loss cannot be calibrated across heads. On a representative head, pointwise regression mis-ranks the many low-importance tokens, while KVP does not show this bias.

A.5 Why supervised ranking fails: per-head importance heterogeneity

  • Oracle-importance distributions span sharp-Zipfian to near-flat across heads, so a fixed magnitude-sensitive loss calibrated to one regime fails on others.The distributions are shown for eight representative heads on the same input.Source: §A.5 · p. 17
  • At 10k context length, single-layer KVP compression takes 0.71ms versus 404ms for full-model prefill, a 570× difference.The latency comparison isolates scoring and eviction-index computation for one KV head; TOVA and SnapKV attention-score recomputation is excluded.Source: Figure 7 · p. 18
  • On head L10.H0, pointwise regression mis-ranks the right tail, where most tokens are unimportant, while KVP (RL) does not exhibit this bias.The passage attributes regression's error pattern to heavy-tailed targets: left-tail outliers dominate the loss, leaving right-tail ordering approximately random.Source: §A.5 · p. 18

A.6 Do the per-head agents learn distinct rankings?

The 112 KVP agents produce markedly different rankings, supporting per-head specialization; mid-layer heads correlate more strongly than early and late heads. Additional evaluations report strong KVP performance across model and task settings, alongside a qualitative failure case.

A.6 Do the per-head agents learn distinct rankings?

  • Across 112 Qwen2.5-7B heads, token rankings show head-specific reordering, with mean rank variance of 294.8 and variance of 150–500 across positions.Many heads consistently rank a small band of early attention-sink positions first.Source: §A.6 · p. 18, Figure 10 · p. 20, Figure 11 · p. 22
  • KVP heads in mid-layers L8–L15 correlate more strongly, while early and late layers are more independent.The pairwise comparison uses Spearman rank correlation.Source: §A.6 · pp. 18–19, Figure 12 · p. 23
  • KVP outperforms JudgeQ at every cache budget across three long-context benchmarks and is on par with or above it on BoolQ.Source: Figure 8 · p. 19
  • The ranking differences support per-head specialization, while explicit modeling of cross-head dependencies is left for future work.The paper relates head specialization to differing functions, including attention sinks, local syntactic patterns, and global semantic retrieval.Source: §A.6 · pp. 18–19

A.7 Additional results

  • On Phi-4 14B, KVP remains strong across RULER accuracy and RULER and OASST2-4k perplexity, while heuristic baselines change relative performance between Phi-4 and Qwen 2.5.The figure attributes the RULER result to adaptation where attention-based SnapKV and TOVA fail due to the task's nature.Source: Figure 13 · p. 24
  • KVP struggles to fully capture the ground-truth future-attention pattern in one qualitative failure case at layer 10, head 0.Other qualitative comparisons cover layer 12, head 0, and a long example at layer 19, head 0.Source: Figure 14 · p. 24, Figure 15 · p. 25, Figure 19 · p. 27
  • Across cache sizes, KVP has perplexity on par with or better than leading baselines at most sizes, despite RULER's structure favoring methods that isolate tokens relevant to its final question.Source: Figure 16 · p. 25
  • The additional plots report MMLU and Hellaswag accuracy, ARC-Challenge accuracy, and per-task RULER accuracy as functions of cache size; KVPR and KVPS remain competitive with the uncompressed model on ARC-Challenge.Source: Figure 18 · p. 26, Figure 17 · p. 26, Figure 20 · p. 28

Related papers

Loading 2602.10238v2…