受 TurboQuant 启发的 KV 缓存量化方案的统计推断与质量评估

# 受 TurboQuant 启发的 KV Cache 量化的统计推断与质量度量

来源: https://arxiv.org/html/2605.08114
Paolo D’Alberto paolodalberto@amd\.com

###### 摘要

我们在公平的位预算下分析了三种 KV cache 量化方案：KV（标量 MSE 基线）、KQV（在 $K$ 上使用 WHT + MSE；在 $V$ 上使用 WHT + MSE + QJL），以及 QKQV（在两者上均使用 WHT + MSE + QJL）。从超球面上的 Beta 分布出发，我们追踪了 $K$ 上的 QJL 如何通过 $\sqrt{\pi/2}$ 放大内积方差，进而通过琴生不等式被 softmax 非线性地放大。

研究得出了三项实证发现。（1）在 $n=4$（实践中占主导地位的预算）时，KQV 在所有分布和秩的测试中，在 KL 散度、几何 $K$ 误差和 6D 距离等所有指标上均胜出。（2）$K$-$V$ 不对称性是绝对的：在所有预算和分布下，QKQV 在 KL 散度方面始终劣于 KQV。（3）存在一种依赖于预算的交叉现象：QKQV 在 $n \in \{2, 3, 5\}$ 时实现更好的几何 $K$ 重建，而 KQV 在 $n \in \{4, 6\}$ 时表现更好，这与秩和尾部权重无关——这是一个开放的速率-失真问题。

KL$(p_{\text{ref}} \| p_{\text{quant}})$ 仅由 $K$ 决定，它架起了 $K$ 方向误差与路由腐败及输出崩溃之间的桥梁。在 $n=4$ 时，KQV 实现的 KL 散度比 QKQV 低 $2.25\times$ ($MW r=-0.983, p<10^{-23}$)，这解释了上述交叉现象。我们引入了一个 6D 误差框架，提供了超越标量指标的分布判别能力和几何分辨率。

## 1 超球面上的 Beta 分布

设 $\mathbf{x}$ 在 $\mathcal{S}^{d-1}$ 上均匀分布。每个坐标 $x_j$ 的密度为

$f_X(t)=\frac{\Gamma(d/2)}{\sqrt{\pi}\,\Gamma\!\left(\tfrac{d-1}{2}\right)}\left(1-t^2\right)^{\frac{d-3}{2}},\quad t\in[-1,1], \tag{1}$

即在 $[-1,1]$ 上的 $\text{Beta}\!\left(\tfrac{d-1}{2},\tfrac{d-1}{2}\right)$ 分布。当 $d \to \infty$ 时，这收敛于 $\mathcal{N}(0,1/d)$。

几何证明如下。固定 $x_1=t$。剩余坐标满足 $x_2^2+\cdots+x_d^2=1-t^2$，在 $\mathbb{R}^{d-1}$ 中形成一个半径为 $\sqrt{1-t^2}$ 的 $(d-2)$-维球面。在值 $t$ 处的表面积元素为

\begin{split} \|dS\|_t &\equiv \|\mathcal{S}^{d-2}\|\cdot\left(1-t^2\right)^{\frac{d-2}{2}}\cdot\frac{dt}{\sqrt{1-t^2}}\\ &=\|\mathcal{S}^{d-2}\|\cdot\left(1-t^2\right)^{\frac{d-3}{2}}\,dt, \end{split} \tag{2}$

