PaperlayerView exact revision on arXiv

Source-linked AI summary

KV Cache Compression Through the Lens of Transform Coding

Hannah Laus, Claudio Mayrink Verdun, Hao Wang, Flavio du Pin Calmon, Felix Krahmer

arXiv:2608.14191v1cs.LGcs.CLeess.SP

TL;DR

Long-context inference makes the KV cache a dominant memory cost, while existing quantization methods minimize cache reconstruction error without modeling its propagation through attention. AATC uses an attention-aware distortion measure with transform coding and reverse water-filling, achieving near-lossless accuracy at approximately 5.8× compression across evaluated benchmarks.

  • Problem

    The KV cache becomes the dominant memory cost at long contexts, while existing quantization methods optimize cache reconstruction error rather than its effect on attention outputs.

  • Method

    AATC derives an attention-aware distortion measure separating key and value errors across tokens and channels, then allocates bits using whitening and reverse water-filling.

  • Results

    AATC achieves near-lossless accuracy across all benchmarks at approximately 5.8× compression, while KIVI, KVQuant, and PALU each degrade on at least one benchmark.

  • Takeaways & Limitations

    The attention-aware distortion perspective unifies several KV cache compression methods and connects their effects to transform coding.

  • Takeaways & Limitations

    The analysis assumes independent tokens, the implementation lacks a dedicated CUDA optimization, and evaluation covers only two model families.

Abstract

from arXiv · show

The key-value (KV) cache stores information from past tokens and is a major memory bottleneck in long-context inference. Existing quantization methods address this bottleneck by representing the KV cache uniformly with lower-precision data types and designing quantization schemes to minimize reconstruction error in the cache itself, without accounting for how that error propagates through attention mechanisms. We prove that, under a white-noise quantization model, the expected attention-aware distortion decomposes into additive key and value contributions that factor across tokens and channels. Building on transform coding and reverse water-filling, which are classical tools from signal processing and rate-distortion theory, we introduce Attention-Aware Transform Coding (AATC), which allocates bits over a calibration set to minimize attention-aware distortion. On Llama-3.1-8B-Instruct and Qwen-2.5-7B-Instruct, evaluated across LongBench, RULER, GSM8K, MMLU-Pro, and MATH-500, our method achieves near-lossless accuracy at approximately $5.8\times$ compression, whereas each baseline degrades in at least some settings.

I. Introduction

KV cache compression is reframed as an attention-aware transform-coding problem: whiten key and value features, then allocate bits to minimize output-relevant distortion. AATC achieves near-lossless accuracy at approximately 5.8× compression across the reported benchmarks, outperforming baselines that degrade in some settings.

  • Motivation: KV caches become the dominant inference memory cost as context length grows, motivating compression for long-context LLM deployment.The cache grows linearly with context length because keys and values from past tokens are retained.
  • Contribution: The paper derives an attention-aware distortion measure that separates key and value quantization errors into token- and channel-dependent factors.This decomposition unifies effects previously studied separately in token eviction, quantization, and low-rank projection.
  • Method: AATC whitens key and value features on calibration data and applies reverse waterfilling to allocate bits across transformed channels and tokens.Whitening approximately decorrelates features and produces components of varying strength, while reverse waterfilling assigns more bits to larger components.
  • Prior approaches: Existing compression methods include scalar and vector quantization, low-rank projection, eviction, merging, rotation, and transform coding.These families differ in whether they reduce precision, jointly encode coefficients, delete dimensions or tokens, or decorrelate before allocation.

B. Attention and the KV cache

Decoder-only transformers use self-attention to mix information across tokens, with keys determining attention weights and values supplying the combined content. During decoding, the keys and values for all prior tokens remain in the KV cache, whose size grows with sequence length.

  • Attention: Self-attention is the transformer operation that introduces interdependence between token representations.Tokens are mixed across the sequence before each per-token nonlinear operation.
  • Attention: For each query, self-attention computes scaled similarities with keys, applies a row-wise softmax, and returns a weighted average of values.The attention weights are nonnegative and sum to one.
  • KV cache: The KV cache grows with sequence length because past keys and values are retained to avoid recomputing them.This growth makes cache size a direct function of the expanding context.
  • Assumptions: The paper primarily analyzes single decode-step attention with a causal mask that allows each position to attend only to itself and earlier tokens.The method is described for this setting before extension to variants such as group query attention.
  • KV cache: During decoding, each new token appends a key and value, and its query attends to all cached keys and values.The cache stores keys and values from the prompt and previously generated tokens at every layer.

