Technical Summary: HeadQ: Model-Visible Distortion and Score-Space Correction for KV-Cache Quantization

1. Problem Statement

Current KV-cache quantization methods predominantly optimize for storage-space reconstruction, typically minimizing Mean Squared Error (MSE) between original and quantized keys and values. However, the authors argue that this objective misaligns with how attention mechanisms actually consume cache data.

• Keys are read through logits (query-key scores), meaning only specific linear projections of the keys affect the attention distribution.
• Values are read through attention-weighted readout ($\sum p_j v_j$).

The paper posits that persistent cache error should be measured in model-visible coordinates rather than raw Euclidean storage coordinates. A perturbation can have a large storage MSE but zero impact on attention (if it lies in a "softmax-null" direction), or a small storage MSE but catastrophic attention damage (if it affects the "score-visible" direction). This mismatch leads to suboptimal quantization strategies that waste bits on irrelevant components while failing to protect behaviorally critical ones.

2. Theoretical Foundation: Model-Visible Geometry

The paper establishes a geometric framework distinguishing between storage error and behavioral error:

2.1 Key Score Space (Fisher Geometry)

For a fixed query $q$, a key perturbation $\delta k_j$ induces a score perturbation $\delta s_j = q^\top \delta k_j / \sqrt{d}$.

• Softmax-Null Directions: A token-shared translation (adding a constant vector $\mu$ to all keys) adds a constant to all scores. Since softmax is invariant to constant shifts, this perturbation is softmax-null and causes zero attention damage, despite potentially large storage MSE.
• Score-Visible Space: Attention depends only on the centered active coefficients of the key perturbation relative to the query.
• Fisher Information: The local second-order geometry of the softmax distribution is captured by the Fisher Information Matrix. The paper derives that the KL divergence between the original and perturbed attention distributions is locally proportional to a Fisher-weighted variance of score errors (Eq. 6), not raw key MSE.

2.2 Value Readout Geometry

Values are consumed via a weighted sum. For fixed attention weights $p$, the readout error is $\Delta o = \sum p_j \Delta v_j$.

• The paper proposes an $A^2$-weighted token distortion as the appropriate surrogate for value error. Unlike raw value MSE, this metric weights the distortion of each token by the square of its attention probability ($p_j^2$), reflecting the fact that errors in highly attended tokens contribute more to the output error.

3. Methodology: HeadQ

HeadQ is a key-side quantization method designed to preserve the "score-visible" residual rather than the raw key vector.

Algorithm Overview

1. Calibration: For each layer and head, a low-rank basis $U_h$ (rank $r$) is learned via PCA from calibration queries (post-RoPE).
2. Cache Write:
   • A standard low-bit quantizer $B_b$ generates a base key $\hat{k}^{base}_j$.
   • The residual $e_j = k_j - \hat{k}^{base}_j$ is computed.
   • The residual is projected onto the learned basis: $z_j = Q_q(U_h^\top e_j)$, where $Q_q$ is a low-bit side quantizer.
   • The cache stores the base key and the side code $z_j$.
3. Cache Read (Inference):
   • The current query $q$ is used to dequantize the side code.
   • An additive logit correction is applied:
     $$\hat{s}_j = \frac{q^\top \hat{k}^{base}_j}{\sqrt{d}} + \frac{q^\top U_h \hat{z}_j}{\sqrt{d}}$$
   • This correction directly targets the score error, bypassing the need to reconstruct the full key vector.

Key Distinction: HeadQ does not store the current query; it stores a residual code in a calibration-learned basis that the runtime query can decode into a logit correction.

4. Key Contributions

1. Formulation of Model-Visible Distortion: The paper defines key distortion as "score-visible error modulo constant shifts" and value distortion as "readout-visible error with an $A^2$ diagonal surrogate." This explains why raw MSE is a poor proxy for behavioral performance.
2. HeadQ Algorithm: A practical cache-write/read algorithm that stores a low-rank residual side code in a query basis and applies it as an additive logit correction.
3. Falsifiable Evidence Chain:
   • Metric Superiority: Fisher/score-space error predicts attention KL divergence significantly better (Spearman $\rho \approx 0.62\text{--}0.69$ gain) than raw key MSE across six models.
   • Counterexamples: Paired experiments show that a quantizer with lower storage MSE can have higher attention KL and output error, proving storage MSE is not a sufficient proxy.
   • Null-Space Interventions: Restoring a "softmax-null" mean component (which dense attention ignores) sharply increases Fisher damage for non-equivariant fixed-grid quantizers, demonstrating that storage-coordinate waste breaks behavioral performance.
   • Deployability: Query-PCA bases learned from calibration data successfully approximate the active residual, whereas random bases of equal budget fail.
   • Directionality: Reversing the sign of the HeadQ correction ("wrong-sign" control) causes catastrophic failure, confirming the method works by correcting the specific score object in the correct direction.

5. Experimental Results

Experiments were conducted on WikiText-103 using six models (GPT-2, Pythia, Qwen2, Qwen2.5, TinyLlama, Mistral).

• K-Only Experiments (Dense Values):
  • On the most damaged 2-bit rows, HeadQ (rank 4) removes approximately 84–94% of the excess perplexity (PPL) compared to a KIVI-like baseline for models like Qwen2, Pythia, and GPT-2.
  • TinyLlama shows material improvement; Mistral is already near-dense performance in this protocol.
• Full-KV Composition (2-bit):
  • Combining HeadQ (key-side) with an $A^2$-weighted value policy improves performance across all six models at 2-bit precision.
• Anomaly Analysis:
  • Pythia-160M showed a weaker correlation between Fisher and KL in specific low-entropy "route-flip" regimes (where small perturbations flip the top-1 token). This anomaly was localized to a specific family-specific boundary and weakened in larger Pythia checkpoints (410M, 1B).

6. Significance and Claims

The paper claims that operator matching is the critical design lesson for KV-cache quantization:

• Keys are score objects; compression must preserve query-visible residuals and avoid wasting resolution on softmax-null shifts.
• Values are readout objects; compression should prioritize $A^2$-weighted error.

The authors explicitly state that this work is mechanistic and diagnostic, not a claim of a production-ready serving system. They do not claim to have optimized packed kernels, latency, or throughput. The significance lies in:

1. Correcting the theoretical objective of KV-cache quantization from storage reconstruction to model-visible distortion.
2. Providing a deployable method (HeadQ) that achieves substantial PPL recovery in low-bit regimes by targeting the correct error geometry.
3. Demonstrating that standard storage metrics (MSE) can be behaviorally misleading, necessitating a shift toward score-space proxies.

The paper concludes that the next step is to extend this operator-matched view into a fully optimized full-KV codec, rather than treating storage MSE as the default target.