Summary (Overview)

  • Core contribution: This paper analyzes the gating mechanism for KV cache eviction, showing that naive learned gates fail because softmax normalizers absorb per-query constants, making dense attention always reachable and preferred by a frozen backbone's loss.
  • Key insight: The paper distinguishes redundancy (whether information is recoverable from remaining cache entries) from query relevance, formulating KV eviction as query-independent redundancy estimation.
  • Method (IndexMem++): A two-branch design where a selection branch evicts redundant KV entries and a recovery branch summarizes evicted content, both operating inside a single softmax. The gate is trained via reverse-KL distillation against the dense model rather than language-modeling loss.
  • Results: At 75% KV cache compression, IndexMem++ achieves 90.14 on RULER (8K) and 46.03 on LongBench with Qwen3-8B, within 3.72 and 1.36 points of full-cache inference respectively.
  • Theoretical finding: Recovery of evicted information depends more on where the summary is read (inside the softmax) than on memory capacity, because the missing mass fraction ρ is a ratio of softmax denominators only accessible within the normalizer.

Introduction and Theoretical Foundation

Background

Large language models are evolving from chatbots into long-horizon agents, generating long interaction histories that make KV cache the dominant bottleneck in autoregressive decoding. As cache grows linearly with sequence length, storage and data movement become critical. KV eviction is particularly valuable because it reduces both memory consumption and attention computation by shortening effective sequence length.

The Redundancy vs. Relevance Distinction

The paper's central theoretical contribution is distinguishing two concepts:

  • Relevance: How much a query reads an entry (what H₂O and SnapKV measure via attention from observed queries)
  • Redundancy: Whether an entry's information can be recovered from other remaining KV entries

Existing methods treat relevance as a proxy for redundancy, but the paper argues this is fundamentally inadequate: an entry may be highly relevant yet redundant, or irrelevant yet irreplaceable.

The Softmax Normalizer Invariance

The foundational theoretical observation is that attention weights are shares of one softmax normalizer:

ot=∑j≤tatjvj,atj=exp⁡(ℓtj)Zt,ℓtj=qtkj⊤d,Zt=∑ℓ≤texp⁡(ℓtℓ)o_{t} = \sum_{j \leq t} a_{tj} v_{j}, \qquad a_{tj} = \frac{\exp(\ell_{tj})}{Z_{t}}, \qquad \ell_{tj} = \frac{q_{t} k_{j}^{\top}}{\sqrt{d}}, \qquad Z_{t} = \sum_{\ell \leq t} \exp(\ell_{t\ell})

The key property: a softmax absorbs any per-query constant. This means a constant gate cancels inside the normalizer and reproduces dense attention exactly. Consequently, a frozen backbone's loss prefers the no-op (constant scores), making direct gate training fail structurally, not just from poor tuning.


Methodology

The Gating Paradigm

The direct instantiation of the framework is:

ai=softmax⁡(qiK≤i⊤d+log⁡g≤i),sj=fθ(hj),log⁡gj=sj\mathbf{a}_{i} = \operatorname{softmax}\left(\frac{q_{i} K_{\leq i}^{\top}}{\sqrt{d}} + \log \mathbf{g}_{\leq i}\right), \quad s_{j} = f_{\theta}(h_{j}), \quad \log g_{j} = s_{j}

where fθf_{\theta} maps the contextualized hidden state hjh_j to a per-KV-head retention score. The gate is placed inside the softmax as a per-key additive logit bias, which in the binary limit recovers attention over the retained set.

Three Key Fixes for Training

1. Removing the no-op solution (protected entries + fixed mass):

