Technical Summary: KVpop — Key-Value Cache Compression with Predictive Online Pruning

1. Problem Statement

Autoregressive decoding in Transformer-based Large Language Models (LLMs) relies on a Key-Value (KV) cache to avoid recomputing past token representations. However, the KV cache grows linearly with context length, creating a significant bottleneck for memory and bandwidth during long-context inference.

Existing solutions to reduce this cost generally fall into two categories:

• Static/Heuristic Eviction: Methods like Sliding Window or Sink-Token preservation (e.g., StreamingLLM) rely on fixed rules, while score-based methods use online proxies (e.g., cumulative attention). These often fail to track future token utility, leading to brittle eviction when relevance shifts.
• Learned Eviction: Approaches like Dynamic Memory Compression (DMC) or Dynamic Memory Sparsification (DMS) train policies to merge or drop tokens. However, many of these methods make eviction decisions at token insertion time or rely on differentiable relaxations that do not directly supervise the specific "keep-or-drop" decision based on future utility. Furthermore, methods with protected windows delay the removal of tokens but often decide their fate immediately upon entry, ignoring evidence that accumulates while the token remains in the protected window.

The core challenge is predicting which tokens will be useful for future queries, as locally salient tokens may become irrelevant, while tokens with low immediate attention may become critical later (e.g., in reasoning traces).

2. Methodology: KVpop

KVpop (KV compression with Predictive Online Pruning) is a sparse-attention retrofit designed to enforce a fixed per-head KV budget without altering the base model architecture. It addresses the limitations of prior work through two primary mechanisms: Future-Attention Supervision and Delayed Stateful Scoring.

2.1 Architecture and Budget

KVpop maintains a bounded cache per KV head defined by the budget $B = s + w + k$:

• $s$ Sink Tokens: The first few tokens are always retained.
• $w$ Protected Window: The most recent $w$ tokens are protected from eviction.
• $k$ Long-Range Top-K: Older tokens compete for a fixed budget of $k$ slots based on learned importance scores.

2.2 Future-Attention Target Supervision

Unlike methods that score tokens based on past attention or local statistics, KVpop supervises the scorer using the future attention mass a token receives after it leaves the protected window.

• Target Definition: For a token $t$, the target is the aggregated log-mass of attention it receives from future queries (positions $d \ge t + w$).
• Efficient Computation: To avoid the $O(S^2)$ cost of materializing dense attention maps during training, KVpop uses a transposed-attention pass. By swapping the roles of queries and keys and reusing the log-sum-exp (LSE) normalizers returned by sparse attention kernels, the system computes the future-attention target efficiently without dense matrix materialization.
• Loss Function: The training objective is a boundary-aware pairwise logistic loss. It compares the predicted score of a newly eligible token ($t_{new}$) against the score of the teacher's cutoff token ($t_{bnd}$). The loss pushes the student to correctly rank $t_{new}$ relative to the boundary, ensuring the top-$k$ selection matches the teacher's future-utility ranking.

2.3 Stateful Scoring with Delayed Decision

KVpop introduces a delayed memory-based scorer (implemented via an mLSTM) that defers the scoring decision until the token reaches the eviction boundary.

• Mechanism: Instead of scoring a token immediately upon insertion, the scorer updates its internal state with the token's representation and subsequent context. When the token exits the protected window ($q - w$), the scorer reads the accumulated memory to generate the importance score.
• Advantage: This allows the scorer to incorporate "near-future" context (evidence that accumulates while the token is in the protected window) before making the keep-or-drop decision, addressing the limitation of insertion-time scoring.

2.4 Inference Efficiency

During inference, KVpop enforces the fixed budget $B$. The selection of the top-$k$ long-range tokens is optimized using a running threshold over static token priorities and a Fenwick tree, allowing the sparse mask to be generated on-the-fly within the attention kernel (e.g., FlexAttention) without storing dense $S \times S$ masks.

