Retention-Gated TRIM-KV: Efficient Cache Eviction

• Retention-Gated TRIM-KV is a memory management technique that uses learned gating mechanisms to dynamically retain essential token representations in transformer models.
• It implements diverse methods such as learned soft queries and token-specific retention gates to optimize key-value cache eviction under tight memory budgets.
• Empirical results demonstrate that these approaches reduce memory footprint and computational cost while maintaining or even improving task accuracy across various benchmarks.

Retention-Gated TRIM-KV refers to a class of methods for efficient key–value (KV) cache eviction in transformer-based LLMs and vision-LLMs (VLMs), where an explicit or learned gating mechanism determines which token representations are retained or evicted under a memory budget. These approaches optimize memory and computation during long-context inference by training, tuning, or dynamically optimizing the retention criterion, with strong empirical results across a variety of tasks, models, and budget regimes.

1. Conceptual Foundations and Motivation

Retention-Gated TRIM-KV is motivated by the need to control the unbounded growth of the KV cache during autoregressive sequence generation. Storing all intermediate keys and values across all layers results in memory and data movement cost that scales as $\mathcal{O}(L N d)$ with $L$ layers, $N$ prompt length, and $d$ hidden dimension (Wang et al., 2024). This becomes prohibitive for extended contexts or large-batch inference on modern LLMs and VLMs.

Traditional KV cache eviction strategies rely on static, local, or heuristic rules (such as retaining the last window of tokens), often ignoring global or semantic importance. Retention-gated approaches instead aim to learn, estimate, or optimize per-token or per-segment retention scores, using this information to perform budgeted trimming (“TRIM-KV”) that preserves essential information for downstream computation. The key technical innovation is the use of a gate—implemented as learned queries, neural scoring modules, or attention-based importance functions—ensuring that the KV cache preferentially retains the most valuable tokens at each layer or head (Liu et al., 13 Sep 2025, Bui et al., 3 Dec 2025).

2. Technical Mechanisms: Forms of Retention Gating

Retention-gated TRIM-KV methods instantiate the gating criterion through a spectrum of architectural and algorithmic choices. The major variants include:

• Learned Soft Queries (Judge Q): Retention gating is realized by augmenting the model with a small list of “soft” tokens whose embeddings are trainable. Their aggregated attention over the prefilling context provides global importance scores for each KV entry, enabling informed trimming under a budget. Only the embeddings are tuned, with all other model weights frozen (Liu et al., 13 Sep 2025).
• Token-specific Retention Gates (TRIM-KV): Each attention head and layer is augmented with a lightweight MLP that predicts, at token creation, a scalar retention score $r_t^{(\ell, h)} \in [0,1]$. This retention is assigned to the token and decays exponentially over time, reflecting intrinsic utility. Decayed scores serve as gating weights for attention aggregation and as an eviction criterion (Bui et al., 3 Dec 2025).
• Attention-based Top-K Gates (DynamicKV, FastKV, PrefixKV): Per-token attention-based saliency or importance scores are used to select the top-$k$ tokens for retention, either within each layer independently or via a coordinated global budget. Methods differ in whether they use pooled past attention weights, task-dependent dynamic allocation, or binary search optimization for configuring per-layer budgets (Zhou et al., 2024, Jo et al., 3 Feb 2025, Wang et al., 2024).

Across these mechanisms, the unifying principle is that retention is no longer governed by a static heuristic, but is adaptively or trainably determined by a learned or dynamically optimized gate.

3. Mathematical Formalism

Retention-gated TRIM-KV typically builds on the standard attention mechanism by introducing retention-related modifications:

• In TRIM-KV (Bui et al., 3 Dec 2025), for each token at position $i$ and current step $t$, the decayed retention is $\alpha_{t,i}^{(\ell, h)} = (r_i^{(\ell, h)})^{t-i}$, and the attention logit for that token is modulated by this quantity. The output is:

$$o_t^{(\ell, h)} = \sum_{i=1}^t \left[ \frac{ \exp(\alpha_{t,i}^{(\ell, h)}\, q_t^{(\ell, h)} \cdot k_i^{(\ell, h)}) }{\sum_{j=1}^t \exp(\alpha_{t,j}^{(\ell, h)}\, q_t^{(\ell, h)} \cdot k_j^{(\ell, h)}) } v_i^{(\ell, h)} \right]$$

At inference, tokens with the smallest decayed retention are evicted to maintain the budget.

• In Judge Q (Liu et al., 13 Sep 2025), aggregate importance for each KV pair $j$ is:

$$s_j = \frac{1}{n H} \sum_{h=1}^H \sum_{i=1}^n \alpha^{(\ell, h)}_{i j}$$

where $\alpha^{(\ell, h)}_{i j}$ is the attention from soft query $i$ to key $j$ at layer $\ell$ and head $h$. The top-$K$ keys are retained.

• In DynamicKV (Zhou et al., 2024), the attention-based gating mask at layer $\ell$ is:

$g_\ell = \mathbf{1}\{A_\ell \geq \tau_\ell\}$

where $A_\ell$ is the average pooled attention score and $\tau_\ell$ is the top-$B_\ell$ threshold.

Optimization for budget satisfaction and information preservation frequently relies on binary search or Lagrangian techniques to match global and per-layer or per-head compression constraints (Wang et al., 2024). Theoretical justifications sometimes invoke information coverage metrics or Gini coefficients to explain budget adaptivity.

