PaperlayerView exact revision on arXiv

Source-linked AI summary

KVpop -- Key-Value Cache Compression with Predictive Online Pruning

Lukas Hauzenberger, Niklas Schmidinger, Anamaria-Roberta Hartl, David Stap, Thomas Schmied, Sebastian Böck, Günter Klambauer, Sepp Hochreiter

arXiv:2607.05061v2cs.LGcs.AI

TL;DR

KV cache growth limits long-context inference, while existing eviction methods struggle to predict which tokens will remain useful. KVpop learns fixed-budget retention from future-attention supervision and retains 95% of dense-attention performance at 75% compression and 94% at 88% on Qwen3-4B.

  • Problem

    KV caches grow linearly with context length, while token utility is difficult to predict because relevance can shift over time.

  • Method

    KVpop learns fixed-budget KV eviction at the eviction boundary using future-attention supervision and delayed memory-based scoring.

  • Results

    KVpop retains 95% and 94% of dense-attention performance on Qwen3-4B at 75% and 88% KV cache compression, respectively, and outperforms eviction baselines.

  • Takeaways & Limitations

    KVpop improves the quality-efficiency trade-off of bounded-memory inference across mathematical reasoning benchmarks.

  • Takeaways & Limitations

    KVpop is a post-training retrofit for dense-attention Transformers with homogeneous per-head cache budgets, and alternative memory-based scorers remain unexplored.

Abstract

from arXiv · show

Key-value (KV) cache growth is a major bottleneck in autoregressive decoding, as memory and bandwidth scale linearly with context length. Existing KV eviction methods often rely on static heuristics or proxy scores, which poorly track future token utility and cause brittle eviction as relevance shifts. To address this, we introduce KVpop, which learns a fixed-budget KV eviction policy by directly supervising the keep-or-drop decision. The scorer is trained against a novel future-attention target, computed efficiently without materializing dense attention maps. We further introduce a delayed memory-based scorer that, uniquely among learned eviction methods, defers scoring for a fixed number of steps to exploit near-future context. On AIME and HMMT mathematical reasoning, KVpop retains 98% of full-attention performance on Qwen3-4B at 75% KV cache compression and 97% at 88% compression, consistently outperforming established eviction baselines. Qwen3-8B shows even stronger results, reaching near-full teacher performance. These results show that supervising eviction with future-attention signals cuts memory costs while maintaining quality.

1 Introduction

KVpop addresses the linear growth and imperfect prediction of KV-cache eviction by learning fixed-budget retention decisions from future-attention supervision. It combines bounded sink, recent-window, and learned long-range top-k retention with efficient target computation and optional delayed scoring that uses near-future context.

  • KV caches grow linearly with context length, creating a bottleneck for long-context autoregressive decoding.
  • Existing KV-reduction methods use fixed token-preservation heuristics, online importance scores, or learned retention policies to predict which tokens remain useful.
  • KVpop: KVpop retains sink tokens, a protected recent-token window, and a learned long-range top-k cache within a bounded per-head KV budget.Lightweight head-wise scoring modules rank older tokens for the long-range cache.
  • KVpop: A delayed memory-based scorer defers token scoring until the eviction boundary, allowing KVpop to exploit near-future context.The future-attention target is computed without forming a dense S × S attention map by reusing attention-kernel LSE normalizers.
  • Contributions: KVpop trains fixed-budget eviction at the eviction boundary using a future-attention target and boundary-aware loss for long-range top-k decisions.The target is computed with a training-only transposed-attention pass, avoiding dense attention-map materialization and inference-time overhead.

2 Related Work

