PaperlayerView exact revision on arXiv

Source-linked AI summary

Make Each Token Count: Towards Improving Long-Context Performance with KV Cache Eviction

Ngoc Bui, Hieu Trung Nguyen, Arman Cohan, Rex Ying

arXiv:2605.09649v1cs.LG

TL;DR

Long-context inference is bottlenecked by growing KV-cache memory and attention computation. DBTrimKV learns token retention scores and allocates a global cache budget across layers, heads, and modalities. Across language, vision-language, and dialogue benchmarks, it reduces KV memory while matching or exceeding full-cache inference.

  • Problem

    KV-cache memory and attention computation grow with context length, creating a bottleneck for long-context and multimodal inference.

    Source: §1 · p. 1

  • Method

    DBTrimKV learns cached-token utility scores and uses one global memory budget to allocate retention across layers, heads, and modalities.

    Source: §2.2 · p. 3, §7 · p. 14

  • Results

    Across language, vision-language, and multi-turn dialogue benchmarks, DBTrimKV substantially reduces KV memory while matching or exceeding full-cache inference.

    Source: §1 · p. 2, §7 · p. 14

  • Takeaways & Limitations

    Selective forgetting can act as attention regularization that suppresses distractors and sharpens reasoning over relevant context.

    Source: §1 · p. 2, §7 · p. 14

    The geometric-retention analysis assumes stable future query-state dynamics and a positive probability of exiting a token’s relaxed top-K region within bounded decoding steps.

    Source: §3.2 · p. 6

Abstract

from arXiv · show

The key-value (KV) cache is a major bottleneck in long-context inference, where memory and computation grow with sequence length. Existing KV eviction methods reduce this cost but typically degrade performance relative to full-cache inference. Our key insight is that full-cache attention is not always optimal: in long contexts, irrelevant tokens can dilute attention away from useful evidence, so selective, learnable eviction can improve generation rather than merely approximate the full cache. We introduce a global retention-based KV eviction method that learns each token's future utility under a unified memory budget. Lightweight retention gates assign utility scores to cached KV entries, and a shared final scoring projection calibrates these scores across all layers and heads. This enables a single global eviction policy in which tokens from different layers, heads, and modalities compete directly for cache capacity. We further provide theoretical analysis showing that preferentially retaining useful tokens reduces attention dilution, and we justify geometric retention as a query-agnostic proxy for future utility. Across diverse long-context language and vision-language reasoning, and multi-turn dialogue benchmarks, our method substantially reduces KV memory while matching or surpassing full-cache inference. These results suggest that learned, globally calibrated KV eviction is not only a compression technique, but also a mechanism for improving long-context reasoning.

1. Introduction

The introduction argues that long-context KV caches create substantial memory and computation costs, while selective learned retention can do more than approximate full-cache inference: it can suppress distractors and improve generation. The method assigns future-utility scores to cached entries and uses a shared global budget to select tokens across layers, heads, and modalities.

1. Introduction

  • KV-cache memory and attention computation grow with retained context, making cache management a major bottleneck for long-context and multimodal inference.Prompts may contain tens of thousands of text tokens or hundreds to thousands of visual tokens.Source: §1 · p. 1
  • Full-cache attention can dilute useful evidence among distractors, so removing the right tokens can sharpen attention and improve generation rather than merely approximate full-cache inference.The introduction says the paper theoretically analyzes how preferentially retaining useful tokens reduces attention dilution.Source: §1 · pp. 1–2
  • Lightweight retention gates predict each cached entry’s future utility, and a shared final scoring projection calibrates scores across layers and heads.The scores are trained under a memory constraint to reflect persistent usefulness for future decoding rather than attention at the current step.Source: §1 · p. 2
  • A single global KV budget ranks entries across layers, heads, and modalities, jointly selecting tokens and allocating capacity dynamically.The introduction also justifies geometric retention as a query-agnostic proxy for future utility, supported by observed survival patterns of attended tokens.Source: §1 · p. 2
  • Across long-context reasoning, multi-turn dialogue, and multimodal understanding, the method reduces KV memory while matching or surpassing full-cache inference, often improving accuracy.The introduction reports evaluations on long-context language and vision-language benchmarks.Source: §1 · p. 2

