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 KV caches make memory and attention increasingly costly, while existing eviction methods typically lose performance relative to full-cache inference. DBTrimKV learns globally calibrated token-retention scores under one memory budget, matching or exceeding full-cache inference across language, vision-language, and dialogue benchmarks while using much less KV memory. The results indicate that selective eviction can suppress distractors and sharpen long-context reasoning.

  • Problem

    Long-context and multimodal inference faces KV-cache memory and attention costs that grow with retained sequence length, while existing eviction methods typically degrade performance relative to full-cache inference.

  • Method

    DBTrimKV learns globally calibrated retention scores for cached tokens and dynamically allocates a shared memory budget across layers, heads, and modalities.

  • Results

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

  • Takeaways & Limitations

    Selective KV eviction can function as attention regularization by suppressing distractors and sharpening reasoning over relevant context, beyond improving inference efficiency.

  • Takeaways & Limitations

    The geometric-retention analysis assumes stable future query dynamics and a nonzero probability of exiting each token's relaxed top-K region within a bounded number of decoding steps.

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

Long-context KV caches create memory and attention bottlenecks, while conventional eviction treats removal as compression toward full-cache inference. This work instead learns globally calibrated future-utility scores to suppress distractors, dynamically allocate one shared budget, and improve or preserve reasoning quality.

  • Problem: KV caches grow linearly with sequence length, making memory and cached-token attention major bottlenecks in long-context and multimodal generation.Caching past keys and values avoids recomputation, but retained context increases resource demands.
  • Related work: Existing KV eviction methods typically approximate full-cache inference using heuristics such as recency, accumulated attention, or local attention magnitude [Li et al., 2024b; Xiao et al., 2023; Li et al., 2024c; Cai et al., 2025].These approaches frame eviction primarily as compression under reduced memory and computation.
  • Motivation: Full-cache attention can dilute useful evidence among irrelevant tokens, so removing distractors can sharpen attention and improve generation [Bansal et al., 2026].The paper identifies attention dilution as a mechanism by which full-cache inference may degrade in long contexts and theoretically analyzes preferential distractor eviction.
  • Method: Lightweight retention gates learn future-utility scores under a memory constraint, while shared final projections calibrate scores across layers and heads for one global budget.The resulting ranking lets entries from different layers, heads, and modalities compete directly, dynamically allocating capacity toward useful long-range information.
  • Empirical scope: Globally calibrated retention substantially reduces KV memory while matching or surpassing full-cache inference across long-context language, vision-language, dialogue, and multimodal tasks.Selective eviction improves accuracy over the full cache in many cases, supporting eviction as a reasoning improvement rather than only a cost reduction.

2. Preliminaries

In autoregressive transformers, the KV cache grows with context length, causing linear memory and quadratic cumulative attention cost. KV eviction limits the cache to M tokens, but selecting an optimal monotone subset is combinatorial and therefore typically approximated with heuristics or learned retention scores.

  • KV-cache memory grows linearly with sequence length, while cumulative attention computation grows quadratically during generation.
  • Eviction restricts attention to at most M retained key-value pairs, with monotonicity preventing an evicted token from re-entering the cache.
  • The ideal policy selects a size-M cache whose attention output stays closest to full attention, but solving this combinatorial problem at every step is infeasible.
  • Learned retention relaxes binary eviction into continuous scores β_i, where larger β_i indicates greater long-term importance and β_i = 1 recovers standard attention.
  • Retention gates predict β_i from token embeddings and are trained so a budgeted student matches the full-cache teacher distribution.

3. Can KV Eviction Improve Long-Context Performance?

