Blog Tags GitHub
Back to blog
kv-cache quantization llm-inference long-context

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

RateQuant用率失真理论把KV cache按注意力头分配bit宽度：假设量化MSE满足$D(b)=\alpha\beta^{-b}$，用reverse waterfilling闭式求最优分配，并识别distortion model mismatch这一失效模式——不同量化器$\beta$从3.6到5.3变化，套错模型比uniform更差。校准per-quantizer $\beta$并对K/V单独分配预算后，Qwen3-8B在KIVI 2.5bit的PPL从49.3降到14.9（−70%），QuaRot改善6.6 PPL；全流程校准仅1.6秒，推理零开销。

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

一、论文概述

项目内容
arXiv ID2605.06675
标题RateQuant: Optimal Mixed-Precision KV Cache Quantization via Rate-Distortion Theory
作者Fei Zuo, Zikang Zhou, Hao Cong, Xiaoyan Xi, Ho Fai Leung
提交日期2026-04-22（v1），2026-06-26（v2）
学科分类cs.LG, cs.CL, cs.IT
评测Qwen3-4B/8B/32B，WikiText-2 PPL + ARC/HellaSwag/PIQA/WinoGrande

二、核心思想

问题定义：LLM 生成时缓存所有 KV pair，规模随 seq len × batch × 层数线性增长；一个 32B 模型 4k seq FP16 已占 >1GB，往往超过权重本身。KV cache 量化能降低此成本，但现有量化器（KIVI、QuaRot、TurboQuant、Hooper 2024 等）都给每个注意力头分配相同 bit-width，忽略头重要性异质性。自然想法是按头分配不同 bit，但作者发现一个隐藏陷阱：不同 quantizer 的失真曲线 D(b)=αβ−bD(b)=\alpha\beta^{-b} 的衰减率 β\beta 差异从 3.6 到 5.3，把一个 quantizer 的 distortion model 套到另一个 quantizer，会翻转 marginal gain 排序，让 mixed-precision 比 uniform 更差——distortion model mismatch。

图1 Distortion model mismatch

解决方案：RateQuant 把 per-head KV cache bit 分配形式化成率失真优化问题，用 reverse waterfilling 闭式求最优连续解，用贪心 求最优整数解，并通过per-quantizer 校准 β\beta 消除 mismatch。进一步把 K/V 视为 2N2N 独立组件、允许它们拿到不同预算，得到”quantizer-agnostic”分配层。

三、技术架构/方法

RateQuant 分四阶段（图 2）：

图2 RateQuant pipeline

3.1 问题建模

LL 层 × HH KV heads = N=LHN = LH 组。每组 sensitivity wi>0w_i>0。假设：

Assumption 1（指数失真–率）：D(b)=α⋅β−bD(b)=\alpha\cdot\beta^{-b}。作者用 TurboQuant Lloyd-Max MSE（d=128d=128）拟合 α≈1.36,β≈3.48,R2>0.99\alpha\approx 1.36, \beta\approx 3.48, R^2>0.99。

优化：

min⁡b∈RNJ(b)≜∑iwiD(bi)s.t.∑ibi=B, bmin⁡≤bi≤bmax⁡\min_{\mathbf{b}\in\mathbb{R}^N} \mathcal{J}(\mathbf{b}) \triangleq \sum_i w_i D(b_i) \quad \text{s.t.}\quad \sum_i b_i = B,\ b_{\min}\le b_i\le b_{\max}

3.2 Reverse Waterfilling（Theorem 2）

连续解：

bi∗=bˉ+ln⁡wi−ln⁡w‾ln⁡βb_i^* = \bar{b} + \frac{\ln w_i - \overline{\ln w}}{\ln\beta}

其中 bˉ=B/N\bar b = B/N。关键洞察：β\beta 越小，同样的 sensitivity ratio 分到越多额外 bit。TurboQuant β=3.48\beta=3.48 时，某头 sensitivity 是均值 ee 倍会分到 1/ln⁡3.48≈0.801/\ln 3.48 \approx 0.80 额外 bit。

Gain ratio（Theorem 3）：

JuJ∗=wˉw~≥1\frac{\mathcal{J}_u}{\mathcal{J}^*} = \frac{\bar w}{\widetilde w} \ge 1

wˉ\bar w 是算术均值，w~=(∏wi)1/N\widetilde w=(\prod w_i)^{1/N} 是几何均值。AM/GM 比可在不量化的情况下预测混合精度潜在收益；Qwen3 上经验 wˉ/w~≈2.0\bar w/\widetilde w\approx 2.0，说明有可观空间。