4. Training Protocols and Implementation Strategies

Retention-gated TRIM-KV frameworks exhibit substantial architectural and training diversity:

• Embedding-only tuning: In Judge Q, only a small number of new soft-token embeddings are trained, with the alignment loss targeting similarity between their attention maps and those of actual response tokens, resulting in low training cost, preservation of the pretrained model, and rapid integration (Liu et al., 13 Sep 2025).
• Gate-parameter training: TRIM-KV inserts one or more MLP gates per head per layer, all pretrained transformer parameters are frozen, and only the gates are trained, typically using a combination of next-token prediction, distillation from a fixed teacher, and capacity-constrained losses which penalize exceeding the target cache size (Bui et al., 3 Dec 2025).
• Heuristic attention scoring: DynamicKV and FastKV infer saliency/importance using attention scores, with top-$k$ selection performed at periodic intervals or per layer, and budget allocation periodically rebalanced based on empirical activation patterns (Zhou et al., 2024, Jo et al., 3 Feb 2025).
• Task-adaptive budgets: DynamicKV can reallocate token retention across layers mid-inference, reflecting observed importance signal distributions, in contrast to purely static or input-agnostic allocations (Zhou et al., 2024).

Implementation generally involves minimal changes at the interface level: for instance, soft-embedding appends, gate-insertion “sandwiching” projection layers, and hooks for extracting attention patterns during prefill. At inference, cache management routines (e.g., per-step eviction based on retention score, dynamic resizing of layer-wise buffers) ensure strict adherence to memory constraints.

5. Empirical Performance and Benchmark Results

Comparable methods have been rigorously evaluated on established long-context and memory-bounded benchmarks. Key findings include:

Method	Cache Regime	Benchmark	Metric	Retention-Gated Gain	Source
Judge Q	512 tokens	LongBench (Llama-3)	avg. match score	+0.86 over SnapKV	(Liu et al., 13 Sep 2025)
Judge Q	1024 tokens	RULER	avg. match score	+5.91 over SnapKV	(Liu et al., 13 Sep 2025)
TRIM-KV	128 tokens	GSM8K (Qwen3-4B)	pass@1 (%)	27% (vs 10% SnapKV)	(Bui et al., 3 Dec 2025)
TRIM-KV	0.5k–8k tokens	LongProc	task acc.	Consistently > heuristics	(Bui et al., 3 Dec 2025)
DynamicKV	1.7% cache	LongBench (16 tasks)	avg. F1/Rouge/Acc.	85–90% of full cache	(Zhou et al., 2024)

These results demonstrate that retention-gated approaches deliver substantial memory and speed benefits with only marginal (sometimes negative) accuracy loss, compared to static or last-window baselines. Notably, some methods (TRIM-KV, Judge Q) can slightly exceed full-cache performance under certain settings (“regularization via hard eviction”) (Bui et al., 3 Dec 2025).

Qualitative analysis reveals that retention scores often align with emergent information structures—such as sliding-windows, quick adaptivity to topic shifts, and layer-specific “gist” or “sink” token behaviors—suggesting a degree of interpretability and flexibility absent in heuristic gating (Bui et al., 3 Dec 2025).

6. Comparative Analysis and Limitations

Retention-Gated TRIM-KV unifies a spectrum of strategies under a gating-and-eviction abstraction. The key technical and practical distinctions are:

• Learned queries (e.g., Judge Q): Trainable, capture global context, require low-overhead fine-tuning, but operate as probes rather than intrinsic attributes (Liu et al., 13 Sep 2025).
• Per-token retention gates (TRIM-KV): Intrinsic, trainable, reveal head/layer specialization, offer end-to-end control and interpretability, but require insertion of additional parameters and dedicated fine-tuning (Bui et al., 3 Dec 2025).
• Attention-based top-$k$ gating (DynamicKV, PrefixKV): Heuristic, adaptive, efficient, budget-flexible, can dynamically match task idiosyncrasies, but may underweight crucial tokens if attention aligns poorly with semantic relevance (Wang et al., 2024, Zhou et al., 2024).

Common limitations include the need for pre-set or tuned memory budgets, reliance on attention-score proxies (which may miss “silent” contributors), some training overhead for learnable variants, and in a few cases, possible inapplicability to settings where activation reordering is not feasible (e.g., strictly sequential, non-reorderable contexts) (Wang et al., 2024).

7. Connections, Interpretability, and Future Directions

Retention-gated TRIM-KV sits at the intersection of memory-efficient model inference, retrieval-augmented generation, and transformer interpretability. Several studies have highlighted that retention gates, when trained in data-driven fashion, rediscover classic heuristics such as sliding windows, content sinks, and gist compression, but also adapt dynamically to context or task (Bui et al., 3 Dec 2025, Zhou et al., 2024).

A plausible implication is that future work may further unify retention scoring with retrieval or planning heads, or leverage retention patterns for model diagnostic and interpretability purposes. There is also a trend toward leveraging layer-wise adaptation, cross-task calibration, and hybridization with other explicit memory mechanisms.

Retention-gated TRIM-KV currently represents state-of-the-art for practical, memory-bounded inference in autoregressive transformers, demonstrating robust trade-offs between efficiency and accuracy with interpretable, theoretically motivated control over cache retention.