KV eviction can improve long-context performance by suppressing distractors that dilute attention away from useful evidence. Geometric retention provides a query-agnostic proxy for future token utility by modeling how long tokens remain useful.

  • Attention dilution: Selective eviction improves attention allocation when it suppresses distractor mass more than useful-token mass, rather than merely approximating full-cache attention.Attention dilution occurs because irrelevant tokens compete with useful evidence in the softmax denominator; preferential retention reduces this dilution.
  • Attention dilution: Many near-tie distractors can force severe attention dilution because their cumulative softmax contribution absorbs substantial attention mass from useful tokens.This is the mechanism formalized by Proposition 3.1.
  • Geometric retention: Geometric retention models future persistence: tokens likely to leave a relaxed top-K survival region sooner receive faster decay, while persistent tokens retain weights closer to one.Theorem 3.3 derives geometric decay under stable future-query dynamics and a positive probability of exiting the survival region.
  • Empirical support: Empirical survival under full-cache inference drops rapidly with horizon, supporting geometric retention as a surrogate for future token persistence.The study uses Qwen3-VL-4B on 98 long MMDU multimodal multi-turn dialogues prefilling 16,384 tokens.

4. Global Token Retention via Weight-Tied Gates

The method uses weight-tied retention gates to place KV entries from all layers and heads on a common score scale, enabling global eviction under one cache budget. At inference, it retains entries with the highest horizon-aware utility scores, while paged attention supports resulting variable-length caches.

  • Global retention design: The method replaces fixed or myopically adaptive per-layer and per-head allocations with a unified ranking motivated by heterogeneous long-range token retention across heads.A small number of heads preserve long-range tokens, while many quickly lose them, making uniform local budgets poorly suited to global capacity allocation.
  • Global retention design: Weight-tied retention gates globally calibrate scores across layers and heads, allowing all KV entries to compete under a single cache budget.Layer- and head-specific projections produce embeddings, while a shared final projection makes scores comparable across locations.
  • Training: Training updates only the retention gates, freezes LLM weights, and replaces the local capacity constraint with a global KV-budget constraint.The global target budget is denoted Mglobal.
  • Inference-time eviction: At each decoding step, the system retains the Mglobal cached entries with the largest horizon-aggregated retention scores across all layers and heads.The lookahead horizon sums predicted future utility over remaining decoding steps; longer horizons favor larger retention parameters and keep older, slowly decaying tokens competitive.
  • Implementation: Paged attention stores variable-length per-head KV sequences in fixed-size pages and computes attention with variable-length kernels, supporting dynamic head-specific cache sizes.Each head maintains a block table for its active pages, avoiding the need for equal logical sequence lengths across heads.

5. Experiments

Across visual question answering, long-horizon reasoning, and dialogue, DBTrimKV consistently outperforms competing eviction methods while substantially reducing KV-cache memory. Dynamic allocation is especially effective under tight budgets, where DBTrimKV can surpass full-cache inference and retain task-relevant evidence selectively.

  • Short-form visual question answering: DBTrimKV and TrimKV consistently outperform prefilling-only visual token pruning and match vanilla performance across memory budgets by jointly compressing visual and textual KV entries.The comparison uses equal total KV budgets, combining all prompt text tokens with each pruning baseline’s visual-token budget.
  • Long-horizon reasoning: DBTrimKV achieves state-of-the-art performance across image, video, and pure-text reasoning settings, with particularly large gains in low-budget regimes.Table 2 reports image and video reasoning results, while Figure 5 covers pure-text reasoning across varying KV budgets.
  • Long-horizon reasoning: Up to 3.75%, DBTrimKV with KV budgets of 1024, 512, and 256 surpasses vanilla full-cache inference, while dynamic allocation increasingly outperforms fixed allocation as budgets shrink.The comparison is reported against TrimKV and against fixed-budget baselines including AdaKV/AdaPyramid-KV versus SnapKV.
  • Qualitative analysis: DBTrimKV dynamically concentrates cache capacity across specialized heads, retaining early visual context as structural anchors while discarding background distractors.Figure 6 visualizes aggressive total-token compression, sparse mid-layer head usage, and selective retention of visual tokens.
  • Multi-turn dialogue: At M = 512, DBTrimKV exceeds full-cache vanilla inference by 14% and consistently outperforms heuristic and fixed-budget eviction methods in interactive dialogue.The advantage over heuristic baselines widens as the KV budget decreases, and AdaKV similarly improves over SnapKV.
  • Ablations and runtime: Weight tying improves low-budget retention, contextual embeddings outperform key-value inputs, and DBTrimKV remains faster than vanilla inference despite overhead from variable-length head-specific cache handling.These findings support cross-layer/head score comparability, stronger long-horizon utility estimation, and practical large-scale decoding.

