PaperlayerView exact revision on arXiv

Source-linked AI summary

A Probabilistic Interpretation of KV Cache Eviction

Renato Geh, Alex Chen, Daniel Israel, Aditya Grover, Guy Van den Broeck

arXiv:2608.28293v1cs.CLcs.AI

TL;DR

KV cache eviction lacks a formal foundation despite its practical promise for reducing inference cost. The paper proves the problem is NP-complete, recasts it as expectation estimation, and uses sampling with decode-time correction; the resulting approach is competitive and more robust across tasks.

  • Problem

    KV eviction is practically effective but formally underdeveloped, with existing methods relying on heuristic entry-selection strategies.

  • Method

    The paper formalizes KV eviction probabilistically as expectation estimation and uses self-normalized importance sampling to correct for evicted entries during decoding.

  • Results

    The paper shows KV eviction is NP-complete and reports that probabilistic eviction with decode-time correction is competitive at the same compression budget and more robust across tasks.

  • Takeaways & Limitations

    Existing zero-variance eviction methods can have unbounded bias, while principled strategies should manage bias and variance when approximating future expectations.

  • Takeaways & Limitations

    The usual post-eviction treatment ignores missing entries, conditioning the distribution on kept entries and severely changing it.

Abstract

from arXiv · show

The premise and promise of KV (cache) eviction is simple: higher throughput can be achieved by evicting some entries from the KV cache, at a negligible cost to quality. This holds empirically for many existing methods, though most rely on creative heuristics for selecting which entries to drop. Despite recent advances, the problem of KV eviction has remained informal in the literature. This paper aims to properly formalize this problem through the lens of probabilistic reasoning and reveal what can be learned from this perspective. Concretely, we (1) formalize the problem of KV eviction and, unfortunately, prove that it is computationally hard, (2) show that by framing it probabilistically, KV eviction reduces to the problem of expectation estimation, which can be approximated through sampling, (3) show that through this probabilistic interpretation, correcting for evicted entries during decoding---a previously ignored problem---becomes feasible, and (4) reveal that existing methods in the literature are zero-variance biased estimators that can be easily adapted in order to enable decode time correction. In practice, we show that this probabilistic version of KV eviction coupled with decode time correction is more robust to different tasks compared to existing eviction methods and achieves competitive performance at the same compression budget.

1 Introduction

KV cache eviction can improve throughput and reduce memory, but its informal heuristic foundations motivate a probabilistic formalization. The paper shows eviction is computationally hard, reframes it as expectation estimation, and proposes correction through sampling-based methods.

  • Motivation: KV cache eviction shrinks the cache to improve throughput and reduce memory while aiming to preserve generation quality.Existing strategies use sophisticated heuristics to choose entries for eviction, despite limited formal treatment.
  • Probabilistic interpretation: Eviction can be framed as expectation estimation, but removing entries distorts the distribution and biases later expectations.
  • Probabilistic interpretation: Existing top-k methods are deterministic zero-variance estimators that may have arbitrarily large bias.The paper instead interprets scores as proposal distributions for self-normalized importance sampling, enabling decode-time correction.
  • Formalization: KV eviction is computationally hard, with the paper showing that it is NP-complete.
  • Correction: The paper formalizes an asymptotically correct decode-time estimator with bounded bias and variance.
  • Empirical validation: Probabilistic eviction with decode-time correction is competitive at the same compression budget and more robust across different tasks.

2 Preliminaries

The preliminaries define attention, KV caching, and eviction. KV caching avoids recomputing the full attention matrix, while eviction further restricts attention to a selected subset of cached entries.

  • Attention: Attention uses key, query, and value embedding matrices to produce attention values for each sequence entry.The embeddings are denoted K, Q, and V, with sequence length n and embedding dimension d.
  • KV caching: KV caching computes and stores the relevant keys and values so generation can reuse them in the next forward pass.This avoids fully recomputing the prohibitively large n × n softmax matrix at longer sequence lengths.
  • KV cache eviction: KV cache eviction reduces attention computation to a subset of entries by retaining an index set I.The objective is to choose I so downstream performance is unaffected or only slightly impacted.
  • Transition: The paper proceeds from this setting by formalizing KV eviction and studying its computational hardness.

3 KV Cache Eviction is Hard

The paper formalizes KV eviction as a constrained selection problem and proves it is NP-complete. The hardness result applies even to the easier eviction-time task of preserving attention values when all ground-truth values remain available.

  • Problem definition: KVEVICTION asks whether a fixed-size subset of entries satisfies a specified error tolerance under given matrices, error, and compression ratio.
  • Reduction: In the PARTITION formulation, the selected set is constrained to a prescribed size, while v_I denotes values not indexed by I.
  • Hardness proof: The NP-completeness proof reduces PARTITION to KVEVICTION using an even sequence length, d = 1, r = 1/2, ε = 0, and all-ones K and Q.
  • Hardness proof: The reduction establishes that solving KVEVICTION for this constructed case would solve any PARTITION instance.
  • Scope of hardness: Preserving attention values at eviction time is already hard, even though real-world eviction must preserve values at future decode steps without evicted entries.
  • Beyond hardness: The paper seeks principled approximations by exploiting the structure of KVEVICTION through a probabilistic interpretation.