Prior KV-cache reduction methods comprise sparse retrieval, heuristic eviction, and learned eviction. Sparse retrieval preserves the full cache while eviction enforces a bounded cache, and KVpop learns explicit future utility at the eviction boundary.

  • Taxonomy: Prior KV-cache reduction methods fall into sparse retrieval, heuristic eviction, and learned eviction.Sparse retrieval reduces tokens read per query, heuristic eviction removes tokens using fixed or online rules, and learned methods train eviction policies during retrofitting.
  • Sparse retrieval over full KV cache: Sparse retrieval selects relevant subsets from the full KV cache, reducing attention compute and bandwidth without imposing a memory bound.Examples include retrieving memory pages, routing attention to blocks, and selecting token groups likely to matter for the current step.
  • Learned eviction methods: Learned eviction replaces hand-designed rules with retrofit-trained policies that can adapt to train–inference mismatch from KV-cache sparsification.Dynamic Memory Sparsification trains binary eviction gates with differentiable relaxation and delays removal through a sliding window.
  • Heuristic eviction methods: Eviction permanently drops cached tokens to enforce a hard, fixed-size cache and directly shrink the KV footprint.Heuristic policies keep recent tokens or attention sinks, or score cached tokens online using attention and query-similarity signals.
  • KVpop: KVpop explicitly predicts token future utility using attention mass received after a token exits the protected window, rather than relying on differentiable relaxation.Its target is computed during training without materializing the dense attention map, while decoding enforces a bounded KV cache and permits delayed scoring at the eviction boundary.

3 KVpop

KVpop retrofits pretrained language models with fixed-budget, per-head KV pruning guided by learned online importance scorers. Its future-attention supervision, boundary-aware training, and delayed stateful scoring target token utility while avoiding dense attention materialization at inference.

  • Fixed-budget pruning: KVpop enforces a fixed per-head KV budget by retaining sink tokens, a protected recent window, and top-k long-range tokens.Attention is computed only over the retained set.
  • Importance scoring: Lightweight per-head scorers rank eligible tokens online, while training supervises them with future-attention targets; inference enforces only budget B.The target is independent of scorer architecture and supports stateless or stateful policies.
  • Teacher policy: The teacher applies shared per-head decay factors γ_h to effective scores, adding recency bias so newer tokens can compete for retention over time.Teacher and student use the same decay, and the teacher retains the top-k eligible tokens.
  • Boundary supervision: KVpop trains at the retention boundary by comparing the newly evictable token with the teacher’s cutoff token using a pairwise logistic loss.The loss focuses on the single comparison that changes cache membership and costs O(1) per sampled query position once the cutoff is known.
  • Efficient target computation: The future-attention target uses a transposed-attention pass and sparse log-normalizers, avoiding S × S probability matrices and adding no inference-time overhead.The sparse-LSE approximation empirically matches the dense-LSE target in downstream performance.
  • Stateful scorers: Stateful scorers maintain objective-shaped memory and exploit the protected window to delay scoring until a token becomes eligible for eviction.A token need not receive a score when it enters the KV cache; scoring begins when it leaves the window.

4 Experiments

Experiments evaluate KVpop on Qwen3 models across mathematical reasoning, out-of-domain reasoning, and long-generation efficiency. Results show strong retention under compression, lower memory growth, faster decoding than DMS, and benefits from delayed scoring.

  • Experimental setup: KVpop distillation trains Qwen3-4B-Instruct-2507 and Qwen3-8B for 2,000 steps using Nemotron-Math v2 sequences packed to length S = 16384.Both stateless KVpopmlp and stateful KVpop variants are trained; stateful budgets are reduced to match stateless memory footprints.
  • Baselines: DMS is trained with KVpop’s cache budget and matched scorer parameter count, isolating retention objectives rather than scorer capacity.Comparisons include training-free methods and full-attention teachers; sparse-retrieval methods are excluded because they retain the full KV cache.
  • Mathematical Reasoning: 97% and 100%: at CR=88%, KVpop preserves these relative teacher scores for Qwen3-4B and Qwen3-8B, while KVpopmlp reaches 93% and 98%.At CR=75%, all methods improve, but KVpop outperforms all sparse baselines on average.
  • Sparsification generalization beyond mathematical reasoning: GPQA Diamond accuracy and LiveCodeBench v6 pass@1 evaluate whether the learned sparsification policy generalizes beyond mathematical reasoning and its training distribution.These out-of-domain evaluations use Qwen3-4B and follow DMS’s general-purpose evaluation protocol.
  • Inference efficiency: 19%: DMS and KVpop increase peak VRAM only to 19GB across 16k-to-131k-token generations, versus dense attention rising from 18GB to 36GB.At 75% compression, KVpop also reduces latency and sustains higher long-generation throughput than DMS.
  • Ablations and eviction patterns: 0.2-point increase: delayed mLSTM scoring improves token accuracy over immediate scoring after 2,000 steps by integrating near-future context before eviction.KVpop also drops purely numeric tokens more often, retains reasoning-structural tokens, and achieves 81% global mean recall of teacher top-k decisions.