6. Related Work

Related work spans KV-cache compression, dynamic budget allocation, and visual token pruning. The paper addresses the limited understanding of long-horizon KV-cache methods in multimodal settings through extensive VLM evaluation.

  • KV cache compression: KV-cache compression methods reduce memory through vector compression, quantization, or token eviction, but compression and selection overhead can limit long-horizon generation.Simple eviction heuristics [Xiao et al., 2023; Li et al., 2024c; Cai et al., 2025] reduce overhead but may degrade performance, while token retention offers a learnable alternative [Bui et al., 2025].
  • KV cache compression: Most long-horizon efficiency studies focus on text-only LLMs, leaving multimodal behavior and effectiveness underexplored.The paper bridges this gap by extensively evaluating KV-cache compression strategies, including retention-based eviction, on VLM benchmarks.
  • Dynamic budget allocation for KV cache: Dynamic budget allocation work distributes KV-cache capacity across layers, heads, or modalities using methods such as AdaKV [Feng et al., 2024], PyramidKV [Cai et al., 2024], and CAKE [Qin et al., 2025].ZigZagKV [Zhong et al., 2025] is also identified as a text-only approach in this line of work.
  • Visual token pruning: Visual token pruning selects vision tokens using pretrained attention maps, diversity constraints, or text-to-vision attention.Representative methods include VisionZip [Yang et al., 2025b], FastV, DivPrune [Alvar et al., 2025], CDPruner [Zhang et al., 2025], and SparseVLM [Zhang et al., 2024b].

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

The conclusion presents DBTrimKV as a globally calibrated KV-eviction method that dynamically allocates memory and improves long-context efficiency and reasoning. Appendix A.1 supplies proofs for the attention-dilution results referenced by the paper.

  • 7. Conclusion: DBTrimKV is introduced as a globally calibrated KV-eviction method for long-context language and vision-language models.
  • 7. Conclusion: DBTrimKV learns future-utility retention scores and uses one global budget to allocate memory across layers, heads, and modalities.
  • 7. Conclusion: Across language, vision-language, and multi-turn dialogue benchmarks, DBTrimKV consistently outperforms prior eviction baselines while matching or exceeding full-cache inference.
  • 7. Conclusion: Selective eviction is reported to reduce attention dilution, improve reasoning over relevant context, and substantially lower KV memory.
  • A.1. Proofs for Attention Dilution: Appendix A.1 provides proofs for the attention-dilution results stated in Section A.1.
  • A.1. Proofs for Attention Dilution: The proof includes an expression for the attention-related quantity δ_t involving |D′_t|/|U_t| and an exponential term e^-∆|D′_t|/|U_t|.
  • A.1. Proofs for Attention Dilution: The proof of Proposition 3.2 proceeds from the definition and analyzes the case ρ_D ≤ ρ_U using monotonicity of a displayed function.
  • A.1. Proofs for Attention Dilution: It further considers the limit ρ_D/ρ_U →0 when 1−δ_t>0, deriving a consequence from the displayed formula.

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