4 A Probabilistic Interpretation of KV Cache Eviction

The paper recasts attention as expectation computation, making KV eviction an expectation-estimation problem with a bias–variance objective. This framing shows why deterministic top-k methods are biased and motivates probabilistic sampling with importance correction.

  • Attention as expectation: Attention computes expectations over categorical distributions on KV-cache entry indices, not over the value-vector space itself.Each attention row induces a distribution over entries whose values are averaged.
  • Eviction objective: KV eviction seeks a subset of entries that estimates future expectations with low bias and variance.The probabilistic interpretation assigns each entry a distributional role in keeping decisions.
  • Existing methods: Existing top-k eviction is deterministic and zero-variance but can be arbitrarily or unboundedly biased.Selecting the highest-scored entries skews the expectation even at eviction time.
  • Probabilistic eviction: Sampling unique entries from a proposal distribution produces an unbiased eviction-time estimator, with importance sampling available for later adjustment.The proposal must cover all entries to support correction.
  • Decode-time problem: Decode-time correction is needed because eviction removes entries and distorts later distributions and expectations.The paper explicitly asks whether missing entries after eviction can be corrected at future steps.

5 Correcting for Eviction

After eviction, missing entries change the attention distribution, so the paper uses sampled proposal counts and self-normalized importance sampling to estimate corrected decode-time attention. The resulting estimator is asymptotically correct and is implemented by combining sampled pre-eviction contributions with exact post-eviction contributions.

  • The correction problem: Ignoring evicted entries at decode time conditions attention on the retained support and severely changes the distribution.This distributional shift can alter future expectations.
  • Practical motivation: The correction pipeline addresses errors that can otherwise propagate and amplify through later layers and future generations.This propagation occurs because subsequent K, V, and Q representations depend on previous attention values.
  • Correction method: Self-normalized importance sampling uses the proposal and sampled counts to correct probabilities during decoding.Entries are sampled at eviction time, unsampled entries are evicted, and the recorded counts provide correction weights later.
  • Correction method: The corrected attention combines an estimated pre-eviction expectation with an exact decode-time expectation before normalization.Algorithm 1 separately computes the normalizing constants and contributions for the two intervals.
  • Estimator properties: The attention estimator is asymptotically correct, while fixed-sample self-normalization retains bias and mean-squared error bounded by O(1/m).Increasing the number of samples reduces this error.
  • Implementation: The implementation samples a proposal per attention head at eviction time and applies the correction algorithm at decode time.Grouped-query attention requires a shared mixture proposal across heads that use the same cache.

6 Bias vs Variance

The bias–variance study finds that correction enables a tunable trade-off, with temperature scaling reducing variance at the cost of bias. Probabilistic eviction with correction is most effective in lower-to-middle compression ranges and is more robust than existing biased top-k strategies.

  • Bias–variance trade-off: Temperature scaling reduces correction variance by moving importance weights toward one, while temperatures below one increase pre-eviction probability mass.As τ approaches infinity, correction is suppressed and the estimator becomes lower variance but more biased.
  • Bias–variance trade-off: In middle compression ranges, increasing bias when variance is already low harms estimation quality.The reported comparison uses corrected and uncorrected probabilistic eviction alongside top-k H2O.
  • Comparison with top-k: Probabilistic eviction alone can yield less attention error than greedy top-k eviction, although the curves meet as compression approaches one.Stochastic eviction and greedy top-k produce different uncorrected error curves except at the extreme compression limit.
  • Practical conclusion: Corrected probabilistic eviction tends to be more robust than existing zero-variance, unbounded-bias strategies because it reduces bias.The paper concludes that bias plays a larger practical role in attention error than variance.

7 Robustness

The evaluation compares probabilistic eviction with established score-based methods across models, datasets, and compression ratios. Probabilistic eviction, especially πmin-h, is generally more robust while remaining competitive in aggregate performance.

  • Probabilistic eviction and decode-time correction achieve better or at least competitive average performance against state-of-the-art eviction methods.
  • Win scores compare four probabilistic proposals with five top-k eviction estimators using normalized average ranks across tasks.
  • Individual splits expose instability in existing methods: StreamingLLM and K-norm rank highly on some tasks but last on others.
  • πmin-h consistently ranks at the top in win score, indicating the strongest robustness across methods, datasets, and models.
  • Probabilistic eviction tends to perform better at low-to-middle compression ratios, where enough samples reduce bias and variance.
  • At extremely high compression, all methods perform equally poorly because most of the prompt has been evicted.

8 Conclusion

The conclusion frames KV eviction as a computationally hard expectation-estimation problem and proposes probabilistic estimators with decode-time correction. These estimators match existing methods competitively while improving robustness across tasks.

  • A probabilistic interpretation reduces KV eviction to expectation estimation, enabling sampling-based approximation and correction for eviction-induced distributional distortion.
  • Existing score-based eviction methods are zero-variance estimators with potentially unbounded bias.
  • The proposed self-normalized importance-sampling estimator has reasonable upper bounds on both bias and variance.
  • Empirically, probabilistic estimators achieve competitive performance with state-of-the-art eviction methods and greater robustness across tasks.
  • KV eviction is NP-complete, even though the paper’s formal problem is easier than preserving attention values at future decode time.