2. Preliminaries

These preliminaries frame KV eviction as a way to control the growing cost of autoregressive attention and describe learnable token retention as a surrogate for selecting cache entries. They introduce a fixed-capacity eviction objective and a teacher–student approach to learning retention scores.

2.1. Self-Attention and KV Eviction

2.2. Token Retention as a Learnable Surrogate

  • Bui et al. [2025] relax discrete retention into exponentially decaying continuous factors, with embedding-predicted gates trained for a student to match a full-cache teacher under a memory budget.The learnable score β_i represents long-term importance; β_i = 1 recovers standard attention. The budget M applies to each attention head.Source: §2.2 · p. 3

3. Can KV Eviction Improve Long-Context Performance?

Full-cache attention can dilute useful evidence with distractor tokens, while selective eviction can reduce that dilution. The section also motivates geometric retention as a query-agnostic proxy for how long tokens remain useful.

3.1. Attention Dilution as Loss of Useful Mass

  • Many distractors with logits close to those of useful tokens can force severe dilution because their combined attention mass absorbs useful-token mass.The proposition attributes this effect to the cumulative contribution of competitive distractors, even when each receives little attention individually.Source: §3.1 · p. 4
  • Selective eviction reduces distractor-induced dilution when useful tokens have a higher logit-weighted average retention rate than distractors.Under that condition, eviction is a corrective mechanism for attention dilution, not merely an approximation to full-cache attention.Source: §3.1 · p. 5

3.2. Geometric Retention as Query-Agnostic Future Utility

  • Under stable future-query dynamics and a bounded-step chance of leaving a token’s relaxed top-K region, its retention decays geometrically, with β_i reflecting persistence.Short-lived local tokens have smaller β_i, while globally useful or structural tokens have β_i closer to one.Source: §3.2 · p. 6
  • Only 13.8% of tokens survive 4,096 positions and 3.2% survive 16k under top-K=256; vision tokens fade at a similar rate to text tokens.Under the 99%-mass criterion, the shown head loses 60% of mass-relevant tokens by 4,096 positions.Source: §3.2 · p. 6, Figure 2 · p. 6
  • The retention gate predicts each token’s decay once when it enters the cache, and a shared final projection globally ranks tokens across layers, heads, and modalities.The resulting geometric weight is interpreted as the predicted probability that the token remains useful, given information available when it was cached.Source: §3.2 · p. 7, Figure 3 · p. 7

4. Global Token Retention via Weight-Tied Gates

The method globally ranks KV entries using retention scores calibrated by a projection shared across layers and heads, allocating a single cache budget without predefined per-layer or per-head budgets.

4. Global Token Retention via Weight-Tied Gates

  • A shared final scoring projection puts retention scores from different layers and heads on a common scale, enabling global ranking under one cache budget.This replaces hand-designed per-layer or per-head allocations with a unified eviction rule.Source: §4 · p. 7, §4.1 · pp. 7–8

4.1. Architecture and Training

  • Each layer–head gate computes a head-specific embedding, then uses a shared scalar readout so equal scores have the same meaning across heads.Without the shared readout, scores are only locally meaningful and cannot reliably support global eviction.Source: §4.1 · pp. 7–8
  • Training follows [Bui et al., 2025] but uses a global capacity constraint, trains only the retention gates, and freezes the LLM weights.The target global KV budget is denoted Mglobal.Source: §4.1 · p. 8