The paper models future token utility as persistence within relaxed top-K regions under stable, low-dimensional query dynamics. Under block-exit conditions, this persistence decays geometrically, motivating learned retention scores that summarize future usefulness without observing future queries.

  • A.2. Geometric Retention as Query-Agnostic Future Utility: Under Assumption A.1, future query states evolve as a stable, ergodic Markov process in a low-dimensional subspace, making future top-K membership regions query-state dependent.The assumption requires dimension m ≪ d and convergence toward an invariant distribution.
  • A.2. Geometric Retention as Query-Agnostic Future Utility: Empirical analysis of Qwen3-8B query states supports the stable-dynamics assumption, with fitted VAR(1) eigenvalues inside the unit circle and median spectral radius ˜ρ = 0.76.The analysis uses 128 rollouts on each of 10 AIME-24 prompts and reports ergodic convergence and sub-exponential innovations.
  • A.2. Geometric Retention as Query-Agnostic Future Utility: Geometric retention assigns value to tokens that remain useful across future query trajectories, rather than only tokens that are immediately in the old-cache top-K.Relaxed survival regions include near-boundary tokens that may become top-K again as query states evolve, which suits monotone eviction because removed tokens cannot re-enter.
  • A.2. Geometric Retention as Query-Agnostic Future Utility: The block-exit horizon is token dependent: locally relevant tokens may have short persistence blocks, whereas sink, delimiter, and other structural tokens may persist much longer.This allows the model to represent both rapidly changing contextual relevance and globally persistent utility.
  • A.2. Geometric Retention as Query-Agnostic Future Utility: Theorem A.3 shows that, with stable query dynamics and top-K block exit, each token’s relaxed survival probability decays exponentially over future steps.The bound follows by applying the Markov property across blocks in which each token has a uniformly positive probability of exiting its relaxed survival region.
  • A.2. Geometric Retention as Query-Agnostic Future Utility: The survival parameter β_i summarizes persistence: tokens that quickly leave relaxed top-K regions have smaller β_i, while tokens compatible with many future query states have β_i closer to one.This gives geometric retention a probabilistic interpretation as a proxy for expected future usefulness.
  • A.2. Geometric Retention as Query-Agnostic Future Utility: The retention gate learns β_i directly from token representations or compatibility features, avoiding explicit estimation of block-exit parameters.Thus, the learned score summarizes the geometry of each token’s relaxed survival region under future query dynamics.

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

The experiments use long-context multimodal training with specified optimization and data-mixture settings, compare against several KV-eviction baselines, and evaluate distractor mitigation through a harder combined Needle-in-a-Haystack setup.

  • B.1. General Experiment Settings: Training uses a 32,768-token maximum sequence length, learning rate of 2 × 10−4, weight decay of 10−6, cosine scheduling, and a per-GPU batch size of 1 with four-step gradient accumulation.Other hyperparameters follow the default HuggingFace Trainer settings.
  • B.1. General Experiment Settings: The training mixture samples 40% of M4Instruct [Li et al., 2024a], 30% of Academic Video [Zhang et al., 2024c], 30% of R1-OneVision [Yang et al., 2025c], and 20% of OpenR1-Math-220k [Hugging Face].
  • B.2. Baselines: SnapKV [Li et al., 2024c] serves as an attention-score-based eviction heuristic, while R-KV extends this approach to reasoning models using key-vector similarity to identify redundant tokens.
  • B.2. Baselines: AdaKV [Feng et al., 2024] extends SnapKV with dynamic head-wise budget allocation driven by attention statistics.
  • B.2. Baselines: R-KV [Cai et al., 2025] is a training-free reasoning-model compression method that jointly scores tokens by attention importance and semantic redundancy.
  • B.3. Needle in a Haystack Experiment Settings: The NIAH evaluation combines MK-NIAH, MV-NIAH, and MQ-NIAH because recent LLMs are near-perfect on single-NIAH tasks, using 5,117 training samples spanning 4k–200k contexts.The samples randomize the numbers of keys, values, and queries to test whether DBTrimKV mitigates attention dilution by evicting distractors.

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