D Score-based Eviction Methods as Probabilistic Eviction Policies

Score-based KV eviction methods can be represented as probabilistic proposal distributions without changing their top-k selections. This connects existing heuristics to probabilistic eviction and permits bias–variance adjustments.

  • Any score-based eviction method can be subsumed by probabilistic eviction because infinitely many proposal distributions can share its top-k modes.
  • StreamingLLM corresponds to proposals whose modes are the first four and last k cache entries.
  • H2O proposal distributions are obtained by normalizing its score values.
  • K-norm converts negative L2-norm scores into a proposal distribution using a softmax because value entries may be negative.
  • SnapKV forms proposals by normalizing scores produced through column averaging and one-dimensional pooling of p(V).
  • Temperature scaling can reduce variance, while mixing a proposal with a uniform distribution can reduce bias.

F On the Relationship Between Number of Samples and Compression Ratio

The number of samples controls the effective compression ratio in probabilistic eviction. Binary search efficiently selects samples for a target ratio, while global soft constraints avoid unnatural per-head compression choices.

  • The compression ratio r specifies the fraction of KV-cache entries evicted: r = 0 means none, r = 0.5 means half, and r = 1 means all.
  • Algorithm 2 uses binary search to choose a non-negative sample count m whose expected unique-value count approximates j∗.
  • Algorithm 2 achieves |Eπ[J(m)] − j∗| ≤ 0.5 under full support and has complexity O(n log(α)).
  • The target number of unique sampled entries is j∗ = (1 − r∗) · n, and J(m) counts unique values after m samples with replacement.
  • Highly peaked proposal distributions can require impractically many samples to achieve low compression ratios, although this issue does not arise empirically at r∗ ≥ 0.1.
  • Global soft constrainting fixes samples across heads to match an average target compression ratio, avoiding potentially unnatural per-head ratios.
  • Table 1 shows that target compression ratios produce effective ratios with minimal error.

G Experimental Details

The experiments compare probabilistic proposals with top-k eviction methods across LongBench and RULER, examining per-split scores, temperature scaling, and layerwise compression patterns. Probabilistic eviction with correction generally performs better, while bias–variance trade-offs and proposal choice shape results.

  • Experimental setup: Experiments use SnapKV, TOVA, H2O, StreamingLLM, and K-norm implementations from KVPress v0.5.3 on NVIDIA RTX GPUs.Each example can run on a 24GB GPU in approximately 3 to 5 seconds.
  • Compression control: The effective compression ratio differs insignificantly from the target ratio in the reported Llama3 RULER experiment, with similar results for Qwen3 and LongBench.This supports practical control of the requested global compression budget.
  • Evaluation scope: Four probabilistic proposals and five top-k methods are evaluated across LongBench splits and six of RULER’s 13 splits.The comparison is reported for both Llama3 and Qwen3.
  • Per-dataset results: Probabilistic eviction with correction tends to achieve better score performance, while individual baselines vary sharply across dataset splits.StreamingLLM ranks first on Llama3 QuestionAnswer but last on MultiKey-NIAH; K-norm is top-three on Qwen3 MultiKey-NIAH but last on CommonWords.
  • Bias–variance trade-off: Increasing the importance-sampling temperature usually reduces variance less than it increases bias, suggesting bias is the more important issue to resolve.Figures 7 and 8 evaluate τ ∈ {1, 2, 3, 4, 5} using pairwise win scores and average total scores.
  • Compression patterns: Layerwise compression is uneven: lower layers tend to evict less, whereas higher layers tend to evict more.πmin-h is more conservative, while πH2O spikes rapidly at certain layers.

J Limitations

The proposed estimator is not ideal: self-normalized importance sampling remains biased and has nonzero variance, while implementation requires correction overhead. Its bias has an O(m^-1) upper bound, and overhead is asymptotically smaller than the KV cache.

  • Estimator limitations: Self-normalized importance sampling is biased and has variance greater than zero, unlike the unattainable ideal estimator.The ideal would be simultaneously zero-bias and zero-variance, requiring perfect correction.
  • Estimator limitations: Deterministic eviction without correction has unbounded bias, whereas the proposed estimator provides reasonable upper bounds on bias and variance.This contrasts the trade-offs of corrected probabilistic and uncorrected deterministic methods.
  • Estimator limitations: The bias of self-normalized importance sampling is upper bounded by O(m^-1), where m is the number of samples.Lower compression generally improves accuracy, but sparse attention distributions can require fewer effective samples at a given precision.
  • Implementation boundary: Correction adds O(k · h · b) bytes of memory overhead, compared with O(d · k · h · b) for the KV cache itself.The paper states that the correction overhead is asymptotically dwarfed by the cache size.
Loading 2608.28293v1…