4.2. Global KV Eviction

  • At inference, the method retains the Mglobal entries with the largest scores across layers and heads, assigning capacity automatically rather than setting separate head or layer budgets.The scores aggregate predicted future utility over the remaining decoding horizon; longer horizons favor tokens with larger retention parameters, allowing older, slowly decaying tokens to remain competitive.Source: §4.2 · p. 8
  • Paged attention stores KV entries in fixed-size pages with per-head block tables, supporting variable-length logical caches and variable-length attention kernels.This avoids requiring contiguous physical storage or kernel specialization for head-specific cache sizes.Source: §4.2 · p. 8, Figure 4 · p. 8

5. Experiments

Experiments evaluate DBTrimKV across short-form visual question answering, long-horizon reasoning, multi-turn visual dialogue, and runtime and gate ablations. Results show gains over eviction baselines and, in several settings, full-cache inference.

5.2. Long-Horizon Reasoning

  • DBTrimKV achieves state-of-the-art across reported long-horizon settings, especially at low budgets; budgets of 1024, 512, and 256 surpass full-cache inference by up to 3.75%.Dynamic allocation generally outperforms fixed allocation, with gaps widening as budgets shrink; gains are strongest on long-form tasks, while shorter-reasoning tasks show smaller differences.Source: Table 2 · p. 10, Figure 5 · p. 11, §5.2 · pp. 11–12

5.1. Short-Form Question Answering

  • DBTrimKV and TrimKV consistently outperform prefilling-only visual token pruning baselines and match vanilla performance under all memory budgets.Joint eviction across text and visual tokens, layers, and heads lets the methods preserve visual tokens where useful and trim less-contributory tokens.Source: §5.1 · p. 10

5.2. Long-Horizon Reasoning

  • On MathVisionmini, DBTrimKV concentrates cache tokens in a sparse subset of heads, retains visual context early, and later preserves vision boundary tokens as anchors.Fine-grained visual patches shrink during text-centric reasoning, while <|vision_start|> and <|vision_end|> persist across numerous heads.Source: Figure 6 · p. 11, §5.2 · p. 12

5.3. Multi-turn Dialogue

  • At all tested budgets, DBTrimKV outperforms full-cache inference in multi-turn visual dialogue, including by 14% at M = 512, and its advantage over heuristic baselines widens as budgets shrink.DBTrimKV also improves over fixed-budget TrimKV, while AdaKV improves over SnapKV, underscoring the value of adaptive allocation for interactive dialogue.Source: §5.3 · pp. 12–13

5.4. Runtime and Ablations

  • Weight tying substantially improves retention-gate performance, especially at low budgets, while token contextual embeddings outperform concatenated key-value vectors, particularly under tight budgets.The contextual embedding provides a stronger long-horizon utility signal than projected KV representations alone.Source: §5.4 · p. 13, Table 4 · p. 13, Table 5 · p. 13
  • With PagedAttention, DBTrimKV remains efficient and outperforms vanilla inference at scale, despite small overhead versus TrimKV from variable-length, head-specific cache handling.The runtime comparison fixes batch size at 16, context length at 16384, and KV budget at 256.Source: §5.4 · p. 13

6. Related Work

Related work addresses KV compression through multiple strategies, dynamic capacity allocation, and visual-token pruning, but often relies on added overhead, heuristics, or local signals. This work instead uses learned retention scores to manage text and vision KV entries under one global budget and evaluates compression on VLM benchmarks.

KV cache compression

  • Compression methods span quantization, retrieval, and eviction or merging, while low-overhead eviction heuristics can degrade performance and multimodal effectiveness remains underexplored.Learnable token retention has been proposed as an eviction alternative [Bui et al., 2025]; this work evaluates retention-based compression on VLM benchmarks.Source: §6 · pp. 13–14

Dynamic budget allocation for KV cache

  • Dynamic-budget methods typically use heuristics or local budgets; our learned retention scores jointly drive eviction and allocation under a single global KV constraint.Prior approaches include AdaKV [Feng et al., 2024] and PyramidKV [Cai et al., 2024] for text-only models, and VL-KV [Tu et al., 2024] for VLMs.Source: §6 · p. 14

