🤖 machine learning

Output-Aware Rotation for INT2 KV-Cache Quantization

This paper proposes OptR, an output-aware rotation method that minimizes post-projection attention-output errors through per-head orthogonal corrections and key reparameterization, thereby significantly improving the performance of INT2 KV-cache quantization for large language models while maintaining negligible inference overhead.

Original authors: Vincent-Daniel Yun, Woosang Lim, Minsoo Cheong, Sunwoo Lee, Murali Annavaram, Sai Praneeth Karimireddy, Sungjoo Yoo

Published 2026-08-05

📖 7 min read🧠 Deep dive

CC0 1.0

Original authors: Vincent-Daniel Yun, Woosang Lim, Minsoo Cheong, Sunwoo Lee, Murali Annavaram, Sai Praneeth Karimireddy, Sungjoo Yoo

Original paper dedicated to the public domain under CC0 1.0 (http://creativecommons.org/publicdomain/zero/1.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to remember a massive story to tell a friend. You have a notebook where you write down every important detail you've heard so far. The bigger the story, the more pages you need. Now, imagine your notebook is running out of space, and you can't afford to buy a bigger one. To solve this, you decide to shrink your handwriting. Instead of writing full, clear letters, you scribble them down using only four tiny symbols: a dot, a dash, a circle, and a cross. This is like compressing a huge amount of information into a very small space.

In the world of artificial intelligence, specifically Large Language Models (LLMs), these notebooks are called "KV caches" (Key-Value caches). They store the context of a conversation so the AI can remember what was said earlier. As conversations get longer, these caches get huge, eating up memory and slowing things down. To fix this, scientists try to shrink the data to just 2 bits (using only four levels, like our four symbols). But here's the catch: when you squish data that tightly, the "outliers"—the really important, unusual details—get squashed into the wrong shape, and the AI starts making mistakes. It's like trying to fit a giant, lumpy potato into a tiny box; the potato gets bruised, and when you take it out, it doesn't look like the potato you put in.

For a while, researchers tried to fix this by rotating the potato before squishing it, hoping to spread the lumps out evenly. But they were measuring the success by how well the potato looked inside the box, not by how well the AI could use it to tell the story later. This paper introduces a new method called OptR (Output-Aware Rotation) that changes the game. Instead of just trying to make the squished data look nice, OptR checks if the AI can still understand the story after the data has been squished and un-squished. It turns out that by paying attention to the final result, the AI can remember much more accurately, even with that tiny, four-symbol notebook.

The Problem: The "Squished Potato" Effect

When an AI reads a long text, it builds a mental map of the words it has seen. This map is stored in the KV cache. To save space, researchers use a technique called quantization to shrink this map. The most extreme version is INT2 quantization, which reduces the data to just four possible values. It's incredibly efficient—using only 1/8th of the memory of standard formats—but it's risky.

Think of the data in the cache as a group of students standing in a line. Most students are average height, but a few are giants. If you try to fit them all into a small room (the INT2 range), the giants get crushed, and the average students get squished too because the room is too small for the giants. This causes "quantization error," where the AI misremembers the giants and the average students alike.

To fix this, previous methods used rotation. Imagine spinning the line of students so the giants are no longer standing straight up but are leaning sideways. This spreads their height out across the room, making it easier to fit everyone without crushing them. However, the paper argues that existing methods were making a mistake. They were spinning the students just to make them fit nicely in the room (minimizing the error of the stored data), but they weren't checking if the students could still run a race correctly once they left the room.

The authors point out a mismatch: the rotation that makes the data look best inside the box isn't necessarily the rotation that helps the AI perform best after it uses that data. The AI doesn't care if the data looks perfect in the cache; it cares if the final answer is correct.

The Solution: OptR (Output-Aware Rotation)

The paper proposes OptR, a method that optimizes the rotation based on the final output of the AI, not just the storage.

Here is how OptR works, step-by-step:

Centering the Data (The "Leveling" Trick): Before squishing the data, OptR subtracts the average value from the keys (the "Key Reparameterization"). Imagine if the giants in our line were actually just standing on a high platform. OptR lowers the platform so everyone is standing on the same ground level. This doesn't change who is tall or short relative to each other, but it stops the giants from hitting the ceiling of the small room. Crucially, this step is "attention-equivalent," meaning it changes the numbers but keeps the AI's attention focus exactly the same. It narrows the range of numbers, making the INT2 squishing much less painful.
Learning the Perfect Spin (The "Output-Aware" Step): Instead of using a fixed spin (like a standard mathematical rotation), OptR learns a custom spin for every single "head" (a specific part of the AI's brain) in the model. It does this by simulating the whole process: squishing the data, un-squishing it, running it through the AI's attention mechanism, and then projecting it to the final answer.
- It looks at the error in the final answer (the "post-WO attention-output error").
- It breaks this error down into two parts: errors caused by the keys (which determine what the AI pays attention to) and errors caused by the values (which determine what information is actually retrieved).
- It then tweaks the rotation angles slightly to minimize the error in the final answer, not just the error in the storage.

Think of it like tuning a radio. Old methods tried to make the static sound as quiet as possible inside the speaker box. OptR listens to the music coming out of the speaker and adjusts the knobs until the song sounds perfect, even if the static inside the box isn't perfectly silent.

What They Found

The researchers tested OptR on three different AI models (Qwen3-4B, Qwen3-8B, and Phi4-14B) and five different challenging benchmarks, including math problems (AIME25), coding tasks (LiveCodeBench), and long-context retrieval (finding a needle in a haystack of 64,000 words).

The results were striking:

Massive Accuracy Gains: On the Qwen3-8B model, using standard INT2 quantization with the previous best method (QuaRot) resulted in an accuracy of only 17.33% on a difficult math test. When they added OptR, the accuracy skyrocketed to 66.67%. That's nearly a four-fold improvement.
Beating the Baseline: Even compared to the state-of-the-art method (OSCAR), which was already quite good, OptR pushed the accuracy from 54.67% to 66.00%.
Long-Context Superpower: The most impressive finding was in long-context tasks. As the story got longer (up to 64,000 tokens), standard INT2 methods failed miserably, dropping to near-zero accuracy. OptR kept the retrieval accuracy high, dropping only slightly from 99.83% (at 4k tokens) to 70.02% (at 64k tokens), whereas the standard method crashed to 0.04%.
No Speed Penalty: The authors confirmed that this magic doesn't come with a heavy cost. They integrated OptR into the system and found that it added negligible overhead. The speed of the AI (latency) and the amount of data it could process (throughput) remained almost identical to the standard methods. It's like getting a supercharged engine without adding any extra weight to the car.

Why This Matters

The paper suggests that for ultra-low-bit quantization (like INT2) to work, we can't just look at the data in isolation. We have to look at how that data flows through the entire system to the final answer. By optimizing for the output rather than the storage, OptR bridges the gap between extreme compression and high performance.

The authors emphasize that this isn't just a theoretical win; it's a practical one. It allows AI models to handle much longer conversations and larger batches of users without running out of memory, all while keeping the model smart enough to solve hard math problems and write code. They didn't just suggest this might work; they measured it across multiple models and found consistent, significant improvements, proving that "output-aware" optimization is the key to unlocking the full potential of 2-bit AI.

Technical Summary: Output-Aware Rotation for INT2 KV-Cache Quantization

Problem Statement

As Large Language Models (LLMs) scale in size and context length, the Key-Value (KV) cache has become a primary bottleneck for memory and bandwidth during long-context inference. While ultra-low-bit quantization (specifically INT2) offers a path to reduce storage to 1/8 of BF16 requirements, it introduces significant quantization errors. In INT2, where only four levels exist, a few extreme outlier values can expand the quantization range, causing large rounding errors for the majority of values.

Existing rotation-based INT2 methods (e.g., QuaRot, OSCAR) attempt to mitigate this by rotating cache vectors into a quantization-friendly basis. However, these methods optimize rotation matrices based on proxy objectives defined before the complete attention readout, such as minimizing cache reconstruction error or matching pre-attention statistics. The paper argues that this creates an objective mismatch: the rotation that best reconstructs cached keys and values does not necessarily minimize the error propagated through the attention mechanism and the output projection matrix ( $W_O$ ) into the model's hidden state. Consequently, the model's final prediction may suffer despite a well-reconstructed cache.

Methodology: OptR

The authors propose OptR, an output-aware rotation method designed to minimize the post- $W_O$ attention-output error directly. The method operates through a one-time offline calibration process where model weights remain frozen, and only rotation corrections are learned.

1. Output-Space Error Decomposition

OptR formulates the quantization problem in the output space. It decomposes the error between the BF16 attention output and the INT2 quantized output into two distinct terms:

Key-induced error: Arising from quantization errors in keys ( $\Delta k$ ) that perturb the attention logits and the resulting softmax distribution ( $\Delta p$ ).
Value-induced error: Arising from quantization errors in values ( $\Delta v$ ) aggregated by the (potentially perturbed) attention weights.

The optimization target is the squared error of the final output contribution ( $\Delta y$ ) after projection by $W_O$ , rather than the raw reconstruction error of the cache.

2. Attention-Equivalent Key Reparameterization

Before rotation, OptR applies a per-channel mean shift (centering) to the keys. By subtracting the calibration mean $\mu$ from all keys, the method shifts the attention logits by a query-dependent constant. Crucially, this operation leaves the softmax attention distribution unchanged (attention-equivalent) but significantly narrows the dynamic range of the keys. This reduces the impact of outliers on the limited INT2 quantization levels without altering the attention mechanism's behavior.

3. Per-Head Orthogonal Corrections

OptR learns specific orthogonal corrections for each KV head to refine a base rotation (e.g., Hadamard or covariance-based).

Parameterization: Corrections are learned as unconstrained matrices $A_K, A_V$ , which are converted to skew-symmetric generators ( $S = A - A^\top$ ) and exponentiated to ensure the resulting rotations remain orthogonal: $R = R_0 \exp(S)$ .
Two-Stage Calibration:
1. Key Calibration: The key rotation is optimized while keeping values in BF16. The loss function combines a KL-divergence term (to preserve the attention distribution) and a term measuring the post- $W_O$ output error induced by key quantization.
2. Value Calibration: Once the key rotation is fixed, the value rotation is optimized to minimize the output error resulting from value quantization under the selected attention distribution.
Differentiability: Since INT2 rounding is non-differentiable, gradients are estimated using a Straight-Through Estimator (STE).

4. System Integration

OptR is integrated into an SGLang-based INT2 KV-cache pipeline. It maintains compatibility with paged KV-cache systems and fused decoding kernels by preserving the tensor shape and regular memory layout. The key reparameterization and rotation are implemented via optimized Triton kernels that fuse operations (rotation, centering, clipping, quantization) to minimize overhead.

Key Contributions

Output-Space Optimization Formulation: The paper reframes INT2 KV-cache quantization as an optimization problem targeting the post- $W_O$ attention-output error, decomposing it into key- and value-induced terms.
OptR Algorithm: A novel method combining attention-equivalent key reparameterization (to reduce channel-wise offsets) with per-head orthogonal corrections learned through the full INT2 quantization and attention path.
Efficiency and Compatibility: The method retains the paged KV-cache format and introduces negligible inference overhead, making it practical for deployment in existing serving systems.

Experimental Results

The authors evaluated OptR on three models (Qwen3-4B, Qwen3-8B, and Phi4-14B) across five benchmarks (AIME24/25, GPQA-Diamond, MBPP+, LiveCodeBench v6) and long-context retrieval (RULER-NIAH).

Reasoning and Coding Accuracy: OptR significantly improves performance over base rotation methods. For Qwen3-8B on AIME25, accuracy increased from 17.33% (QuaRot-INT2) to 66.67% (QuaRot-INT2 + OptR), and from 54.67% (OSCAR) to 66.00% (OSCAR + OptR). These results approach the BF16 baseline of 68.00%.
Long-Context Retrieval: OptR demonstrates superior robustness as context length increases. On RULER-NIAH with 64K context, QuaRot-INT2 dropped to near-zero accuracy, whereas QuaRot-INT2 + OptR maintained 70.02% accuracy.
Error Analysis: Per-layer analysis shows that OptR reduces both immediate post- $W_O$ output error and propagated residual-stream error across most layers and KV heads, not just a subset.
Efficiency: Integration of OptR adds negligible latency (within 2% of the base pipeline) and minimal memory overhead (e.g., +18 MiB for a batch of 128 with 2K/4K context).

Significance

The paper claims that OptR addresses a fundamental mismatch in existing low-bit quantization strategies: optimizing for cache reconstruction does not guarantee optimal model performance. By directly optimizing for the error observed by subsequent layers (post- $W_O$ ), OptR enables ultra-low-bit (INT2) KV caching to achieve performance comparable to BF16 in complex reasoning and long-context tasks. The method is presented as a practical, drop-in enhancement for existing rotation-based pipelines that preserves system efficiency while drastically reducing quantization-induced degradation.