内积感知量化：可证明快速、准确且自适应的算法

# 内积感知量化：可证明快速、准确且自适应的算法

**来源：** https://arxiv.org/html/2606.00289

**Nathan White**  
宾夕法尼亚大学  
[email protected]

**Krish Singal¹**  
宾夕法尼亚大学  
[email protected]

###### 摘要

量化是一种基础工具，用于压缩数据集、神经网络权重以及在多种计算任务中减少内存使用。许多向量量化的下游应用需要与任意输入执行内积运算。这推动了*内积感知*量化方案的研究，这种方案能够近似保持与未见向量之间的内积——其目标不仅仅是单纯最小化均方误差。在本工作中，我们定义了能够捕捉自然需求的目标函数，并开发了*自适应*且*无偏*的量化方法，这些方法能够近似保持与最坏情况和平均情况输入的内积。对这些目标函数的分析展示了其与广为人知的“自适应随机量化”（ASQ）概念之间的紧密联系。我们针对所提出的目标函数，开发了可证明快速的精确算法和近似算法。我们的理论结果启发了高效的实用算法，这些算法在各种工作负载分布下均表现良好。同时，它们还为标准ASQ问题带来了实用算法，在保持质量的同时，相比此前最先进的方法速度快了2-10倍。这些理论和实验成果有助于使自适应量化技术在实际环境中更加高效和易于处理。

## 1 引言

向量量化是一种极其重要的工具，在众多计算和机器学习应用的空间与运行时间优化中扮演着核心角色。例如，量化被用于向量搜索的数据集压缩 [GHKS 13; GGX+ 25]、大模型权重和键值缓存的压缩 [SZY+ 23]，以及量化感知训练 [MNA+ 17] 和训练后量化 [FAHA 23; JLPK 23]。

形式化地，我们如下定义向量量化。考虑 \(w \in \mathbb{R}^d\)，并设 \(Q \subset \mathbb{R}\) 为一组*量化值*，满足 \(|Q| = s \ll d\)。那么，向量 \(w\) 可以通过某种*舍入分布* \(\mathcal{D}\) 被*量化*，其中 \(\text{Supp}(\mathcal{D}) \subseteq Q^d\)。在本工作中，我们考虑**无偏**和**自适应**的量化方案。若 \(\mathop{\mathbf{E}}\limits_{\widehat{\boldsymbol{w}} \sim \mathcal{D}} [\widehat{\boldsymbol{w}}] = w\)，则该量化方案是无偏的¹；而自适应方案则指 \(Q\) 和 \(\mathcal{D}\) 都可能依赖于 \(w\)。在最一般的情况下，我们希望能够根据某个合适的目标函数，联合优化元组 \((Q, \mathcal{D})\)。

¹ 我们用粗体字表示随机变量。

此前的许多量化工作考虑的是非自适应/弱自适应或有偏的方案；例如，QSGD [AGL+ 17] 和 NUQ SGD [RKFM+ 21] 等方法仅使用向量的浅层属性，如范数或长宽比。其他工作 [FHH+ 20; FTM+ 20] 则针对来自常见、预设分布的向量进行定制，但并未适应具体的输入向量。例如，有证据表明，机器学习应用中重要且相关的数据遵循对数正态 [CBUS+ 20]、正态 [BNS 19] 或次韦布尔分布 [VAM 18]。另一方面，像“四舍五入到最近”（RTN）这样有偏且非自适应的技术，在训练后量化中已被证明不如自适应方法 [NAVB+ 20]。

在[BBBIMV 24; ZLK+ 17; FTM+ 20]等工作中，有一系列关于*自适应随机量化*（ASQ）的研究探索了自适应且无偏的量化方案，并指出相对于其他类型的量化有显著改进（参见[BBBIMV 24]的图1）。在ASQ中，目标是构建一个量化集，以最小化均方误差 \(\operatorname{\mathsf{MSE}}(\widehat{\boldsymbol{w}}, w) \triangleq \mathop{\mathbf{E}}\left[~\mathinner{\!\left\lVert\widehat{\boldsymbol{w}}-w\right\rVert}_2^2\right]\)。

从中采样 \(\widehat{\boldsymbol{w}}\) 的舍入分布简单而自然：对每个 \(w_i\) 独立地以唯一无偏的方式舍入到最接近 \(w_i\) 的较大或较小的量化点。我们将这种舍入分布称为**标准随机量化**，并将得到的分布记作 \(\operatorname{\smash{\mathcal{D}}^{\mathsf{SSQ}}}(w, Q)\)²；我们在1.1节给出了正式定义。标准随机量化是一种自然的舍入分布选择，因为对于每个坐标 \(i\)，它在所有无偏分布中最小化了 \(\mathop{\mathbf{Var}}[\widehat{\boldsymbol{w}}_i]\)。³