Visual token pruning

  • Visual-token pruning methods use attention or diversity signals, whereas our KV-cache approach jointly manages vision and text tokens across heads and layers by downstream utility.Examples include VisionZip [Yang et al., 2025b], FastV, DivPrune [Alvar et al., 2025], CDPruner [Zhang et al., 2025], and SparseVLM [Zhang et al., 2024b].Source: §6 · p. 14

7. Conclusion … A.1. Proofs for Attention Dilution

DBTrimKV learns token-utility scores and uses one global budget to allocate KV memory across layers, heads, and modalities. The paper reports that selective eviction reduces attention dilution and KV memory while improving reasoning, matching or surpassing full-cache inference across language, vision-language, and dialogue benchmarks.

7. Conclusion

  • Across language, vision-language, and multi-turn dialogue benchmarks, DBTrimKV outperforms prior eviction baselines and can match or exceed full-cache inference with much less KV memory.The conclusion attributes these results to selective eviction that reduces attention dilution and improves reasoning over relevant context.Source: §7 · p. 14
  • DBTrimKV estimates cached-token future utility with learned retention scores and allocates memory dynamically under a single budget across layers, heads, and modalities.The method is described as globally calibrated.Source: §7 · p. 14
  • The conclusion frames selective forgetting as attention regularization that suppresses distractors and sharpens reasoning over relevant context.The paper presents this as a broader interpretation of KV eviction beyond inference efficiency.Source: §7 · p. 14
  • The authors identify adaptive retention objectives, retrieval-aware cache policies, and training-time integration as directions for future learned memory management.These are proposed as future research for long-context models.Source: §7 · p. 14

A.1. Proofs for Attention Dilution

  • The appendix proof states that the attention mass over U_t equals 1 − δ_t and tends to zero.This is given in the supplied proof passage as ∑_{i∈U_t} α_{t,i} = 1 − δ_t → 0.Source: §A.1 · p. 22
  • The proof relates the exponential-score sum over U_t to δ_t/(1 − δ_t).The supplied expression is ∑_{i∈U_t} e^{z_{t,i}} = δ_t/(1 − δ_t).Source: §A.1 · p. 22
  • For ρ_D ≤ ρ_U, the proof uses monotonicity of aδ_t/((1 − δ_t) + aδ_t) for a ≥ 0 to establish its stated bound.It then considers the limiting case ρ_D/ρ_U → 0 when 1 − δ_t > 0.Source: §A.1 · p. 22

A.2. Geometric Retention as Query-Agnostic Future Utility · B. Experiments

The section frames geometric retention as a query-agnostic proxy for how persistently cached tokens remain useful under future query dynamics. Under its stated assumptions, relaxed top-K persistence and expected future utility decay exponentially, motivating a learned retention score.

A.2. Geometric Retention as Query-Agnostic Future Utility

  • The analysis assumes future query states follow stable, low-dimensional Markov dynamics; fitted Qwen3-8B query-state models were stable across all 10 contexts, with median spectral radius 0.76.The caption also reports convergence consistent with ergodicity and sub-exponential innovation tails.Source: §A.2 · pp. 23–24, Figure 8 · p. 24
  • A token can merit retention even when it is not currently in the old-cache top-K, if it remains near the boundary or could become top-K under plausible future queries.The relaxed survival region includes hard top-K states and may also include states where the token is sufficiently compatible to justify retention.Source: §A.2 · pp. 24–25
  • Because monotone eviction prevents removed tokens from re-entering the cache, retention should reflect usefulness across intervening query trajectories, not just one future query.The formulation measures whether a token remains inside its relaxed survival region for consecutive future steps.Source: §A.2 · pp. 25, 27
  • The block-exit assumption requires a token inside its relaxed survival region to have probability at least ε_i of exiting within b_i steps; these parameters can vary by token.Locally relevant tokens may have small b_i, while persistent structural tokens such as sinks or delimiters may have much larger b_i.Source: §A.2 · pp. 25–26
  • At horizon L = 4096, representative heads lose most tokens even under the 99% mass criterion, while top-K survival drops sharply across text and multimodal data.The caption identifies the gap between hard top-K and τ-mass survival curves as attention-dilution mass.Source: Figure 9 · p. 26
  • The old-cache future utility in Eq. (10) is bounded above by Pr(S_n | F_t), and the passage says its bound follows by substituting the persistence bound.Source: §A.2 · p. 27
  • The learned survival parameter β_i summarizes future usefulness: it is smaller for tokens that quickly leave their relaxed regions and closer to one for tokens compatible with many future queries.Rather than estimating block-exit parameters explicitly, the retention gate predicts β_i from token representations or features determining compatibility.Source: §A.2 · p. 27

