PaperlayerView exact revision on arXiv

Source-linked AI summary

Cache What Lasts: Token Retention for Memory-Bounded KV Cache in LLMs

Ngoc Bui, Shubham Sharma, Simran Lamba, Saumitra Mishra, Rex Ying

arXiv:2512.03324v2cs.LGcs.AI

TL;DR

Long-context LLM inference is bottlenecked by quadratic attention and growing KV-cache memory, while existing memory-bounded methods can add overhead or use unreliable importance proxies. TRIM-KV learns decaying, layer- and head-specific retention scores with lightweight gates trained by distillation and a capacity loss, then evicts the lowest-scoring tokens. Across math, procedural, and conversational long-memory benchmarks, it consistently outperforms strong eviction and retrieval baselines, sometimes exceeding full-cache models.

  • Problem

    Long-context inference faces quadratic self-attention costs and rapidly growing KV-cache memory, while existing bounded-memory methods rely on overhead-heavy or unreliable strategies.

  • Method

    TRIM-KV uses lightweight layer- and head-specific retention gates that assign exponentially decaying token scores and evict the lowest-scoring tokens under budget M.

  • Results

    TRIM-KV consistently outperforms strong eviction and retrieval baselines across math reasoning, procedural generation, and long-memory benchmarks, sometimes surpassing full-cache models.

  • Takeaways & Limitations

    Learned retention scores recover human-intuitive behaviors and provide a lightweight probe of layer- and head-specific dynamics.

  • Takeaways & Limitations

    The current method keeps backbone parameters frozen during training and still relies on standard attention at inference time.

Abstract

from arXiv · show

Memory and computation remain core bottlenecks in long-horizon LLM inference due to the quadratic cost of self-attention and the ever-growing key-value (KV) cache. Existing strategies for memory-bounded inference, such as quantization, offloading, or heuristic KV eviction, either incur high orchestration costs or rely on unreliable attention-based proxies of importance. We propose TRIM-KV, a novel approach that learns each token's intrinsic importance at creation time via a lightweight retention gate. Each gate predicts a scalar retention score that decays over time, reflecting the long-term utility of the token for a specific layer and head. Tokens with low scores are evicted when the memory budget is exceeded, ensuring that the cache always contains the most critical tokens. TRIM-KV is trained efficiently through distillation from a frozen LLM combined with a capacity loss, requiring only gate fine-tuning and adding negligible inference overhead. Across mathematical reasoning (GSM8K, MATH-500, AIME24), procedural generation (LongProc), conversational long-memory benchmarks (LongMemEval), and long-context understanding (LongBenchV2 and SCBench), TRIM-KV consistently outperforms strong eviction and learnable retrieval baselines, especially in low-memory regimes. Remarkably, it even surpasses full-cache models in some settings, showing that selective retention can serve as a form of regularization, suppressing noise from uninformative tokens. Qualitative analyses further reveal that learned retention scores align with human intuition, naturally recovering heuristics such as sink tokens, sliding windows, and gist compression without explicit design. Beyond efficiency, retention scores provide insights into layer- and head-specific roles, suggesting a new path toward LLM interpretability.

1 INTRODUCTION

Long-context inference is constrained by quadratic attention and growing KV-cache memory, while existing compression, offloading, and attention-guided eviction methods have important limitations. TRIM-KV learns decaying, layer- and head-specific token-retention scores, trains lightweight gates with distillation and capacity losses, and improves memory-bounded performance across several benchmarks.

  • Motivation: The method addresses long-context bottlenecks caused by quadratic self-attention and KV-cache growth, which increase memory and latency costs.Compression and quantization reduce memory but scale poorly with generation length, while offloading introduces orchestration overhead.
  • Method: TRIM-KV predicts each token’s intrinsic importance when created using a layer- and head-specific retention gate whose score decays exponentially over time.When the cache exceeds budget M, the token with the smallest current score is evicted, favoring tokens judged most important and more recent tokens.
  • Method: The gates are trained while the pretrained backbone remains fixed, using distillation to preserve outputs and a capacity loss to encourage memory-bounded sparsity.Inference computes scores online and uses simple score comparisons for eviction, adding minimal overhead.
  • Results and Contributions: TRIM-KV consistently outperforms eviction baselines across mathematical reasoning, procedural generation, and long-memory tasks, including cases where baselines use 4× more KV budget.It also delivers a 58.9% relative pass@1 gain over the SOTA learnable KV retrieval baseline and sometimes surpasses full-cache models.
  • Interpretability: Learned retention scores align with human intuition and recover sink-token, sliding-window, and gist-compression behaviors without explicit hard-coding.The scores also provide a lightweight probe of layer- and head-specific attention dynamics.

