Attention mass is a proxy that can misrank evictions
KV caches store key-value states for every processed token during autoregressive decoding, so their memory and attention cost grow linearly with sequence length. Training-free compression keeps only a subset of cached entries, and the eviction rule is computed from inference-time statistics rather than learned from task traces or optimized through extra cache-policy training.
Existing training-free methods decide what to keep using indirect signals. Attention-based rules retain entries that receive high attention in the current context or are expected to receive high attention during future decoding; other methods favor entries that help reconstruct the original context from the compressed cache. The paper argues these proxies need not match the predictive effect of eviction, because attention can overvalue redundant entries and reconstruction utility can favor context recovery rather than preserving the next-token distribution.
Next-token KL is the observable part of behavior loss
The paper defines behavior preservation against the same model before eviction. For a future continuation of length tau, the discrepancy between the full-cache and compressed-cache conditional distributions factorizes by the chain rule into the next-token KL at the current prefix plus an expectation over future tokens of the remaining divergence.
Only the first term is observable when compression decisions are made, because later terms depend on future tokens and query states that are not yet available. The method therefore uses the next-token distribution at available query positions as the compression-time behavioral target, and scores evictions by their estimated output-distribution perturbation.
Figure 1 supports this choice. Using LLaMA3.1-8B, the authors group sampled eviction masks by each method's own score percentile and measure mean KL over subsequent positions along the same reference continuation. The behavior-preserving score shows the clearest association with future-position KL (Spearman rs = 0.90 on LongBench, 0.86 on MATH500, 0.89 on AIME25), while proxy-based scores are less pronounced (for example, SnapKV at 0.22, 0.45, and 0.32).
Logit shifts are estimated from pre-eviction statistics
The ideal objective minimizes the KL from the full-cache prediction to the compressed-cache prediction over a retention mask, but obtaining the exact logit perturbation for many candidate masks would require candidate-specific masked-cache forward passes. The method instead estimates the perturbation from full-cache forward statistics.
For each layer, head, and scoring position, it computes a deletion counterfactual that holds the pre-eviction query, keys, and values fixed, removes masked entries, and renormalizes attention over the retained entries. Summing the resulting head-output perturbations over heads gives a residual-stream perturbation at that layer.
The final-state perturbation is approximated by the direct residual-path contribution, summing layer perturbations and treating the responses of later attention, MLP, and normalization layers as approximation error. A first-order expansion through the final normalization and LM head converts this into an estimated logit shift, which is plugged into the KL objective. For RMSNorm-based models the final-normalization Jacobian-vector product is computed in closed form; for models without final normalization it is the identity.
Top-K KL keeps scoring cheap enough to be practical
Full-vocabulary KL evaluation is expensive, so the implemented score evaluates the KL on the top-K support of the full-cache distribution, using renormalized top-K probabilities and a numerically stable log-sum-exp. When K is the full vocabulary, this recovers the estimated-logit KL.
Selecting B entries to evict is intractable by direct search because eviction costs are non-additive: removing multiple entries changes both the renormalized attention outputs and the resulting predictive distribution. The method uses a two-stage greedy procedure. It first builds a candidate pool from the eligible entries with the lowest singleton-eviction scores, then refines the selection by repeatedly adding the candidate that causes the smallest increase in the mask-level estimated-logit KL, updating cached mask statistics after each step.
During refinement, only the removed-mass and removed-contribution statistics of the affected layer change, so statistics for other layers are reused. A stability rule rejects any candidate whose addition would make retained attention mass fall below a threshold epsilon, avoiding unstable attention renormalization.
Gains over attention heuristics grow under tight budgets
The method is evaluated in two inference settings. In prefill-time compression, the prompt cache is compressed once after long-context encoding, with scoring positions drawn from the last m_pre prompt positions. In generation-time compression, prompt entries are always protected, compression is triggered every I_comp generated tokens, and only generated-token entries outside a recent protected window are eligible.
Across diverse architectures and both settings, the method delivers substantial gains in downstream task quality over lightweight attention-based heuristics at matched retained-KV budgets, with the largest gains under aggressive compression. Budgets are matched by the final retained layer-KV-head-token budget, and protected entries are counted as retained entries, so no method receives uncounted protected cache.
The behavior-preserving scoring incurs additional compression-time computation, but the resulting compressed inference remains faster than full-cache inference in the evaluated settings. Default hyperparameters are fixed within each regime and are not tuned per benchmark.
The guarantee covers compression-time queries only
Because eviction decisions must be made before the future decoding queries that will read the cache are known, the method preserves behavior only with respect to the available compression-time queries. It does not guarantee behavior for arbitrary unseen future continuations, where later queries may amplify or dampen the effect of eviction.
The residual-path approximation is a second boundary: it keeps the component computable from pre-eviction statistics and treats the responses of later attention, MLP, and normalization layers as approximation error. The paper validates this ranking signal in an appendix rather than claiming an exact equivalence to masked-cache forward passes.
Prefer behavior-preserving scoring when budgets are tight
The decision-relevant split in this paper is budget pressure. At matched retained-KV budgets the behavior-preserving score beats lightweight attention-based heuristics, and the margin widens as compression becomes more aggressive, so the tighter the cache budget, the stronger the case for paying the extra compression-time computation.
The second split is the compression regime. Prefill-time compression scores prompt-end positions once after encoding, while generation-time compression protects prompt entries and scores recent generated positions at intervals; the two regimes use separate implementation hyperparameters because their eligible candidates and schedules differ. Readers comparing methods should match the retained layer-KV-head-token budget and count protected entries, since those choices determine whether the comparison is fair.
Sources used in this answer
Behavior-Preserving KV Cache Compression
A training-free framework scores candidate KV evictions by the estimated KL between compressed-cache and full-cache next-token distributions, computed from pre-eviction forward statistics, and outperforms lightweight attention-based heuristics at matched retained-KV budgets, with the largest gains under aggressive compression.