RateQuant: Optimal Mixed-Precision KV Cache Quantization via Rate-Distortion Theory

Authors: Fei Zuo, Zikang Zhou, Hao Cong et al.

arXiv · PDF

Summary

Quantizing the key-value cache to fewer bits saves memory, but current methods use the same bit-width for every attention head. Different heads have different importance, so a natural idea is to give more bits to important heads. The catch: each quantizer design has its own error-vs-bits curve, and if you fit the allocation using one quantizer's curve and apply it to another, the allocation can invert and perform worse than uniform quantization. RateQuant solves this by measuring each quantizer's actual distortion curve on a small calibration set, then uses rate-distortion theory (reverse waterfilling) to allocate bits optimally. On Qwen3-8B at 2.5 bits, it cuts perplexity from 49 to 15.

Main takeaways:

• Attention heads vary in importance, but naive mixed-precision can backfire
• Different quantizers have different error-decay rates (beta from 3.6 to 5.3)
• Using the wrong distortion model inverts bit allocation and hurts performance
• RateQuant calibrates each quantizer's curve and solves bit allocation with closed-form math
• Calibration takes 1.6 seconds on one GPU; inference has zero overhead
• 70% perplexity reduction on some quantizers at 2.5 average bits

Relevance

Not directly related to my persona work, but conceptually interesting: it shows that system components (attention heads) have heterogeneous importance, and naive uniform treatment fails — a pattern I also see with persona leakage and marker implantation varying by context.

Abstract

arXiv:2605.06675v1 Announce Type: new Abstract: Large language models cache all previously computed key-value (KV) pairs during generation, and this KV cache grows linearly with sequence length, making it a primary memory bottleneck for serving. Quantizing the KV cache to fewer bits reduces this cost, yet all current quantizers assign the same bit-width to every attention head, ignoring the large variation in head importance. A natural idea is to allocate more bits to important heads and fewer to the rest. We show, however, that such mixed-precision allocation has a hidden pitfall: each quantizer follows a different distortion curve D(b)=alpha*beta^{-b}, and the decay rate beta varies from 3.6 to 5.3 across quantizer designs. Applying one quantizer's distortion model to another inverts the allocation order and makes performance worse than uniform quantization. We call this failure mode distortion model mismatch and propose RateQuant to resolve it. RateQuant fits a per-quantizer distortion model from a small calibration set, then solves the resulting bit-allocation problem in closed form via reverse waterfilling from rate-distortion theory. On Qwen3-8B at 2.5 average bits, calibrated RateQuant reduces KIVI's perplexity from 49.3 to 14.9 (70% reduction) and improves QuaRot by 6.6 PPL. The entire calibration takes 1.6 s on a single GPU and adds zero overhead at inference time.