C. Quantization of the KV cache

KV cache quantization replaces keys and values with lower-precision reconstructions, but attention-output distortion captures effects that cache reconstruction error alone misses. Under white-noise and bounded-error assumptions, the expected distortion separates into key and value terms, with a higher-order remainder.

  • Quantization setup: Uniform scalar quantization maps each cache element to a quantized integer and reconstructs it using a scale and zero-point.Within the clipping range, reconstruction error is bounded by half the quantization step.
  • Quantization setup: KV cache methods can use per-tensor, per-token, per-channel, or group quantization, while transform-coding methods may optimize cache reconstruction rather than attention-output distortion.The attention output is computed from the quantized keys and values, so the paper defines distortion at the output level.
  • Attention-aware distortion: Three effects distinguish attention-output distortion from cache reconstruction error: key-query alignment, attention weight, and output-projection shaping of value errors.Tokens with negligible attention weight can be quantized more coarsely without substantially increasing output distortion.
  • Assumptions and approximation: The white-noise model assumes zero-mean, mutually independent, symmetrically distributed key and value errors, while cross-token independence is acknowledged as an idealization.The model is exact only in the high-rate limit, although experiments suggest the distortion metric remains predictive at 2–4 bits.
  • Theoretical decomposition: Under the white-noise model and bounded key errors, expected attention-output distortion decomposes into additive key and value contributions that factor across tokens and channels.The derivation expands output error to second order, takes expectations, and bounds the remaining terms.
  • Assumptions and approximation: The distortion takes the form D = DK + DV + R, where R is higher order and becomes negligible as the effective quantization noise decreases.The experiments suggest the approximation remains predictive even in low-bit regimes.

B. Interpretation of the Distortion Measure

The key and value components of attention-aware distortion assign different importance weights to cached dimensions. This factorization identifies which tokens and channels can receive lower precision with limited effect on the attention output.

  • Key distortion: Key distortion weights channel c of token i by a_i^2, the residual ||v_iW_O − o||^2, and the query component q_c^2.These factors capture attention participation, output change from switching attention, and query sensitivity to the key channel.
  • Key distortion: Channels with small q_c^2 and tokens with small a_i^2 or ||v_iW_O − o||^2 contribute little to key distortion.Such channels or tokens can therefore be compressed more aggressively without substantially changing the distortion.
  • Value distortion: Value distortion weights each value channel by a_i^2 and further scales it through the output projection W_O.Only tokens receiving significant attention contribute strongly to value distortion.
  • Factorization and design implications: Because quantization-error variance depends on channel allocation but is shared across tokens, the total distortion factorizes into token-dependent and channel-dependent terms.This factorization supports channel-wise or token-wise compression designs.
  • Factorization and design implications: The separate DK and DV terms allow keys and values to be quantized independently without cross-term coupling between their errors.The proposed method synthesizes channel-level factors, while token-wise allocation is outside its scope.

V. AATC: Attention-Aware Bit Allocation via Transform Coding

