TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate

• 作者/机构: Amir Zandieh (Google Research), Majid Daliri (New York University), Majid Hadian (Google DeepMind), Vahab Mirrokni (Google Research)

A1 主要贡献

本文研究高维欧几里得向量的量化问题，旨在在最小化失真（均方误差MSE和内积误差）的同时，将高维向量压缩为低比特宽度的整数。这一问题源于香农的信源编码理论，在当前的大型AI模型训练与部署、向量数据库等领域至关重要。

核心问题与挑战：
现有的向量量化（VQ）算法存在权衡：一些算法不兼容加速器（如GPU），计算缓慢，不适用于如KV缓存量化等实时AI应用；另一些算法虽然速度快，但在不同比特宽度下的失真界限并非最优。特别是在大型语言模型（LLM）中，推理延迟主要受限于内存与计算单元间的通信瓶颈，而KV缓存的大小随模型规模和上下文长度线性增长，成为主要瓶颈。因此，迫切需要一种既高效又能保持模型性能的KV缓存压缩方法。同样，在向量数据库的近邻搜索（NN search）中，高效的向量压缩和精确的内积估计也是核心需求。

研究目标与创新点：
本文旨在设计一种轻量级、可在线应用（无需数据预处理）、且高度兼容加速器的向量量化算法TurboQuant，以解决上述挑战。其核心贡献如下：
1. 提出两种优化的TurboQuant算法：
* MSE优化的TurboQuant：通过随机旋转输入向量，使其坐标服从集中的Beta分布。利用高维空间中不同坐标的近似独立性，对每个坐标独立应用最优的标量量化器（通过求解一维k-means问题得到）。该方法实现了接近最优的MSE失真率。
* 内积优化的TurboQuant：认识到MSE最优量化器在内积估计上存在偏差，本文提出一个两阶段方法。首先应用一个(b-1)比特的MSE量化器，然后对产生的残差向量应用一个1比特的量化Johnson-Lindenstrauss（QJL）变换。这种组合方法产生了一个无偏且低失真的内积估计器。

1. 提供严谨的理论保证：

   • 证明了TurboQuant在MSE和内积失真方面均能达到近乎最优的失真率，与现有方法相比，在比特宽度依赖性上实现了指数级的改进。
   • 利用香农下界和姚氏minimax原理，形式化地证明了任何向量量化器可达到的最佳失真率的信息论下界。理论分析表明，TurboQuant的性能仅与该理论下界相差一个很小的常数因子（约2.7）。

2. 广泛的实验验证：

   • 在KV缓存量化任务中，使用3.5比特/通道实现了与全精度模型完全相同的质量，使用2.5比特/通道也仅有微小的质量下降，实现了超过5倍的压缩。
   • 在近邻搜索任务中，TurboQuant的召回率优于现有的乘积量化（PQ）技术，同时将索引构建时间几乎降为零。

问题定义：
本文的目标是设计一个量化映射 $Q : R^d \to {0, 1}^B$，将d维向量转换为B比特的二进制字符串（其中 $B = b \cdot d$，$b$为比特宽度），以及一个相应的逆映射（反量化） $Q^{-1} : {0, 1}^B \to R^d$。核心目标是最小化两种失真度量：

• 均方误差（MSE）:

  $$D_{\text{mse}} := \underset{Q}{\mathbb{E}} \left[ \| \boldsymbol{x} - Q^{-1} (Q(\boldsymbol{x})) \|_2^2 \right]$$

• 内积误差:

  

其中期望$E_Q$是针对量化器$Q(\cdot)$的随机性计算的。此外，对于内积量化器，要求其估计是无偏的：

本文旨在设计计算高效的量化器 $Q_{mse}$ 和 $Q_{prod}$，以在给定任意比特宽度 $b$ 的情况下，为上述失真度量实现最优界限，并确保 $Q_{prod}$ 提供无偏的内积估计。

A3 背景知识