其中因子 $1/\sqrt{1-t^2}$ 是弧长修正（球面从垂直轴弯曲）。通过 $\|\mathcal{S}^{d-1}\|$ 归一化并代入 $\|\mathcal{S}^k\|=2\pi^{(k+1)/2}/\Gamma\!\left(\tfrac{k+1}{2}\right)$ 可得方程 (1 (https://arxiv.org/html/2605.08114#S1.E1))。

两个表面积均通过 Gamma 函数表示：$\|\mathcal{S}^{d-1}\|=2\pi^{d/2}/\Gamma(d/2)$ 和 $\|\mathcal{S}^{d-2}\|=2\pi^{(d-1)/2}/\Gamma\!\left(\tfrac{d-1}{2}\right)$。它们的比值即为方程 (1 (https://arxiv.org/html/2605.08114#S1.E1)) 中的归一化常数。Gamma 函数自然地出现在高斯径向积分 $\int_0^\infty r^{d-1}e^{-r^2}dr=\Gamma(d/2)/2$ 中，这是 $d$ 维球面几何的通用 size 单位。每次固定一个坐标并将球面投影时，维度减少一，Gamma 的参数偏移 $1/2$：从 $d/2$ 变为 $(d-1)/2$。密度 (1 (https://arxiv.org/html/2605.08114#S1.E1)) 正是这种维度计量的体现。

对于较大的 $d$ 和较小的 $t$，$\left(1-t^2\right)^{(d-3)/2}\approx e^{-dt^2/2}$，这是一个方差为 $1/d$ 的高斯分布。这种收敛在经验上是可见的：坐标直方图在 $d=2$ 时呈 U 形（反正弦分布），在 $d=4$ 时呈轻微拱形，在 $d=8$ 时呈钟形，在 $d=1024$ 时与高斯分布无法区分。图 1 (https://arxiv.org/html/2605.08114#S1.F1)、2 (https://arxiv.org/html/2605.08114#S2.F2)、4 (https://arxiv.org/html/2605.08114#S6.F4) 和 7 (https://arxiv.org/html/2605.08114#S8.F7) 展示了这四个代表性维度下的这一进展。注意，Beta 分布定义在区间 $[-1,1]$ 上，而对于任何有限的 $d$，高斯分布并非如此，只能对其进行近似。

Refer to caption Figure 1: $d=2$, Beta$(0.5,0.5)$——反正弦（U 形）。顶部：旋转坐标直方图（蓝色）与 Beta 密度（红色）对比，数据来自 gaussian、heavy_tail、low_rank。底部：KL 散度箱线图，MSE 4-bit 与 MSE 3-bit++QJL 1-bit（等预算）。MSE 在所有分布中胜出。圆圈表示超出 $1.5\times\mathrm{IQR}$ ($\approx 2.7\sigma$ for a Gaussian) 的异常值。

## 2 TurboQuant 流水线

在 Transformer 注意力机制中，每个解码步骤计算出的键（Key）和值（Value）张量会被存储并复用于所有后续步骤；这个内存存储区即为 *KV cache*。其大小随上下文长度线性增长，使其成为推理阶段主要的内存瓶颈，也是压缩的主要目标。

TurboQuant\[1 (https://arxiv.org/html/2605.08114#bib.bib1)\] 通过两阶段流水线压缩 KV cache 向量。第一阶段应用结构化随机旋转，随后使用 Beta 最优标量量化器；第二阶段通过 1-bit QJL 草图（sketch）纠正残差中的内积偏差。这两个阶段共同解决正交的问题——方差和偏差，下文将逐一描述。

### 2.1 第一阶段：旋转与 Beta 最优量化

在量化之前，应用随机哈达玛-拉德马赫变换 $\mathbf{x}_{\mathrm{rot}}=\tfrac{1}{\sqrt{d}}\,H\,\mathrm{diag}(\mathbf{s})\,\hat{\mathbf{v}}$（其中 $\hat{\mathbf{v}}=\mathbf{v}/\|\mathbf{v}\|$ 且 $\mathbf{s}\in\{-1,+1\}^d$ 为随机符号向量）。如第 1 节所述，$\mathbf{x}_{\mathrm{rot}}$ 的每个坐标遵循 Beta 分布，这使得可以针对此已知密度设计一次标量 Lloyd-Max 量化器\[3 (https://arxiv.org/html/2605.08114#bib.bib3),4 (https://arxiv.org/html/2605.08114#bib.bib4)\] 并按坐标应用。原始范数 $\|\mathbf{v}\|$ 单独存储。这将量化表示的均方误差（MSE）最小化。

局限性。随机旋转是单个固定映射，而非针对每个向量的新旋转。它将数据移出坐标轴，并在 $\mathbf{s}$ 的随机性上提供期望的均匀性，但它不能使各向异性分布变为各向同性。换句话说，带有适当归一化的哈达玛变换是一个正交变换：应用于球面上的单位向量时，它执行旋转。如果随机化独立于生成向量的过程，则引入它是为了破坏病态的轴对齐模式（例如，向量 $(1,0,...,0)$）。

对于高斯输入，Beta 保证成立；对于低秩或子空间集中的输入，边缘坐标可能看似 Beta 分布，但联合结构仍然是简并的。我们重现了这种病态但现实的嵌入分布，并展示了方程 (1 (https://arxiv.org/html/2605.08114#S1.E1)) 在实际中失效的地方。

### 2.2 第二阶段：QJL 残差校正

设 $\mathbf{r}=\mathbf{x}_{\mathrm{rot}}-\hat{\mathbf{x}}_{\beta}$ 为第一阶段残差。使用新鲜符号向量 $\mathbf{s}'$ 的 1-bit QJL 草图\[2 (https://arxiv.org/html/2605.08114#bib.bib2)\] $\mathbf{q}=\mathrm{sign}\!\left(H(\mathbf{s}'\odot\mathbf{r})\right)$ 产生无偏估计量：

$\hat{\mathbf{r}}=\frac{\sqrt{\pi/2}\,\|\mathbf{r}\|}{d}\,\mathbf{s}'\odot H\mathbf{q},\qquad\mathbb{E}[\langle\mathbf{q},\hat{\mathbf{r}}\rangle]=\langle\mathbf{q},\mathbf{r}\rangle\quad\forall\,\mathbf{q}. \tag{3}$

常数 $\sqrt{\pi/2}$ 校正了符号量化偏差：对于 $z\sim\mathcal{N}(0,\sigma^2)$，$\mathbb{E}[\|z\|]=\sigma\sqrt{2/\pi}$，因此 $\mathrm{sign}(z)$ 低估了幅度，系数为 $\sqrt{2/\pi}$。

$\hat{\mathbf{r}}$ 的无偏性是*无条件的*：它对任何残差向量 $\mathbf{r}$ 都成立，无论第一阶段是否最优。第一阶段的最优性影响的是 $\|\mathbf{r}\|$，从而影响 QJL 估计的方差，其正比于 $\|\mathbf{r}\|^2/d$。

Refer to caption Figure 2: $d=8$, Beta$(3.5,3.5)$——轻微钟形。顶部：坐标直方图开始集中在零附近；low_rank (rank==1) 显示近乎平坦的直方图，偏离 Beta 分布——这是联合结构失效的第一个迹象。底部：在等预算下，MSE 在所有分布中胜出。圆圈表示超出 $1.5\times\mathrm{IQR}$ ($\approx 2.7\sigma$ for a Gaussian) 的异常值。注意，第二阶段应用与第一阶段相同的哈达玛结构，但使用独立于 $\mathbf{s}$ 的新随机化 $\mathbf{s}'$，确保草图与量化误差无关。

### 2.3 研究的三种方案

所有实验均在每个 cache 向量的相等有效位预算 $n$ 下比较三种方案。我们在全文中采用以下简写：

表 1: 本文比较的三种量化方案。KV 是标量基线：无旋转，无草图。KQV 对两个 cache 应用 WHT，然后在 $V$ 上仅使用 QJL 作为残差校正器——这是原始 TurboQuant 论文中描述应用于内积的方案。QKQV 将 QJL 草图对称地扩展到 $K$。三者均以 float32 存储精确向量范数；其余所有位用于量化。Shannon 参考值 $2^{2n}$ 在每种预算 $n$ 下对 KQV 和 QKQV 相同，使得比较公平。

成对对比各自隔离一种机制：KQV 与 KV 对比隔离了 WHT + QJL-on-$V$；QKQV 与 KQV 对比隔离了在 $K$ 上添加 QJL 的效果；QKQV 与 KV 对比隔离了完整的 WHT + QJL 包。

## 3 与 Softmax 注意力的交互

在 $\mathrm{softmax}\!\left(QK^T/\sqrt{d}\right)V$ 中，张量 $K$ 和 $V$ 扮演根本不同的角色：$K$ 仅通过内积 $\langle\mathbf{q},\mathbf{k}_i\rangle$ 进入，而 $V$ 作为线性加权和进入。这种不对称性证明了优化 $K$ 量化以针对内积失真而非 MSE 的合理性。

### 3.1 指数中的琴生偏差

即使内积误差均值为零 $\epsilon_i=\langle\mathbf{q},\hat{\mathbf{k}}_i\rangle-\langle\mathbf{q},\mathbf{k}_i\rangle$，$e^x$ 的凸性也会引入偏差：

$\mathbb{E}\!\left[e^{\langle\mathbf{q},\mathbf{k}_i\rangle+\epsilon_i}\right]=e^{\langle\mathbf{q},\mathbf{k}_i\rangle}\cdot e^{\sigma_i^2/2}, \tag{4}$

对于 $\epsilon_i\sim\mathcal{N}(0,\sigma_i^2)$。如果 $\sigma_i^2$ 在所有键上是*均匀*的，则因子 $e^{\sigma^2/2}$ 会在 softmax 比率中抵消。危险的 regime 是当 $\sigma_i^2$ *与注意力分数相关*时：具有大量化方差的高注意力键会产生无法抵消的不对称失真。

### 3.2 QJL 的作用

第二阶段确保每个键的 $\mathbb{E}[\epsilon_i]=0$，防止注意力 logits 发生系统性漂移。第一阶段最小化 $\sigma_i^2$，减少琴生膨胀。两者共同针对会最严重破坏主导注意力权重的不均匀性。

### 3.3 每实例偏差与 Softmax 抵消论证

经过 WHT 和 Lloyd-Max 量化后，特定键-查询对的内积误差为

$\langle\mathbf{q},\mathbf{k}_i\rangle-\langle\mathbf{q},\hat{\mathbf{k}}_i\rangle=\langle\mathbf{q},\mathbf{r}_i\rangle. \tag{5}$

在随机键的期望下，由 Lloyd-Max 量化器的对称性可知 $\mathbb{E}[\mathbf{r}_i]=0$；但对于*特定*键 $\mathbf{k}_i$，值 $\langle\mathbf{q},\mathbf{r}_i\rangle$ 是一个固定的非零数。这就是 QJL 校正的每实例偏差。

人们可能会问这种偏差在 softmax 中是否无关紧要，因为它出现在分子和分母中：

$\frac{e^{a_i+b_i}}{\sum_j e^{a_j+b_j}},\quad b_i=\langle\mathbf{q},\mathbf{r}_i\rangle. \tag{6}$

如果 $b_i=b$ 在所有键上是常数，它将完全抵消。但 $b_i$ 是*键特定*的：每个键都有其自己的量化残差 $\mathbf{r}_i$，因此偏差不抵消。具有高 $|b_i|$ 的高注意力键会被不公平地重新加权。

QJL 校正了这种键特定偏差。然而，如第 4 节所示，校正的代价是内积估计中方差增加了 $2\pi$ 倍，这通过琴生不等式比其移除的每实例偏差造成了更多的系统性 softmax 失真。

对 $V$ 的影响。值张量*线性*地进入输出：$\sum_i w_i\hat{\mathbf{v}}_i$。没有指数，因此方差不会被放大。$V$ 量化中的每实例偏差直接且加性地传播到输出，使得 QJL 的无偏校正在这里真正有用，而不会带来使其对 $K$ 有害的 $2\pi$ 方差惩罚。

### 3.4 KL 散度作为机制桥梁

设 $p_{\mathrm{ref}}=\mathrm{softmax}(QK^T/\sqrt{d})$ 和 $p_{\mathrm{quant}}=\mathrm{softmax}(Q\hat{K}^T/\sqrt{d})$ 表示固定查询的参考和量化注意力权重分布。KL 散度

$\mathrm{KL}(p_{\mathrm{ref}}\|p_{\mathrm{quant}})=\sum_i p_{\mathrm{ref},i}\log\frac{p_{\mathrm{ref},i}}{p_{\mathrm{quant},i}} \tag{7}$

由构造*仅取决于 $K$*：$V$ 从未进入 softmax，因此 $V$ 量化对该量具有恒为零的影响。这使得 KL 成为 $K$-cache 误差路径的清晰机制探针，与 $V$ 解耦。

$e^x$ 的凸性（琴生不等式）将 $K$ 方向误差超线性地放大到 KL 中。对于得分扰动为 $\epsilon_i=\langle\mathbf{q},\hat{\mathbf{k}}_i-\mathbf{k}_i\rangle$ 的键，即使均值为零的扰动也会放大 softmax 分子：$\mathbb{E}[e^{s_i+\epsilon_i}]=e^{s_i}\cdot e^{\sigma_i^2/2}$。当 $\sigma_i^2$ 在键之间不均匀时（在任何量化方案下的通用情况），膨胀是不对称的，KL 随着 $K$ 方向误差 $\epsilon_K^{\mathrm{dir}}$ 超线性增长。

这种 KL 膨胀具有直接的路由后果，由*top-5 召回率* $\mathrm{topk5}\in[0,1]$ 测量：量化分布中恢复的参考 top-5 注意 token 的比例。路由误差 $1-\mathrm{topk5}$ 增加