gij={1,j∈Pi,βexp⁡(sj)∑ℓ∈Eiexp⁡(sℓ),j∈Ei,aij∝exp⁡(qikj⊤d)gijg_{ij} = \begin{cases} 1, & j \in P_{i}, \\ \beta \frac{\exp(s_{j})}{\sum_{\ell \in E_{i}} \exp(s_{\ell})}, & j \in E_{i}, \end{cases} \quad a_{ij} \propto \exp\left(\frac{q_{i} k_{j}^{\top}}{\sqrt{d}}\right) g_{ij}

Dense attention becomes impossible because it requires gij=1g_{ij} = 1 throughout EiE_i, hence β=∣Ei∣\beta = |E_i|, whereas β≪∣Ei∣\beta \ll |E_i|. The evictable entries must compete for a fixed total mass β\beta, making sparsity penalties and schedules unnecessary.

2. Distillation instead of LM loss: The cross-entropy gradient mixes eviction errors with frozen-backbone errors:

ptgated−eyt=(ptgated−ptdense)⏟caused by eviction+(ptdense−eyt)⏟frozen-backbone errorp_{t}^{\text{gated}} - e_{y_{t}} = \underbrace{(p_{t}^{\text{gated}} - p_{t}^{\text{dense}})}_{\text{caused by eviction}} + \underbrace{(p_{t}^{\text{dense}} - e_{y_{t}})}_{\text{frozen-backbone error}}

Distillation against the dense model's predictions removes the second term entirely.

3. Head budget allocation via retained mass: Scores are not comparable across heads (adding a constant to all scores of one head shifts its normalizer). The allocation is read from the attention mass each head retains, computed online from the final prefill query.

Recovery via Centroid KV

The paper shows that eviction error grows with ρ(oE−oS)\rho(o_E - o_S) where ρ=ZE/(ZS+ZE)\rho = Z_E/(Z_S + Z_E) is the mass on evicted entries. Recovery requires both direction (what was evicted) and the mass fraction ρ, which is a ratio of softmax denominators only available inside the normalizer. The solution: cluster evicted entries, average each cluster into a centroid key/value, and append R centroids to the retained cache where qikc⊤/d+log⁡nrq_i k_c^{\top}/\sqrt{d} + \log n_r sets a centroid's weight.


Empirical Validation / Results

Ablation Progression (Qwen3-0.6B, RULER 8K, 75% compression)

ConfigurationComponentRULER Avg.NIAH
Direct instantiationnone4.332.00
+ always-attended entriessinks + window44.86 (+40.53)53.73
+ distillationKL58.86 (+14.00)72.78
+ head budgetmass67.60 (+8.74)86.17
Dense backbone (no eviction)—70.0186.95

RULER Results (Qwen3-8B, 75% compression)

MethodQwen3-8B (8K)Qwen3-8B (16K)
Full-KV93.8692.96
SnapKV35.1341.54
KeyDiff67.1466.87
DMA41.7636.19
TrimKV80.2679.09
KVpop72.7171.81
IndexMem++90.1480.93

LongBench Results (Qwen3-8B)

Method50% compression75% compression
Full-KV47.3947.39
SnapKV45.3942.02
KeyDiff36.9128.19
DMA42.2532.72
TrimKV46.5445.41
IndexMem++47.5346.03

Mathematical Reasoning (Qwen3-8B, pass@1)

MethodMath500 (K=1024)Math500 (K=512)AIME25 (K=2048)AIME25 (K=1024)
Full-KV0.8320.8320.7000.700
TriAttention0.8630.7580.3830.192
Random Attention0.8690.7840.3880.221
IndexMem++0.8240.7940.6000.450

Key finding: IndexMem++ stays close to Full-KV in generation length (5,226 vs. 5,258 tokens on Math500), while other methods generate substantially longer outputs, suggesting it preserves reasoning states rather than forcing compensation with longer chains of thought.

Key Ablation Findings

  • Scorer capacity matters: Linear readout (0.5M) costs 14.07 points vs. MLP (7.5M)
  • History variants don't help: Convolution, recurrent state, and attention variants all stay within 1.2 points of the reference, suggesting hjh_j already contains needed historical information
  • Joint backbone training decalibrates the scorer: Full-KV improves (70.01→71.65) but eviction performance falls (67.60→62.83)
  • Centroid recovery helps: At matched budgets, 64 slots improve Math500 by +4.00 and RULER 8K by +0.49

Theoretical and Practical Implications

Theoretical Implications

  1. The normalizer is the mechanism: The paper establishes that the softmax normalizer is not just a computational detail but the fundamental quantity governing both eviction (mass redistribution) and recovery (mass fraction ρ is a ratio of denominators).

  2. Why learnable eviction fails naively: The invariance of softmax to per-query constants means a constant gate is always a no-op, and a frozen backbone's loss prefers exactly that. This explains why many existing learnable eviction methods underperform.

  3. Readout location > memory capacity: Recovering evicted information requires both direction and mass fraction; the mass fraction is only available inside the softmax, making the readout location more important than memory size.

  4. Scores cannot compare across heads: The same invariance means scores are not identifiable across heads; budget allocation must be read from retained attention mass.

Practical Implications

  1. Decoding efficiency: IndexMem++ largely decouples decoding throughput from input length. At batch size 32, increasing input from 8K to 256K (32×) reduces throughput by only ~28%.

  2. Not a sliding window: Analysis shows only 52.5% (AIME25) and 48.4% (SWE) of retained entries fall inside a recency window of the same size, and all 39 early facts needed later in SWE trajectories remain cached at point of use.

  3. Training recipe: The paper provides a complete recipe: frozen backbone, protected entries (4 sinks + 128-token window), reverse-KL distillation, per-key retention rate, and online head budget allocation.


Conclusion

The paper's central message is that KV cache eviction must be understood through the lens of the softmax normalizer. A gate trained on the model's own loss will find that keeping everything is preferable and return constant scores. Making the question answerable requires:

  1. Denying the no-op: Protected entries fix the scale; evictable entries compete for fixed mass β
  2. Distillation over LM loss: Supervise against the dense model's predictions, not the next token
  3. Mass-based head allocation: Read budget from retained attention mass, not scores

For recovery, the lesson is not to build larger memories but to read small summaries through the same softmax as the cache they supplement, since the missing mass fraction is only available there.

Future Work

The paper identifies on-policy distillation as the key next step: training the gate on prefixes generated under hard eviction rather than dense teacher-forced prefixes, to expose it to the state distribution it actually meets at inference.

Related papers