AATC combines calibration-based transform coding with attention-aware reverse waterfilling to allocate bits across transformed key and value channels. Whitening decorrelates and orders dimensions, while separate budgets determine precision for keys and values.

  • AATC procedure: AATC uses a linear transform to decorrelate and importance-order cache dimensions, followed by global bit allocation minimizing attention-aware distortion.The allocation operates per token across channels because calibration captures channel importance but not individual-token importance.
  • Transform coding: Whitening-based SVD transforms key and value projections using calibration activations, producing decorrelated representations with ordered dimensions.The transformed representations are cached during inference and mapped back to the original keys and values through the corresponding factors.
  • Transform coding: Small singular-value dimensions contribute little to the output and are natural candidates for low-precision or zero-bit allocation.These dimensions reflect the combined effect of input variance and projection magnitude.
  • Bit allocation: Reverse waterfilling assigns bits to transformed dimensions under a global budget, with zero bits when a channel’s weight is too small to justify one bit.The continuous relaxation uses a water level λ, while the practical allocation supports discrete bit counts.
  • Bit allocation: AATC solves separate allocation problems for keys and values because they play asymmetric roles in attention.The respective budgets are denoted B_K and B_V.
  • Rank reduction: Zero-bit dimensions reduce effective cache rank adaptively according to attention-aware weights and the available budget, unlike fixed-rank SVD truncation.The effective dimensionality becomes the number of dimensions receiving at least one bit.
  • AATC procedure: AATC performs offline calibration to determine channel-wise allocation and online inference quantization, using per-layer transforms under a global budget.This differs from KVTC’s global cross-layer transform and reconstruction-error allocation.

VI. Experiments

The experiments compare AATC with KIVI, PALU, and KVQuant across long-context, reasoning, and knowledge benchmarks on Llama-3.1-8B-Instruct and Qwen-2.5-7B-Instruct. The evaluation also includes a variance-only ablation and a per-layer KVTC-related comparison.

  • Evaluation setup: AATC is evaluated on LongBench, RULER, GSM8K, MATH-500, and MMLU-Pro using Llama-3.1-8B-Instruct and Qwen-2.5-7B-Instruct.The MMLU-Pro evaluation uses math and computer science subtasks.
  • Evaluation setup: The main comparisons use KIVI, PALU, and KVQuant, with KIVI and KVQuant reimplemented and a recent-token window added for fairness.The standard full-precision retention uses the first 4 tokens and most recent 128 tokens.
  • Ablation and related baseline: AATC var-only retains the transform and reverse-waterfilling procedure but removes the attention-aware q_c^2 and W_Oc weights.Its allocation is based on quantization-error variance alone.
  • Ablation and related baseline: KVTC is closest to AATC but uses a single global cross-layer transform and reconstruction-error allocation rather than per-layer output-aware allocation.In the per-layer setting, the transforms match and the variance-only allocation corresponds to the reconstruction-error criterion.
  • Implementation details: The implementations use grouped-query attention, with 2 head-groups for Llama and 1 head-group for Qwen.Llama has 8 KV-heads, 32 query heads, and 32 layers; Qwen has 4 KV-heads, 28 query heads, and 28 layers.

A. Main Results

AATC remains statistically indistinguishable from FP16 across all evaluated cells at 5.82× compression, while attention-aware allocation becomes especially valuable under more aggressive compression and on harder-to-compress Qwen models.

  • Across all 18 cells, AATC is never statistically distinguishable from the uncompressed FP16 model.Statistical indistinguishability is defined as falling within 2σ of FP16.
  • At 5.82× compression, AATC is lossless to within evaluation noise on every benchmark and both backbones.The Llama KV-cache memory budget falls from 1.07 GB at FP16 to 184 MB.
  • On Qwen, AATC exceeds the strongest quantization baseline by 7.2 points on RULER-32k and 16.8 points on MMLU-Pro math.These are the settings where AATC is the unique best compressed method.
  • The evaluation spans LongBench, RULER, GSM8K, MMLU-Pro, and MATH-500, with further analyses covering bit sweeps, ablations, and implementation details.Table II uses LongBench’s task-averaged score over seven English subsets.
  • AATC var-only is statistically indistinguishable from full AATC for Llama at 5.82× compression, while full AATC gains on Qwen under stronger compression.At 2 bits, the full method’s advantage becomes visible as context length increases.
  • At approximately 7× compression and 2 bits, full AATC retains near-baseline RULER accuracy across context lengths, whereas AATC var-only degrades sharply as context grows.At 2.5 bits, the variants are indistinguishable on Llama but already show a small AATC advantage on Qwen.

Appendix