向量与矩阵表示法。本文使用粗体小写字母（如$\mathbf{x}$和$\mathbf{y}$）表示向量，粗体大写字母（如$\mathbf{M}$）表示矩阵。向量$\mathbf{x}$从索引$i$到$j$（含端点）的分片表示为$\mathbf{x}{i:j}$。矩阵$\mathbf{M}$的第$i$行向量表示为$\mathbf{M}_i$。}$，简称为$\mathbf{M

基本符号定义。使用$S^{d-1}$表示$R^d$中半径为1的超球面。随机变量$x$的微分熵记为$h(x)$。随机变量$x$和$y$之间的互信息记为$I(x; y) = h(x) - h(x|y)$。

超球面上随机点的坐标分布。由于TurboQuant采用随机旋转来处理最坏情况的输入，理解超球面上随机点的统计特性至关重要。以下引理概述了其中一个我们分析和设计所必需的性质：

引理1（超球面上随机点的坐标分布）。对于任意正整数$d$，如果$\mathbf{x} \in S^{d-1}$是一个在单位超球面上均匀分布的随机变量，那么对于任意$j \in [d]$，其坐标$x_j$服从以下（经过缩放/平移的）Beta分布：

$$\boldsymbol{x}_{j} \sim f_{X}(x):=\frac{\Gamma(d / 2)}{\sqrt{\pi} \cdot \Gamma((d-1) / 2)}\left(1-x^{2}\right)^{(d-3) / 2}$$

在高维情况下，该Beta分布收敛于正态分布$f_X(\cdot) \to N(0, 1/d)$。

引理1的证明。$f_X(x)$等于$d-1$维空间中半径为$\sqrt{1 - x^2}$的球的面积与$d$维空间中单位球的体积之比，再按$1/\sqrt{1 - x^2}$进行缩放（根据勾股定理）。因此，

$$f_X(x) = \frac{\frac{2\pi^{(d-1)/2}}{\Gamma((d-1)/2)} \cdot (1-x^2)^{(d-2)/2}}{\frac{2\pi^{d/2}}{\Gamma(d/2)}} \cdot 1/\sqrt{1-x^2} = \frac{\Gamma(d/2)}{\sqrt{\pi} \cdot \Gamma((d-1)/2)} (1-x^2)^{(d-3)/2}.$$

2.1 香农失真下界

香农下界（SLB）的定义。香农下界（SLB）是一个源自香农有损信源编码定理【【索引编号，49，Coding theorems for a discrete source with a fidelity criterion，1959，IRE Nat. Conv. Rec】】的强大工具，它为任何有损压缩方案的最优可达失真率提供了一个普适的下界。我们特别使用一个适用于一般d维信源的均方误差（MSE）失真度量的SLB版本。

引理2（SLB）。设$\mathbf{x} \in R^d$是一个具有任意概率分布$p_X$和有限微分熵$h(\mathbf{x})$的随机向量。定义总比特复杂度为$B \ge 0$的MSE失真率函数$D(B)$为：

其中，下确界取自$\mathbf{x}$和重建随机向量$\mathbf{y} \in R^d$的所有联合分布，使得互信息$I(\mathbf{x}; \mathbf{y})$至多为$B$，且$E[|\mathbf{x} - \mathbf{y}|_2^2]$是期望MSE失真，该期望是关于$\mathbf{x}$和$\mathbf{y}$的联合分布计算的。那么，对于任意比特复杂度$B \ge 0$，以下香农下界成立：

SLB的证明依据。这是一个经典结果，可通过反向高斯测试信道证明（证明见【【索引编号，14，Elements of information theory，1999，John Wiley & Sons】】）。我们的下界结果使用了SLB的一个推论，该推论对应于单位超球面上均匀分布的随机点。我们在以下引理中呈现这一点：

引理3（超球面上随机点的SLB）。设$\mathbf{x} \in S^{d-1}$是一个在单位超球面上均匀分布的随机变量，并根据引理2定义总比特复杂度为$B$的MSE失真率函数$D(B)$。那么，对于任意比特复杂度$B \ge 0$，以下失真下界成立：

$$D(B) \geq 2^{-2 B / d}.$$

引理3的证明。若设$A_d$为超球面$S^{d-1}$的面积，则超球面上均匀分布的熵为$h(\mathbf{x}) = \log_2 A_d$。将此代入引理2的SLB，我们得到$D(B) \ge \frac{d}{2\pi e} \cdot A_d^{2/d} \cdot 2^{-2B/d}$。使用斯特林近似公式计算伽马函数，我们有$A_d = \frac{2\pi^{d/2}}{\Gamma(d/2)} \ge (\frac{2\pi e}{d})^{d/2} \cdot \sqrt{2d\pi} \cdot (1 - O(1/d))$。将此代入从引理2得到的不等式，即可得到所需的下界。

2.2 QJL: 1比特内积量化

QJL在TurboQuant中的作用。如前所述，我们设计了两种VQ算法：一种为最小化MSE而优化，另一种为最小化内积误差而优化。我们证明，MSE最优的量化器不一定能提供无偏的内积估计，尤其在较低比特宽度时表现出显著偏差。我们针对内积量化的解决方案是一个两阶段算法。首先，我们使用比目标比特宽度预算少一个比特的MSE最优量化器，从而最小化残差的L2范数。接下来，我们对残差应用一个无偏且最优的单位比特量化器。对于单位比特内积量化器，我们利用了最近提出的量化Johnson-Lindenstrauss（QJL）算法【【索引编号，62，Qjl: 1-bit quantized jl transform for kv cache quantization with zero overhead，2024，https://arxiv.org/abs/2406.03482】】，它是一个比特宽度为1的最优内积量化器。这里，我们介绍QJL算法及其关键的理论保证 。

定义1 (QJL)。对于任意正整数$d$，QJL映射$Q_{qjl} : R^d \to {-1, +1}^d$定义为：

其中$S \in R^{d \times d}$是一个随机矩阵，其元素是独立同分布地从正态分布$N(0, 1)$中采样的，符号函数（sign）逐元素地应用于其向量输入。逆/反量化映射$Q^{-1}_{qjl} : {-1, +1}^d \to R^d$定义为：

$$Q_{\mathrm{qj1}}^{-1}(\boldsymbol{z}) := \frac{\sqrt{\pi/2}}{d} \cdot \boldsymbol{S}^{\top} \cdot \boldsymbol{z} \quad \text { for any } \boldsymbol{z} \in\{-1,+1\}^{d} .$$

QJL的性能保证。在下一个引理中，我们重申【【索引编号，62，Qjl: 1-bit quantized jl transform for kv cache quantization with zero overhead，2024，https://arxiv.org/abs/2406.03482】】中的结果，这些结果表明QJL是无偏的，并且内积失真很小 ：

引理4 (性能保证: QJL)。设$Q_{qjl}$和$Q^{-1}{qjl}$如定义1所定义。对于任意向量$\mathbf{x} \in S^{d-1}$和任意$\mathbf{y} \in R^d$，我们有：
- 无偏性: $E[\langle\mathbf{y}, Q^{-1}\rangle$.}(Q_{qjl}(\mathbf{x}))\rangle] = \langle\mathbf{y}, \mathbf{x
- 方差界: $\text{Var}[\langle\mathbf{y}, Q^{-1}{qjl}(Q|_2^2$}(\mathbf{x}))\rangle] \le \frac{\pi}{2d} \cdot |\mathbf{y

引理4的证明。无偏性直接源于【【索引编号，62，Qjl: 1-bit quantized jl transform for kv cache quantization with zero overhead，2024，https://arxiv.org/abs/2406.03482】】的引理3.2。为证明方差界， 设$s_1, s_2, \ldots, s_m$表示定义1中随机矩阵$S$的行。我们有：

由于$s_i$是独立同分布的，上式实际上是$d$个独立同分布随机样本$z_i := \sqrt{\pi/2} \cdot s_i^\top \mathbf{y} \cdot \text{sign}(s_i^\top \mathbf{x})$（对于$i \in [d]$）的平均值。现在我们使用【【索引编号，62，Qjl: 1-bit quantized jl transform for kv cache quantization with zero overhead，2024，https://arxiv.org/abs/2406.03482】】中的事实3.4来对单个$z_i$的方差进行上界限定 ：

上式最后一个等号成立是因为$s^\top\mathbf{y}$是一个均值为零、方差为$|\mathbf{y}|_2^2$的高斯随机变量。现在，$d$个独立同分布随机样本$z_1, z_2, \ldots, z_d$的平均值的方差是：

$$\operatorname{Var}\left(\left\langle\boldsymbol{y}, Q_{\mathrm{qjl}}^{-1}\left(Q_{\mathrm{qjl}}(\boldsymbol{x})\right)\right\rangle\right)=\frac{1}{d^{2}} \sum_{i \in[d]} \operatorname{Var}\left(z_{i}\right) \leq \frac{\pi}{2 d} \cdot\|\boldsymbol{y}\|_{2}^{2} .$$

A2 方法细节

我们开发了两种向量量化（VQ）算法，每种都针对一个特定目标。第一种算法旨在最小化原始向量与量化重建后向量之间的均方误差（MSE）。第二种算法则为无偏内积估计进行了优化，解决了MSE最优量化器固有的偏差问题。这些算法将在以下小节中详细介绍。

此外，在3.3节中，我们为任何向量量化器建立了可达到的最佳失真率的信息论下界。这一分析表明，TurboQuant 达到了近乎最优的性能，在所有比特宽度上仅与下界相差一个很小的常数因子。

3.1 MSE最优的TurboQuant

核心思想：随机旋转与标量量化。设$\mathbf{x} \in S^{d-1}$是维度为$d$的单位球面上的一个（最坏情况）向量。我们的目标是每个坐标用$b$比特来量化$\mathbf{x}$，同时最小化公式(1)中定义的重建MSE。我们首先通过将该向量与一个随机旋转矩阵$\Pi \in R^{d \times d}$相乘来进行随机化。我们可以通过对一个具有独立同分布正态分布元素的随机矩阵进行QR分解来生成$\Pi$。

旋转后向量的统计特性。由此产生的旋转向量$\Pi \cdot \mathbf{x}$在单位球面$S^{d-1}$上均匀分布。如引理1所示，$\Pi \cdot \mathbf{x}$的每个坐标都服从一个Beta分布，该分布在高维时收敛于正态分布。此外，在高维情况下，$\Pi \cdot \mathbf{x}$的不同坐标变得几乎不相关【【索引编号，55，High-dimensional probability: An introduction with applications in data science，2018，Cambridge university press】】，这使我们能够独立地对每个坐标应用最优的标量量化器。因此，根据引理1，我们的任务简化为为服从分布$f_X(x) = \frac{\Gamma(d/2)}{\sqrt{\pi} \cdot \Gamma((d-1)/2)} (1 - x^2)^{(d-3)/2}$（其中$x \in [-1, 1]$）的随机变量设计一个标量量化器。

优化问题：一维k-means。在已知概率分布的情况下，最优标量量化问题可以被构建为一个一维的连续k-means问题。具体来说，我们的目标是将区间$[-1, 1]$划分为$2^b$个簇/桶。最优解遵循Voronoi剖分【【索引编号，42，Least squares quantization in pcm，1982，IEEE transactions on information theory】】，这意味着当按排序顺序排列时，区间边界是连续质心之间的中点。因此，用$c_i$表示升序排列的质心，我们可以将标量量化问题表述为以下k-means优化问题：

$$\mathcal{C}(f_X, b) := \min_{-1 \leq c_1 \leq c_2 \leq \ldots \leq c_{2^b} \leq 1} \sum_{i=1}^{2^b} \int_{\frac{c_{i-1}+c_i}{2}}^{\frac{c_i+c_{i+1}}{2}} |x-c_i|^2 \cdot f_X(x) \ dx.$$

最优码本的预计算。请注意，公式(4)中的$C(f_X, b)$表示比特宽度为$b$时的最优MSE代价函数，我们将对这个量进行界定，以证明TurboQuant端到端MSE的上界。公式(4)中的问题可以使用迭代数值方法求解，以达到任何期望的精度。我们为一系列实际相关的比特宽度$b$一次性求解公式(4)，并存储结果以供量化器未来使用。

高维近似与码本示例。例如，在中等到高维度$d$的情况下，分布$f_X(x)$非常接近于正态分布，对于比特宽度$b=1, 2$，最优量化质心分别为$\pm\sqrt{2/\pi d}$和${\pm 0.45/\sqrt{d}, \pm 1.45/\sqrt{d}}$。

算法1: TurboQuant_mse

1: 输入: 维度 d 和比特宽度 b
// 设置TurboQuant_mse的全局参数
2: 生成一个随机旋转矩阵 Π ∈ R^(d×d)
3: 通过寻找最小化公式(4)中MSE代价的质心 c_1, c_2, ... c_{2^b} ∈ [-1, 1] 来构建码本
4: 过程 Quant_mse(x)
5:   y ← Π · x
6:   对每个 j ∈ [d]，idx_j ← arg min_{k∈[2^b]} |y_j − c_k|  {idx_j是b比特整数}
7:   输出: idx
8: 过程 DeQuant_mse(idx)
9:   对每个 j ∈ [d]，ỹ_j ← c_{idx_j}
10:  x̃ ← Π^⊤ · ỹ
11:  输出: x̃

量化与反量化流程。因此，量化器$Q_{mse} : R^d \to {0, 1}^{b \cdot d}$首先计算$\Pi \cdot \mathbf{x}$，然后计算并存储该向量每个坐标的最近质心的索引。反量化映射$Q^{-1}_{mse} : {0, 1}^{b \cdot d} \to R^d$通过检索存储索引对应的质心来重建向量，然后通过与$\Pi^\top$相乘将结果旋转回原始基底。这些过程的伪代码在算法1中给出。

性能保证：MSE失真上界。我们现在准备证明关于TurboQuant_mse的主要定理。

定理1 (性能保证: TurboQuant_mse)。对于任意比特宽度$b \ge 1$和任意向量$\mathbf{x} \in S^{d-1}$，算法1中的Quant_mse(x)过程输出一个索引向量$\mathbf{idx} \in [2^b]^d$。当此索引向量传递给DeQuant_mse(idx)原语时，它会产生一个重建向量$\tilde{\mathbf{x}} \in R^d$，该向量满足以下失真界限：
- 定义为$D_{mse} := E_{\tilde{\mathbf{x}}}[|\mathbf{x} - \tilde{\mathbf{x}}|2^2]$的MSE，对于任意$b \ge 0$，其上界为$D$。} \le \frac{\sqrt{3}\pi}{2} \cdot \frac{1}{4^b
- 对于较小的比特宽度，特别是$b = 1, 2, 3, 4$，MSE表现出更精细的失真值：$D_{mse}$分别约为$0.36, 0.117, 0.03, 0.009$。

定理1的证明。我们首先证明$D_{mse} = d \cdot C(f_X, b)$，其中$C(f_X, b)$是在公式(4)中为标量量化器定义的最优MSE代价。设$\tilde{\mathbf{y}}$如算法1的第9行所定义。由于$\Pi$是旋转矩阵，我们可以写出：$|\mathbf{x} - \tilde{\mathbf{x}}|^2 = |\Pi \cdot \mathbf{x} - \tilde{\mathbf{y}}|^2$。使用算法1第5行的记法$\mathbf{y} = \Pi \cdot \mathbf{x}$并代入$D_{mse}$的定义，我们可以写出：

证明过程解释。上面的第三个等式源于算法1第9行中$\tilde{y}j$的定义，第四行则是因为所有$y_j$都具有相同的分布$y_j \sim f_X(\cdot)$，如引理1所示。最后两行是因为在第6行中，$c$被选为每个坐标$y_j$的最近质心。

代价函数的界定。现在我们必须对最优k-means代价$C(f_X, b)$进行界定。对于中等大小的$d$，$f_X \to N(0, 1/d)$。通过数值求解公式(4)中的优化问题，对于$b=1, 2, 3, 4$，我们得到$C(f_X, b)$分别约为$0.36/d, 0.117/d, 0.03/d, 0.009/d$。对于更大的比特宽度$b > 4$，我们可以应用Panter-Dite【【索引编号，44，Quantization distortion in pulse-count modulation with nonuniform spacing of levels，1951，Proceedings of the IRE】】关于固定速率标量量化器失真的高分辨率公式，得到以下界限：

$$ \mathcal{C}(f_{X}, b) \leq \frac{1}{12} \cdot \left( \int f_{X}(x)^{1/3} \ dx \right)^{3} \cdot \frac{1}{4^{b}} = \frac{\sqrt{3}\pi}{2d} \cdot \frac{1}{4^{b}}. $$

这就完成了证明。

熵编码优化。通过对指向最近码本元素的索引应用熵编码，可以进一步提高TurboQuant的效率。具体来说，每个码字索引在量化向量中出现的概率可以计算为$p_\ell := \int_{(c_{\ell-1}+c_\ell)/2}^{(c_\ell+c_{\ell+1})/2} f_X(x) dx$。对这些索引进行最优编码，可以将平均比特宽度降低到接近分布${p_i}{i \in [2^b]}$的熵。这种无损压缩不影响失真，并能免费提供比特宽度的减少。最显著的减少发生在$b=4$时，此时${p_i}$的熵约为3.8。对最优前缀码的详细计算表明，平均比特宽度可以减少5%。然而，考虑到增益有限，为了保持简单和速度，我们选择不将此技术整合到TurboQuant中。

3.2 内积最优的TurboQuant

问题：MSE优化量化器的偏差。对于像近邻搜索这样的重要应用，拥有一个无偏的内积估计器至关重要。然而，3.1节中介绍的TurboQuant_mse在与查询向量进行内积估计时并不能提供无偏的结果。为了说明这一点，考虑比特宽度$b=1$的情况。在这种情况下，对于足够大的$d$，解决公式(4)中优化问题的最优码本是${\pm\sqrt{2/\pi d}}$。这意味着对于任何$\mathbf{x} \in R^d$，TurboQuant_mse的量化映射是$Q_{mse}(\mathbf{x}) = \text{sign}(\Pi \cdot \mathbf{x})$，而反量化映射是$Q^{-1}{mse}(\mathbf{z}) = \sqrt{2/\pi d} \cdot \Pi^\top \cdot \mathbf{z}$，对于任何$\mathbf{z} \in {-1, +1}^d$。因此，对于足够大的$d$，根据引理4，我们有$E[\langle \mathbf{y}, Q^{-1} \rangle$，这存在一个$2/\pi$的乘法偏差。这个偏差会随着比特宽度$b$的增加而减小，正如我们在4.1节的实验中所展示的。}(Q_{mse}(\mathbf{x})) \rangle] = \sqrt{2/\pi} \cdot \langle \mathbf{y}, \mathbf{x

解决方案：两阶段方法。为了解决这个偏差问题，我们提出了一个结合TurboQuant_mse和QJL实例的解决方案【【索引编号，62，Qjl: 1-bit quantized jl transform for kv cache quantization with zero overhead，2024，https://arxiv.org/abs/2406.03482】】。具体来说，设$Q_{mse}$是对应于比特宽度为$b-1$的TurboQuant_mse的量化映射。对于任 何$\mathbf{x} \in S^{d-1}$，残差向量定义为$\mathbf{r} := \mathbf{x} - Q^{-1}{mse}(Q$，从而得到总比特宽度为$b$的量化，并提供以下无偏的内积估计器：}(\mathbf{x}))$，其L2范数很小，即期望$E[|\mathbf{r}|] = \sqrt{C(f_X, b-1)}$（根据公式(4)）。然后我们可以对这个残差向量应用QJL量化映射$Q_{qjl

$$\left\langle\boldsymbol{y}, Q_{\mathrm{mse}}^{-1}\left(Q_{\mathrm{mse}}(\boldsymbol{x})\right)\right\rangle+\|\boldsymbol{r}\|_{2} \cdot\left\langle\boldsymbol{y}, Q_{\mathrm{qj} 1}^{-1}\left(Q_{\mathrm{qj} 1}(\boldsymbol{r})\right)\right\rangle .$$

形式化定义与伪代码。更形式化地，量化映射$Q_{prod} : S^{d-1} \to [2^{b-1}]^d \times {-1, 1}^d \times R$定义为：

该过程的伪代码在算法2中给出。
算法2: TurboQuant_prod

1: 输入: 维度 d 和比特宽度 b
// 设置TurboQuant_prod的全局参数
2: 从算法1中获取比特宽度为 b-1 的 Q_mse 和 Q⁻¹_mse
3: 从定义1中获取 Q_qjl 和 Q⁻¹_qjl
4: 过程 Quant_prod(x)
5:   idx ← Quant_mse(x) {使用 b-1 比特}
6:   r ← x - DeQuant_mse(idx)
7:   qjl ← Quant_qjl(r)
8:   γ ← ||r||_2
9:   输出: (idx, qjl, γ)
10: 过程 DeQuant_prod(idx, qjl, γ)
11:  x̃_qjl ← γ · Q⁻¹_qjl(qjl)
12:  x̃ ← DeQuant_mse(idx) + x̃_qjl
13:  输出: x̃

性能保证：无偏内积与失真上界。我们在以下定理中证明了TurboQuant_prod的主要结果。

定理2 (性能保证: TurboQuant_prod)。对于任意比特宽度$b \ge 1$和任意向量$\mathbf{x} \in S^{d-1}$，算法2中的Quant_prod(x)过程输出一个索引向量$\mathbf{idx} \in [2^{b-1}]^d$、一个符号向量$\mathbf{qjl} \in {-1, 1}^d$和一个正数$\gamma \ge 0$。当这些向量和标量值传递给DeQuant_prod(idx, qjl, γ)原语时，它产生一个重建向量$\tilde{\mathbf{x}} \in R^d$，该向量对于任何向量$\mathbf{y} \in R^d$满足以下性质：
- 期望内积: $E_{\tilde{\mathbf{x}}}[\langle\mathbf{y}, \tilde{\mathbf{x}}\rangle] = \langle\mathbf{y}, \mathbf{x}\rangle$
- 内积失真: 定义为$D_{prod} := E_{\tilde{\mathbf{x}}}[|\langle\mathbf{y}, \mathbf{x}\rangle - \langle\mathbf{y}, \tilde{\mathbf{x}}\rangle|^2]$，对于任意$b \ge 0$，其上界为$D_{prod} \le \frac{\sqrt{3}\pi^2 \cdot |\mathbf{y}|2^2}{d} \cdot \frac{1}{4^b}$。
- 小比特宽度下的失真: 特别是对于$b=1, 2, 3, 4$，$D$分别约为$1.57/d, 0.56/d, 0.18/d, 0.047/d$。}$表现出更精细的失真值：$D_{prod

定理2的证明（无偏性）。首先，我们计算以内积估计$\langle\mathbf{y}, \tilde{\mathbf{x}}\rangle$在给定$\tilde{\mathbf{x}}_{mse}$的条件下的条件期望，如下所示：

$$\begin{aligned} \begin{aligned} \mathbb{E}\left[\langle\boldsymbol{y}, \tilde{\boldsymbol{x}}\rangle \mid \tilde{\boldsymbol{x}}_{\mathrm{mse}}\right] & =\underset{\tilde{\boldsymbol{x}}_{\mathrm{qj} 1}}{\mathbb{E}}\left[\left\langle\boldsymbol{y}, \tilde{\boldsymbol{x}}_{\mathrm{mse}}+\tilde{\boldsymbol{x}}_{\mathrm{qj1}}\right\rangle \mid \tilde{\boldsymbol{x}}_{\mathrm{mse}}\right] \\ & =\left\langle\boldsymbol{y}, \tilde{\boldsymbol{x}}_{\mathrm{mse}}\right\rangle+\underset{\tilde{\boldsymbol{x}}_{\mathrm{qj} 1}}{\mathbb{E}}\left[\left\langle\boldsymbol{y}, \tilde{\boldsymbol{x}}_{\mathrm{qj} 1}\right\rangle \mid \tilde{\boldsymbol{x}}_{\mathrm{mse}}\right] \\ & =\left\langle\boldsymbol{y}, \tilde{\boldsymbol{x}}_{\mathrm{mse}}\right\rangle+\langle\boldsymbol{y}, \boldsymbol{r}\rangle \\ & =\langle\boldsymbol{y}, \boldsymbol{x}\rangle, \end{aligned} \end{aligned}$$

其中第一个等式源于算法第12行中$\tilde{\mathbf{x}}$的定义。上面的第三个等式源于引理4，最后一行源于第6行中残差向量$\mathbf{r} = \mathbf{x} - \tilde{\mathbf{x}}{mse}$的定义。现在我们可以使用全期望定律计算无条件期望：$E}}}[\langle\mathbf{y}, \tilde{\mathbf{x}}\rangle] = E_{\tilde{\mathbf{x}{mse}}[E[\langle\mathbf{y}, \tilde{\mathbf{x}}\rangle|\tilde{\mathbf{x}}\rangle$，这证明了定理的第一个论断。}]] = E[\langle\mathbf{y}, \mathbf{x}\rangle] = \langle\mathbf{y}, \mathbf{x

定理2的证明（失真界）。在计算失真时，我们对$\tilde{\mathbf{x}}_{mse}$应用相同的条件，然后计算由此产生的条件失真：

$$\begin{aligned} \begin{aligned} \mathbb{E}\left[|\langle\boldsymbol{y}, \boldsymbol{x}\rangle-\langle\boldsymbol{y}, \tilde{\boldsymbol{x}}\rangle|^{2} \mid \tilde{\boldsymbol{x}}_{\mathrm{mse}}\right] & =\mathbb{E}_{\tilde{\boldsymbol{x}}_{\mathrm{qj} 1}}\left[\left|\langle\boldsymbol{y}, \boldsymbol{x}\rangle-\left\langle\boldsymbol{y}, \tilde{\boldsymbol{x}}_{\mathrm{mse}}+\tilde{\boldsymbol{x}}_{\mathrm{qj} 1}\right\rangle\right|^{2} \mid \tilde{\boldsymbol{x}}_{\mathrm{mse}}\right] \\ & =\mathbb{E}_{\tilde{\boldsymbol{x}}_{\mathrm{qj} 1}}\left[\left|\langle\boldsymbol{y}, \boldsymbol{r}\rangle-\left\langle\boldsymbol{y}, \tilde{\boldsymbol{x}}_{\mathrm{qj} 1}\right\rangle\right|^{2} \mid \tilde{\boldsymbol{x}}_{\mathrm{mse}}\right] \\ & =\operatorname{Var}\left(\left\langle\boldsymbol{y}, \tilde{\boldsymbol{x}}_{\mathrm{qj} 1}\right\rangle \mid \tilde{\boldsymbol{x}}_{\mathrm{mse}}\right) \\ & \leq \frac{\pi}{2 d} \cdot\|\boldsymbol{r}\|_{2}^{2}\|\boldsymbol{y}\|_{2}^{2}, \end{aligned} \end{aligned}$$

其中上面的第二个等式源于算法2第6行和第10行中$\mathbf{r}$和$\tilde{\mathbf{x}}{mse}$的定义。上面的第三行是因为根据引理4，$E[\langle\mathbf{y}, \tilde{\mathbf{x}}|$重新缩放的事实。}\rangle] = \langle\mathbf{y}, \mathbf{r}\rangle$。最后一行源于引理4中所示的QJL估计器的方差界，并利用了第11行中$\tilde{\mathbf{x}}_{qjl}$被$\gamma = |\mathbf{r

最终失真界的推导。现在，通过全期望定律以及$\mathbf{r} = \mathbf{x} - \tilde{\mathbf{x}}_{mse}$的事实，我们可以如下界定内积失真：

$$\begin{aligned} \begin{aligned} D_{\text {prod }} & =\underset{\tilde{\boldsymbol{x}}_{\mathrm{mse}}}{\mathbb{E}}\left[\mathbb{E}\left[|\langle\boldsymbol{y}, \boldsymbol{x}\rangle-\langle\boldsymbol{y}, \tilde{\boldsymbol{x}}\rangle|^{2} \mid \tilde{\boldsymbol{x}}_{\mathrm{mse}}\right]\right] \\ & \leq \frac{\pi}{2 d} \cdot\|\boldsymbol{y}\|_{2}^{2} \cdot \mathbb{E}\left[\left\|\boldsymbol{x}-\tilde{\boldsymbol{x}}_{\mathrm{mse}}\right\|_{2}^{2}\right] \\ & =\frac{\pi}{2 d} \cdot\|\boldsymbol{y}\|_{2}^{2} \cdot D_{\mathrm{mse}} \cdot \end{aligned} \end{aligned}$$

通过引用定理1中比特宽度为$b-1$的MSE界限，定理得证。

3.3 下界

方法：Yao的minimax原理与SLB。我们通过证明任何压缩算法可达到的最佳失真的下界，来表明TurboQuant在任意比特宽度下都能达到一个最优的失真率（最多相差一个小的常数因子）。我们的下界证明利用了Yao的minimax原理。该原理允许我们将随机算法在最坏情况确定性输入下的下界与确定性算法在随机输入下的下界联系起来。随后，我们利用2.1节中介绍的香农下界（SLB）为后者推导出可达到的失真率的下界。

理论下界。形式上，我们证明以下定理。

定理3 (最佳可达压缩失真下界)。对于任何比特宽度为$b$的随机量化算法$Q : S^{d-1} \to {0, 1}^{b \cdot d}$和任何重建映射$Q^{-1} : {0, 1}^{b \cdot d} \to R^d$，存在一个硬输入实例$\mathbf{x} \in S^{d-1}$，使得：

$$D_{\mathrm{mse}}(Q):=\mathbb{E}\left[\left\|\boldsymbol{x}-Q^{-1}(Q(\boldsymbol{x}))\right\|_{2}^{2}\right] \geq \frac{1}{4^{b}} .$$

此外，存在一个$\mathbf{y} \in S^{d-1}$，使得：

定理3的证明。根据Yao的minimax原理，最优随机压缩算法在最坏情况输入下的期望MSE（$D_{mse}$）等于最优确定性压缩算法应用于从最大难度随机分布中抽取的输入时的期望MSE。根据定义，后一种情况的MSE由单位超球面上均匀分布的输入可达到的最佳MSE所下界。

证明过程解释。对于比特宽度为$b$、作用于从球面$S^{d-1}$均匀分布的输入的压缩算法，其可达到的最佳MSE在引理3中给出了下界。因此，通过援引引理3，我们得出$D_{mse} \ge \frac{1}{4^b}$。

内积失真下界的推导。此外，从$D_{mse} \ge \frac{1}{4^b}$以及$D_{mse}$的定义，我们得出：

$$\begin{aligned} \begin{aligned} D_{\mathrm{mse}} & =\sum_{j=1}^d \mathbb{E}\left[\left|\boldsymbol{x}_j-\left[Q^{-1}(Q(\boldsymbol{x}))\right]_j\right|^2\right] \\ & =\sum_{j=1}^d \mathbb{E}\left[\left|\left\langle\boldsymbol{e}_j, \boldsymbol{x}\right\rangle-\left\langle\boldsymbol{e}_j, Q^{-1}(Q(\boldsymbol{x}))\right\rangle\right|^2\right] \\ & \geq \frac{1}{4^b} . \end{aligned} \end{aligned}$$

根据鸽巢原理，存在一个索引$j \in [d]$，使得$E[|\langle \mathbf{e}_j, \mathbf{x} \rangle - \langle \mathbf{e}_j, Q^{-1}(Q(\mathbf{x})) \rangle|^2] \ge \frac{1}{d} \cdot \frac{1}{4^b}$，这完成了证明。

与球堆积界的对比。我们注意到，一个可比较的向量量化最坏情况失真下界可以通过“球堆积”论证得出（实际上，由于这是一个更难的问题，其常数会更大）【【索引编号，26，On the structure of vector quantizers，1982，IEEE Transactions on Information Theory】】。然而，定理3为我们的分析提供了一个更稳健和相关的下界。这是因为它建立的是期望失真的下界，而非最坏情况误差，并且与我们在定理1和定理2中提出的上界无缝对接。

A4 实验环境

• 硬件配置: 所有实验均在单块NVIDIA A100 GPU上进行。

• 数据集:

  • DBpedia Entities: 使用OpenAI3嵌入编码为1536维和3072维的向量。用于理论验证和近邻搜索实验。
  • GloVe: 标准的低维（200维）词嵌入数据集，用于近邻搜索实验。
  • LongBench: 包含多种长文本场景的数据集，用于端到端生成任务评估。使用其子集LongBench-E以获得更均衡的上下文长度分布。

• 模型:

  • Llama-3.1-8B-Instruct: 用于"Needle-In-A-Haystack"测试和LongBench评估。
  • Ministral-7B-Instruct: 用于LongBench评估。

A4 实验结果

4.1 理论验证

本节使用DBpedia数据集（10万训练点，1000查询点）对TurboQuant_prod和TurboQuant_mse两种方法的理论结果进行验证。

• 内积估计的偏差与方差：

  • Fig. 1 显示了两种方法在内积估计任务中的误差分布。TurboQuant_prod在所有比特宽度下都保持无偏。而TurboQuant_mse在低比特宽度时存在明显偏差，该偏差随着比特宽度的增加而减小并最终收敛到零。

  

  • Fig. 2 (比特宽度b=2) 表明，TurboQuant_prod的内积误差方差不随原始向量内积值的变化而改变。相反，TurboQuant_mse的偏差与平均内积值正相关，平均内积越大，偏差越大。

  

• 与理论界的比较：

  • Fig. 3 将两种方法的平均内积误差和MSE与理论推导的上下界进行了比较。实验结果与理论预测高度一致。具体而言，对于内积估计，TurboQuant_prod在低比特率下表现更优；而随着比特数增加，TurboQuant_mse的偏差减小，最终在内积估计上取得更优性能。

  

4.2 Needle-In-A-Haystack 测试

该测试评估模型在长文档中检索特定信息（“针”）的能力。

• 实验设置：使用Llama-3.1-8B-Instruct模型，上下文长度从4k到104k。所有压缩方法均在0.25的内存压缩率（即4倍压缩）下进行测试。基线方法包括PolarQuant, SnapKV, PyramidKV, KIVI。

• 实验结果：

  • Fig. 4 表明，具有理论保证的量化方法（如PolarQuant和TurboQuant）优于基于令牌的压缩技术（SnapKV, PyramidKV）和缺乏理论保证的标量量化方法（KIVI）。
  • TurboQuant 表现出色，即使在4倍压缩下，其性能也与全精度基线完全相同（得分0.997），展示了其在长上下文处理中的鲁棒性。

  

4.3 LongBench 上的端到端生成

在LongBench-E数据集上评估各种KV缓存压缩算法。

• 实验设置：比较TurboQuant与KIVI, PolarQuant等基线方法在Llama-3.1-8B-Instruct和Ministral-7B-Instruct模型上的性能。TurboQuant在流式生成过程中也对生成的token进行量化，并采用了2.5比特和3.5比特的混合精度策略（对离群值通道使用更高精度）。

• 实验结果：

  • Table 1 显示，TurboQuant在两个模型上均取得了最高的平均分。
  • 使用3.5比特的TurboQuant在Llama-3.1-8B-Instruct上达到了与全缓存（16比特）几乎相同的平均分（50.06 vs 50.06）。
  • 使用2.5比特的TurboQuant也表现优异（平均分49.44），优于使用更多比特的KIVI和PolarQuant。这表明TurboQuant在至少4.5倍压缩下，仍能保持与未量化模型相当的性能。

  表1: Llama-3.1-8B-Instruct上各种KV缓存压缩方法在LongBench-V1上的结果。
  

4.4 近邻搜索实验

评估TurboQuant在近邻搜索任务中的性能。

• 实验设置：使用DBpedia（1536d, 3072d）和GloVe（200d）数据集。比较TurboQuant与Product Quantization (PQ)和RabitQ。评估指标为Recall@k (1@k)。

• 实验结果：

  • 量化时间 (Table 2)：TurboQuant的量化速度比PQ和RabitQ快几个数量级。例如，在d=3072时，TurboQuant用时0.0021秒，而PQ和RabitQ分别需要494.42秒和3957.19秒。

  表2: 不同方法在不同维度下使用4比特量化的量化时间（秒）。
  

  • 召回率 (Fig. 5)：在所有数据集和比特宽度设置下，TurboQuant的召回率始终优于PQ和RabitQ。例如，在OpenAI3 (d=1536)数据集上，4比特的TurboQuant在Top-k=1时召回率接近1.0，显著高于PQ和RabitQ。这证明了TurboQuant在高维量化搜索任务中的鲁棒性和高效性。