B.1. General Experiment Settings … C. Additional Experimental Results

The appendix specifies training and baseline settings, then describes a combined Needle-in-a-Haystack evaluation of TrimKV and DBTrimKV; no additional-results findings are included in the supplied passages.

B.1. General Experiment Settings

  • Training uses a 32,768-token maximum sequence length, a 2 × 10−4 learning rate, 10−6 weight decay, cosine scheduling, and four H100 GPUs.Per-GPU batch size is 1 with gradient accumulation of 4; other hyperparameters follow HuggingFace Trainer defaults.Source: §B.1 · p. 28
  • The training mixture randomly samples 40% M4Instruct [Li et al., 2024a], 30% Academic Video [Zhang et al., 2024c], 30% R1-OneVision [Yang et al., 2025c], and 20% OpenR1-Math-220k [Hugging Face].Source: §B.1 · p. 28

B.2. Baselines

B.3. Needle in a Haystack Experiment Settings

  • The NIAH evaluation combines MK-NIAH, MV-NIAH, and MQ-NIAH because recent LLMs perform near-perfectly on single-NIAH tasks.Training uses 5,117 samples with contexts from 4k to 200k; 200 held-out tests fix keys, values, and queries at 8, 4, and 4.Source: §B.3 · p. 28
  • Using Qwen3-4B-Instruct-2507 with its 256k context window, the experiment compares trainable TrimKV and DBTrimKV while freezing the base model and training only the retention head.Source: §B.3 · p. 28

C.1. MMDU · C.2. Long Context Benchmarks · C.2.1. LongBench-V2

On MMDU, DBTrimKV outperforms Vanilla and competing eviction methods across the tested budgets, including under the extreme 128-token constraint. Table 7 reports LongBench-V2 evaluations for Phi3-mini-128K at a 1024-token budget, but the supplied passage gives no results.

C.1. MMDU

  • At a 128-token budget, DBTrimKV reaches 104.10% of Vanilla’s overall performance, while SnapKV, R-KV, and AdaKV fall to 60–64%.Table 6 presents the complete MMDU metric comparison across budgets.Source: §C.1 · p. 28, Table 6 · p. 29
  • DBTrimKV consistently outperforms competing eviction methods and Vanilla across the tested KV budgets of 512, 256, and 128 tokens.The passage reports this trend across all three tested budgets.Source: §C.1 · p. 28, Table 6 · p. 29
  • As budgets shrink, baseline eviction methods sharply decline on Visual Perception and Image Relative, indicating loss of visually grounded context.Both TrimKV and DBTrimKV preserve these critical tokens.Source: §C.1 · pp. 28–29
  • At 128 tokens, DBTrimKV scores 3.54 vs. 3.40 on Visual Perception, exceeding Vanilla.The passage attributes this to preserving critical visual tokens.Source: §C.1 · p. 29, Table 6 · p. 29
  • At 128 tokens, DBTrimKV scores 4.11 vs. 3.90 on logical coherence, exceeding Vanilla.The comparison is reported for the lowest tested budget.Source: §C.1 · p. 29, Table 6 · p. 29

