ICLR2026AlgoTurboQuant: Online Vector Quantization with Near-optimal Distortion Rate

接近最优失真率的在线向量量化方法。

个人评价:这篇文章主要是针对大模型的KV-Cache的压缩,虽然是同一作者不同方法的浓缩(PolarQuant+QJL),但是补充了在结合方法下的量化误差上下界。目前实验停留在纯语言模型阶段,也许可以拓广到多模态阶段。

总结:本文主要介绍了一种针对高维向量量化的创新方法,旨在通过大幅度压缩数据规模来优化AI模型推理,KV Cache管理以及向量数据库检索的效率。其核心在于结合了随机旋转技术和最优标量量化器,能在极低的比特位宽下实现接近理论极限的MSE。针对内积检索中的偏置问题,作者设计了一个两阶段架构,利用 1-比特 QJL 变换补偿余数,从而确保了内积估算的无偏性。实验数据表明,该算法在 Llama-3.1 等大语言模型的长文本测试中,仅需 2.5 至 3.5 比特即可保持与全精度近乎一致的性能。此外,相较于传统的乘积量化 (PQ) 技术,TurboQuant在保持高召回率的同时,将索引构建时间降低至接近于零,展现出卓越的加速器友好性。

简单来说,进行向量量化(VQ)的目的是最小化下列两个误差:

$$ D_{MSE} = \mathbb{E}_{Q}[||x-Q^{-1}(Q(x))||_2^2]\tag{1} $$$$ D_{prod} = \mathbb{E}_Q[|\langle y,x\rangle-\langle y,Q^{-1}(Q(x))\rangle|^2]\tag{2} $$

此外,对于内积量化,在大模型推理中我们更希望向量的内积是无偏的,即满足:

$$ \mathbb{E}_Q[\langle y,Q^{-1}(Q(x))\rangle]=\langle y,x\rangle\tag{3} $$

而这两个优化目标是难以兼顾的,因此在VQ中通常会设计两个Quantizer,分别是$Q_{MSE}$,$Q_{prod}$。对应到KV-Cache量化场景下就是对K向量用$Q_{prod}$,对V向量用$Q_{MSE}$。

对于$Q_{MSE}$,我们的目标就是最小化公式(1)。在此之前,我们有如下假设:

待量化向量满足$||x||^2 = 1$,即$x\in \mathbb{S}^{d-1}$,即分布在d维球面上。若不满足这个条件,在实际中可以通过存储L2范数进行Scale使向量满足条件。

我们有如下引理(Lemma 1):

若$x\in \mathbb{S}^{d-1}$,是在单位超球面上均匀分布的随机变量,那么对任意$j\in [d]$,坐标$x_j$服从(缩放/平移后的)Beta型分布:

$$ x_j ~ f_X(x)=\frac{\Gamma(\frac{d}{2})}{\sqrt{\pi}\Gamma(\frac{d-1}{2})}(1-x^2)^{\frac{d-3}{2}},\quad x\in[-1,1] $$

在高维情况下, 该分布收敛到正态分布:

$$ f_X(.)\to N(0,\frac{1}{d}) $$

证明略.

High-Level层面上可以理解为固定球上一点的一个坐标相当于用一个平面去截这个高维球面,那么此时其余坐标构成的截面是一个维度为d-2,半径为$\sqrt{1-x^2}$的球面。(可以想象一下三维球面被平面截得到圆),那么这个分布自然就是中间多(x=0),两边少($x=\pm 1$).

在这个引理的支持下,我们对原始的向量x乘以一个随机的旋转矩阵$\Pi $(这里相当于做了一个极坐标变换),使其成为在单位超球面上均匀分布的随机变量$z = \Pi x$。那么此时根据Lemma 1,z的每个坐标都可以认为符合上述Beta型分布,且在高维情况下收敛为正态分布。此外,在高维下,z不同坐标之间会变得近似独立,因此我们可以对每个坐标独立地应用最优标量量化器。于是我们的问题转变为:

为服从如下分布的随机变量设计一个标量量化器。

$$ x_j ~ f_X(x)=\frac{\Gamma(\frac{d}{2})}{\sqrt{\pi}\Gamma(\frac{d-1}{2})}(1-x^2)^{\frac{d-3}{2}},\quad x\in[-1,1] $$

在随机变量分布给定的情况下的最优标量量化问题可以表述为一个一维连续K-means问题。更具体而言,我们希望把区间[-1,1]划分为$2^b$个簇。最优解需要满足:当所有质心按照升序排序时,区间边界应当是相邻质心的中间。因此,若记这些升序排序的质心为$c_i$那么,该标量量化问题可以描述为如下k-means优化问题:

$$ C(f_x,b) = \min_{-1 \leq c_1\leq c_2\leq \dots \leq c_{2^b}\leq 1}\sum_{1}^{2^b}|x-c_i|^2f_x(x)dx\tag{4} $$

该问题可以通过迭代数值方法进行求解(Lloyd-Max量化器,本质和K-means相似,可以看作一维的K-means)。此外,我们只需要针对一组实际有效的bit-width b离线求解一次,然后把结果存储下来,供量化器之后重复使用。

至此,$Q_{MSE}$的做法很明确:先计算$z=\Pi x$,然后对z的每个坐标找到最近的质心,并存储该质心的索引。对应的反量化流程则通过读取这些索引对应的质心来重建旋转后的向量,再乘以$\Pi^\top$从而得到原始向量。

论文中还给出了该量化器损失的上界,理解起来并不困难,这里不做过多赘述。

之所以$Q_{MSE}$需要与$Q_{prod}$不能统一,是因为前者不满足内积无偏性的特性,即式(3)。(论文中给出了证明)

为了保证$Q_{prod}$的内积无偏性,作者提出了将$Q_{MSE}$与QJL相结合的方案。具体而言,设$Q_{MSE}$是对应于位宽 b−1 的$ Q_{mse} $的量化映射。对于任意$x\in \mathbb{S}^{d-1}$,我们定义残差向量:

$$ r := x - Q_{mse}^{-1}(Q_{mse}(x)) $$

其L2范数很小,即在期望意义下(见式(4))

$$ \mathbb{E}[||r||]=\sqrt{C(f_X,b-1)} $$

随后,我们可将QJL 的量化映射 $Q_{QJL}$ 应用于该残差向量,从而使总体位宽达到 b,并得到如下无偏内积估计器:

$$ \langle y,Q^{-1}_{MSE}(Q_{MSE}(x))\rangle + ||r||^2\cdot\langle y,Q^{-1}_{qjl}(Q_{qjl}(r))\rangle $$

更形式化的来说,我们可以定义:

$$ Q_{prod}(x) = [Q_{MSE}(x),Q_{qjl}(x-Q^{-1}_{MSE}(Q_{MSE}(x))),||x-Q^{-1}_{MSE}(Q_{MSE}(x))||_2] $$

论文还对该方法的误差下界进行了估计,具体的参考原文。