3. Key Contributions

1. Supervised Eviction Policy: The authors formulate fixed-budget KV eviction as a supervised prediction problem at the eviction boundary, using a future-attention target to train long-range top-$k$ decisions.
2. Efficient Target Computation: They introduce a training-only transposed-attention pass to compute future-attention targets without materializing dense attention maps, ensuring no inference-time overhead.
3. Delayed Scoring: They demonstrate that a stateful scorer can exploit near-future context by deferring the scoring decision until the token becomes evictable, outperforming immediate scoring methods.
4. Empirical Validation: The method is evaluated on mathematical reasoning (AIME, HMMT) and out-of-domain tasks (GPQA-D, LiveCodeBench), showing that learned eviction policies can maintain high performance under aggressive compression.

4. Experimental Results

The authors evaluated KVpop on Qwen3-4B and Qwen3-8B models, trained on the Nemotron-Math v2 dataset with a sequence length of 16,384.

4.1 Mathematical Reasoning (AIME & HMMT)

• Compression Ratio 75% (Budget 2048):
  • Qwen3-4B: KVpop retained 98% of the full-attention teacher's performance.
  • Qwen3-8B: KVpop retained 100% (near-full teacher performance).
  • Comparison: KVpop significantly outperformed training-free baselines (StreamingLLM, TOVA) and the learned baseline DMS. For instance, at 75% compression, StreamingLLM dropped to ~76% relative performance on Qwen3-4B, while KVpop maintained 98%.
• Compression Ratio 88% (Budget 4096):
  • Qwen3-4B: KVpop retained 97% of teacher performance.
  • Qwen3-8B: KVpop retained 100% of teacher performance.
  • Comparison: Training-free methods collapsed significantly (e.g., StreamingLLM dropped to ~47% relative performance), while KVpop remained robust.

4.2 Out-of-Domain Generalization

Evaluated on GPQA-Diamond (scientific reasoning) and LiveCodeBench (code generation):

• KVpop models trained only on mathematical reasoning data generalized well to these domains.
• At 75% compression, KVpop on Qwen3-4B achieved accuracy comparable to the dense teacher on GPQA-D (0.59 vs 0.59) and competitive results on LCB.

4.3 Inference Efficiency

• Memory: Unlike dense attention, where VRAM usage scales linearly with generation length (reaching 36GB at 131k tokens), KVpop and DMS maintained a nearly constant memory footprint (19GB) regardless of sequence length.
• Latency: KVpop demonstrated lower end-to-end decoding latency compared to DMS and dense attention at long generation lengths (up to 131k tokens). The authors attribute this to KVpop's fixed per-head budget, which avoids the "ragged" cache structures of DMS that complicate compilation and execution.

4.4 Ablation Studies

• Delayed Scoring: Removing the delay (scoring immediately upon insertion) resulted in a 0.2-point drop in token accuracy, confirming the benefit of integrating near-future context.
• Stateless vs. Stateful: The stateful variant (KVpop) generally matched or slightly outperformed the stateless variant (KVpopmlp), particularly at higher compression ratios.

5. Significance and Claims

The paper claims that supervising eviction with future-attention signals is a superior strategy for KV cache compression compared to static heuristics or proxy scores. By directly training a policy to predict which tokens will matter for future queries, KVpop achieves:

1. High Fidelity: It maintains near-full teacher performance even at high compression ratios (up to 88%), a regime where other methods fail.
2. Robustness: The learned policy transfers effectively to out-of-domain tasks, suggesting the model learns general principles of token utility rather than overfitting to specific training distributions.
3. Efficiency: It enforces a hard memory bound, enabling stable, long-context decoding with constant memory usage and improved throughput compared to dynamic gating methods like DMS.

The authors conclude that KVpop demonstrates that learned eviction policies, when properly supervised by future utility and allowed to utilize near-future context, can effectively solve the memory bottleneck in autoregressive generation without sacrificing model quality.