The proof expands attention-output error under key and value perturbations, shows that cross-terms vanish or are higher-order in expectation, and derives additive key and value distortion terms with a controlled remainder.

  • Proof strategy: The proof first expresses output error through attention-weight and value perturbations, then controls cross-terms before evaluating key and value contributions.The softmax response is expanded along the perturbation direction using first- and second-order Taylor terms.
  • Error decomposition: Key quantization perturbs attention weights through logit errors α_i, while value quantization directly perturbs the attention output.The key term captures perturbed attention weights; the value term captures value quantization.
  • Assumptions: Under token independence and zero-mean symmetric perturbations, off-diagonal expectations vanish and odd moments do not contribute to the leading key distortion.These assumptions remove mixed-token and cubic contributions in the proof.
  • Cross-term control: The leading expected squared output error retains only diagonal key and value terms plus a controlled remainder.Two cross-terms vanish exactly, while the remaining cross-terms are bounded as higher-order contributions.
  • Conclusion: The theorem concludes that attention-aware distortion decomposes into additive key and value terms together with bounded higher-order remainders.This decomposition is the basis for independently quantizing keys and values.

B. Derivation of AATC via Waterfilling

AATC whitens key and value features on calibration data, quantizes transformed channels independently, and allocates bits using an attention-aware reverse-waterfilling surrogate.

  • Transform coding: AATC quantizes each transformed key and value channel independently with a non-overloading uniform quantizer.The transform basis makes the theorem’s independence conditions hold by construction.
  • Distortion model: The transformed-channel standard deviation equals the square root of its variance, and the quantization distortion at b_l,c bits scales as 2^-2b_l,c.The range-to-deviation ratio is absorbed into a common factor that does not affect allocation.
  • Calibration: Bit allocation uses calibration queries to form an attention-agnostic surrogate that averages query-dependent factors over the calibration set.Attention weights are unavailable at compression time, so token-dependent factors are dropped and averaged.
  • Allocation weights: The allocation weight combines query alignment with the reconstruction projection for each channel.For grouped query attention, weights sum over all queries attending to a shared key head.
  • Global allocation: The analogous formulation is applied to value distortion, and per-layer distortions are combined into global distortions.The resulting global formulation leads to the waterfilling problem.

C. Mapping Related Work onto the Distortion Formula

Attention-derived importance appears in the token-dependent factor of the distortion formula, linking AATC’s analysis to token eviction and related cache-reduction methods.

  • Tokens with vanishing attention weights contribute negligibly to distortion regardless of their key and value quantization.The factor a_i^2 therefore identifies a token-dependent importance signal.
  • The distortion factorization connects attention-based token eviction criteria to projections of the paper’s decomposition onto its token-dependent factor.This places eviction and transform coding as different ways of exploiting the same decomposition.

1) Token-Level Attention Weight:

Token-dependent attention factors motivate compression methods that prioritize which tokens matter most to attention. Existing eviction methods use attention magnitude, accumulated scores, or recent attention patterns to select retained tokens.

  • 1) Token-Level Attention Weight:: Token eviction methods exploit the token-dependent factor of the attention-aware distortion decomposition.This factor is orthogonal to the channel-wise allocation developed by AATC.
  • 1) Token-Level Attention Weight:: StreamingLLM preserves initial and recent tokens, H2O retains accumulated-attention heavy hitters, and SnapKV selects tokens from recent prefill attention patterns.These methods differ in their attention-derived eviction criteria.
  • 1) Token-Level Attention Weight:: The value-side factor ∥viWO − o∥2 measures how much token i’s projected value deviates from the attention output.When viWO is close to o, key quantization errors at that token are more forgiving because attention shifts change the output little.

2) Token-Level Output Relevance:

The paper connects token-level relevance and channel-level sensitivity to a unified distortion-based allocation framework. It uses explicit query, output-projection, and quantization-error factors while relating its transform to PCA-based constructions.

  • 2) Token-Level Output Relevance:: CAOTE’s token-eviction error and the paper’s distortion formula both emphasize residual value contribution rather than raw token magnitude.CAOTE measures the residual in value space, whereas the paper’s formula uses the residual after output projection.
  • 2) Token-Level Output Relevance:: The key-side factor q_c^2 motivates precision for key directions strongly read by queries, while the value-side factor ||W_Oc|| weights channels by output-projection sensitivity.Prior methods use query relevance or value covariance, but the paper makes these structural factors explicit.
  • 2) Token-Level Output Relevance:: Quantization-error variances E[(δkic)^2] and E[(δvic)^2] capture channel error, while RateQuant instead fits empirical distortion-rate curves and gradient-based sensitivity weights.The paper argues that query weighting, output-projection weights, and attention-score weighting remain explicit in its decomposition.
  • 2) Token-Level Output Relevance:: AATC jointly allocates channel precision using query norms, output-projection weights, and quantization-error variances rather than optimizing one factor independently.Its calibration-derived rotation is shared across the allocation step.
  • 2) Token-Level Output Relevance:: KVTC’s PCA transform and AATC’s whitening-based SVD yield the same decorrelation transform, but AATC allocates using output-aware distortion.The equivalence applies to decorrelation, not to the subsequent allocation objective.
  • 2) Token-Level Output Relevance:: The transform-coding coordinates from right-multiplying K by P equal those obtained from the whitened projection construction.The paper notes that its coordinates differ from a related construction by the singular-value factor S.

E. Experimental Setting for KV Cache Quantization Methods

The experiments compare AATC with several KV-cache compression baselines across long-context, reasoning, and code-oriented evaluations. Results include operating-point comparisons, bit-budget sweeps, and ablations of allocation and calibration choices.

  • E. Experimental Setting for KV Cache Quantization Methods: KIVI uses channel-wise key and token-wise value quantization with a 128-token recent window, 32-token groups, and four full-precision initial tokens.These settings follow the original implementation used in the comparison.
  • E. Experimental Setting for KV Cache Quantization Methods: PALU uses Fisher-uniform decomposition, head groups of four, a 128-token recent window, and Hadamard rotation for its low-rank component.The low-rank window slides by 16 tokens, while the first four tokens remain full precision.
  • E. Experimental Setting for KV Cache Quantization Methods: KVQuant uses pre-RoPE channel-wise key and token-wise value quantization, Fisher-sensitivity-weighted codebooks, dense outlier retention, and matched full-precision windows.The implementation retains the top 1% magnitude outliers and uses a 128-token recent window flushed in blocks of 32.
  • E. Experimental Setting for KV Cache Quantization Methods: AATC allocates 0–16 bits per channel globally across layers, with four initial full-precision tokens and a 128-token recent window sliding by 16.Layer-wise normalization removes activation-scale bias during global allocation.
  • E. Experimental Setting for KV Cache Quantization Methods: The evaluation covers LongBench, RULER, GSM8K, MATH-500, and MMLU-Pro, with LongBench measuring multiple subtasks and GSM8K measuring extracted-answer accuracy.LongBench uses 500 samples per task, while GSM8K uses 8-shot chain-of-thought prompting.
  • E. Experimental Setting for KV Cache Quantization Methods: PALU’s Qwen degradation is concentrated in code completion: RepoBench-P falls from 0.67 to 0.34 and LCC from 0.61 to 0.48, while QMSum remains near FP16 at 0.234 versus 0.237.The reported pattern is attributed to truncation of fine-grained token-level information.
  • E. Experimental Setting for KV Cache Quantization Methods: AATC remains within FP16 noise from 4-bit to 2-bit on LongBench, while GSM8K declines smoothly and drops sharply only at 1.5-bit and below.At 1-bit on Llama, AATC scores 0.33 versus KIVI’s 0.20 on LongBench and 0.50 versus 0.03 on GSM8K.
  • E. Experimental Setting for KV Cache Quantization Methods: At approximately 7× compression, AATC stays close to FP16 on Qwen across context lengths, whereas its variance-only variant declines on long-context RULER.This comparison isolates the effect of omitting structural weighting factors from allocation.
Loading 2608.14191v1…