² 当上下文明确时，我们直接简写为 \(\operatorname{\smash{\mathcal{D}}^{\mathsf{SSQ}}}\)。
³ 参见定理3.2以证明该事实。

虽然MSE是一个自然的目标函数，但许多量化的下游应用不仅仅依赖于低的 \(\ell_2\) 距离。例如，许多现代应用会将量化后的向量与其他向量（可能是任意向量）进行内积运算（这在神经网络权重量化和向量搜索中都是如此）。因此，量化在这些应用中的成功很大程度上取决于其保持内积的能力。在本工作中，我们引入并研究了几种*内积感知*量化目标，旨在近似保持内积。

由于我们考虑无偏量化方案，即 \(\mathop{\mathbf{E}}[\widehat{\boldsymbol{w}}] = w\)，根据期望的线性性，对于任意向量 \(x \in \mathbb{R}^d\)，有 \(\mathop{\mathbf{E}}[\langle \widehat{\boldsymbol{w}}, x \rangle] = \langle w, x \rangle\)。因此，最小化 \(\langle \widehat{\boldsymbol{w}}, x \rangle\) 的期望（平方）误差就简化为最小化 \(\mathop{\mathbf{Var}}[\langle \widehat{\boldsymbol{w}}, x \rangle]\)。在我们的第一个目标中，我们针对最坏情况下的 \(x\) 选择来最小化方差。

###### 定义 1.1（最大方向方差 \(\operatorname{\mathsf{MDV}}\)）。
对于向量 \(w \in \mathbb{R}^d\) 和目标量化集大小 \(s \in \mathbb{N}\)，我们定义

\[
\operatorname{\mathsf{MDV}}(w,s) \triangleq \min_{Q \subset \mathbb{R} : |Q| \leq s} \max_{x \in \mathbb{R}^d : \|x\|_2 \leq 1} \mathop{\mathbf{Var}}_{\widehat{\boldsymbol{w}} \sim \operatorname{\smash{\mathcal{D}}^{\mathsf{SSQ}}}(w,Q)} [\langle \widehat{\boldsymbol{w}}, x \rangle]
\]

在实践中，输入向量 \(x\) 通常并非最坏情况，而是来自某个已知（或估计的）分布。因此，我们还考虑针对某个（已知的）向量 \(x\) 分布来优化，以最小化平均方差。

###### 定义 1.2（平均方向方差 \(\operatorname{\mathsf{ADV}}\)）。
对于向量 \(w \in \mathbb{R}^d\)、目标量化集大小 \(s \in \mathbb{N}\) 以及 \(\mathbb{R}^d\) 上的输入分布 \(\mathcal{X}\)，定义

\[
\operatorname{\mathsf{ADV}}_{\mathcal{X}}(w,s) \triangleq \min_{Q \subset \mathbb{R} : |Q| \leq s} \mathop{\mathbf{E}}_{\boldsymbol{x} \sim \mathcal{X}} \left[ \mathop{\mathbf{Var}}_{\widehat{\boldsymbol{w}} \sim \operatorname{\smash{\mathcal{D}}^{\mathsf{SSQ}}}(w,Q)} \left[ \langle \widehat{\boldsymbol{w}}, \boldsymbol{x} \rangle \right] \right] \tag{1}
\]

我们注意到，这些目标函数将标准随机量化固定为舍入分布，并优化量化集。虽然这是一个自然且标准的分布选择，但人们可能想知道是否有更好的舍入分布选择能够获得更低的内积方差。我们证明，这可能不是一个值得探索的方向：对于 \(\operatorname{\mathsf{MDV}}\)，我们证明标准随机量化实际上是最优的舍入分布（定理3.2）；对于 \(\operatorname{\mathsf{ADV}}\)，我们证明计算最优分布是NP难的（定理3.2）。

针对这两个目标，我们设计了可证明快速的算法，该算法能够在用户指定的 \(\varepsilon > 0\) 下，获得目标最优成本的一个 \((1+\varepsilon)\) 因子内的解。参见第3节对我们算法结果的讨论。

我们还实现了这些算法的实用版本进行评估，并展示了它们优于先前的自适应量化算法。特别地，由于 \(\operatorname{\mathsf{ADV}}\) 与传统ASQ之间存在紧密联系（见1.1节），我们采用我们的技术开发了显著优于当前最先进QUIVER算法 [BBBIMV 24] 的算法；参见图1。

