KVarN: Variance-Normalized KV-Cache Quantization Mitigates Error Accumulation in Reasoning Tasks
Lorenz K. Muller, Philippe Bich, Chiara Boretti, Hyun-Min Chang, Jiawei Zhuang, Lukas Cavigelli
cs.LG
2026-06-02
用 Hadamard 旋转加双向方差归一化修掉 token 尺度离群,把 KV cache 压到 2.3 bit;Qwen3-4B 在 AIME24 上 60.0%,接近 FP16 的 61.1%,远超 KIVI。
test-time scaling(推理时多花算力、让模型多想几步)能让大模型推理更强,但代价是 KV cache 随解码越积越大,撑爆显存。把 KV cache 量化到 2 bit 能省显存,可现有方法是在「类 prefill」设置下评测的:把一段固定长上下文一次性并行量化。在真正的自回归解码里,token 是边生成边压缩的,量化误差会逐 token 累积,因为早期 token 的误差会顺着后续 transformer block 的输出滚雪球。现有方法没保住每个 token 的尺度,误差就这么叠上去。
作者专门设计了一个「伪解码」(pseudo-decode)评测来量化这种累积:把序列切成块(通常 128 token),每生成一块就量化一次,再看注意力输出的重建误差怎么随块数涨。正是在这个设置下,KIVI 这类现有方法和 KVarN 的差距才被拉开。用老的 prefill 评测,差距看不出来。
KVarN 的起点是一个观察:把量化误差拆成幅度误差(‖K‖−‖Kdq‖)² 和方向误差 2‖K‖‖Kdq‖(1−cosθ),会发现误差最大的那 5% 离群 token 主要是幅度上的(单个 token 的尺度乘子算错),方向上其实没偏多少。而这 5% 离群只占整体 MSE 的一小部分,却贡献了端到端掉点的大头。所以关键是把 token 尺度修对,不让那个错误乘子逐时间步滚雪球。
流程(每元素 2.3 bit)三步:通道维做 Hadamard 旋转把通道间离群摊平 → 双向方差归一化(Sinkhorn-Knopp 式,token 维和通道维轮流归一,从 SINQ 改造)→ RTN 2-bit 量化。方差归一化产出的归一系数吸收进一个高精度按组尺度里,所以 token 尺度离群被修掉了。存储是 2-bit 值加 FP8 尺度加每 128 元素一个 FP16 零点,块大小 128×128;两个 Hadamard 变换分别融进 WV 和 WO 权重,RoPE 之后再做两次在线变换。
| 方法 | Qwen3-4B AIME24 | Phi-4-14B AIME24 | bit/元素 |
| FP16 | 61.1 | 62.2 | 16.0 |
| KIVI | 55.5 | 57.8 | 2.3 |
| QuaRot | 56.7 | 58.9 | 2.3 |
| KVQuant-1% | 40.0 | 55.6 | 2.4 |
| KVarN | 60.0 | 61.7 | 2.3 |
同样 2.3 bit,KVarN 在 AIME24 上几乎贴着 FP16,把 KIVI、QuaRot、KVQuant 都甩开一截。MATH500 上 KVarN 在 Qwen3-4B 到 79.2(FP16 82.6),在 Phi-4-14B 到 84.8(FP16 84.9,基本无损)。HumanEval 两边都在 88% 上下,和 FP16 持平。600 行长检索里 KVarN 85%,KIVI 只有 74%,FP16 是 90%。显存从 16 bit 压到 2.3 bit,降 85.6%,方差归一化每 128 token 只花 1.9ms,相对 vLLM 基线约 0.18% 开销。整组实验约花了 50 GPU 天。
长链推理模型在 KV cache 上烧显存;KVarN 把它压到原来的约七分之一,且在数学(AIME24、MATH500)和长检索这种过去 2-bit 方法崩盘的任务上近乎不掉点,代码集成进了 vLLM。对要部署长推理又卡显存的人,这是个直接可用的方案。但收益有一部分是「理论上的」,要等推理引擎支持端到端 2-bit KV 才完全兑现。
KVarN 只适用于带 KV cache 的架构,不适用于 SSM 这类没有 KV cache 的模型。对 MLA(DeepSeek 那种多头潜在注意力)适用性不明。最尴尬的是,当前还没有任何服务框架支持端到端 2-bit KV,所以完整好处得等推理引擎跟上。另外它最亮的数字集中在 AIME24 这类数学任务,generative 长文本任务的累积行为有没有这么干净,论文没充分覆盖。