2 RELATED WORK

Related work addresses KV-cache growth through token eviction or merging, vector compression or quantization, retrieval, and explicit forgetting mechanisms. These approaches trade memory efficiency against long-range memory preservation, altered attention dynamics, or additional training requirements.

  • KV Cache Compression: KV-cache compression research primarily follows token eviction or merging, vector compression or quantization, and token retrieval.The cited work spans methods that remove, compact, or selectively recover cached information.
  • Forgetting in Language Models: Linearized and recurrent attention variants summarize history into fixed-size states, reducing computation but potentially degrading tasks requiring long-range memory.Later methods increase state capacity, while other approaches modify attention logits or impose trainable sparsity.
  • Forgetting in Language Models: Methods that alter attention dynamics substantially typically require model training rather than serving as drop-in inference-time cache policies.This distinguishes them from approaches focused on selectively managing an existing KV cache.

3 PRELIMINARIES

Autoregressive self-attention caches past key-value pairs to avoid recomputation, but memory grows linearly with sequence length while computation grows quadratically. KV eviction introduces binary retention decisions under a capacity constraint, yet exact optimization is combinatorial and practical systems therefore use heuristics or learnable policies.

  • Preliminaries: Caching running key-value pairs avoids recomputation during autoregressive decoding, but memory grows linearly with sequence length and computation grows quadratically.This creates inefficiency for long-context inputs and extended generation.
  • Revisiting KV Cache Eviction: KV eviction prunes less-important cached pairs as new tokens arrive, appending new tokens while enforcing a fixed memory budget.The cache-management procedure directly trades retained history against bounded memory.
  • Revisiting KV Cache Eviction: The eviction formulation uses binary variables α_ti to indicate whether a key-value pair remains available, with monotonicity preventing retrieval after eviction.The objective seeks attention outputs close to those produced by the full KV cache.
  • Revisiting KV Cache Eviction: The capacity constraint keeps at most M tokens at each inference step, while the loss penalizes differences between attention with and without eviction.The displayed formulation expresses the memory bound and approximation goal for eviction.
  • Revisiting KV Cache Eviction: Exact constrained optimization is impractical at every decoding step because the decision problem is combinatorial, motivating heuristic and learnable eviction methods.This practical difficulty motivates replacing exact decisions with efficient policies.

4 METHODOLOGY

TRIM-KV learns token retention from creation-time embeddings, using exponentially decaying scores to guide memory-bounded KV eviction. Jointly trained retention gates preserve model quality while enforcing capacity constraints, enabling efficient inference under fixed budgets.

  • 4.1 SELECTIVE IN-CONTEXT MEMORY VIA RETENTION-GATED ATTENTION: The smooth exponential formulation replaces discrete eviction signals with differentiable decay, avoiding the optimization difficulties of hard binary retention.The method adopts β_i^(t−i) after identifying sigmoid-based decay as poorly normalized and prone to vanishing gradients.
  • 4.1 SELECTIVE IN-CONTEXT MEMORY VIA RETENTION-GATED ATTENTION: Retention-gated attention assigns each token a layer- and head-specific score β_i from its embedding, with effective influence decaying exponentially over time.The score lies in [0, 1], where larger values indicate slower decay and stronger long-term retention.
  • 4.1 SELECTIVE IN-CONTEXT MEMORY VIA RETENTION-GATED ATTENTION: Unlike attention scores, which measure query-dependent short-term utility, retention scores encode intrinsic long-term utility available when each token is created.This reframes eviction as deciding how long a token should remain useful rather than how much it contributes to the current query.
  • 4.2 TRAINING: The gates are trained jointly with frozen model weights using quality losses and a capacity penalty that discourages exceeding the target KV budget.Quality training combines distillation from the standard-attention model with next-token prediction, while capacity regularization controls aggregate retention.
  • 4.3 INFERENCE: At inference, gates operate as eviction decision-makers: when a new token exceeds capacity, the cache removes the token with the lowest retention score.The policy maintains the predefined memory budget while adapting retention as new context arrives.
  • 4.3 INFERENCE: At 32K context, TRIM-KV achieves approximately 2× higher decoding throughput than full-cache decoding and is faster than SnapKV.The implementation is parallelizable and avoids materializing full attention or retention matrices during training.

