Output-Aware Rotation for INT2 KV-Cache Quantization

2026-08-03

AI Summary

This work addresses the memory and bandwidth bottlenecks imposed by KV caching in long-context large language model inference, where existing INT2 quantization methods are limited due to a mismatch between their optimization objectives and actual output errors. The paper proposes OptR, the first output-aware orthogonal rotation strategy that jointly reparameterizes keys to mitigate channel-wise bias by minimizing attention errors after output projection. OptR preserves the paged KV cache format and incurs negligible inference overhead while significantly improving quantization accuracy. Experimental results demonstrate that OptR consistently outperforms QuaRot and OSCAR across three models and five reasoning and code benchmarks, with notable gains in long-context retrieval capabilities.

Abstract

The key-value (KV) cache has become a major memory and bandwidth bottleneck in long-context large language model inference, making ultra-low-bit quantization increasingly important. However, existing rotation-based INT2 methods optimize cache statistics or proxy errors before the complete attention readout, even though the model is ultimately affected by the error propagated through attention and the output projection $W_O$. To address this mismatch, we propose \textit{OptR}, an output-aware rotation method that minimizes post-$W_O$ attention-output error. OptR decomposes the post-$W_O$ attention-output error into key- and value-induced terms and learns per-head orthogonal corrections through the full INT2 quantization and attention path. OptR further applies an attention-equivalent key reparameterization to reduce large channel-wise offsets without changing the softmax distribution. Across three models and five reasoning and coding benchmarks, OptR consistently improves both QuaRot and OSCAR and strengthens long-context retrieval, while preserving the paged KV-cache format with negligible inference overhead.

Problem

• KV-cache quantization
• INT2
• output-aware
• attention mechanism
• large language models

Innovation

• output-aware rotation
• INT2 quantization
• KV-cache
• attention output error
• orthogonal correction