Corollary 4：若 ln⁡wi∼N(μ,σ2)\ln w_i \sim \mathcal{N}(\mu,\sigma^2)，则 Ju/J∗=exp⁡(σ2/2)\mathcal{J}_u/\mathcal{J}^* = \exp(\sigma^2/2)。

3.3 Sensitivity Estimation

用 KV 投影输出的平方梯度范数：

wl,hK=Ex∼D[1T∑t∥∂L∂Kl,h,t∥2]w_{l,h}^K = \mathbb{E}_{\mathbf{x}\sim\mathcal{D}}\left[\frac{1}{T}\sum_t \left\|\frac{\partial\mathcal{L}}{\partial\mathbf{K}_{l,h,t}}\right\|^2\right]

wl,hVw_{l,h}^V 类似。D\mathcal{D}：16 sequences × 512 tokens。

Proposition 5（Loss–distortion connection，2 阶 Taylor + diagonal Fisher）：

E[L(θ^)−L(θ)]≈∑l,h[wl,hKD(bl,hK)+wl,hVD(bl,hV)]\mathbb{E}[\mathcal{L}(\hat\theta) - \mathcal{L}(\theta)] \approx \sum_{l,h}[w_{l,h}^K D(b_{l,h}^K) + w_{l,h}^V D(b_{l,h}^V)]

梯度型 sensitivity 直接出现在 loss 展开中，activation-based 只 bound 前向误差放大——这解释了后者 1.07 PPL 劣化。

3.4 整数贪心分配（Algorithm 1）

对每个组从 bmin⁡b_{\min} 起，每次选 arg⁡max⁡iwi[Di(bi)−Di(bi+1)]\arg\max_i w_i[D_i(b_i)-D_i(b_i+1)]，bi∗+=1b_{i^*}\mathrel{+}=1，直到消耗完预算 B−Nbmin⁡B-N b_{\min}。Proposition 6：D(b)D(b) 凸时贪心即最优（polymatroid 结构）。

3.5 Quantizer-Agnostic 扩展

Empirical distortion calibration：TurboQuant β≈3.6\beta\approx 3.6，KIVI/QuaRot β≈5.0\beta\approx 5.0–5.3。测 b∈{2,3,4,5,6}b\in\{2,3,4,5,6\} MSE，ln⁡D\ln D vs bb 做最小二乘拟合 (αq,βq)(\alpha_q, \beta_q)。至关重要：错的 β\beta 反转 marginal gain 排序（图 3），KIVI 2.5bit 从 49.3 恶化到 87.0。

图3 Marginal gain mismatch

Separate K/V allocation：K 和 V 用不同量化 schema（per-channel key vs per-token value）时，令它们成为 2N2N 独立组件，允许 bˉK≠bˉV\bar b_K\ne \bar b_V。

Pipeline summary：(1) 16 forward+backward 做 sensitivity（8B on H200 ~1.6s）；(2) 5 个 bit 上 distortion 建模（<0.1s）；(3) 贪心分配 2N2N 组件（<0.01s）；(4) 推理侧使用静态 2KB lookup table，零 runtime 开销。

3.6 Per-head 分配可视化

图7 Per-head bit 分配图

Qwen3-8B bˉ=4.0\bar b=4.0（bmin⁡=3,bmax⁡=6b_{\min}=3, b_{\max}=6）：高敏感度头（早/晚层）拿 5–6 bit，低敏感度中间层拿 3 bit。梯度型 sensitivity 呈 U 型（早+晚层高），activation-based 单调递增（图4）。

图4 Per-head sensitivity

四、核心创新

创新点说明
率失真优化框架首次把 per-head KV bit 分配写成 min⁡∑wiD(bi)\min\sum w_i D(b_i)，闭式 reverse waterfilling 求解
Distortion model mismatch 概念识别并解决”用错 quantizer 的 β\beta 反而更差”的失效模式
Per-quantizer β\beta 校准5 bit-width MSE + 最小二乘拟合，让方法 quantizer-agnostic
Separate K/V allocationK/V 独立预算，可差异化到 2.85/2.15 bit
Gradient-based sensitivity + Prop 5理论证明梯度型是 loss-preserving 分配的正确 proxy
AM/GM 预算收益预测器wˉ/w~\bar w/\widetilde w 无需量化即可预判混合精度潜在收益

五、实验结果

5.1 主结果：Uniform vs RateQuant（TurboQuant）

WikiText-2 PPL（seed 42；8B 3 seeds 均值），Δ\Delta = PPL 改善：

Modelbˉ\bar bUniformRateQuantΔ\DeltaFP16
Qwen3-4B3.513.8913.70+0.2013.19
4.013.9213.47+0.45
Qwen3-8B3.510.009.76+0.249.53
4.09.949.67+0.27
Qwen3-32B3.57.707.64+0.067.50

