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

Abstract: 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 (LongBench 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.

• The paper presents TRIM-KV, a novel method that dynamically computes token retention scores at creation and evicts low-importance tokens under memory constraints.
• It employs retention-gated attention, distillation, and capacity loss to align with LLM outputs while reducing computational overhead.
• Experimental results demonstrate that TRIM-KV outperforms traditional KV eviction strategies across diverse benchmarks, enhancing both efficiency and interpretability.

Token Retention for Memory-Bounded KV Cache in LLMs

Introduction

The paper "Cache What Lasts: Token Retention for Memory-Bounded KV Cache in LLMs" introduces TRIM-KV, an innovative method designed to address the challenges of memory and computation constraints in long-horizon LLM inference. The intrinsic quadratic cost of self-attention and the growing demand for key-value (KV) cache storage are identified as bottlenecks in efficient model deployment. Traditional methods like quantization and heuristic KV eviction tend to impose either high orchestration costs or rely on unreliable proxies of token importance based on attention.

TRIM-KV proposes a novel approach where each token's importance is evaluated at the time of its creation via a lightweight retention gate. The retention gate assigns a scalar retention score that reflects the long-term utility of a token for specific layers and heads. The score decays over time, leading to the eviction of low-importance tokens when memory limits are reached. This method ensures the KV cache maintains only the most crucial tokens, optimizing memory utilization and increasing efficiency especially in memory-constrained environments.

Methodology

Retention-Gated Attention: The paper introduces a mechanism called retention-gated attention. This involves computing a retention score $\beta \in [0,1]$ for each token that decays exponentially. This decay mimics human memory by gradually reducing the weight of older tokens. A high $\beta$ score suggests the token has critical importance and needs longer retention, whereas a lower score indicates that the token will be quickly discarded.

Training Framework: The retention gates are trained using distillation from a frozen LLM combined with a capacity loss. This minimizes the inference overhead and allows the model to optimally manage the KV cache. By incorporating a distillation loss and next-token prediction loss, TRIM-KV ensures alignment with the original LLM's output distribution while exploring sparsity patterns.

Eviction Policy: The scoring system is integral to determining eviction policy, where tokens with minimal score $\beta$ are evicted when the cache exceeds the preset memory budget. This is computed using a simple score comparison, making the inference process straightforward and efficient.

Figure 1: Patero frontiers of competing algorithms with different budgets on math benchmarks.

Experimental Results

The experimental results demonstrate significant improvements over competing KV eviction strategies in both controlled and less constrained memory settings. TRIM-KV exhibits superior performance across various tasks including mathematical reasoning datasets like GSM8K, MATH-500, and AIME24, as well as procedural generation and conversational benchmarks. Remarkably, it often surpasses full-cache models, revealing that selective retention can suppress noise from uninformative tokens.

TRIM-KV's designs naturally recover several heuristic strategies like sliding windows and gist compression without explicit instructions, indicating that learned retention scores align closely with human cognitive processes.

Qualitative Analysis

TRIM-KV's qualitative analysis shows that token retention scores serve as an insightful probe into LLM interpretability, reflecting layer- and head-specific dynamics. The model tends to assign high retention scores to task-critical tokens, while punctuation and filler tokens generally receive lower scores.

Figure 2: a) Average retention scores across all layers and heads of Qwen3-4B on tokens of an AIME24 example. b) Top 10 tokens with the highest (left table) and lowest (right table) average retention. c) The layer- and head-wise sparsity level estimated by token retentions.

Implications for AI Development

The approach presented by TRIM-KV suggests practical and theoretical advancements in optimizing LLM memory use, especially for applications requiring prolonged inference under strict memory constraints. The notion of treating retention scores as proxies for token importance shifts the focus from attention-based dynamics to intrinsic attributes of token utility over long contexts.

These advancements could influence future developments in AI, offering pathways to more efficient architectures that seamlessly balance performance and computational resource use. Additionally, the insights into head-specific dynamics may foster enhancements in model interpretability and diagnostics.

Conclusion

The paper showcases TRIM-KV as a compelling strategy for memory-bounded LLM operations. Its core contribution lies in leveraging retention scores for adaptive token caching, ensuring models remain efficient without retraining from scratch. This method potentially paves the way for future explorations of intrinsic token importance and memory-efficient model design tailored for complex reasoning and long-context scenarios.