**图1:** 比较此前传统ASQ的最先进算法与我们更快的算法（带热身二分的Wilber算法）在从 \(\text{LogNormal}(0,1)\) 采样的向量上的运行时间。左侧固定 \(s=64\) 并变化向量大小 \(d\)，右侧固定 \(d=500,000\) 并变化 \(s\)。阴影区域显示了在采样向量上运行时间的10-90%范围。图3类似，但额外比较了更多方法。

由于精确算法是[BBBIMV 24]中近似QUIVER算法的核心子程序，我们立即也可以对该算法实现加速。更一般地，任何将QUIVER作为子程序的下游任务（例如，[BBBIMV 25]）都可以通过我们的实现得到加速。参见第4节我们快速算法的具体细节。

由于自适应方法存在固有的权衡（更高的量化质量以更慢的预处理时间为代价），我们工作的一个主要影响是利用理论洞见推动自适应技术在实际环境中走向高效。

### 1.1 预备知识

##### 标准随机量化。
对于给定向量 \(w \in \mathbb{R}^d\) 和量化集 \(Q\)，定义**标准随机量化**舍入分布 \(\operatorname{\smash{\mathcal{D}}^{\mathsf{SSQ}}}(w,Q)\)，该分布独立地舍入每个 \(w_i\)，使得对于 \(\widehat{\boldsymbol{w}} \sim \operatorname{\smash{\mathcal{D}}^{\mathsf{SSQ}}}(w,Q)\)，有

\[
\widehat{\boldsymbol{w}}_i = 
\begin{cases}
w_i^\uparrow & \text{以概率 } (w_i^\uparrow - w_i)/(w_i^\uparrow - w_i^\downarrow) \\
w_i^\downarrow & \text{以概率 } (w_i - w_i^\downarrow)/(w_i^\uparrow - w_i^\downarrow)
\end{cases}
\]

其中，\(w_i^\uparrow := \min\{q \in Q : q \geq w_i\}\) 和 \(w_i^\downarrow(Q) := \max\{q \in Q : q \leq w_i\}\) 分别是 \(Q\) 中从上方和下方最接近 \(w_i\) 的值。注意，\(\mathop{\mathbf{E}}[\widehat{\boldsymbol{w}}] = w\)，且对于所有 \(i \in [d]\)，有 \(\mathop{\mathbf{Var}}[\widehat{\boldsymbol{w}}_i] = (w_i^\uparrow - w_i)(w_i - w_i^\downarrow)\)。为简化符号，我们通常记 \(\smash{\operatorname{\mathsf{VarSSQ}}}(w_i, Q) := \mathop{\mathbf{Var}}[\widehat{\boldsymbol{w}}_i] = (w_i^\uparrow - w_i)(w_i - w_i^\downarrow)\)。

##### \(\operatorname{\mathsf{MDV}}\) 与 \(\operatorname{\mathsf{ADV}}\) 的性质。
下面，我们针对这两个目标函数的结构进行一些简单但有用的观察。我们首先为固定量化集下的每个目标函数引入简写符号。我们记

\[
\operatorname{\mathsf{MDV}}(w,Q) \triangleq \max_{x \in \mathbb{R}^d : \|x\|_2 \leq 1} \mathop{\mathbf{Var}}_{\widehat{\boldsymbol{w}} \sim \operatorname{\smash{\mathcal{D}}^{\mathsf{SSQ}}}(w,Q)} [\langle \widehat{\boldsymbol{w}}, x \rangle]
\]
和
\[
\operatorname{\mathsf{ADV}}_{\mathcal{X}}(w,Q) \triangleq \mathop{\mathbf{E}}_{\boldsymbol{x} \sim \mathcal{X}} \left[ \mathop{\mathbf{Var}}_{\widehat{\boldsymbol{w}} \sim \operatorname{\smash{\mathcal{D}}^{\mathsf{SSQ}}}(w,Q)} [\langle \widehat{\boldsymbol{w}}, \boldsymbol{x} \rangle] \right]
\]

\(\operatorname{\mathsf{MDV}}\) 具有一个简单的组合结构；具体而言，

\[
\begin{aligned}
\operatorname{\mathsf{MDV}}(w,Q) &= \max_{x \in \mathbb{R}^d : \|x\|_2 \leq 1} \left[ \mathop{\mathbf{Var}}_{\widehat{\boldsymbol{w}} \sim \operatorname{\smash{\mathcal{D}}^{\mathsf{SSQ}}}(w,Q)} [\langle \widehat{\boldsymbol{w}}, x \rangle] \right] \\
&= \max_{x \in \mathbb{R}^d : \|x\|_2 \leq 1} \left[ \sum_{i=1}^d x_i^2 \mathop{\mathbf{Var}}_{\widehat{\boldsymbol{w}} \sim \operatorname{\smash{\mathcal{D}}^{\mathsf{SSQ}}}} \right]
\end{aligned}
\]