DBTrimKV outperforms competing eviction methods and Vanilla across MMDU KV budgets, including 104.10% relative performance at 128 tokens. On LongBench-V2, the study compares KV eviction methods under a fixed 1024-token budget for deep long-context reasoning tasks.

  • C.1. MMDU: As KV budgets shrink, traditional eviction methods sharply lose Visual Perception and Image Relative performance, indicating weaker retention of visually grounded context.The metric breakdown identifies multimodal-specific degradation as a principal failure mode of baseline eviction strategies.
  • C.1. MMDU: At 128 tokens, DBTrimKV exceeds Vanilla on visual perception, 3.54 vs. 3.40, and logical coherence, 4.11 vs. 3.90.TrimKV and DBTrimKV preserve critical tokens associated with these metrics.
  • C.1. MMDU: 104.10% relative performance to Vanilla at 128 tokens, with DBTrimKV consistently outperforming competing eviction methods and full-cache inference across tested MMDU budgets.The evaluated budgets are 512, 256, and 128 tokens.
  • C.2. Long Context Benchmarks: LongBench-V2 evaluates deep reasoning in realistic long-context scenarios including multi-document QA and code repository comprehension.The comparison includes Full KV Cache, TrimKV, Locret [Huang et al., 2025b], and the proposed method.
  • C.2.1. LongBench-V2: The LongBench-V2 comparison uses Phi-3-mini-128k-instruct with chunk-prefill settings and a 1024-token KV budget.Table 7 reports performance and average relative change compared with Full KV.

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

DBTrimKV performs best with a moderate lookahead horizon under typical KV budgets, while longer horizons help only under extreme memory constraints. It also improves efficiency across scaling regimes and progressively retains visual tokens that align with the evidence used in generation.

  • C.3. Ablation Study on Lookahead Step (T-t): A lookahead horizon of T −t = 2 generally achieves the highest MathVisionmini accuracy at moderate-to-high KV budgets, while T −t = 1 consistently degrades performance.At a budget of 1024, T −t = 5 reduces accuracy from 52.63% to 46.05%.
  • C.3. Ablation Study on Lookahead Step (T-t): Under extremely constrained KV budgets of 64 and 128, the longer T −t = 5 horizon outperforms both shorter lookahead settings.
  • C.4. Efficiency Scaling: Efficiency is evaluated across context length, generation length, and KV budget using throughput and total decoding time for Vanilla, DBTrimKV, and TrimKV.The scaling experiments use fixed batch size and task-specific context, generation, and KV-budget settings, with results averaged over five random seeds where specified.
  • C.4. Efficiency Scaling: DBTrimKV maintains a flatter efficiency profile than Vanilla as context length increases, avoiding Vanilla’s rising decoding time and falling throughput.The comparison uses Vanilla (Full Cache), DBTrimKV, and TrimKV, measuring decoding time and generated tokens per second after warmup.
  • C.4. Efficiency Scaling: DBTrimKV scales more efficiently with generation length than Vanilla, which exhibits a severe quadratic increase in decoding time and substantially lower throughput.Vanilla performance is not reported at 32k generated tokens because of an OOM error.
  • C.5. Visualization of DBTrimKV: As generation begins, DBTrimKV retains a high density of visual tokens to establish structural and semantic understanding.This behavior is illustrated in Figures 13 and 14.
  • C.5. Visualization of DBTrimKV: During decoding, DBTrimKV aggressively evicts background distractors and concentrates retention on informative image regions aligned with the model’s textual output.In Figure 13, retained tokens focus on the scythe, clock, and “RATHER TIME” banner used to infer the final answer.

D. Limitations and Future works

DBTrimKV’s current approach trains only lightweight retention gates while freezing the base language model, preserving efficiency but leaving the effects of joint end-to-end training unknown. Future work should investigate whether jointly training the base LLM and retention network improves compatibility with KV-cache eviction.

  • Limitations and Future Work: The method freezes base language-model weights and trains only lightweight retention gates, which is computationally efficient but leaves joint training’s impact unknown.Future work should test end-to-end training of the base LLM alongside the retention network to determine whether internal representations become more amenable to KV-cache eviction.
Loading 2605.09649v1…