Qwen3-8B @ 4.0 bit：恢复 66% headroom；Qwen3-4B @ 4.0 bit：恢复 62%。sweet spot 在 3.5–4.0 bit。

5.2 跨量化器校准（Qwen3-8B）

Table 3（bmin⁡=2b_{\min}=2）四策略对比——Theo=用 TurboQuant 的 D(b)D(b)；Cal=校准；+Sep=K/V 分离：

Quantbˉ\bar bUniformTheoCalCal+SepbˉK/bˉV\bar b_K/\bar b_V
KIVI2.549.3286.9573.1214.86 (−70%)2.85/2.15
KIVI3.010.8112.4311.3010.973.07/2.93
QuaRot2.534.88271.950.5228.33 (+6.6)2.61/2.39
QuaRot3.011.9012.2710.8410.58 (+1.3)3.04/2.96
TurboQuant3.010.929.969.969.88 (+1.0)3.06/2.94
FP169.53
  • mismatch 会灾难化：KIVI 从 49.3 → 87.0，QuaRot 从 34.9 → 271.9；
  • calibration 部分恢复，但 K/V separation 是关键：KIVI 2.5bit 从 73.1 降到 14.9，靠”K 2.85bit, V 2.15bit” 的差异化预算（per-channel key 更 error-prone，per-token value 较容忍）；
  • TurboQuant + RateQuant 3.0bit（9.88） 击败 KIVI uniform 3.0（10.81）与 QuaRot uniform 3.0（11.90）——在弱 quantizer 上做好分配可以超过强 quantizer + uniform。

图5 PPL vs bits 图6 KIVI per-layer MSE

图 6：KIVI key cache 失真远高于 value cache（~4× MSE），驱动 K/V 拆分（2.85/2.15 @ 2.5bit）。

5.3 Sensitivity proxy 消融（Qwen3-8B, bmin⁡=3b_{\min}=3）

bˉ\bar bUniformGradientActivationSwing
3.59.959.76 (+0.19)10.83 (−0.88)1.07
4.09.949.59 (+0.35)10.02 (−0.08)0.43
4.59.629.58 (+0.05)9.85 (−0.22)0.27

3.5bit 处 swing 1.07 PPL 超过 uniform→FP16 的 gap，说明 proxy 选择比分配算法更关键。梯度型呈 U 形（早/晚层敏感），activation-based 单调递增（图 4）。

5.4 下游 & 吞吐（Qwen3-8B @ 4.0 bit, TurboQuant）

指标FP16UniformRateQuant
ARC-C55.852.554.1
HellaSwag57.155.256.6
PIQA76.974.476.4
WinoGrande67.666.969.5
Avg64.462.264.2 (恢复 89.8% gap)
tok/s37.738.138.0

吞吐等价（38.0 vs 38.1 tok/s），推理零开销。

5.5 与 mixed-precision 基线对比（KIVI 2.5bit）

  • Layer-level 方法：PPL ↓ ~25%；
  • 全局 K>V split：↓ 37%；
  • RateQuant：↓ 70%。

六、总结

核心贡献

  1. 首次把 mixed-precision KV cache 分配放入率失真理论框架，reverse waterfilling 给出最优解与 gain ratio；
  2. 识别并解决 distortion model mismatch——per-quantizer β\beta 校准是必需的；
  3. Separate K/V allocation 把 K/V 视为 2N2N 独立组件，为不同 schema（per-channel key vs per-token value）差异化预算；
  4. 校准仅 1.6s（Qwen3-8B on H200），推理零开销，wˉ/w~\bar w/\widetilde w 提供先验收益预测器。

技术影响

  • KIVI 2.5bit PPL 从 49.3 → 14.9（−70%），QuaRot 2.5bit −6.6 PPL；
  • Qwen3-8B @ 4.0bit 达到 FP16 89.8% 精度差恢复率、平价吞吐；
  • 为长上下文 / 高 batch 服务提供了理论基础扎实、部署零成本的 KV 压缩层；
  • 建立”AM/GM 预测 mixed-precision 是否值得做”的通用判据。

局限性

  • 静态 per-head 分配，未捕获 input-dependent sensitivity 变化（长上下文里可能进一步波动）；
  • bˉ≥3.5\bar b\ge 3.5 时头room 快速消失，主要收益在激进 bit 区间（≤3.0\le 3.0）；
  • 需要基础 quantizer 支持 per-head 或 per-group bit-width（非所有工程实现都开放此接口）；
  • 未与 KV eviction / token merging 等 orthogonal 压缩协同评测。

七、参考资源