Technical Summary: MOMENTKV

Problem Statement

Autoregressive decoding in Transformer-based language models relies on the Key-Value (KV) cache, which stores previously computed projections to reuse them across decoding steps. As sequence lengths scale to tens or hundreds of thousands of tokens, the KV cache's memory footprint grows linearly, becoming the primary bottleneck for long-context inference.

KV cache eviction addresses this by retaining a fixed-size subset of key-value pairs and discarding the rest. While existing methods have improved token selection strategies (e.g., based on cumulative attention weights or value-aware scoring), they share a fundamental limitation: post-eviction inference renormalizes attention exclusively over the retained tokens, completely discarding the evicted set.

The authors identify that the primary source of output degradation is not merely the residual attention mass on evicted tokens (which existing methods already minimize), but a directional mismatch between the retained and evicted token sets. Empirical analysis reveals that in practice, the evicted tokens are often near-orthogonal to the retained ones. Consequently, even a small amount of evicted attention mass can produce a substantial directional error that standard renormalization cannot recover. This creates a "directional gap" where improving token selection alone yields diminishing returns because more selective retention concentrates the retained set on specific query-aligned directions, further widening the gap with the evicted set.

Methodology: MOMENTKV

To address this directional gap, the authors propose MOMENTKV, a framework that maintains compact, small-size moment statistics over the evicted token set. These statistics serve a dual purpose: guiding eviction decisions and correcting the attention output during inference.

1. Compact Moment Statistics

Instead of storing evicted tokens, MOMENTKV maintains four running sums per attention head, updated incrementally upon each eviction:

• Count ($n_e$): The number of evicted tokens.
• Key Sum ($s_k$): $\sum_{i \in E} k_i$.
• Value Sum ($s_v$): $\sum_{i \in E} v_i$.
• Outer-Product Sum ($S$): $\sum_{i \in E} v_i k_i^\top$.

These statistics allow for the computation of the key mean ($\bar{k}$), value mean ($\bar{v}$), and value-key covariance ($\tilde{S}$) without revisiting the evicted tokens. The total storage overhead is $O(d^2)$ per head, independent of context length.

2. First-Order Approximation of Evicted Output

During inference, MOMENTKV derives a closed-form first-order approximation of the evicted attention output, $\hat{f}_E(q)$. By recentering the logits around the mean key and applying a first-order Taylor expansion to the softmax function, the authors derive:
$$\hat{f}_E(q) = \bar{v} + \frac{1}{n_e \sqrt{d}} \tilde{S} q$$

• The $\bar{v}$ term acts as a query-independent baseline, recovering the mean direction of the evicted sub-output.
• The $\tilde{S} q$ term acts as a query-adaptive correction, capturing the perpendicular variation (directional content) that the baseline alone misses.

This approximation allows the model to reconstruct an estimate of the lost evicted contribution, effectively restoring the directional information that standard eviction obliterates.

3. Moment-Informed Eviction

The moment statistics also guide the eviction process itself. For each retained token $j$, the algorithm computes a moment residual:
$$r_j = v_j - \bar{v} - \frac{1}{n_e \sqrt{d}} \tilde{S} k_j$$
This residual measures how much novel directional information token $j$ carries beyond what the accumulated summary already captures. The eviction score is defined as the product of the attention weight and the residual norm: $\text{score}(j) = \alpha_j \cdot \|r_j\|$.

• Tokens with low residuals are well-predicted by the moment model and are prioritized for eviction.
• Tokens with high residuals (carrying unique directional information) are retained even if their attention weights are lower.

This creates a mutually reinforcing loop: moment-informed eviction keeps the evicted set "geometrically regular" (low logit spread $\sigma$), which tightens the approximation bound for the inference correction. In turn, the accurate correction provides better residuals for future eviction decisions.

4. Normalization-Corrected Inference

The final attention output is reconstructed as a convex combination of the retained output ($f_R$) and the approximated evicted output ($\hat{f}_E$):
$$\hat{f}(q) = \hat{w}_R f_R(q) + (1 - \hat{w}_R) \hat{f}_E(q)$$
The mixing weight $\hat{w}_R$ is estimated using a lower bound on the evicted partition function derived via Jensen's inequality. This ensures the model accounts for the magnitude of the evicted mass while correcting its direction.

Key Contributions

1. Identification of Directional Gap: The paper establishes that the dominant error source in KV cache eviction is the near-orthogonality between retained and evicted sub-outputs, not just the magnitude of the evicted attention mass.
2. Compact Moment Statistics: Introduction of a memory-efficient method to track the geometric properties (mean and covariance) of evicted tokens using only $O(d^2)$ storage per head.
3. Closed-Form Correction: Derivation of a first-order approximation for the evicted attention output that recovers both the mean direction and query-dependent variations.
4. Dual-Role Mechanism: A unified framework where moment statistics simultaneously optimize eviction selection (by identifying redundant tokens) and correct inference outputs (by restoring lost directional information).

Experimental Results

The authors evaluated MOMENTKV on LongBench and RULER benchmarks using LLaMA-3.1-8B-Instruct and Qwen3-4B-Instruct.

• Performance: MOMENTKV consistently outperformed all baselines (H2O, SnapKV, PyramidKV, Ada-KV) across all cache budgets ($L \in \{128, 256, 512, 1024\}$).
• Aggressive Compression: The largest gains were observed under aggressive compression. On LongBench with $L=128$ and LLaMA-3.1-8B, MOMENTKV achieved a +1.35 average score improvement over the strongest baseline (Ada-KV), retaining 94.1% of the full-cache performance.
• Retrieval Tasks: On RULER, which tests precise token retrieval, MOMENTKV showed proportional gains of +3.4 and +3.5 over Ada-KV for LLaMA and Qwen respectively, validating its effectiveness in retrieval-intensive settings.
• Ablation Studies:
  • Normalization-Corrected Inference (NC) contributed significantly more (+2.41) than Moment-Informed Eviction (MI) alone (+0.57), confirming that recovering the evicted direction is the primary driver of performance.
  • The combination of both mechanisms yielded a synergistic effect (+3.17 total gain), exceeding the sum of individual gains.
• Efficiency: The method adds only 1.6 ms of latency per token and requires 4.1 MB of auxiliary memory (6.4% of the retained cache), while reducing per-token latency by over 65% compared to a full cache.

Significance

The paper argues that the path to further improvement in long-context inference lies not solely in better token selection, but in recovering the directional information that current eviction strategies obliviate. MOMENTKV demonstrates that by maintaining compact statistics over the evicted set, it is possible to close the directional gap, transforming the "irreducible cost" of eviction into a correctable approximation error. This approach offers a practical, scalable solution for deploying long-context language models without sacrificing the efficiency of fixed-budget KV caches.