PN-Mix Uses Query–Key Pair Norms for Low-Bit KV Cache Quantization

Why KV-cache quantization remains difficult

In long-output inference, the key-value cache grows with the sequence and must be read repeatedly during decoding. This creates both a memory-capacity problem and a bandwidth bottleneck. Quantization reduces the footprint, but aggressive settings such as 2-bit storage can introduce enough error to damage attention and, ultimately, model accuracy. Mixed precision offers a compromise: reserve higher precision for the parts of the cache that matter most.

A saliency signal tied to attention

Earlier mixed-precision key-cache methods generally estimate channel importance from statistics of the keys themselves. PN-Mix takes a more direct route. Its analysis examines how quantization errors perturb attention logits at the level of RoPE channel pairs. Because RoPE operates on paired channels, the method evaluates the interaction between the corresponding query and key components.

The central observation is that channel pairs with a large product of query and key norms can make a larger contribution to logit variation when quantized. PN-Mix therefore assigns higher precision to these salient pairs, instead of treating every channel equally or ranking channels solely by key-side statistics.

The allocation is determined during the prefill stage, using the current input, and then reused throughout decoding. This makes the strategy input-adaptive without requiring importance estimation at every generated token. For long-output workloads, avoiding repeated allocation decisions is important because the cache is accessed continuously after prefill.

Reported results

Across standard reasoning, code-generation, and AIME benchmarks, the paper reports that PN-Mix reaches comparable or better average accuracy at 2.26 average bits than prior mixed-precision approaches using 2.32–2.39 bits. Its Pareto sweep also shows a more gradual accuracy decline as the budget is reduced to 1.5 bits. The compared baselines exhibit sharper degradation below 2.25 bits.

The result highlights an important design choice in cache compression: the bit budget is only part of the problem. The location of the high-precision budget matters as well. A saliency measure connected to the attention computation can be more useful than a generic measure of activation magnitude or a cache-only statistic.

Implications and limitations

PN-Mix suggests a practical direction for reducing the cost of long-context and long-output serving. By concentrating precision on input-relevant RoPE channel pairs, it may lower cache storage and data movement while retaining more of the model’s behavior than uniform low-bit quantization.

The scope is also clear. The method is designed for RoPE-based decoder-only LLMs and focuses on the key cache. Its behavior with other positional-encoding schemes, attention architectures, or joint key-value quantization is not established by the supplied results. Still, the work offers a useful principle: mixed-precision cache allocation should reflect the query–key interactions that actually drive attention, not just the standalone statistics of stored keys.

Source: Hugging Face Daily Papers