5 Conclusion · Appendix · A Background: Linear Attention and mLSTM

KVpop is a fixed-budget, future-attention-supervised KV cache compression method that makes predictive online retention decisions through delayed stateful scoring. Its mLSTM scorer provides a lightweight stateful retrofit, while the paper identifies broader memory scorers, compressed-cache training, and hybrid dense-sparse layers as future directions.

  • 5 Conclusion: KVpop’s mLSTM scoring is a strong stateful KV pruning mechanism, but alternative memory-based scorers remain unexplored.The authors also leave compressed-cache architectures trained from scratch and hybrid dense-sparse layers for future work.
  • 5 Conclusion: KVpop learns fixed-budget KV cache retention from future-attention supervision and targets the eviction-boundary decision determining cache membership.The method optionally delays stateful scoring until near-future context is available, bounding inference memory while improving the quality-efficiency trade-off over prior eviction methods.
  • A Background: Linear Attention and mLSTM: Linear attention replaces softmax similarity with a feature-map kernel κϕ(q, k) = ϕ(q)⊤ϕ(k).The explicit feature representation enables the linear-attention formulation.
  • A Background: Linear Attention and mLSTM: Associativity yields equivalent chunkwise parallel and online recurrent implementations, enabling linear-time prefill and training with constant-memory autoregressive generation.The parallel view suits training and prefill, while the recurrent view supports step-by-step decoding.
  • A Background: Linear Attention and mLSTM: In recurrent linear attention, each head maintains a running key–value summary Ct and optionally a normalizer zt, updated through rank-one outer products.Queries retrieve outputs from the current summary with normalization.
  • A Background: Linear Attention and mLSTM: mLSTM augments linear-attention outer-product updates with expressive gates and is selected as KVpop’s stateful importance scorer.The three gates control the input write, state attenuation, and readout modulation.
  • A Background: Linear Attention and mLSTM: For retrofit efficiency, the mLSTM scorer removes its output gate, reuses Transformer attention projections, and constructs inputs solely from cached quantities.For each KV head, the scorer forms xt = [kt; vt] ∈R2dqkv and applies a small head-specific linear projection.
  • A Background: Linear Attention and mLSTM: The lightweight scorer adds only one small matrix for each query, key, and value head after projection and Hedgehog activation.The adaptations use the softmax-stabilized Hedgehog feature map and are based on Hauzenberger et al., 2026.

B Target Variants and Boundary-Loss Details … D.3 Sparse Normalizer Approximation

The appendix specifies target aggregation, normalization, temporal and boundary-loss refinements, fixed-budget sparse attention, and efficient future-target computation. It also shows that sparse normalizer reuse provides a practical approximate target while exact dense supervision remains available.

  • B.1 Future-Attention Target Variants: Max aggregation is the default GQA target, retaining a token when any sharing query head strongly relies on it; probability-space mean instead rewards broad shared utility.Mean aggregation can underweight tokens crucial to only one query head.
  • B.2 Count Normalization: Count normalization converts total future attention into average future attention and clips the future-query count to N_t ≤ L + 1 under finite lookahead.Unnormalized mass favors early tokens because they have more future query positions.
  • B.3 Temporal Decay and Static Priorities; B.4 Margin Weighting; B.5 Keep/Drop Balancing: Temporal decay uses learned per-KV-head factors constrained by 0 < γ_min < γ_max < 1, while margin weighting and headwise balancing adjust boundary-loss example weights.The final weight is the product of margin and balancing weights; margin weighting reduces near-tie influence.
  • B.6 Sampling Query Positions: Sampling a small set of post-saturation query positions mixes uniform and later-biased samples, evaluating the loss with only a gather and pairwise logistic term.This covers both early saturation and the long-range regime without computing the boundary loss everywhere.
  • C Efficient Running Top-k Sparse Attention; C.1 Mask Definition: The fixed-budget mask retains sink tokens, the recent window, and top-k eligible long-range tokens, enforcing B = s + w + k once the sequence is sufficiently long.During decoding, attention runs directly over retained cache entries; training uses a parallel sparse mask over the full sequence.
  • C.2 Fenwick-Tree Cutoff Computation; C.3 Use with FlexAttention: Running cutoffs are computed in O(S log S) time and O(S) space per head with a Fenwick tree, then evaluated on the fly in FlexAttention without storing an S × S mask.The same thresholds instantiate the sparse mask and construct teacher boundary tokens.
  • D Efficient Future-Target Computation; D.1 Dense Target Identity; D.2 Transposed Attention: The future-attention target is a per-key log-sum-exp over future queries, computed exactly by transposed attention with swapped query/key roles and a mask enforcing d ≥ t + w.The transposed pass subtracts the original query normalizer, and its value output is unused.
  • D.3 Sparse Normalizer Approximation: Reusing the sparse-attention log-normalizer yields an empirically sufficient approximate target that tracks most softmax mass, avoids an additional dense causal pass, while exact dense normalization remains optional.The approximation reuses quantities already produced by the forward sparse-attention call.