5 EXPERIMENTS

TRIM-KV is evaluated across long-generation, long-context, qualitative, and ablation settings, consistently retaining strong performance under constrained KV budgets. The analyses show that learned retention scores recover task-relevant tokens and heterogeneous layer- and head-specific sparsity.

  • Math Reasoning Tasks: 198.4% relative improvement over attention-guided eviction baselines at the same budget, while TRIM-KV also delivers a 58.9% pass@1 gain over SeerAttn-R.On AIME24, GSM8K, and MATH-500, it can outperform eviction methods even with 4× larger KV budgets and sometimes surpass full-cache inference.
  • Long Procedural Generation Tasks: TRIM-KV consistently outperforms other eviction baselines on LongProc and surpasses the full-cache model in several settings.Retention gates trained on math-reasoning data generalize to this non-math procedural-generation benchmark.
  • Qualitative Result: Retention scores prioritize task-relevant tokens and attention sinks while discarding whitespace and punctuation early.The learned policy naturally recovers attention sinks, sliding windows, and other heuristics, with behavior adapting across layers and heads.
  • Qualitative Result: KV heads retain different token types and often preserve a small number of dispersed, high-context tokens rather than contiguous chunks.This head specialization makes selective token retention more budget-effective than chunk- or block-based caching in the reported analysis.
  • Qualitative Result: Later layers are typically sparser than earlier layers, but efficient per-head variable-length caches remain future work because current implementations assume uniform sequence lengths.Retention scores can support heterogeneous budgets across heads under a global constraint.
  • Long-context Decoding: On LongMemEval, TRIM-KV maintains good performance with only 25% of the budget, while all eviction methods struggle on incompressible SCBench retrieval tasks.It remains competitive across most SCBench tasks and performs strongly across both long-context and long-generation settings.
  • Ablation Studies: Removing the memory capacity loss causes a sharp drop in AIME24 pass@1, showing that this loss is essential for compression.Forward KL and next-token prediction each perform well alone, while their combination further improves accuracy.

6 CONCLUSION AND FUTURE WORK

TRIM-KV uses retention gates and distillation-based training to manage KV caches by intrinsic token importance, outperforming strong baselines and sometimes full-cache models. Future work proposes jointly training retention with attention and extending it to multimodal, tool-calling, and adaptive-budget settings.

  • TRIM-KV prioritizes tokens by intrinsic importance through lightweight retention gates, enforcing strict memory budgets with efficient eviction.The gates are trained using distillation and a capacity loss.
  • Across math reasoning, procedural generation, and conversational long-memory benchmarks, TRIM-KV outperforms strong eviction and retrieval baselines and sometimes surpasses full-cache models.
  • Retention scores align with human intuitions and expose layer- and head-specific dynamics, providing a probe for interpretability.
  • Future work would jointly train retention gating with attention layers instead of optimizing it over a frozen backbone and standard attention stack.This could support explicit task-performance and memory-use trade-offs during training.
  • Planned extensions include multimodal inputs, tool-calling applications, and adaptive memory budgets across layers, heads, and tasks.

DISCLAIMER

The document is presented for informational purposes by JPMorgan Chase & Co. and its affiliates, not as a Research Department product. It disclaims warranties and liability and is not investment research or advice.

  • The document is informational material from JPMorgan Chase & Co. and its affiliates, not a product of the Research Department.
  • It disclaims representations, warranties, and liability for the information’s completeness, accuracy, or reliability.
  • The document is not investment research, investment advice, or a recommendation, offer, or solicitation.

ETHICS STATEMENT

