Cache What Lasts: Token Retention for Memory-Bounded KV Cache in LLMs
Authors: Ngoc Bui
Shubham Sharma
Simran Lamba
Saumitra Mishra
Rex Ying
Venue: ICLR 2026 Poster (2026)
Keywords: KV cache
eviction
large language models
llm
long-context generation
Primary area: foundation or frontier models, including LLMs
TL;DR: We propose a learnable KV evicton method for long-context and long-horizon generation 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 (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.

Published as a conference paper at ICLR 2026
CACHE WHAT LASTS: TOKEN RETENTION FOR MEMORY-BOUNDEL KV CACHE IN LLMS
Ngoc Bui -Shuhham Sharma', Simran Lamba', Saumitra Mishra', Rex Ying X
Depaarment Computer Science, Yale University. 'JPMorganChase AI Research (ngoc. bui, rex. rex.ying) 8ya laalleedd
(shubham, x2. .2aaraa,, simran. lamba, swwmi henncaaa.ee.aa.isas........e.e.. jpmchase. com
O Source Code: https://gitim:: com com/nggocbh/trimk)
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 mneoorybbouneed 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 'ss intrinsic importance at creation time via a lightweight retention gate. Each gate predicts 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 frozen LLM combined with 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, even surpasses full-cache models some settings, showing that selective retention can serve 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- -pec-specific roles, suggesting a new path toward LLM interpretability.
INTRODUCTION
Modern large language models (LLMs) can, in principle, handle extremely long input contexts some recent models support context windows of 128k tokens or more (Yang et al.
2025 Gao et al.
2024). Yet, extending context length comes with steep computational costs. The self aattention mechanism has quadratic time complexity in sequence length, and storing the key-value (KV) cache
for thousands of tokens can quickly exhaust GPU memory (Wang et al. 2025 Li et al. 2024a). In practical deployments, the KV cache, which saves past key and value vectors to avoid re-computation, becomes major memory and latency bottleneck for long-context inference. Decoupling resource usage from context length is therefore critical for enabling efficient and scalable applications such as long-horizon reasoning (Chen et al. 2025) and lifelong agents (Zheng et al. 2025 Li et al. 2024d).
To address this challenge, recent work has explored memory-bounded LLMs that can operate effec- tively under constrained KV budgets (Li et al. (2024a). One line of research focuses on compression and quantization, aiming to reduce memory footprint by learning compact representations of past tokens rather than storing all keys and values explicitly (Hooper et al. 2024 Saxena et al. 2024).
Work done during an internship JPMorganChase AT Research

---

Published as a conference paper at ICLR 2026
These techniques mostly effective during the prefill phase but scale poorly with generation length. Another line leverages attention sparsity to offload most of the cache to CPU or secondary storage, and retrieve only relevant segments on demand via similarity search (Tang al. 2024) or learned indices (Gao et al. 2025). While offloading lowers the on-GPU footprint, it incurs nontrivial orchestration overhead that accumulates over long generations, undermining end-to-end throughput.
A more common and direct approach to enforce a fixed memory budget is KV cache eviction, which directly drops certain tokens from the KV cache (Xiao et al. (023). Many KV eviction strategies have been proposed to decide which tokens to remove. However, most of them are attention-guided heuristics: they track attention from new queries to cached tokens and retain those that are recently
or frequently attended, adapting the cache to the current focus (Zhang et al.. 2023 Li et al.. [024cc]
Wang et al., 2025 Liu et al., 2025 Ghadia al. 2025 Cai et al. (025). While being efficient, these methods assume that recent attention is reliable proxy for future importance. This assumption often breaks for long-horizon generation and reasoning tasks: token might be crucial much later, even it has not been attended to in the recent past (Jiang al. 2024). Moreover, attention- -baoodd----dddd-- eviction can suffer from attention bias, e.g., the model might temporarily overlook a needed token due to a distracting context (Shi et al. 2023), causing it to be evicted prematurely. While some recent studies have attempted to learn better eviction decisions (Chen et al. 2024 Zeng et al.. (024), these methods typically scale poorly with sequence length and therefore limited to the prefilling stage.
In this work, we take a new perspective on the KV eviction problem. Rather than relying on the attention-guider importance, we propose learn each token's 's intrinsic importance at the time of its creation and use that as the basis for eviction. Intuitively, not all tokens are created equal: some carry significant semantic or task-related weight (e.g. critical fact, a question being answered, or the first few "sink" tokens that often encode the topic or instructions), while others are relatively inconsequential (e.g. filler words, stopwords, or trivial arithmetic steps in a chain-of- thain-of-thought). Moreover, the importance of tokens is not uniform across the network, but it varies systematically by layers and heads, reflecting their functional specializations (Voita et al. 2019 Wu et al., 2024b).
We posit that the contextual embedding of a token already encodes much of its long-term utility. We therefore introduce a retention gate that maps the token's 's embedding and produces a scalar retention score 3 E [0, IJ reflecting the token's 's inherent importance for specific layer and head. Especially, we design this retention score to decay exponentially as the context grows, mimicking the gradual forgetting of old information in human brains (Ebbinghaus) 2013). Thus, a highly important token will have se and retain a high score for a long time, whereas a token deemed unimportant will have closer to 0 and its influence will vanish quickly. We leverage this score to drive a simple eviction policy: whenever the number of cached tokens exceeds the budget M, we evict the token with the smallest current retention score. This approach, which we call Token RetentIon for Memory-bounder KV Cache (TRIM-KV), ensures that at all times, the cache is filled with the M tokens judged most intrinsically important, with a preference toward more recently generated tokens.
Implementing retention-based caching in an existing LLM only requires adding a few lightweight components. We integrate the retention gates into each self-attention layer of pretrained model to modulate attention weights by token importance during training. We then train only the gates with a two-part loss: a distillation loss that compels the modified model to mimic the original model's outputs, thus preserving quality, and a capacity loss that penalizes exceeding the target memory budget, thus encouraging sparseness in attention via eviction. Importantly, by training the gates across all layers jointly, the model can learn a coordinated, globally optimal caching policy rather than greedy layer-wise decisions. At inference time, the learned retention gates produce per-token scores on the fly, and eviction is implemented with simple score comparison, adding minimal overhead.
Results and Contributions. Through extensive experiments on long-context and generennration benchmarks, demonstrate that our learnable token retention approach substantially improves the performance of memory-bounded L.LMs. On challenging mathematical reasoning datasets, GSM8K, MATH, AIME, long procedural generation benchmark, LongProe, and long-memory chat assistant benchmark, Memval our method consistently outperforms eviction baselines, even when those baselines use 4x more KV budget, and deliver 58.9% relative gain on pass@) compared to the SOTA learaalle KV retrieval baseline (Gao et al. 2025). especially in low-memory regimes. Remarkably. in several settings, TRIM-KV even surpasses a full-cache model, suggesting that selective retention can function as an effective regularizer by suppressing noise from uninformative tokens.
2

---

Published as a conference paper at ICLR 2026
We also present qualitative evidence that learned retention scores align with human intuition: the model tends to assign high scores to initial tokens and problem descriptions, and low scores to less meaningful punctuation. Notably, many behaviors reminiscent of common heuristics, such keeping sink tokens, sliding windows, and gist compression (Mu et al. 2023), emerge naturally and adaptively from our learned policy, without being hard-coded. Finally, we show that these learned retention scores can also act as a diagnostic tool for probing layer- and head-specific dynamics, providing a lightweight means to analyze and ultimately improve the interpretability of attention patterns.
RELATED WORK
KV Cache Compression. As model sizes and context windows grow, optimizing KV-cache memory
is increasingly critical. Prior work largely falls into three directions: (i) token eviction/merging (Xiao
al.
2023 Li al. 2024c Zhang al.
2023; Nawrot al.) 2024, Zhang et al.
2024 Qin al.
2025
Wang et al. 2025; Liu ct al. 2025
Park et al. Cai et al.
2025 Park et al.
2025 Kim
al.
2024). (ii) vector compression/quaniizzn
(Hooper et al. 2204
Liu al. [024b]
Yue et al.
2024
Sun et al.) 2024a). and (iii) token retrieval (Tang et al. 2024 Liu et al. 2024a Gao et al.
2025).
While effective in many settings, vector compression and retrieval either discard fine-grained
information introduce nontrivial systems overhead (e.g., coordination and data movement) (Li
ct al. 2024a). Moreover, their memory and computation still scale with sequence length, making them inefficient for long-horizon generation applications. Token evietion offers simple, memory- bounded alternative; however, most existing policies are heuristic and can significantly degrade performance, especially on long reasoning trajectories. Recent work has introduced learnable eviction policies (Chen et al.) 2024 Zeng et al. 2024 Huang et al. 2024), but these are primarily designed for the pre-filling stage and thus not well suited to sustained long-horizon generation. We bridge this gap by introducing a learnable and efficient eviction policy designed for long-horizon LLM inference under fixed memory budgets.
Forgetting in Language Models. A key limitation of vanilla self-attention is the lack of an explicit forgetting mechanism, forcing the model earry potentially irrelevant information and making long-context processing inefficient. Early work tackled this by replacing quadratic attention with linearized and recurrent variants (Katharopoulos
al. 2020 Wang al.) 2020 [un et al. 2023 Yang et al. 2023 2024) that summarize the past into a fixed-size state, often a single vector. While computationally attractive, such heavy compression can degrade performance on tasks requiring long-range memory. Follow-up studies (Behrouz et al. 2024| Sun et al. 2024bb Karami et al. Karami and Mirrokni 2025) increase memory capacity by replacing this hidden vector with a more expressive neural state. Complementary lines of work retain softmax attention but enforce forgetting by modifying attention logits (Lin et al., (025) or imposing trainable sparsity patterns (Yuan et al. 2025). However, these approaches typically alter attention dynamics substantially and thus require training models from scratch. This incurs significant training cost and leaves their scalability to contemporary LLM sizes uncertain. In contrast, we introduce a plug-in forgetting mechanism for pretrained LLMs that converts them into memory- bouuddd models, providing long-context efficiency without retraining from scratch.
PRELIMINARIES
3.1 TRANSFORMERS WITH SELF-ATTENTION
Given a sequence of d-dimensional input vectors Xi...., Xy, (causal) self-attention layer attends only to past positions. For each = 1,...,,7 ,,, the attention output o, computed as
exp (q. ks)
q Wox = WKX, V2 Wvx, O,
2j=1 exp (q. kj)
where q, k, v are query. key, and value states, respectively, and Wo. WK, Wv E Rddd are lincar transformation weights. Here, assume single-head attention layer and omit the scaling factor 1/vd for simplicity. The sequence of key-value pairs {(k, vz)}, is the in-context memory of the LLM. During the autoregressive decoding, we typically generate one token at a time and cache the running key-value pair (k,, vt) to our in-context memory avoid recomputation. However, this
3

---

Published as a conference paper at ICLR 2026
vanilla caching approach leads to a linear increase in memory footprint with the sequence length, while computation grows quadratically (Keles et al. 2023). This reduces efficiency when handling long-context inputs and extended generation tasks.
3.2 REVISITING KV CACHE EVICTION
common method to address the linear growth in the memory is to prune or compress the running key-value pairs into fixed-size (slot) memory. As new tokens arrive, we evict un-for less- Jimportnn tokens from our memory and append the new ones, To understand this procedure, we revisit and rewrite the attention computation with eviction at inference step as follows:
Oie exp (q k;)
0 Ej=1 Oty exp (q, k,)
Vi where On E {0,1) and on 2 0441 Vi,
In Equation ① introduce binary variable du 0, indicating whether key-value pair i has been evicted at time & and the monotonicity constraint Oi
04+11,1 ensures that cannot retrieve token once evicted (Figure I. The goal is to choose a decision variable so that the attention output deviates little as possible from the full KV cache (all Cty 1, ti,
0
min Lbase (0/i Oz) 8. t.
Ces S M.
(2)
0
1=1
Figure 1: Attention
Here. L penalizes differences between attention with and without eviction, and the constraint enforces keeping at most M tokens at any inference step L.
w/ eviction (M 3)-
Solving the above constrained optimization at every time step is impractical due to its combinatorial nature and efficiency requirements of LLM inference in worldw applications. Most existing
approaches (Xiao al. 2023 Han al. 2023 Zhang al. 2023 Li 2024c] Cai al. 2025
Ghadia et al. 2025) opt to determine a heuristically while we focus on a learnable eviction method.
METHODOLOGY
In this section, we propose a learning-basee eviction policy that prunes the KV cache based on the intrinsic importance of the tokens at each layer and head. The policy ranks tokens by relative importance to decide which should be evicted from the KV memory. To learn token importance, we introduce small neural network that takes token embeddings input and produces scalar retention score. We then integrate this retention score into the attention computation to modulate the attention weights. We term this proxy attention mechanism a retention- attention. We train the LLM with retention-gate -ated attention against baseline model with standard attention, using a combination of distillation and hinge-liks like regularization losses enforce memory capacity constraints while preserving response quality. A visualization is shown in Figure
4.1 SELECTIVE IN-CONTEXT MEMORY VIA RETENTION-GATER ATTENTION
We introduce retention-gan gatedd attention, trainable mechanism that mimics the information loss induced by inference-time eviction. From the formulation (T). the sequence Oziz (((+ Otti represents how token retained in the attention computation over time. Retention begins 1 and then abruptly drops to 0 once the token is evicted. While this binary behavior matches the inference stage, poses challenges for learning: the signal is discrete, non-differentiable, thus providing no gradients for optimization. To remedy this, we replace the hard binary variable with smooth, monotonically decreasing function that models the gradual decay of importance while enabling gradient-basee training. A natural candidate is the sigmoid function, ãy = 1/(1 + ex, t))). which predicts the time at which the token is evicted. However, this design suffers from two drawbacks: (i) the domain of is unnormalized since the sequence length is unknown during decoding, and (ii) the sigmoid flattens across most of its range, producing negligible variation between steps and leading to vanishing gradients during training.
To overcome these limitations, adopt exponential decay formulation, āu 8 where 3, E [0, 1]. to model the retention rate of token over time. Larger values of 3; correspond to higher intrinsic importance, implying slower decay and stronger memory retention. Substituting this design
4

---

Published as a conference paper at ICLR 2026
for in Equation ① yields our proposed retention-gated attention:
8 8 exp (q. k,)
q: = Wox,, = WKX, Vi = Wvx,, 3r - g(xx), O, 2j=1 exp (q. k2)
Vi. (3)
Here, we propose a retention gate g. which is a lightweight network, to parametrize the token importance 8,. The retention gate can be a linear projection, i.e., g(x) = a y W E R Rixd or a simple MLP, i.e., g(x) a (MLPP(x) + b). The sigmoid function squashes the output of g to the range [0, 1]. while b is a learnable bias. When all 3, 1, Vt, our retention-gate recovers the vanilla attention. By incorporating the retention score into the exponential term, the attention weight can be written as exp(q, k, + (t i) log 3,). which reveals that the retention score acts as an additive bias on the attention logits (Press et al. 2021).
attention
Token Retention vs. Attention Scores. In standard self-attention, the importance of past token i at decoding step t is given by au & exp(q: k;), which depends explicitly on the current query qr- These scores capture short-ferm utility for predicting the next token and are recomputed at every step, making them local, myopic, and highly dependent on the transient decoding state.
KV cache eviction, in contrast, is long- horizon decision: once token is dropped, it cannot influence any future prediction. An effective eviction policy should depend on a token's is intrinsic long-term utility that reflects how useful it is expected be over the remainder of the sequence and how long it has already stayed in the cache, rather than on the current query alone.
Token retention provides a more suitable abstraction. Instead of asking "how much should token i contribute now?" it asks "how important is token i for the long run, and for how long should it stay in the cache?" Concretely. each token receives scalar retention score Br [0. 1] based only on information available at creation time (its representation, layer, and head), and its effective contribution at future step t decays as , This yields smooth, exponentially decaying retention curve that is aligned naturally with long-term utility under memory constraints.
Brain-inspired Interpretation. Our intentioo-gated attention bears a natural connection to the Ebbinghaus's 's forgetting curve (Ebbinghaus) 2013), which models human memory retention as exponential decay, R = exp( -t),, where S denotes memory strength (Wožniak et al. 1995).
Similarly, our retention- gatedd attention models each token's 's contribution as exponentially decaying with temporal distance, exp((t - i) log 3,). Tokens start with full weight and gradually lose influence as new tokens arrive. The parameter 33 acts analogously to memory strength: larger values yield more persistent, durable memories, while smaller values indicate weaker memories that fade quickly. This design embeds a forgetting mechanism into attention, enabling dynamic prioritization of recent or salient tokens while discarding less useful context, akin to human memory. Unlike prior work that applies the forgetting curve to directly into the attention mechanism.
generation (Zhong et al. 2024), we integrate
4.2 TRAINING
Cqualiy
Cypp
Our goal is to train the retention gate g so that the LLM can preserve response quality under memory constraint, thereby bridging the gap with the inference stage. Instead of training separate gate for each layer and head, formu- lated in Problem (2). we optimize all retention gates jointly an end-to-end fashion. This holis- tic approach mitigates error propagation. allow- ing the model to learn coordinated, globally optimal caching policy rather than greedy layer- wise decisions. Starting from a pretrained LLM, we replace every standard attention block with our proposed retention gated attention. Each block is equipped with a lightweight retention gate g that maps token representations retention scores B, E [0, 1), which are then used to modulate attention weights according to Equation (3) We call this proxy L.L a retention-gated L.M..
Transfomerr Blook
Standard Attention
Retention Attention
o
KK
Linoer
Unser
X
Figure 2: Training architecture.
5

---

Published as a conference paper at ICLR 2026
Objectives. To train these retention gates, we formulate the training objective that balances two goals: (i) preserving the predictive quality of the original pretrained LLM, and (ii) enforcing memory capacity constraints by controlling the sum of retention scores at each step.
For the first objective, we use a combination of the distillation and standard next-token prediction losses. The distillation loss encourages the proxy LLM to align its output distribution with that of the baseline LLM using standard attention. In parallel, the next-token prediction loss enables the model uncover sparsity patterns directly from the data, extending beyond the constraints of the pretrained LLM. Let p( ) and qa(-) be the output distribution of the pretrained LLM and retention-gan -gatedd LLM, respectively, where 8 denotes the parameters of all retention gates. The quality loss is given by
Lquality = LKL + LNP = DRu(p( (z)) ll go((((z)) + E log go (y|z)]
(4)
Here, DKL is the standard forward Kullback- Lulbber divergence (Kullback and Leibler 1951).
For the second objective, we impose a hinge-like like regularization penalty, which discourages the model from exceeding the available KV memory slots at each step. For a retention gate within a given layer and KV head, the memory capacity loss is defined as:
1 T
E max {0, D0T- 1=1
Leap
(5)
[=1
where T is the sequence length and M is the predefined memory capacity. Here, M acts as a soft hyperparameter, primarily intended prevent over-optimization during the early decoding stage when the sequence remains short. Training is performed with a fixed value of M, while inference can flexibly accommodate different KV budgets. This regularization is applied uniformly across all layers and KV heads of the transformer. The combined training objective is then:
min Lquality + Ceap
where Aeap hyperparameter balancing between quality and capacity loss. Note that during training, only the retention gate parameters are updated, while all other model weights remain frozen.
Hardware-aware Computation. Retention gated attention is fully parallelizable and compatible with FlashAttention-sttle kernels (Dao 2023). We implement it with FlexAttention (Dong et al.) 2024) plus a custom Triton kernel for the capacity loss Leap. performing forward/backward without materializing the full attention or 3 matrices. This enables long-context training (up to 128K tokens on four H100 GPUs) with minor overhead versus standard parameter-efficien fine-tuning.
4.3 INFERENCE
At inference time, the base LLM is augmented with the retention gates learned during training (Section 4.2). These gates provide token-level intrinsic importance scores 3,, which quantify how strongly each past token should be retained for future computations. Unlike training, where the retention gates are used to modulate the attention weights, at inference, they act purely as decision- makers for eviction, operating alongside but independently of attention computation.
The eviction process is designed to ensure that the KV cache respects a predefined memory budget. Let S, C (1....,t) f) denote the set of tokens currently stored in the KV cache at decoding step t. When a new token f + 1 is generated, it is provisionally added to the cache. If this addition causes the cache size to exceed the memory capacity M, an eviction is triggered. The eviction rule is simple yet principled: we remove the token with the lowest retention score, i.e.,
Jevie arg ( E St). jES.
Intuitively, this favors retaining tokens deemed globally important by the learned retention gates while discarding those with little long-term value. In practice, this makes inference both memory- efficient and adaptive: new context arrives, the model continually re-evaluates the importance of older tokens, enabling long-context generation while keeping memory usage bounded. Algorithm [ illustrates single decoding step, where attention computation is coupled with token eviction.
Complexity. Our inference is simpler and more efficient than existing works, including pure heuristic baselines (Li et al.) (224c)). Table shows that, 32K context, TRIM- KV achieves 2x higher decoding throughput than full-cache decoding and even faster than SnapKV, a purely heuristic method. More details provided in Appendix A.2

---

Published as a conference paper at ICLR 2026
TRIM-KV (ours)
SnapKV
SeorAttn-R
AIME24
GSM8K
MATH-500
80
Owen3 48
(96)
60
75
60
40
50
40
20
0
256
512
1024
2048
4096
32768
64 128
512
1024
16384
128
2022
16384
80
100
80
75
Qwenn
(96
60
50
40
20
25
20
0
512
102 2048 4096
32766
64
128
512
1024
2048
2048
16384
64
512 11288004
26384
100
80
Owwnn3 148
75
60
40
50
40
25
256 512 2048 4096
32768
64
128
512
202
16384
64 128
512
1220820088
16384
KV Budget
KV Budget
KV Budget
Figure 3: Patero frontiers of competing algorithms with different budgets on math benchmarks.
EXPERIMENTS
In this section, we conduct extensive experiments to demonstrate the performance advantages of our method on both long-context and reng-generation tasks.
5.1 LONG GENERATION EVALUATION
Benchmarks. Following prior work (Gao et al. 2025 Cai et al. 2025), evaluate standard math-reasoning
(Art of Problem Solving 2024). GSMSK (Cobbe et al. 2021), and MATH-500 (Hendrycks et al. 2021). To assess performance beyond math reasoning and under
long-context settings, we also report results on LongProe (Ye et al. 2025). Following (Gao et al. 2025), report average pass@ accuracy over 64 samples for AIME24 and 8 samples for GSM8K, MATH-500. We use greedy decoding for LongProc as the default in the benchmark.
Base Models. Following (Gao al. 2025), we mainly use Qwen3's family models (Yang al.
2025 including Qwen3-1.7B. Qwen3-4B, Qwen3-8B, Qwen3-14B and DeepSeek RI Distill (uuo al. 2025) variants including, DeepSeek-RR Bicck----Dsstlll--O --Distill-Qwen-7B and DeepSeck-RI-Disse Dis---illl--ilam----
We report the results with Qwen3 models in the main paper, and the remaining is in Appendix
B
Baselines. We compare our method against SeerAttn-R (Gao al. 2025). R-KV (Cai al. 2025). SnapKV (Li et al. 2024c), H2O (Zhang et al.) 2023), StreamingL.L. (Xiao et al. 2023). R-KV SnapKV, H2O, and StreamingLLMM are heuristic, recency-driven KV eviction policies for long- form generation under a fixed memory budget. SeerAutn-R is learnable KV retrieval approach for reasoning tasks: rather than evicting, offloads the full KV cache host memory and uses recent queries to fetch relevant blocks for attention. KV retrieval methods preserve all past information but require nontrivial CPU-GPU orchestration and incur offloading overhead. We therefore treat SeerAttn-R strong leamable baseline, and R- R-KV/SnapKV representative eviction baselines.
Implementation Details. We train the retention gates using OpenRI-PeA -MAH-- 220k (Hugging Face) dataset, similar to (Gao et al., (025). Note that only train the retention gates' weights while keeping the original model parameters frozen. We set the objective hyperparameters Arap = 1.0 and the memory capacity M 256. Each transformer block has retention gate g, which is single
7

---

Published as a conference paper at ICLR 2026
Layer
Layer Head
Layer
Layer Head
Layer
Layer Head
Figure 4: Visualization of token retention score 8 (top) and eviction decisions oz (bottom).
MLP layer with the hidden dimension of 512, thus having dimensions of d - 512 + h, where is
the number of KV heads. We the activation function as the default activation function in MLP layers of the base model. We initialize the bias in the retention gates to large value (e.g.. == 18.0) begin training with minimal forgetting or compression. All trainings are on 4 H100 80G GPUs.
5.1.1 QUANTITATIVE RESULT
Math Reasoning Tasks. Figure shows our method outperforming baselines by a large margin, especially in low-budget regimes. Notably, TRIM-KV surpasses attention guided methods (R-KV, SnapKV) even when they are given 4x KV budget. Under the same budget, i..511 for AIME24 and 128 for GSMSK/MATH-500 it yields 198.4% relative improvement over these baselines. Against the SOTA learnable KV retrieval baseline, TRIM-KV outperforms SeerAttn-R across all settings, yielding 58.9% pass @1 gain the same budget. Crucially, TRIM-KV operates in pure KV-eviction regime, a stricter setting than the KV retrieval methods such as SeerAttn-R, and thus avoids CPU-GPU offloading overhead. In some settings, like for Qwen3-4B model and AIME24 dataset, our method can even surpass the standard full KV cache. These results suggest that large fraction of KV-cache tokens reasoning models redundant and can be discarded without degrading performance.
Long Procedural Generation Tasks. We evaluate KV-eviction methods on tasks that require both long-
Methodgy budjet
CountDown
Pseudo
context comprehension and extended generation. Ta-
FullKV
50.8
25.0
ble reports results with Qwen3-4B model. Over- all, TRIM-KY consistently outperforms all other evic- tion baselines and, in several settings, even surpasses the full-cache model. Moreover, this result highlights that TRIM-KV with retention gates trained on math- reasoming data generalizes well non-math tasks, Full
StreamingLLLMamns
5.0
20.6
1.5
H2O
12.0
2.5
33.7
0.5
Snapk Vonis
57.0
49.0
13.0
42.7
4.5
R KV
88.5
81.0
63.0
48.2
2.5
TRIM-KV
96.5
93.5
66.3
49.3
18.0
Table 1: Qwen3-4B LongProc. Bold is for
results and analysis are provided in Appendix B
the best, underline is for the best KV eviction.
5.1.2 QUALITATIVE RESULT
To examine the eviction policy learned by retention gates, we run TRIM-KV on Qwen3-4B for the first example in AIME24. Figure 5 show the mean retention score, averaged over all layers and heads, for each token in the example sequence. Aligning with our intuition, retention gates assign high scores to task-relevant tokens (e.g., ometer, shop, wa 1 k, mi nutes) and to the initial token attention sink. In contrast, whitespace and punctuation
< 1lim_sta mm__tarr
which often serves receive low retention scores and are discarded carly, yielding short lifespans in the KV cache. Next, we examine retention scores and eviction decisions at layer-head granularity.
Emergent Eviction Heuristics. Figure visualizes the retention scores 8 and eviction decisions ae: for selected layers and heads. Many eviction heuristics, such as attention sinks (Xiao et al. (2023). sliding windows (7hu et al. 2021). A-shape (Jiang et al. 2024), emerge naturally from our learned
8

---

Published as a conference paper at ICLR 2026
giaa
3.010
34114
52990
4.000
20
3.013
0.7
44779
36
03
100
1000
1.500
2000
2500
3000
35.00
4000
101011
Figure 5: 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. e) The layer- and hese-wise sparsity level estimated by token retentions.
Method
MultiiChhhice
(CL.Ma Math
Codk: MekkRRepoAA
Sum
Mis Sum=NIAH H
Rew MuleHHop
KV
28.78
SceamingLLLMam
5.08
2.20
100.00
13:20
6.84
2.96
29.21
28.25
0.00
H20
4,80
0.00
100.00
8.00
3.70
0.46
8.97
6.51
0.27
SnapKV
10.04
0.00
100.00
18.60
6.29
0.23
27.90
29 29.28
0.31
TRIM KVanoe
23.58
100.00
13:40
13.68
34,02
35.44
49.16
Table 2: Performance long-context tasks from the SCBench benchmark.
policy without being hard-coded, and they adapt to the functional roles of individual layers and heads. For instance, sliding-window window behavior is common in early layers, whereas attention sinks appear more frequently in later layers (see Figure [1 and 12 for a comprehensive view). Moreover, TRIM-KV adapts the window size by layer and head: in Layer I/Head 1, tokens receive nearly uniform retention scores, so the KV cache behaves like a recency-biased window that keeps the most recent tokens; in Layer O/Head 0, multiple sliding windows of varying widths emerge from the learned policy; in Layer 15/Head 2, no sliding window is observed because certain tokens receive substantially higher retention than others, suggesting specialized functional role for this head. The A- A-shaped pattern typically appears in layers that emphasize instruction/proolenn -eteemm--tteeenttt tokens (e.g.. Layer 20/Head and Layer 30/Head 4) or chain-of-thoo thoothh//rrssoonig prompts (e.g.. Layer 35/Head 1). These heads also exhibit context switching, where small, dense lower-triangular blocks emerge and then fade quickly when the context changes or sentence completes. To the best of our knowledge, the absence of sliding window, the presence of multiple coexisting windows, and context switching are newly observed eviction patterns that arise naturally from our learned policy.
Token Retention Enables Interpretability. Beyond guiding eviction policy, token-level retention scores provide a diagnostic tool for analyzing the functional roles of individual KV heads in the base LLM. Visualizing retention scores alongside the tokens preserved in the KV cache after generation reveals distinct specializations: some heads emphasize a recency window (Figure T4 and T7). whereas others preferentially retain mathematical tokens-numbers and operators (Figures 19 as well as problem-description tokens (Figures [6). Even function or filler words, such as pronouns, prepositions, conjunctions, wait and let. tend be kept by specific heads (Figures 13) In particular, some heads tend to retain the period token while discarding others (Figures 18] We
hypothesize that, in these heads, periods act implicit gist tokens Mu al. 2023). summarizing the information in the preceding sentences.
Our analyses indicate that KV heads in LLM develop different functional roles and therefore keep different types of tokens. These tokens are often dispersed across the context rather than forming contiguous chunks, as each already captures contextual information. This observation contrasts with existing approaches (Yuan et al.) 2025) Gao et al.) (025) that advocate chunk- or block-based KV- cache. Instead, we show that keeping a small number of high-context tokens more budget-effective.
Budget Allocation. Figure 5p reports head- and layer-wise sparsity estimated from the retention scores i.e., 1 2
T(T+I) Eiee g.- We observe that later layers are typically sparser than earlier ones, consistent with prior findings in (Cai et al. 2024). Practically, the retention scores enable heterogeneous budgets across KV heads under a global constraint by evicting tokens with low global retention. However, existing KV-cache and FlashAttention implementations assume uniform sequence lengths across heads within a layer; enabling efficient per-head variable-length caches is left to future work.

---

Published as a conference paper at ICLR 2026
5.2 LONG-CONTEXT EVALUATION
Long-context Decoding. We evaluate TRIM-KV on Long Wu
Methodgy sudee
Acc
al. 20244) and SCBench (Li al., 2024b), two 128K-token benchmarks
Full KV Virorr
49.4
for long-term memory and long-context KV compression. We use Qwen3-
Streamingt M.
27.6
4B-Instruct (Qwen3 2025) the base model and train the retention gates
SnapK V 52708
27.8
on SynthLong (Lazarevich et al., 2025), BookSum (Kryšciñski et al.) (201))
TRIM-KV 52708
44.8
and Buddhi (Singhal 2024) with M 1024. All other experimental settings
Streaming Lisss8
19.0
follow S.c As shown in Table 3 and T TRIM-KV significantly outperforms baselines on Long maintaining good performance with only 25% of the budget, while others degrade sharply. On SCBench, it re-
SnapKV 16384
18.2
TRIM-KV L6381
38.6
StrcamingLLMete2
13.0
SnapKV K112
15.8
mains competitive across most tasks, though all KV eviction methods struggle on incompressible retrieval tasks (e.g., Retr.KV, Code. Repo,), consistent Table 3: Results on with (Li al. 2024b). Overall, TRIM-KV shows strong performance in both long-context and long-generation settings, unlike prior work that targets only one stage. More details of this experiment are Appendix B.2
TRIM-KV
34,0
LongMemEvals.
Chunk Prefilling. We follow Huang et al. (2024) to evaluate TRIM-KV
Method
Acc
A (%)
the chunk-prefill settings. We train Phi3-mini-1 1-iK-iiii-KKKK
128K (Abdin et al. 2024) on Full KV
28.79
0.00
LongAlpaca (Chen et al., (023) with M 512. All experimental settings
LocRet
28.03
-2.64
TRIM-KV
34.09
+18.41
follow (Huang et al.) (024)) for fair comparison, detailed in Section B.3
Table T0 shows that TRIM-KV outperforms the LocRet, trainable KV
Table 4:
Results
eviction baseline, and even full KV inference, by a significant margin.
LongBench-V2.
5.3 ABLATION STUDIES
We ablate the objective by training the Qwen3-4B retention gates with different loss combinations and report AIME24 pass@) 1 at a 4096- token budget in Table 5 Both forward KL and next-token prediction perform well on their own, and their combination further improves accuracy. The memory capacity loss is essential for compression, and removing it leads to a sharp drop. We provide comprehensive ablations with reversed KL, generalization with different training datasets, gate's s architecture, and other hyperparameters such as M in Appendix B.4
Methodgy balge
pass
Full K252008
65.5
TRIM-KV s008
75.8
(TRIM-KV -LiI.)on
72.1
(TRIM-KV -an
72.5
(TRIM-KV KV -Copphuus
42.9
Table 5: Objective ablation.
6 CONCLUSION AND FUTURE WORK
We introduced TRIM-KV, a learnable, retention gaten approach to KV cache management that prioritizes tokens by intrinsic importance rather than recent attention. By training lightweight gates with distillation and a capacity loss, our method enforces strict memory budgets with a simple and efficient eviction policy. Extensive experiments across math reasoning, aroopprceeurrll generation, and conversational long-memory benchmarks demonstrate that our method outperforms strong eviction and retrieval baselines, and sometimes surpasses full-cache models. Analyses show that the learned retention scores align with human intuitions and reveal layer- and head-specific dynamics, offering simple probe for interpretability.
Future Work. Our results indicate that retention-gated attention is an effective learnable proxy for approximating standard attention with eviction during inference. In the current work, however, we keep the backbone parameters frozen during training and still rely on standard attention at inference time. A natural next step is replace standard attention with retention- -ated attention and train the retention mechanism jointly with the attention layers during pretraining post-training. This could allow the retention scores to better cooperate with the learned query, key, and value states, shaping the model's behaviors from the outset rather than optimizing retention top of fixed attention stack. Such a design would enable training objectives that explicitly trade off task performance and memory usage, potentially yielding models that are inherently memory- bounded without requiring any post hoc compression policy. Besides, building on these results, plan extend retention gating multimodal inputs, tool- calling applications, and develop adaptive budgets that allocate memory across layers, heads, and tasks to further improve both performance and efficiency.
10

---

Published as a conference paper at ICLR 2026
DISCLAIMER
This paper was prepared for informational purposes by the Artificial Intelligence Research group of JPMorgan Chase & Co. and its affiliates ("JP Morgan") and is not product of the Research Department of JP Morgan. JP Morgan makes no representation and warranty whatsoever and disclaims all liability, for the completeness, accuracy or reliability of the information contained herein. This document is not intended as investment research or investment advice, or recommendation offer or solicitation for the purchase or sale of any security, financial instrument, financial product or service, or to be used in any way for evaluating the merits of participating in any transaction, and shall not constitute a solicitation under any jurisdiction or to any person, if such solicitation under such jurisdiction or such person would be unlawful.
ETHICS STATEMENT
This work aims to improve the efficiency of large language models by reducing their memory and computational footprint. Our method can make long-context reasoning more accessible by lowering hardware costs, which may democratize access advanced LLM capabilities. However, efficiency improvements may also accelerate the deployment of LLMs in high-stakes or sesource-limite settings where risks around misinformation, bias, misuse persist. We stress that method does not mitigate these broader societal risks and should be paired with ongoing efforts in safety, fairness, and responsible deployment.
REPRODUCIBILITY STATEMENT
We ensure reproducibility by providing detailed descriptions of the model architecture, training objectives, and evaluation benchmarks in the main text and appendix. Hyperparameters, training schedules, and implementation details are included for all experiments. All datasets we use are publicly available, and we will release code, model checkpoints, and scripts for training and evaluation upon publication. Together, these materials allow independent researchers to fully reproduce and verify our results.
The authors used large language models to help refine and polish the writing of this manuscript.
REFERENCES
Marah Abdin, Jyoti Aneja, Harkirat Behl, Sébastien Bubeck, Ronen Eldan, Suriya Gunasckar, Michael Harrison, Russell Hewett. Mojan Javaheripi, Piero Kauffmman, et al. Phi-4 technical report. arXiv preprint arXiv:2412.08905, 2024.
Art of Problem Solving.
AIME problems and solutions, 2024.
URL
https
7/artofproblemsslviing Solutions
phpp/iaei..p.pp...aaaaaap Problems, and
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, et al. Longbench: bilingual, multitask benchmark for long context understanding. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 3119-3137, 2024.
Yushi Bai, Shangqing Tu, Jiajie Zhang, Hao Peng, Xiaozhi Wang, Xin Lv, Shulin Cao, Jiazheng Xu, Lei Hou, Yuxiao Dong, et al. Longbench v2: Towards deeper understanding and reasoning realistic long-context multitasks. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 3639-3664, 2025.
Ali Behrouz, Peilin Zhong, and Vahab Mirrokni. Titans: Learning to memorize at test time. arXiv preprint (0XXx:2501000666 2024.
Zefan Cai, Yichi Zhang, Bofei Gao, Yuliang Liu, Yucheng Li, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong. Junjie Hu, et al. Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling, arXiv preprint arXiv:2406.0200 2024.

---

Published as a conference paper at ICLR 2026
Zefan Cai, Wen Xiao, Hanshi Sun, Cheng Luo. Yikai Zhang, Ke Wan, Yucheng Li, Yeyang Zhou, Li- Wen Chang, Jiuxiang Gu, et al. R-kv: Redundancy-aware kv cache compression for training- reasoning models acceleration. arXiv preprint arXivv2505524113, 2025.
Qiguang Chen, Libo Qin, Jinhao Liu, Dengyun Peng, Jiannan Guan, Peng Wang. Mengkang Hu, Yuhang Zhou, Te Gao, and Wanxiang Che. Towards reasoning era: A survey of long chain-of- thought for reasoning large language models. arXiv preprint M.kiv:2503.09500 2025.
Yilong Chen, Guoxia Wang, Junyuan Shang, Shiyao Cui, Zhenyu Zhang, Tingwen Liu, Shuohuan Wang, Yu Sun, Dianhai Yu, and Hua Wu. Nacl: A general and effective kv cache eviction framework for lims inference time. arXiv preprint arXiv:2408.03667. 2024.
Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Longlora: Efficient fine-tuning of long-context large language models. arXiv preprint arXiv:2309. /2307, 2023.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reicchiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 14168, 2021.
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning, arXiv preprint arXiv:2307.0869) 086911 2023.
Juechu Dong, Boyuan Feng, Driss Guessous, Yanbo Liang, and Horace He. Flex attention: A programming model for generating optimized attention kernels. arXiv preprint arXiv:2412.05496 2024.
Hermann Ebbinghaus. [image] memory: A contribution to experimental psychology. Annals of neurosciences, 20(4):155,2011 2013.
Tianyu Gao, Alexander Wettig, Howard Yen, and Danqi Chen. How to train long-context language models (effectively). arXiv preprint arXiv:2410.02660, 2024.
Yizhao Gao, Shuming Guo, Shijie Cao, Yuqing Xia, Yu Cheng, Lei Wang, Lingxiao Ma, Yutao Sun, Tianzhu Ye, Li Dong, et al. Seerattention-r: Sparse attention adaptation for long reasoning. arXiv preprint arXiv:2506. 250.08299000.889,,, 08889, 2025.
Ravi Ghadia, Avinash Kumar, Gaurav Jain, Prashant Nair, and Poulami Das. Dialogue without limits: Constant-sized kv caches for extended responses in llms, arXiv preprint arXiv:2503.00979, 2025.
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang. Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability 1lms via reinforcement learning, arXiv preprint arXiv:2501. 12948 2025.
Chi Han, Qifan Wang, Hao Peng, Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang. Lm- infinite: Zero-shot extreme length generalization for large language models. arXiv preprint arXiv:2308. 16137, 2023.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. arXiv preprint ar arXiv:2103.038// 2021.
Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh. Michael Mahoney, Yakun Shao, Kurt Keutzer, and Amir Gholami. Kvquant: Towards 10 million context length Ilm inference with kv cache guantization. Advances in Neural Information Processing Systems, 37: 1270-1303, 2024.
Yuxiang Huang, Binhang Yuan, Xu Han, Chaojun Xiao, and Zhiyuan Liu. Locret: Enhancing eviction in long-context Ilm inference with trained retaining heads on consumer-gran grade devices. arXiv preprint arXiv:2410.01885. 2024.
Hugging Face. Open rl: A fully open reproduction of deepseck-rl. URL com homm/huugiingfaac/ope---
https: github
12

---

Published as a conference paper at ICLR 2026
Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir Abdi, Dongsheng Li, Chin- Yew Lin, et al. Minference 1.0: Accelerating pre-filling for long-context Iims via dynamic sparse attention. Advances in Neural Information Processing
Systems, 77-52481-52515 2024.
Mahdi Karami and Vahab Mirrokni. Lattice: Learning to efficiently compress the memory. arXiv preprint urXiv:2504.05640 2025.
Mahdi Karami, Ali Behrouz, Praneeth Kacham, and Vahab Mirrokni. Trellis:
Learning to compress key-value memory in attention models. In Second Conference on Language Modeling.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference machine learning, pages 5156-5165. PMLR, 2020.
Feyza Duman Keles, Pruthuvi Mahesakya Wijewardena, and Chinmay Hegde. On the computational complexity of self-attention. In International conference on algorithmic learning theory, pages 597-619. PMLR, 2023.
Minsoo Kim, Kyuhong Shim, Jungwook Choi, and Simyung Chang. Infinipot: Infinite context processing on memory-constrainee lims. arXiv preprint arXiv:2410.01518, 2024.
Wojciech Kryšcińskii, Nazneen Rajani, Divyansh Agarwal, Caiming Xiong, and Dragomir Radev. Booksum: A collection of datasets for long-form narrative summarization. 2021.
Solomon Kullback and Richard A Leibler. On information and sufficiency. The annals of mathemati- cal statistics, 22(1):79-86, 1951.
Ivan Lazarevich, David Bick, Harsh Gupta, Srinjoy Neema, Gokul Ramakrishnan, and Ganesh Venkatesh. Extending 11m con- text with 99% less training tokens.
Mukherjee,
Nishit
bttps://cerebrat ai/bi /log/
February 2025.
Haoyang Li, Yiming Li. Anxin Tian, Tianhao Tang, Zhanchao Xu, Xuejia Chen, Nicole Hu, Wei Dong, Qing Li, and Lei Chen. A survey on large language model acceleration based on kv cache management. arXiv preprint 1.. 2024a.
Yucheng Li, Huiqiang Jiang, Qianhui Wu, Xufang Luo, Surin Ahn, Chengruidong Zhang, Amir H Abdi, Dongsheng Li, Jianfeng Gao, Yuqing Yang, al. Schench: A kv cache-centric analysis of long-context methods. arXiv preprint arXiv:2412.103319 2024b.
Yuhong Li, Yingbing Huang, Bowen Yang. Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapkv: Lim knows what you are looking for before generation. Advances in Neural Information Processing Systems, 37:22947-22970, 2024c.
Zhuoling Li, Xiaogang Xu, Zhenhua Xu, SerNam Lim, and Hengshuang Zhao. Larm: Large auto-regressive model for long-horizon embodied intelligence. arXiv preprint arXiv:2405. 17424 2024d.
Zhixuan Lin, Exgenii Nikishin, Xu Owen He, and Aaron Courville. Forgetting transformer: Softmax attention with a forget gate. arXiv preprint arXiv:2503.02 2025.
Di Liu, Meng Chen, Baotong Lu, Huiqiang Jiang, Zhenhua Han, Qianxi Zhang, Qi Chen, Chen- gruidong Zhang, Bailu Ding, Kai Zhang, et al. Retrievalattention: Accelerating long-context llm inference via vector retrieval. arXiv preprint arXiv:2409. 2024a.
Xiang Liu, Zhenheng Tang, Hong Chen, Peijie Dong, Zeyu Li, Xiuze Zhou, Bo Li, Xuming Hu, and Xiaowen Chu. Can lims maintain fundamental abilities under kv cache compression? arXiv preprint arXiv:2502.011 01944,,,,,, 2025.
Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu. Kivi: A tuning-free asymmetric 2bit quantization for kv cache. arXiv preprint arXiv:2402.02755 02750, 2024b.
13

---

Published as a conference paper at ICLR 2026
Jesse Mu, Xiang Li, and Noah Goodman. Learning to compress prompts with gist tokens, Advances in Neural Information Processing Systems, 06:19327-19352, 2023.
Piotr Nawrot, Adrian Lańcucki, Marcin Chochowski, David Tarjan, and Edoardo M Ponti. Dynamic memory compression: Retrofitting Ilms for accelerated inference. arXiv preprint arXiv:4403.09636, 2024.
Junyoung Park, Dalton Jones, Matthew Morse, Raghavv Goel, Mingu Lee, and Chris Lott. Keydiff: Key similarity-basee kv cache eviction for long-context Ilm inference in resource-constrained environments. arXiv preprint arXiv:2504. 15364, 2025.
Ofir Press, Noah A Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2 / 12409, 2021.
Ziran Qin, Yuchen Cao, Mingbao Lin, Wen Hu, Shixuan Fan, Ke Cheng, Weiyao Lin, and Jianguo
Li. Cake: Cascading and adaptive kv cache eviction with layer preferences. arXiv preprint arXiv:2503.12491, 2025.
Qwen3. Qwenn//wen3-4--insstuu-----
2025. URL https:/hhuginggfccccoomm
Qwen Instruct 2507
Utkarsh Saxena, Gobinda Saha, Sakshi Choudhary, and Kaushik Roy. Eigen attention: Attention in low-rank space for kv cache compression. arXiv preprint urXiv:2408.05646. 2024.
Freda Shi, Xinyun Chen, Kanishka Misra, Nathan Scales, David Dohan, Ed H Chi, Nathanael Schärli, and Denny Zhou. Large language models can be easily distracted by irrelevant context. In International Conference on Machine Learning, pages 31 31210-31227. PMLR, 2023.
Chaitanya Singhal. Introducing buddhi: Open-source chat model with 128k context window. AI Planet (Medium), April 2024. htt.htts::// https:://uugg.gggaaaae co wtt.aatttaaapplaeeetttt. buddhi-dat aset
Accessed: 2025-08-23.
Hanshi Sun, Li- Len Chang, Wenlei Bao, Size Zheng, Ningxin Zheng, Xin Liu, Harry Dong, Yuejie Chi, and Beidi Chen. Shadowkv: Kv cache in shadows for thghuthroughput long-context Ilm inference. arXiv preprint arXiv:2410.21465, 2024a.
Yu Sun, Xinhao Li, Karan Dalal, Jiarui Xu, Arjun Vikram, Genghan Zhang, Yann Dubois, Xinlei Chen, Xiaolong Wang, Sanmi Koyejo, et al. Learning to (lear at test time): Rnns with expressive hidden states. arXiv preprint arXiv:2407.04620 2024b.
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models. arXiv preprint arXiv:2307.08621, 2023.
Jiaming Tang, Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikci, and Song Han. Quest: Query-aware sparsity for efficient long-context Ilm inference. arXiv preprint arXiv:2406. 10774, 2024.
Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi-heal self-attention: Specialized heads do the heavy lifting, the rest can be pruned. arXiv preprint urXiv:1905.099 2019.
Guangtao Wang, Shubhangi Upasani, Chen Wu, Darshan Gandhi, Jonathan Li, Changran Hu, Bo Li, and Urmish Thakker. Lims know what to drop: Self-attention guided kv cache eviction for efficient long-context inference. arXiv preprint arXiv:2503. 15892750.008879 2025.
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.0477 2020.
Piotr Wožniak, Edward Gorzelańczyk, and Janusz Murakowski. Two components of long-term memory. Acta neurobiologiae experimentalis, 55(4):301-305, 1995.
Di Wu, Hongwei Wang, Wenhao Yu, Yuwei Zhang, Kai- Wei Chang. and Dong Yu. Longmemeval: Benchmarking chat assistants on long-term interactive memory. arXiv preprint arXiv:2410.10813 2024a.
14

---

Published as a conference paper at ICLR 2026
Wenhao Wu, Yizhong Wang, Guangxuan Xiao, Hao Peng, and Yao Fu. Retrieval head mechanistically explains long-context factuality, arXiv preprint 2xXiv:2404.15774, 2024b.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309. / /7453, 2023.
An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, et al. Qwen3 technical report. arXiv preprint FrXiv:2505.09386 2025.
Songlin Yang, Bailin Wang. Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficient training. arXiv preprint arXiv:2312.06633 2023.
Songlin Yang. Bailin Wang. Yu Zhang, Yikang Shen, and Yoon Kim. Parallclizing linear transformers with the delta rule over sequence length. arXiv preprint arXiv:2406.0664 06484, 2024.
Xi Ye, Fangeong Yin, Yinghui He, Joie Zhang, Howard Yen, Tianyu Gao, Greg Durrett, and Danqi Chen. Longproc: Benchmarking long-context language models on long procedural generation. arXiv preprint arXiv:2501.0544 .05414, 2025.
Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, YX Wei, Lean Wang, Zhiping Xiao, et al. Native sparse attention: Hardware-aligned and natively trainable sparse attention. arXiv preprint arXiv:2502. 1089, 2025.
Yuxuan Yue, Zhihang Yuan, Haojie Duanmu, Sifan Zhou, Jianlong Wu, and Liqiang Nic. Wkvquant: Quantizing weight and key/value cache for large language models gains more. arXiv preprint arXiv:2402.12065, 2024.
Zihao Zeng. Bokai Lin, Tianqi Hou, Hao Zhang, and Zhijie Deng. In-context kv-cache eviction for Ilms via attention gate. arXiv preprint arXiv:2410.128 128766 2024.
Yuxin Zhang, Yuxuan Du, Gen Luo, Yunshan Zhong, Zhenyu Zhang, Shiwei Liu, and Rongrong Ji. Cam: Cache merging for memory-efficiet llms inference. In Forty-first International Conference Machine Learning, 2024.
Zhenyu Zhang. Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, (6:34661-34714 2023.
Junhao Zheng, Chengming Shi, Xidi Cai, Qiuke Li, Duzhen Zhang, Chenxing Li, Dong Yu, and Qianli Ma. Lifelong learning of large language model based agents: roadmap. arXiv preprint arXiv:2501.07278, 2025.
Wanjun Zhong, Lianghong Guo, Qiqi Gao, He Ye, and Yanlin Wang. Memorybank: Enhancing large language models with long-term memory. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pages 19724-19731, 2024.
Chen Zhu, Wei Ping, Chaowei Xiao, Mohammad Shoeybi, Tom Goldstein, Anima Anandkumar, and Bryan Catanzaro. Long-short transformer: Efficient transformers for language and vision. Advances in neural information processing systems, 34:1233-77777 2021.
A METHODOLOGY
A.1 INFERENCE ALGORITHM.
Algorithm| T Jillustrates the attention computation with KV eviction using retention gates for a single decoding step. We mark the parts that are different from the standard attention computation in blue.
15

---

Published as a conference paper at ICLR 2026
Algorithm 1: Attention computation with KV eviction (single decoding step)
Input :current hidden x,; KV cache K,-1, V,-1 indexed by S retention gate g
Param :capacity M;
Output Output=attention :attention output O,; updated (K,, v,, By); updated index set St
// 1) Project to Q/K/V for the current token
q: < Woxx; k + WKx;; Vt f- WVx; 8 = g(x));
1/ 2). Append current token to the KV cache
K2 <- K-K-1UU
f Vi-1
B St + U
// 3) Compute attention over currently cached keys /values (restricted to Si)
O, FLASHATTN(Q;: K., V.);
// 4) If capacity exceeded, evict the least important token while Stl M do
Jevic & arg lb Si}:
Remove K, Ve[jevie). B, Jevie): S + Se (Jevict:
end
Positional Encoding in KV Cache Eviction. Our retention mechanism is designed to be positional- encoding agnostic and does not add any extra recency bias beyond what is already present in the base model. The exponential decay in the retention-gad gated attention is smooth approximation of the decay process from information. Therefore, regardless of whether the base model uses absolute positions, RoPE, or no positional encoding, that information already folded into x, q, and k in the standard attention (see [.
of the standard attention with eviction, not to encode the positional
Implementation-wise when using RoPE, we follow prior work (R-KV, SnapKV) and cache post- rotated keys, so the eviction is orthogonal to the positional embeddings used.
A.2 COMPLEXITY
Memory efficiency, Like other KV-eviction schemes, TRIM-KV uses fixed-size cache with (.M) slots, independent of sequence length 7' For each token (per head), it stores single scalar retention score 3, adding only z 1/d, overhead, where d, the dimension of the key and vector states, relative to the KV states, which is negligible in practice. Unlike R-KV (Cai et al. 2025 TRIM-KV does not store queries.
Method
Context Length
Gen Length
Batch Size
Throughput (tok/sec)
Decode Time (s)
F
68 .44
84
ScerAttn-R
32786
1024
4
68.93
59.41
SnapKV
124.67
33,00
TRIM-KV
130.48
31.39
FullK
T14.39
35.8
SeerAttn-R
16378
1024
100.45
40.77
SnapKV
153.21
26.73
TRIM-KV
151.04
27.11
Fullk
138.97
58.94
SeerAttn-R
16378
1024
139.34
58.78
SnapKV
244.60
33.49
TRIM-KV
279.90
29.26
Table 6: Throughput and decoding time comparisons of different KV cache methods on a single H200 GPU. The memory budget M 1024.
Computational efficiency. For each generated token, TRIM-KV computes a scalar retention score 8, via lightweight MLP that can be fused with QKV projections; scores are cached and not recomputed each step. During compression, applies a temporal discount (elementwise power) and evicts the argmin; both costs only ((M). This is cheaper than heuristics like R-KV, which require
16

---

Published as a conference paper at ICLR 2026
key similarity scoring over the cache. T throughput and latency: at 32K context, TRIM-KY KV achieves ~2 higher decoding throughput than full-cache decoding and even faster than SnapKV, purely heuristic method. SeerAttn-R does not provide any computation advantage over full cache model.
B ADDITIONAL EXPERIMENTS
B.1 LONG GENERATION EVALUATION
We provide more comprehensive experiment details in this section.
Experiment settings. For the training, we set the maximum training length be 16384. We train the retention gates with a learning rate of 2 x 10-4 and a weight decay of 0.01. For math reasoning tasks, we follow SeerAttn-R (Gao et al. 2025) that uses OpenRI-Math-22 which has 564M tokens for training. During training, we use a batch size of I for each GPU, and gradient accumulation set to 4. Other hyperparameters are set to the default in Huggingface Trainer. All training is conducted on 4 H100 GPUs
(Hugging Face) dataset,
Benchmarks. AIME24 (At Problem Solving 2024), GSM8K (Cobbe et al. 2021), and MATH-
500 (Hendrycks et al., 2021) are standard math reasoning benchmarks. LongProc (Ye et al. 2025) a long-context benchmark of six sereluuralgeneratiion tasks that require integrating dispersed information and producing structured long-form outputs (up tables from HTML to executing multi-step search to build feasible travel itineraries. The suite spans varied access patterns (sequential vs. targeted retrieval), deductive reasoning demands, and search execution, enabling stress tests of long-range coherence and procedure following. Each task includes deterministic solution procedures and structured outputs, allowing rule-based evaluation (row-level F1 for HTML-TTVV unit tests for pseudocode-ccode exact-match traces for Path/ToM, and validators for Countdown/Travel)) To probe generation length, we use three difficulty tiers targeting 0.5K, 2K. and 8K output tokens.
8K tokens) from extracting
Vanilla
TRIM-KV
-
R.KV
SnapkV
Qwen3 1.7B
R1 Distill Owen 7B
R1 Distill Liama 8B
40
40
40
AIMEZ4
30
20
10
0.
0
256
512, 112
2044
32768
256
512 2098
32768
256
512,
1023044006
32768
KV Budget
KV Budget
KV Budget
Figure 6: Patero frontiers of competing algorithms with different budgets on AIME24.
Math reasoning results. F reports AIME24 performance for Qwen3-1.772 1.7B and DeepSeek-R1- Distill variants. Across both families, TRIM-KV consistently outperforms eviction baselines.
Comparison to query-agnostic baseline. We provide comparison to KeyDiff (Park et al. 2025 a query-agnostic baseline that only considers a key diversity for eviction. The result in Figure 7 shows that the performance of KeyDiff is significantly worse than that of other baselines. Note that R-KV can be considered as a generalization of KeyDiff since it considers both key diversity and attention scores for eviction heuristics.
Results on LongProc. Table 7 reports KV-eviction results on long procedure generation tasks. Across tasks and budgets, TRIM-KV achieves the best performance, and it even surpasses the full-cache baseline on COUNTDOWN (0.5K/2K). Under tighter memory budgets, its margin over heuristic baselines widens.
poocedure-generation
17

---

Published as a conference paper at ICLR 2026
Qwen 4B AIME24
60
Acc (%6)
40
Vanille
R-KY
SnapKV
SeerAttn-R
TRIMKV
KeyOift
256
512
1024
2048
4096
32768
KV Budget
Figure 7: Comparison to KeyDiff, a query-agnoste baseline.
Methodgy budaet
HTML, to
TSV
Thought of Mind
Travel Planning
0.5k
2k
8k
0.5k
2k
8k
2k
8k
FallKY
49.0
41.6
13.9
33.0
10.5
0.0
0.0
0.0
SnapK
37.1
9.3
0.1
26.0
7.0
0.0
0.0
0.0
H20g192
0.4
7.0
0.0
0.0
R-KV
38.0
7.1
0.5
26.0
7,5
0.0
0.0
0.0
TRIM-KVaish
48.7
31.2
7.5
32.5
10.5
0.0
0.0
0.0
Streaming
1.2
0.0
0.0
2.0
0.0
0.0
0.0
0.0
SnupKVan.s
1.5
0.2
0.0
15.0
0.0
0.0
0.0
H20geas
0.4
0.8
0.0
7.6
0.0
0.0
010
0.0
R-KV 2048
1.6
0.1
0.0
3.0
0.0
0.0
0.0
0.0
TVIM-KVa0.s
14.7
3,9
0.0
18.0
0.0
0.0
0.0
Table 7: Results of Qwen3-4B across LongProc tasks: F1-score for HTML to TSV task and accuracies (%) for the remaining tasks. Best per task column in bold.
B.2 LONG-CONTEXT EVALUATION
Experimental settings. We adopt Twen3-4B-Instrun (Qwen3) 2025) as the base model, which supports context window of up to 256K tokens. Retention gates are trained on a mixture of SynthLong-322K (Lazarevich al. 2025). BookSum (Krýšcinski al. 2021), and Buddhi (Singhal 2024), covering sequence lengths from 32K to 128K tokens. We shuffle the combined corpus and train for 10,000 steps (i.e., 10,000 randomly sampled examples), with a maximum training sequence length of 128K and memory capacity M = 1024. All other settings follow Section 5.1
Benchmark Dataset. We evaluate chat-assistant capabilities under strict memory budgets using Long (Wu et al. 2024a). This subset provides contexts up to 123k tokens (measured with the Qwen3 tokenizer) and includes six question types that probe long-term memory: single-session- user (SS-User) and wonggle-ession-assiieee assistant within session; seadde-sssioo--refeeeee from stored personal information; multi-session (Multi), which aggregates or compares information across sessions; sewwlelge-uddate (Knowledge), which tracks and answers with the most recent, changed user information; temporal-reasoninn (Temporal), which reasons over timestamps and time references.
(SS-Assist), which test recall of facts stated by the user or (SS-Pref), which requires personalized responses
To evaluate KV-cache eviction methods on this benchmark, we follow the multi-turn, multi-session protocol of (Li et al. 2024b). Specifically, before each query, the eviction-based model must compress the accumulated dialogue into fixed-size, reusable KV cache mirroring real-world assistants that maintain state across turns and sessions under strict memory budgets. We use Qwen3- 4B-Instruce (Qwen3) 2025) to assess whether model outputs match the ground-truti responses.
Results. The results in Table Male@Shhe show that our method outperforms baseline eviction strategies by a significant margin.
18

---

Published as a conference paper at ICLR 2026
Methodky budget
Overall
Multi
Knowledge
SS-User
SS-Pref'
SS-Assist
Temporal
Full KV 131072
49.4
25.6
68.0
62.9
93.3
85.7
30.1
Streamingl ....iinggg......r... M.M2277s
27.8
15.8
50.0
32.9
56.7
33,9
15.0
SnapK 32778
27.6
15.8
42.3
24.3
73.3
28.6
21.8
TRIM-KY VaTTNN
44.8
24.8
69.2
47.1
76.7
73.2
30.1
StramingILLMaan
19.0
12.8
35.9
24.3
26.7
17.9
I[.3
SnapK 16584
18.2
9.0
25.6
17.1
70.0
12.5
14.3
TRIM-KV 18884
38.6
15.0
69.2
44.3
73.3
60.7
27.8
Mere
13.0
9.0
25.6
15.7
16.7
10.7
8.3
SnapKV
15.8
9.8
19.2
12.9
70.0
7.1
12.8
TRIM-K KVeugg
13.5
53.9
73.3
50.0
StreamingL.LMM
10.2
9.8
14.1
14.3
16.7
7.1
6.0
SnapK V s096
13.8
11.3
14.1
14.3
56.7
7.1
9.0
TRIM-KV 1091
26.8
12.0
46.2
25.7
63.3
33,9
19.6
Table 8: Results on LongMemEvals: overall and partial accuracies (%).
B.3 CHUNKED-PRERES EVALUATION
In this section, we evaluate our method the chunked thunked-prefill setting (Huang et al. 2024), which enables long-context inference without exceeding memory limits. In this framework, long prompts are split into multiple chunks; the model computes the KVV cache for each chunk sequentially and compresses the cache whenever it surpasses the memory budget. We compare our method against LocRet (Huang et al. 2024), a Iearnable KV eviction baseline that also assigns voken-importance scores for eviction. Following the experimental setup of LocRet, we evaluate TRIM-KY on the LongBench (Bai et al., 2024) and LongBench- (Bai et al.] (025) benchmarks. For LongBench-Ve we restrict evaluation to examples with context length below 128K tokens, matching the maximum To train TRIM-KV, we use only LongAlpaca (Chen al. 2023), as in LocRet, to ensure that improvements are not attributable to differences training data. In this setting, set M 512, and keep all other hyperparameters identical those in Section
context length advertised for Phi3-mini-1288
For fair comparison, we adopt the hyperparameters used by LocRet: the chunk size is set to 3072 and the budget size to 6000. We evaluate performance on Phi-3-mini-1D 128K (Abdin et al. 2024). reproducing Table 6 from the LocRet paper. Since the original LongBench experiments have not been released, we use the default chat template of Phi-3-mini-
for all runs. Table D and TO report the results for LongBench (Bai et al. 2024) and LongBench-V2 (Bai et al.) (025) benchmarks, respectively. We observe discrepancies in the full-KV performance, likely due to differences in chat templates.
Overall, TRIM-KV remains highly effective in the chunked-prefill setting. On LongBench, TRIM-KV nearly matches full-KV performance, whereas LocRet exhibits 4.82% drop relative to full-KV cache inference. Notably, on more challenging benchmark, LongBench V2, we even surpass the performance of the full KV cache by 6.5%.
Conceptually, LocRet is designed for chunk prefilling in long-context inference, while our method general. The training paradigms, also differ: LocRet predicts attention logits independently for each KV head, whereas we train retention gates end-to-end, allowing importance scores across heads to jointly optimize the eviction strategy. At inference time, LocRet further relies on a hand-crafted sliding window to preserve the most recent tokens from the latest chunk, and they show that removing this heuristic substantially degrades performance. In contrast, our method requires no such manually designed mechanism: sliding- wiihiingwinoowlieee behavior emerges automatically from the learned policy when beneficial, and some heads (c.g., layer 15, head in Figure 4 do not exhibit Biding--wnndww patterns all.
19

---

Published as a conference paper at ICLR 2026
Nethod
hotpong
Iec
repobeadivg
23.15
Lockat
35.99
13,46
48.70
53,11
Ranky
3710
33.33
5336
15.8
34.5
2550
3617
500
96.3
68.35
54.99
50.7
35.50
300
Lacaa
27.34
2487
36.7
36.3
07.91
16.76
37.38
Table 9: Performance on long-context tasks with Phi3-mini- 128K on the LongBench benchmark, including average relative change (Avg. A) compared to Full KV. indicates that the numbers are
reported from (Huang et al. 2024. Table 6).
Method
Acc. Short
Acc. Medium
Acc. Easy
Ace. Hard
Avg. Acc
Avg. (%)
Fall KV
33.71
18.60
34.44
25.86
28.79
0.00
Locket
32.02
26.67
28.74
28.03
-2 -2.64
TRIM-KV
38,39
31,39
37.78
32.18
34.09
+18.41
Table 10: Performance long-context tasks of KV eviction methods with Phi3-mini- minii118KKK-..- on the
LongBench-V V2 benchmark, including average relative change (Avg. A) compared to Full KV.
B.4
ADDITIONAL ABLATION STUDIES
Qwen 8B AIME24
80
60
Acc. (%6
-- Full
20
R-KV
SnapKV
TRIM-KV (Reasoning Dataset 16K]
TRIM-KY (Genoral Dataset
256
512
1024
2048
4096
32768
KV Budget
Figure 8: Generation Ablation.
Ablation on training datasets. In Section 5.1 we train the retention gates on reasoning PaapssttOpeeIIIMMMe (Hugging Face) and evaluate on AIME24. MATH-500, and GSM8K. This follows standard practice and matches the setting of (Gao et al. 2025). ensuring a fair com- parison. To assess cross-domain generalization, we instead train the gates on general long-context datasets (SynthLong. BookSum, Buddhi), similar to Section 5.2 and then evaluate on math reasoning benchmarks to test whether retention scores learned from general data transfer to long chain-of- thought tasks. As shown in Figure 8 gates trained on general datasets generalize well and even surpass math-specific training at a 2048 KV budget. However, their performance degrades more quickly under tighter KV budgets. Overall, these results are promising and suggest that scaling the training of the retention gates by combining all datasets could yield further improvements.
Owen3 1.7B AIME24
20
Acc
Full
10
R-KV
TRIMKV to
TRIMKY (g= Lineer. b=
TRIM-V (g MLP. b- 01
512
1024
2048
4096
16384
KV Budget
Figure 9: Ablating the retention gate's architecture.
20

---

Published as a conference paper at ICLR 2026
Ablation for the retention gate's architecture. We evaluate several retention gate architectures and report the performance of Qwen3 1.7B on AIME24 in Figure Due to computational constraints, this ablation uses greedy decoding. For the MLP gate, we
single-hicden-layer MLP with width 512. We find that the MLP gate outperforms a simple linear projection, and that a large positive initial bias is crucial for stable training by keeping the gate's output nearly at initialization to ensure minimal early forgetting.
Qwen3 4B AIME24
60
(%)
40
Full
TRIM-V 512)
20
TRIM-V 128)
TRIMKV (M m, no Cool
SnapKV
256
512
1024
2048
4096
32768
KV Budget
Figure 10: Ablating the training memory capacity M.
Ablation on training memory capacity M. We evaluate multiple settings of M in Figure 10 With M oc, there is no capacity penalty, which hurts performance due the absence of sparsity pressure. Setting M = 128 outperforms attention-guide guidedd heuristics but shows signs of over-optimizing for sparsity. practice, we recommend choosing M match the expected deployment-ting memory budget.
C ADDITIONAL QUALITATIVE RESULTS
In this section, we provide more qualitative results to illustrate the eviction decisions made by TRIM- KV. All visualizations are from the first example in the AIME24 dataset, Please refer to Section 5.1.2 for discussions.
21

---

Published as a conference paper at ICLR 2026
Ieeed
Heed
Head 2
Nled 3
Iree
Iteed
Ilead
Heed 7
Layer
Layer Layer 1
Leyer 10
Layer
Token Position
6.6
Layer 20
Layer
Layer 30
Layer 34
Layer 35
Tokan Position
Figure 11: A visualization of token retention matrices of Qwen3-4B when answering math question in the AIME24 dataset. Each subplot is a token retention matrix { in specific layer and head. Observations: earlier layers often exhibit sliding-window- window-like patterns, while later layers develop clearer functional specializations.
22

---

Published as a conference paper at ICLR 2026
Head
Itedd
Head
Ieed
Itead
Head
Layer Layer 0
Layer Layer 5
Layer 15
Taken Position
Layer 20
Layer
0.0
Uppe 50
Layer 3
Layer 35
Taken Position
Figure 12: visualization of eviction decisions Qwen3-4B when answering a math question the
AIME24 dataset. Each subplot is matrix of eviction decisions Oz in specific layer and head.
23

---

Published as a conference paper at ICLR 2026
din atartb You are much probiam Problame Every moning
E8 donaten
and saege scether
aped
the
Sappene for maloo at Ba-leuct1in (1221
feel that Your dets which erdiud de Parts sonees Thins apttiituds
prodee
need provide
proide
frns
Find uroven leaved foims, and dirn sterbaaunh
Okay leto
de
Se
Cu bud Elenh So iudi
Therefore
she walle at speed ensh
Pert
pense protlem agait, There
Abenaabeeyy magte
Therstor
corfee So apert
enetii0 slmns by
Walt problem
plaa Therefore
waiding
the unte conainent Phat
c Thererr
sacend Wallnng speed lects
Thwrsfoee
achve
and Then
17x
denora 5G3 an vaciable 190 So
bert
Now Total (i
L-j this
Thenk so
12-
Therefors
100
de Righ edde begote-- that
soffee shop
bith dais
decininami 1004 20 Tharefere
Figure 13: Visualization of tokens retained in the layer 4 head 5 of the KV cache after generation, where the KV budget is 256. The model is Qwen3-4B. Bold blue indicates tokens retained in the KV cache, where gray indicates evicted tokens.
Whith 1S Mariple by
and shops corhee shop offerm What
bid tour. tika ins the
Mut
this applctte 10
than
by
10 bluneeees hour
shop. You
dos Dac 201' Y054 onto 796, rediwas dok
biiied unddli-oie continnniitter ottili Okay, lets sike the Si.
every montnng stops
Thersfore, when fid unch the time plud hours Teki Shoun
again. There acaradis
plus cofice
Danoct do, talt connec Tharlors the antwer 204
she vabla ate karrts lata tree a hours
S- and Aff Sop 1: Duliny
Thereteee vo the cosn waliing trms
lan and and ariiiin corket it da waiding
based thst Wanning
tims
Ays
22 Step
Forms
tim 4
speed lotel tins
56 (nso)7 sice lote 55 Mu buutody both sone d--e 55 ay
21 S.A- 1961 ve two
and
da Matply
Bring
by criplllly
Meybe
bsti Then, escation
thee
thee
express in eeteraadee ic an 114 Thin subeti
wething
compute
compute
the Time walk banh From surtic Wragl
tiee
sides of so 2014 177 20112 an mouut buth Laft
8. makking tieer including corfee shop
repord mumplly YPSA
ane 55 (001) 55-the andi-
atogeeher Tightt sise 42 175h Bie brt
tnc bh
de- edi ds
M
1200s Lot bur 2-
lute buth
maye
(iona) Computt disciiiss
1256 2002
E44 and
Figure 14: Visualization of tokens retained in the layer 6 head 4 of the KV cache after generation, where the KV budget is 256. The model is Qwen3-4B. Bold blue indicates tokens retained in the KV cache, where gray indicates evicted tokens.
24

---

Published as a conference paper at ICLR 2026
seetwwien You enchirttors
Are for 5SS
Den when 1488 (1 oor fo 12.5 bet 12.0
and eisp ship
per hour the moee her hours and b4 nikucing minass spee oh Suppico
1owd Buthens do de
Agy walk Savtteet
thet Thin
the chaiv photogghs than provide the fhns prover Provie the finst asswo' thi fornn Fios anower
codlin Bo par bot
So, web de thut' Thersion
beceelpddry and sturbbenittn
Clay left
walks at apeed of 1GPt knifl, anck clnn lins the coonds coffee shop. Firt,
the 206 But of
prodler upai There
sctearcce given
the 4--0 aj
snop So, time spent walking mig wourd be hours
nuntor niiuuss..
spertis
Wolk
sags "ectuding
the coffee shis So
bened the kamoraa
lrea to
7 yeal
ineecuaton
time
dust dikuris
Thardoes tiee
1901
the lar irrdaad
24000-1 2.4 has 1'le- 11
D 190
by sirpplyy
Mayie seproout
using qundrati
201244 equition 1/9 w214 Exsution the
expasse
the 40
5 (mmalic anl would
an
oI i 22004 (17- book Theellon,
(17 Roly
borh sodes 24 arye-uo (17
So apeed: S -15 3 krsh Walding time Miec 0000 hours 100 minute. Total time
mncsdingg coffee whop: 6 180 24 204 minuuiess... Arawer 55 2500 Siici, andis
34, tul tere with a This 1P The do'' atsgether Hight
beme dee wift die.
4-1-1 Reariinee both
Set will
form 130 sony plynn
tue do Bata 2000 doc a makk serae.
Figure 15: Visualization of tokens retained in the layer 7 head 3 of the KV cache after generation, where the KV budget is 256. The model is Qwen3-4B. Bold blue indicates tokens retained in the KV cache, where gray indicates evicted tokens.
Whith ah Ve
(3) kalmeeeeioo walk at: coffee shop constant spoed
shep
194
but
And
kmm hours
coneat tarta
tractpe4 22) 1igCD i0
wakking 1lg 21
aquatons
trlo lig de o9c
Rearangen
It 2712-4 aa horm we
dy lmwnld
i000 Thee speed 5.16+
Shiff Wiiking
becemes a (17- ddde 033 304- hil beth
AN 60 bowed 204) Ciodees
Then -drt -ids Then sbeved wdal Hight side
48 12
this
Lit moddee 132
(Xght) sinent 1634
ded
19% so iidia cftucheck them volu Pisd Then 8370 34
Figure 16: Visualization of tokens retained in the layer 9 head 7 of the KV cache after generation, where the KV budget is 256. The model is Qwen3-4B. Bold blue indicates tokens retained in the KV cache, where gray indicates evicted tokens.
25

---

Published as a conference paper at ICLR 2026
din interbocee You are given much prosees Probiem Deny moming Age gan lor 505
Which O
1:48 1200 for ap White ber 1Dis
walk toose hor- houm, indluing SS
heer the with lowas her
560 sartin shep
rthat hew, didx
operiit coffee
probirn
first
pay' 1004 d40- sefi744) Solutine
chia thoughs, from prueduelty the fin Proddde
in the kanel
and tuat viceet Theer
alcoria shep,
fins ling
B0rd 200 chook
mulko sp toorl, ieduting live sondia oothee shop mist
again
lest the time (naling shool tut in
shop 3a canee wilking oi heart
Aye.
apentin
days trotuding opem thes cotte E-4
they Nanabies Ea ttnoppo6
lours 26 menutio sin But
bof
bdld fiee do
Thonkoes
bacad fomusa desance agond
d toral tira
lino it tinn aboo ic
hours nood
the
88 lnsttt 85
Thirsors waildng
Simplly
minene a 20000 hours So siultin 2h4
20024 of
20 12
have
Right
2ort
licnt Sos lit mo fire io
agale Rqudtis
bay
to 201244 --3 So, aepartio 1:0 Exuaton a exprice
inwald would negative -xx Vrac/ 4
ja 2ij/t- 40 (17- dook of Tharibess
Mulldyy
wateo new speed 2.5+ time: (loogg))) =3 168
28 winh Prome earier, theal (t 6.4 Vnae fi11 40
Right side seegh
thes Fiet,
shop:
34, 17704 d.l,
dill Thin (7750 TT db124
spalin endie
attoetiien Hight
decell be des 4.8 10
ahat vi0ll
abx i Latire solee
132/1
sprttaa Thanstrs.. 32 abty1o
4090
tue de
sidlree 190, ao dDc Lett check thow scludera rike sorsz.
Figure 17: Visualization of tokens retained in the layer 15 head 5 of the KV cache after generation, where the KV budget is 256. The model is Qwen3-4B. Bold blue indicates tokens retained in the KV cache, where gray indicates evicted tokens.
When Diesy
but
Walk but 2ah 12 Be -4.8 E.A whtch
her inoluding 560 collooo walks
-di Dut
Aye wath Bu- 1220 foar and
loe dee dit a12:0 So that's Then
coffee ship Your ardders thep by daip.
Thardton
tonodjeeis endh eiee
avory cope oct Chiy,, hot
prosoies dys
Thenkorr quly rid t-123 lish de, 3-2.5 oo 100 lou check the 50504,
prodken theee
waloo
1a-2 og
lenit, the time king
Ape bects sharity snatyerng the
oulti Dhartss
2.4
eport corfoe shops the
Si thaol5000 foy wato lex
Thenstone the ucao, malking hours
hour minocee
concittnn, First conenity We spood Toul byy time minutes Dut
do watking
88 wi
Simplly the
makiing brrtt Tiral in
224
vore
lime
30
17s 4e 2er2
thees tio solion fore and This cice
Do,
dde DS Diidd 22 30x
1it5, Thiny dits, So,
Wesid) 048 thus
mae compure tais iimo trdt been the fidt
Coepuee
Trta Tras Aus
100 Totil taso ineluding
Right adde 117- dugg o mopndinn Fiss multory by 2:41 that 177 34- 40.3
corfee stop: &
coooppoo (idde, andil
dae bring all side with
to alrines dedinale topa -lile will Divide
Do mon 1220
equadtons
aqyiTiii 28 Therkor, 130 uhig
0011. (2 day---0..... 411 Sox ana- was cerree 60% theme serus.
Figure 18: Visualization of tokens retained in the layer 27 head 2 of the KV cache after generation, where the KV budget is 256. The model is Qwen3-4B. Bold blue indicates tokens retained in the KV cache, where gray indicates evicted tokens. Observations: This head mostly keeps period tokens. This suggests that this head may implicitly perform gist tokens that summarize information from the previous sentence. This contrasts with recent trends that advocate for saving a chunk of tokens (Yuan et al. 2025 Gao et al. (025). Our results suggest that it can be more budget-efficient to save individual tokens because they already capture contextual information.
26

---

Published as a conference paper at ICLR 2026
starb- You ruch prodees Prodiem Dvery moninn
Whish
wethe shonesds Whan wilk
Mheewtang oor tour
Chiit dot es for taat ecilel win cano
hour wich takoo
menuts, inckd minuto the
"ya decerrac 301
Fied
mome indlding
-dt full than mulichy roa dhee of whic Biee So taas cored Then,
provies the chulril thoughs prowdd fin anpoor Pooily find crower lores Find
anave: tocserp-drr, and
creat Diay,
sod wald So, Lubidde de Theetton bully -qudict quild 1.85 benn sull
tours
woloo a speed of a unuhy ind
the colees thep Firt,
problem again.
whtee give
at
eeoand eouitio Fquation a 10 aad Thansors arover
shop So, the line ipant naiking mould be hour minas lanllcty the
malio at
201 knn sost trolcing
Wast
saye troludin inirndes ip the coher shoo. So loud tne drig time plas
comrore Aya ls winil bous Thandor,
hours
beaed femulic desenne spond's tine convet to hous inta
Thestrre enuaooni 85 wil
hours- I'm mutess Butt fnrindtss
Thenstoo,, willing time in s80 tours didarin-
menuke 2400 144 hour best 2-4
wolk 24 190
124
80 do ceaation
di (oe- chit
but equidirn: ane
tot sawi So, ae with torm agait out Bidd 190 i
seme brone cid 10 6 40-00- 4
simplly taese equations Mayde axpress
Ie 22/50-4-0 So, equation - 2144 Equation 2 Bickk 20 bore empress
12 d0
on volr
aucatiste
Mutryy the DA, 17704 --8. traal be 6.3 488 Then 1770 whe Ther DeTA Thas 201 <bet tanng left do 40.3 dhe
3 d- dopol x0 1CP 2ay00 20 Tharebes equitir 2 beconee E17
lsscti:: 00
Rigmadde di00
equand Pist
244 1772
oothe 150 24 206
Beand 234 chadin andil
Muutyyy conaaattc
taree cat the th 4.00 1ittt
4
to decimate o thte Divide
120 12 qsedrato
don w soluy the auudiaao
nantini- thaeler, adiy1 th Bos-doora w Deto wps 1900 Eicc Lat's chack tthae soluiire male
Figure 19: Visualization of tokens retained in the layer 32 head 2 of the KV cache after generation, where the KV budget is 256. The model is Qwen3-4B, Bold blue indicates tokens retained in the KV cache, where gray indicates evicted tokens.
27


Reviews:
[{"review_id": "JhOJOtXovK", "reviewer": "Reviewer_dB8X", "summary": "This paper proposes TRIM-KV, a trainable approach for KV cache eviction that assigns intrinsic importance scores to tokens at creation time through lightweight retention gates. The method trains these gates using distillation and capacity losses to enable memory-bounded inference while preserving model quality.", "questions": "Please refer to Weaknesses.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The paper introduces a trainable approach to KV cache eviction that learns token-level importance scores through retention gates. This design allows the model to identify and retain intrinsically important tokens \n2. The authors conduct experiments across multiple mathematical reasoning tasks and show competitive results.", "weaknesses": "1. The paper lacks comparison with LocRet, a highly relevant recent work that also computes importance scores for each token and discards low-importance tokens to reduce memory overhead. \n\nLocret: Enhancing eviction in long-context LLM inference with trained retaining heads\n\n2. The paper would benefit from comprehensive evaluation on established long-context understanding benchmarks such as RULER and LongBench-V2 to validate the method's effectiveness.\n\n3. The use of exponential decay for computing token importance scores appears problematic. Even with an initial retention score of 0.999, the importance decays to less than 5% after approximately 3000 tokens. This aggressive decay suggests the model may essentially degenerate into a sliding window approach for longer contexts.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "R459HEgQsx", "reviewer": "Reviewer_gKFH", "summary": "The paper tackles KV cache limits for long context inference and proposes a learned decay gate to keep useful tokens. \nThe idea is easy to follow and fits well with current attention sparsity efforts. \nIn practice, it works reasonably well and the empirical bump over several benchmarks looks consistent.\n\nHowever, I am not fully convinced the method is as conceptually “beyond attention” as the paper suggests. Even though the authors emphasize intrinsic token importance, the approach ultimately still plays through the attention path by reweighing the attention kernel. \nIt is not really removing attention as a relevance signal, but shifting when the weighting is introduced. \nThe learned score still modulates q·k, so the line between “intrinsic importance” and “attention-based importance” is blurrier than the paper presents.\n\nOn the training side, the method is positioned as a plug-in module that sits on top of frozen weights. If retention is so critical, it is natural to know why this is only applied for this scenario instead of making it part of the main training pipeline. Integrating the retention gates into pretraining or post training would likely change the model’s memory habits in a deeper way. \n\nThe model relies on RoPE or similar positional embedding, and the learned gate sits on top. The paper does not clarify whether this interacts with positional encodings. It might be fine in practice, but some comment or ablation on this point would strengthen the clarity of the claim that the model learns “intrinsic” retention.\n\nOverall, the method seems useful and the results are good enough to merit attention.", "questions": "1. Since the method relies on training, why isn’t it integrated into the general training pipeline instead of being applied frozen models?\n2. Does the interaction between the decay gate and existing positional encoding introduce redundancy or bias toward recency, and has this been tested or ablated?\n3. If the model were trained from scratch with the retention mechanism active, would it converge to a different or stronger retention strategy than the post-hoc approach?\n4. How does the system behave under much longer sequences or tool calling contexts?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. Novel idea, fits cleanly into existing models, and works well without heavy changes. \n2. Shows solid gains across tasks, occasionally even beating full-cache runs. \n3. Lightweight enough to feel practical, not just academic.", "weaknesses": "Since it’s still a trained model, it raises the natural question of why this isn’t folded into normal model training, and the paper doesn’t address that integration path. It also doesn’t discuss how the learned decay interacts with positional encoding, which leaves some ambiguity around whether it’s learning true importance or just reinforcing a recency style bias.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "t2ZkybvJQq", "reviewer": "Reviewer_dqKn", "summary": "This paper proposes TRIM-KV, a learnable token retention mechanism for memory-bounded large language model (LLM) inference. Instead of relying on attention-based heuristics for KV-cache eviction, the authors introduce a lightweight retention gate that predicts each token’s intrinsic importance at creation time. The predicted retention score decays exponentially over time, inspired by Ebbinghaus’s forgetting curve, and governs which tokens are evicted when memory exceeds a fixed budget. The gates are trained via distillation from a frozen teacher LLM and a capacity regularization loss, ensuring negligible inference overhead. Experiments across multiple benchmarks — including GSM8K, MATH-500, AIME24, LongProc, and LongMemEval — demonstrate that TRIM-KV outperforms strong heuristic and learnable baselines under the same or even tighter memory budgets. Qualitative analyses further show that the learned retention scores align with intuitive heuristics such as sliding windows, sink tokens, and gist tokens, suggesting interpretability benefits.", "questions": "1. During inference, is the retention gate computed for every generated token, or cached and reused across time steps?\n2. How sensitive is performance to the λ_cap hyperparameter and the initial bias b in the retention gate?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Novel and well-motivated idea:\nThe paper identifies a fundamental limitation of attention-based eviction — that “recent attention ≠ importance” — and replaces it with a predictive, intrinsic importance estimation per token. This shift from reactive to proactive cache management is conceptually elegant and well-justified.\n2. Brain-inspired design:\nModeling token importance decay via an exponential forgetting curve connects the method to cognitive science (Ebbinghaus), offering both theoretical grounding and interpretability.\n3. Method simplicity and practicality:\nTRIM-KV requires only a small MLP per layer/head, trained with frozen base weights. The added computation and memory are negligible, and inference remains fully compatible with Flash/FlexAttention kernels. The implementation is clear and deployable.\n4. Comprehensive empirical validation:\nResults are reported on five diverse benchmarks covering reasoning, long-context generation, and long-memory conversation. TRIM-KV consistently outperforms SeerAttn-R (learnable retrieval) and heuristic eviction methods (R-KV, SnapKV, H2O), sometimes even surpassing the full-cache baseline — a strong empirical claim rarely seen in this line of work.\n5. Interpretability and analysis:\nThe authors visualize layer/head-wise retention scores, revealing emergent behaviors like sliding windows and sink tokens. This provides valuable insight into how LLMs allocate and forget contextual information.", "weaknesses": "Limited exploration of dynamic budgets:\nThe current method assumes a fixed memory budget M. It would be interesting to explore adaptive budgets (e.g., varying by layer, head, or task), especially since the authors mention this as future work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nPKIdFgeLR", "reviewer": "Reviewer_aGD2", "summary": "This paper proposes TRIM-KV, a trainable sparse attention method for KV cache eviction. The core design of TRIM-KV is to learn a forgetting gate $\\beta_i^{t-i}$ for the attention mechanism. Such gate is learned through minimizing the KL loss and language modeling loss while reducing the cap loss $L_{cap}$. A small MLP is applied to predict the gate value $\\beta_i$ at each position, and only the MLP is trained while other parameters are frozen. KV cache eviction is conducted by evicting the KV pair with the minimal $\\beta_i^{t-i}$. This paper trains such architecture on math reasoning datasets, and evaluate the trained model on various reasoning datasets that requires long cot generation. Empirical results show that TRIM-KV is able to outperform all the KV cache eviction baselines, and appears lossless acceleration in long-cot reasoning. The authors conduct detailed ablation studies to show the effectiveness of TRIM-KV's design and the generalizability across various training datasets.", "questions": "See above.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The design of TRIM-KV is reasonable and effective. Learning a forgetting gate $\\beta_i$ for KV cache eviction is a novel and interesting idea.\n\n2. Experimental results show significant speedups, and TRIM-KV outperforms classical baselines such as SnapKV and R-KV. The ablation studies are detailed and comprehensive.\n\n3. The paper is clearly written, making it easy for readers to understand both the training and inference procedures of TRIM-KV.", "weaknesses": "1. Since TRIM-KV evicts the KV cache according to query-agnostic features, i.e., the importance score of each KV pair is determined by itself rather than by specific query tokens, certain types of baselines should be included in the discussion, and, if possible, compared experimentally. InfiniPot [1] provides a training-free and query-agnostic method that aligns with chunked prefill KV cache eviction. Locret [2] also employs a training-based method to learn the importance score of each KV pair. These methods should be discussed and compared.\n\n2. Although this paper focuses on evicting unimportant KV cache units during reasoning, its ability to process general long-context inputs should also be evaluated. Typical long-context benchmarks such as RULER [3] and LongBench [4] should be included. Eviction-based KV cache optimizations are not expected to perform well on difficult retrieval tasks in RULER (as shown in the Locret [2] paper), but this limitation should still be discussed in the paper to highlight potential weaknesses of the method.\n\n3. The description of the training setup is unclear. The volume of training tokens, batch sizes, sequence lengths, and other hyperparameters should be included in the paper (these could be placed in the appendix). Note that another type of attention optimization: trainable sparse attention (e.g., NSA [5]), requires significantly more training data to achieve higher performance. Listing the training cost of TRIM-KV would help clarify the differences between these methods.\n\n4. Is multi-turn conversation supported by TRIM-KV?\n\n5. Writing should be improved. Appendix C is excessively long yet provides limited useful information. Figure 4 is also not providing much information (a smaller figure is preferred). The conclusion should be written in the authors’ own words rather than by including the model’s generated answers. The methods mentioned in Weakness 1 should also be incorporated into the related work section.\n\nOverall, I appreciate this work and plan to increase my overall rating if the weaknesses and questions above are addressed.\n\n---\n\n[1] Infinipot: Infinite context processing on memory-constrained llms.\n\n[2] Locret: Enhancing eviction in long-context llm inference with trained retaining heads.\n\n[3] RULER: What's the Real Context Size of Your Long-Context Language Models?\n\n[4] Longbench: A bilingual, multitask benchmark for long context understanding\n\n[5] Native sparse attention: Hardware-aligned and natively trainable sparse attention", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)