D.4 Implementation Steps

The target is computed through main and transposed attention passes, followed by masking, normalization, and aggregation to recover per-key scores. The transposed pass is training-only and is unnecessary during autoregressive decoding.

  • The main attention pass produces per-query log-normalizers using either dense causal or sparse normalizers, depending on the target variant.
  • A transposed attention pass treats keys as queries and queries as keys, then subtracts original log-normalizers and masks positions with d < t + w.
  • Auxiliary log-sum-exp values yield log ¯m(h,g) t for all keys, which are aggregated across query groups and count-normalized to recover m(h,g) t.
  • The transposed pass is used only during scorer training and is not required for autoregressive decoding.

E Scorer Architecture Details … E.3 mLSTM Scorer

The appendix specifies stateless and stateful KV scorers sharing one target and boundary loss. It details headwise inputs, stateless tradeoffs, and an mLSTM scorer with recurrent memory and stabilized gating.

  • E Scorer Architecture Details: The scorer variants are stateless or stateful, and both use the same target and boundary loss.
  • E.1 Headwise Inputs: For each KV head h and token t, the scorer input is formed from cached quantities.
  • E.1 Headwise Inputs: Using [k; v] aligns scoring with grouped-query KV heads and enables delayed scoring without storing additional hidden states.Keys and values remain in the cache until the eviction decision.
  • E.2 Stateless Linear and MLP Scorers: Stateless scorers predict each token’s score independently using headwise linear layers or small headwise MLPs.They add minimal overhead and are easy to parallelize during training and prefill.
  • E.2 Stateless Linear and MLP Scorers: Stateless scorers cannot use context arriving after token t but before that token becomes evictable.
  • E.3 mLSTM Scorer: The memory-based scorer uses an mLSTM-style recurrent state for each KV head, with projections producing recurrent-write and delayed-read features.The features are derived from x_h,t.
  • E.3 mLSTM Scorer: The feature map may be softmax or Hedgehog; Hedgehog halves the projected feature dimension before positive and negative softmax features.This reduces parameter overhead.
  • E.3 mLSTM Scorer: The recurrent state comprises matrix memory C_h,t and normalizer z_h,t, updated with scalar input and forget gates produced by headwise projections.Optional softcapping improves numerical stability, and implementation uses the numerically stabilized mLSTM formulation.

E.4 Delayed Readout · E.5 Initialization and Practical Variants

The delayed readout is causal yet forward-looking because it scores a newly eligible token using recurrent state summarizing preceding protected-window context. Initialization and practical variants preserve a common training target and boundary loss while varying scorer architecture and stabilization choices.

  • E.4 Delayed Readout: At query position q, the newly eligible token is u = q − w, while the recurrent state has processed tokens through q.
  • E.4 Delayed Readout: The readout uses an optional normalization and SiLU-based transformation with a bias term applied to the recurrent representation.
  • E.4 Delayed Readout: The time-shifted readout remains causal because all information in C_h,q is available at query position q.
  • E.4 Delayed Readout: It is forward-looking relative to token u because the recurrent state summarizes protected-window context observed before u becomes evictable.
  • E.5 Initialization and Practical Variants: Initialization starts recurrent projections near identity-style mappings, derives query/key and value projections from base pathways, and can begin with neutral ranking and stable memory updates.The final score projection may be zero-initialized, while gate biases encourage stable early memory updates.
  • E.5 Initialization and Practical Variants: The evaluated variants include value projection versus direct cached-value use, delayed-token skip connections, headwise normalization, gate softcapping, and stateless MLP scoring.
  • E.5 Initialization and Practical Variants: All practical variants use the same target and boundary loss, isolating the effect of memory-based delayed scoring.