The work targets lower memory and computation costs for long-context reasoning, while acknowledging that efficiency gains do not address broader risks from deployment in high-stakes or resource-limited settings.

  • The method aims to reduce large language models’ memory and computational footprint, potentially lowering hardware costs for long-context reasoning.
  • Efficiency gains may accelerate deployment where misinformation, bias, and misuse risks persist, but the method does not mitigate those societal risks.
  • The authors recommend pairing efficiency improvements with continued efforts in safety, fairness, and responsible deployment.

REPRODUCIBILITY STATEMENT

The authors state that the paper provides detailed methodological and experimental information to support independent reproduction, with public datasets and planned release of code and checkpoints.

  • The paper describes its architecture, training objectives, evaluation benchmarks, hyperparameters, schedules, and implementation details.
  • The datasets are publicly available, and the authors plan to release code, model checkpoints, and evaluation scripts upon publication.

A METHODOLOGY

TRIM-KV assigns each token a retention score, computes attention over the retained cache, and evicts the least important tokens when capacity is exceeded. The design uses fixed-size memory with negligible score overhead and improves decoding throughput over full-cache and SnapKV decoding at 32K context.

  • Attention computation with KV eviction: Each decoding step projects the current token, computes its retention score, appends key-value-score entries, attends over the cache, and evicts tokens when |S_t| exceeds M.The eviction loop removes the least important cached token until the capacity constraint is satisfied.
  • Positional encoding: TRIM-KV is positional-encoding agnostic: exponential retention decay approximates eviction-induced forgetting rather than adding recency bias, and postrotated keys keep eviction orthogonal to RoPE.The mechanism is intended to work with absolute positions, RoPE, or no positional encoding.
  • Efficiency: At 32K context, TRIM-KV achieves ∼2× higher decoding throughput than full-cache decoding and is faster than SnapKV.The cache uses O(M) slots and stores one scalar score per token, adding approximately 1/d_h overhead relative to KV states.

B.1 LONG GENERATION EVALUATION

Across mathematical reasoning, procedural generation, long-memory dialogue, and long-context understanding, TRIM-KV remains effective under constrained KV budgets. It consistently outperforms eviction baselines, reaches the best LongProc results, and can exceed full-cache performance in selected settings.

  • Math reasoning results: On AIME24, TRIM-KV consistently outperforms eviction baselines across Qwen3-1.7B and DeepSeek-R1-Distill model families, while KeyDiff performs significantly worse than other baselines.Figure 6 compares competing algorithms across budgets, and Figure 7 evaluates the query-agnostic KeyDiff baseline.
  • Results on LongProc: TRIM-KV achieves the best performance across LongProc tasks and budgets, surpassing full-cache performance on COUNTDOWN at 0.5K/2K and widening its margin under tighter memory.Table 7 reports F1 for HTML-to-TSV and accuracy for the remaining tasks.
  • Long-memory evaluation: TRIM-KV significantly outperforms baseline eviction strategies on LongMemEvalS under a reusable fixed-size cache protocol for multi-turn, multi-session dialogue.The benchmark tests long-term memory across six question types, and Table 8 reports overall and partial accuracies.
  • Chunked-prefill evaluation: In chunked-prefill evaluation, TRIM-KV nearly matches full-KV performance on LongBench and surpasses full-KV by 6.5% on LongBench-V2.LocRet exhibits a 4.82% drop relative to full-KV on LongBench.
  • Ablation studies: The retention-gate architecture benefits from an MLP with a large positive initial bias, while training capacity M affects sparsity and deployment alignment.M = ∞ removes sparsity pressure and hurts performance; M = 128 outperforms attention-guided heuristics but can over-optimize sparsity, so matching deployment memory is recommended.

C ADDITIONAL QUALITATIVE RESULTS

Qualitative visualizations show that TRIM-KV develops layer- and head-specific retention patterns rather than applying one uniform eviction rule. Earlier layers often resemble sliding windows, while later heads specialize in functions such as gist-token retention.

  • Retention patterns: Earlier layers often exhibit sliding-window-like retention, whereas later layers develop clearer functional specializations.The visualizations show distinct retention and eviction matrices for individual layers and heads.
  • Gist-token specialization: Layer 27 head 2 mostly retains period tokens, suggesting that it may use gist tokens to summarize information from the preceding sentence.The pattern contrasts with strategies that preserve a recent chunk of tokens.
Loading 2512.03324v2…