C.2.1. LongBench-V2

  • Table 7 compares KV eviction methods on LongBench-V2 using Phi3-mini-128K, with a 1024-token budget and average relative change versus Full KV.The supplied passage identifies the table’s scope but does not provide its result values.Source: Table 7 · p. 29

C.3. Ablation Study on Lookahead Step (T-t) · C.4. Efficiency Scaling · C.5. Visualization of DBTrimKV

The ablation shows that the best lookahead depends on KV budget, while DBTrimKV maintains efficiency as context, generation length, and KV budget scale. Its visual-token retention shifts from broad early coverage to task-focused evidence as decoding progresses.

C.3. Ablation Study on Lookahead Step (T-t)

  • At budgets 256, 512, and 1024, T −t = 2 generally performs best; at 1024, T −t = 5 drops accuracy from 52.63% to 46.05%.T −t = 1 consistently degrades performance at moderate-to-high budgets, whereas T −t = 5 outperforms both alternatives at the most constrained budgets, 64 and 128.Source: §C.3 · pp. 29–30, Table 8 · p. 30

C.4. Efficiency Scaling

  • As context length grows, DBTrimKV maintains a flat efficiency profile while Vanilla (Full Cache) decoding time rises and throughput falls.The synthetic benchmark measures total decoding time and throughput during generation; Figure 10 fixes batch size at 16, generated tokens at 8192, and KV budget at 256, averaging five seeds.Source: §C.4 · p. 30, Figure 10 · p. 31
  • As generation length increases, DBTrimKV scales more efficiently than Vanilla, with lower decoding times and higher throughput.Figure 11 fixes batch size at 16, context length at 16384, and KV budget at 256; Vanilla at 32k generation length is omitted due to OOM.Source: §C.4 · p. 30, Figure 11 · p. 31
  • Across KV budgets, DBTrimKV retains its efficiency advantage over Full Cache, with throughput and decoding time largely stable across allocations.Figure 12 reports throughput and total decoding time across budgets; results are averaged over five random seeds.Source: §C.4 · p. 30, Figure 12 · p. 32

C.5. Visualization of DBTrimKV

  • Early in generation, DBTrimKV preserves a high density of visual tokens to establish structural and semantic understanding.Figure 13a illustrates this early-stage retention behavior.Source: §C.5 · p. 30
  • As decoding progresses, DBTrimKV evicts background distractors and retains task-relevant visual evidence aligned with the model’s textual output.In Figure 13, retained tokens concentrate on the scythe, clock, and “RATHER TIME” banner; Figure 14 isolates coordinate data points. This is consistent with Figure 6’s layer-wise pattern of gradual KV-cache pruning.Source: §C.5 · p. 30

D. Limitations and Future works

The limitations and future-work agenda centers on joint training, evaluation at larger model scales, and learning eviction policies to improve downstream reasoning.

D. Limitations and Future works

  • DBTrimKV trains retention gates while freezing base-model weights, leaving open whether joint end-to-end training improves compatibility with KV eviction.The current approach preserves base-model capabilities and is computationally efficient; future work could test joint training of the base LLM and retention network.Source: D. Limitations and Future works · p. 32
  • Evaluations cover models only up to 8B parameters, so scaling retention and attention-dilution mitigation to 70B+ models remains untested.The paper identifies establishing token-retention scaling laws as a critical next step.Source: D. Limitations and Future works · p. 32
  • At specific low budgets, DBTrimKV can outperform full-cache inference by suppressing distractors, motivating reinforcement learning to optimize retention for downstream performance.The proposed future direction treats token retention as an action space for post-training exploration of eviction policies.Source: D. Limitations and Future works · p. 32

Related papers

Loading 2605.09649v1…