F Experiment Details

KVpop experiments used distributed mixed-precision training with numerically stable scoring operations, packed sequences, and truncated teacher supervision. Evaluation used a fixed decoding configuration to isolate retention-policy effects.

  • Training setup: Training used 8 H100 GPUs, PyTorch FSDP, global batch size 128, mixed precision, and gradient clipping at 1.0.Score decay and top-k ranking/selection used float32, while most other activations used bfloat16.
  • Training setup: Packed samples filled sequences to maximum context length, while preserving attention masks across packing boundaries improved hybrid-architecture performance.Attention was not reset at packing boundaries.
  • Teacher supervision: Teacher supervision minimized KL divergence using only the top-256 teacher logits, substantially reducing supervision cost.The passage notes that offline precomputation of teacher logits could avoid loading the teacher during training, but leaves this optimization for future work.
  • Evaluation: Evaluation used identical decoding settings to attribute performance differences to retention policy rather than sampling variability.The configuration used nucleus sampling with temperature 0.6, top-p 0.95, top-k 20, no repetition penalty, and zero-shot prompts.

G KVpopSparsity Patterns

Figure 9 shows that KVpop learns nonuniform retention patterns rather than simply preserving recent tokens. It retains contiguous reasoning-structure blocks while evicting intermediate computations and numerical tokens after their local role, especially in later layers.

  • KVpop Sparsity Patterns: Figure 9 shows that KVpop does not simply preserve the most recent tokens or use a uniform sparsity pattern.The learned retention behavior varies across sequences.
  • KVpop Sparsity Patterns: KVpop often retains contiguous text blocks that carry the structure of the reasoning trace.These blocks support interpreting the evolving solution.
  • KVpop Sparsity Patterns: KVpop evicts many intermediate computation steps and purely numerical tokens once they have served their local role, especially in later layers.It does not treat all operands and intermediate values equally.

H DMS Sparsity Patterns · I KVpopPseudocode

DMS concentrates its long-range budget unevenly across heads, whereas KVpop uses fixed per-head allocations. The KVpop forward pass scores KV entries, selects protected and top-ranked tokens, and trains scores from future-attention estimates derived during sparse attention.

  • H DMS Sparsity Patterns: DMS exhibits heterogeneous eviction, with a few heads near dense attention and many early or late heads collapsing to sliding-window-only attention.This indicates that long-range capacity is concentrated in a small subset of heads.
  • H DMS Sparsity Patterns: KVpop instead assigns each head a fixed long-range token allocation, producing homogeneous and explicitly controlled sparse-attention budgets.The comparison is made under the same 75% compression setting.
  • I KVpopPseudocode: KVpop scores each KV entry and retains sinks, the protected recent window, and highest-ranked long-range entries using a running prefix top-k rule.The resulting sparse mask restricts attention to past tokens that are sinks, recent, eligible, and selected.
  • I KVpopPseudocode: During training, sparse attention supplies log-normalizers that estimate each key’s future attention mass through a transposed FlexAttention call.The future-attention target is computed under the current sparse policy using reused sparse-attention denominators.
  • I KVpopPseudocode: With delayed scoring, KVpop scores key t when it reaches the eviction boundary rather than immediately.The pseudocode identifies delayed mlstm~ scoring as the mechanism for this deferred decision.
  • I KVpopPseudocode: At inference, KVpop skips the training-only future-target computation and uses the sparse student-attention path.The pseudocode explicitly branches on the training flag before computing future-attention quantities.
  • I KVpopPseudocode: The future-attention computation transposes the roles of original keys and queries before constructing the future mask and target.The pseudocode repeats KV heads so original keys become queries and treats original queries as keys.
Loading 2607.05061v2…