Output-Aware Rotation for INT2 KV-Cache Quantization

When a language model processes long inputs, the key‑value (KV) cache that stores intermediate representations grows quickly and begins to dominate memory usage and bandwidth consumption. Researchers have turned to extreme low‑bit quantization—often INT2—to shrink this cache, but most approaches focus on reducing error inside the cache itself, ignoring how that error is later amplified by the attention mechanism and the output projection matrix $W_O$. The paper “Output‑Aware Rotation for INT2 KV‑Cache Quantization” introduces a rotation technique that directly minimizes the error after $W_O$, aligning the quantization step with the part of the model that ultimately determines the generated token.

What You Need to Know

The authors observe that existing INT2 quantization methods apply orthogonal rotations to the KV cache based on statistics such as variance or on proxy metrics computed before the attention readout. While these rotations can make the cache more quant‑friendly, they do not guarantee that the resulting error will stay small once it passes through the scaled dot‑product attention and the subsequent linear layer $W_O$. To close this gap, they propose OptR (Output‑Aware Rotation), which formulates an objective that measures the error of the final attention output after $W_O$.

OptR first derives an analytical expression for the post‑$W_O$ error and shows that it can be split into a component caused by quantizing another caused by the value side. By treating each attention head independently, they learn a per head that minimizes the combined error The optimization is performed through stochastic gradient descent on a small calibration set, using the full INT2 quantization and attention pipeline so that the rotations see the exact quantization noise that will appear at inference time.

Why It Matters

Reducing the KV cache footprint directly translates to lower GPU memory pressure and higher throughput for long‑context generation, which is crucial for applications such as document summarization, code completion, or chatbots that maintain extensive conversation histories. By targeting the error that actually affects the model’s output, OptR achieves higher accuracy for a given bit‑width compared to methods that optimize only cache‑internal statistics.

In practice, this means that developers can deploy INT2‑quantized models with confidence that the loss in quality will be smaller, or alternatively push the quantization even further (e.g., to binary representations) while staying within acceptable error margins. The approach adds only a modest overhead during calibration and does not change the inference kernel, making it easy to integrate into existing serving stacks.

Key Details

• OptR decomposes the post‑$W_O$ attention‑output error into key‑induced and value‑induced terms, allowing separate optimization per head.
• Each head learns an orthogonal rotation matrix that is applied before INT2 quantization and removed after the attention computation.
• The calibration objective uses a mean‑squared‑error loss on the final attention output, measured on a small subset of training data.
• Experiments on LLaMA‑2‑7B and OPT‑6.7B show up to 0.4 perplexity improvement over variance‑based rotations at INT2 KV‑cache quantization.
• Memory savings remain at roughly 75 % reduction compared to FP16 caches, with negligible impact on kernel latency.

What’s Next

Future work could explore extending OptR to other quantization granularities (e.g., mixed‑bit schemes) or to the projection matrices inside the feed‑forward network. Additionally, adapting the method to dynamic KV‑cache eviction policies might further improve efficiency for variable‑length inputs without sacrificing the output‑aware guarantees established here.

📌 Source: Arxiv Ml