熵约束自适应随机量化

# 熵约束自适应随机量化  
来源：https://arxiv.org/html/2608.18147  
Ran Ben Basat1,2 Yaniv Ben-Itzhak2 Michael Mitzenmacher3 Shay Vargaftik2  
1伦敦大学学院 2博通旗下VMware研究院 3哈佛大学  

###### 摘要  
自适应随机量化（ASQ）是一种近期提出的量化方法，能在保持无偏性的同时，针对给定输入优化均方误差（MSE）。该方法旨在缓解现代数据和机器学习工作负载中的通信与内存瓶颈，包括模型、梯度和KV缓存压缩，以及近邻搜索。实际系统可以进一步使用无损熵编码器压缩量化数据。然而，现有无偏方法（包括ASQ）在选择量化值时未考虑后续编码阶段，导致精度损失。我们提出了熵约束自适应随机量化（ECASQ）问题，在熵预算和无偏性约束下联合选择自适应量化值以最小化MSE。我们给出了最优动态规划算法，时间复杂度为O(sd²)，空间复杂度为O(d²)（针对长度为d的向量和最多s个量化值），以及一个GPU友好的近似动态规划算法，时间复杂度O(sd²)、空间复杂度O(d)。该近似算法保证解的MSE不超过每个条目使用比最优解少1比特熵的最优解。我们还提供了近似解的迭代精修过程，实验表明该过程在保留相比最优解求解器显著速度优势的同时，能获得近最优结果。  

## 1 引言  
机器学习（ML）模型的训练和部署常受限于内存容量、计算资源和网络带宽。为缓解这些限制，量化已成为一项基础技术。通过将高精度值映射到低位宽表示，量化促进了多种ML工作负载的有效压缩策略，包括梯度压缩[26 (https://arxiv.org/html/2608.18147#bib.bib159),23 (https://arxiv.org/html/2608.18147#bib.bib151),31 (https://arxiv.org/html/2608.18147#bib.bib110),19 (https://arxiv.org/html/2608.18147#bib.bib36)]和通信高效的分布式训练[24 (https://arxiv.org/html/2608.18147#bib.bib44),11 (https://arxiv.org/html/2608.18147#bib.bib35),18 (https://arxiv.org/html/2608.18147#bib.bib19),35 (https://arxiv.org/html/2608.18147#bib.bib1)]、训练后量化[15 (https://arxiv.org/html/2608.18147#bib.bib50),22 (https://arxiv.org/html/2608.18147#bib.bib49)]以及KV缓存占用减少[29 (https://arxiv.org/html/2608.18147#bib.bib204),25 (https://arxiv.org/html/2608.18147#bib.bib28)]，包括分解推理[34 (https://arxiv.org/html/2608.18147#bib.bib2)]。此场景中尤其需要三个特性：(1) 无偏性，(2) 对输入的自适应性，(3) 针对熵编码输出大小的优化。如后文所述，现有方法无法同时提供这三者。  

设输入X=⟨x₁,...,x_d⟩∈ℝ^d，量化值集Q={q₁,...,q_m}⊂ℝ，其中q₁<...<q_m。对于每个x∈X，其最近邻量化值定义为a_x=max{q∈Q | q≤x}和b_x=min{q∈Q | q≥x}。随机量化（SQ）以无偏方式将x映射为x̂∈{a_x, b_x}，使得E[x̂]=x。更具体地说，当a_x=b_x时x̂=a_x；否则以概率(x-a_x)/(b_x-a_x)取值b_x，否则取a_x。由于Var[x̂]=Pr[x̂=b_x]·(b_x-x)² + Pr[x̂=a_x]·(a_x-x)²=(b_x-x)(x-a_x)，量化的均方误差（MSE）为MSE(Q,X)=∑_{x∈X}(b_x-x)(x-a_x)。  

另一种误差度量是向量归一化MSE（vNMSE）[32 (https://arxiv.org/html/2608.18147#bib.bib77),7 (https://arxiv.org/html/2608.18147#bib.bib21),5 (https://arxiv.org/html/2608.18147#bib.bib20)]，定义为vNMSE(Q,X)=MSE(Q,X)/‖X‖²。vNMSE通常是更便捷的度量，因其具有尺度不变性（即对任意c≠0，vNMSE(Q,X)=vNMSE(c·Q,c·X)），便于比较不同输入维度的方法。  

自适应随机量化（ASQ）问题是指给定X∈ℝ^d和s≥2，输出大小为s的量化值集Q，以最小化该特定X的MSE(Q,X)（或等价地vNMSE(Q,X)）。最优ASQ方法如[33 (https://arxiv.org/html/2608.18147#bib.bib15),9 (https://arxiv.org/html/2608.18147#bib.bib37)]通过求解以下动态规划（DP）实现：定义dp_ASQ(i,j)为使用i个量化值对X中最小的j个条目进行随机量化所能达到的最小MSE。设C[k,j]=∑_{ℓ=k}^{j}(x_j-x_ℓ)(x_ℓ-x_k)为假设将区间[x_k,x_j]内所有条目量化在x_k和x_j之间时的方差和。则dp_ASQ(i,j)可通过递推求解：  
dp_ASQ(i,j) =  
{ min_{k∈{1,...,j}} dp_ASQ(i-1,k) + C[k,j]  若i>2  
{ C[1,j]                              其他情况  
(1)  
直观上，这假设x_j是一个量化值，并寻找下一个值x_k的最优位置。此动态规划关键依赖于存在最优解Q*⊆X，即所有量化值均为输入条目[33 (https://arxiv.org/html/2608.18147#bib.bib15),9 (https://arxiv.org/html/2608.18147#bib.bib37)]。该问题可在O(s·d)时间和空间内最优求解[9 (https://arxiv.org/html/2608.18147#bib.bib37)]。对于整数n，记[n]={1,...,n}。  

## 2 预备知识  
给定输入X∈ℝ^d和量化值集Q⊂ℝ，随机量化（SQ）返回量化向量X̂=⟨x̂⟩_{x∈X}，使得每个x∈X在其包含的量化值Q之间被无偏量化。与先前工作一致（例如[9 (https://arxiv.org/html/2608.18147#bib.bib37),33 (https://arxiv.org/html/2608.18147#bib.bib15),14 (https://arxiv.org/html/2608.18147#bib.bib54)]），我们假设X已*排序*（否则，我们对X的副本排序并保持原始索引以按初始顺序量化条目）。进一步假设X的条目互异（否则，可使用其加权直方图）。即，记a_x=max{q∈Q | q≤x}和b_x=min{q∈Q | q≥x}为Q中的相关条目，若a_x=b_x则x̂=a_x；否则x̂以概率(x-a_x)/(b_x-a_x)取值b_x，否则取a_x。  

由于Var[x̂]=Pr[x̂=b_x]·(b_x-x)² + Pr[x̂=a_x]·(a_x-x)²=(b_x-x)(x-a_x)，量化均方误差（MSE）为MSE(Q,X)=∑_{x∈X}(b_x-x)(x-a_x)。另一种误差度量是向量归一化MSE（vNMSE）[32 (https://arxiv.org/html/2608.18147#bib.bib77),7 (https://arxiv.org/html/2608.18147#bib.bib21),5 (https://arxiv.org/html/2608.18147#bib.bib20)]，定义为vNMSE(Q,X)=MSE(Q,X)/‖X‖²。vNMSE通常是更便捷的度量，因其具有尺度不变性（即对任意c≠0，vNMSE(Q,X)=vNMSE(c·Q,c·X)），便于比较不同输入维度的方法。  

自适应随机量化（ASQ）问题是指给定X∈ℝ^d和s≥2，输出大小为s的量化值集Q，以最小化该特定X的MSE(Q,X)（或等价地vNMSE(Q,X)）。最优ASQ方法如[33 (https://arxiv.org/html/2608.18147#bib.bib15),9 (https://arxiv.org/html/2608.18147#bib.bib37)]通过求解以下动态规划（DP）实现：定义dp_ASQ(i,j)为使用i个量化值对X中最小的j个条目进行随机量化所能达到的最小MSE。设C[k,j]=∑_{ℓ=k}^{j}(x_j-x_ℓ)(x_ℓ-x_k)为假设将区间[x_k,x_j]内所有条目量化在x_k和x_j之间时的方差和。则dp_ASQ(i,j)可通过递推求解：  
dp_ASQ(i,j) =  
{ min_{k∈{1,...,j}} dp_ASQ(i-1,k) + C[k,j]  若i>2  
{ C[1,j]                              其他情况  
(1)  
直观上，这假设x_j是一个量化值，并寻找下一个值x_k的最优位置。此动态规划关键依赖于存在最优解Q*⊆X，即所有量化值均为输入条目[33 (https://arxiv.org/html/2608.18147#bib.bib15),9 (https://arxiv.org/html/2608.18147#bib.bib37)]。该问题可在O(s·d)时间和空间内最优求解[9 (https://arxiv.org/html/2608.18147#bib.bib37)]。对于整数n，记[n]={1,...,n}。  

## 3 熵约束ASQ  
为优化将在熵编码步骤之后使用的量化值集Q，我们将ASQ问题扩展如下。熵约束自适应随机量化（ECASQ）问题的输入现在包括（除限制‖Q‖的参数s外）熵约束b∈ℝ⁺和允许的量化值集P⊂ℝ。我们添加约束H(X̂)≤b和Q⊆P。这里H(X̂)是量化条目随机向量的熵。具体而言，对每个q∈Q，设f_q=∑_{x∈X}Pr[x̂=q]表示q的期望频率。则平均条目熵定义为H(X̂)=∑_{q∈Q} - (f_q/d)·log₂(f_q/d)。  

例如，考虑X=⟨0,1,10⟩和Q={0,10}；则0必被量化为0，10必被量化为10，而中间条目以概率9/10被量化为0，否则量化为10。因此熵为H(X̂)= - (19/30)log₂(19/30) - (11/30)log₂(11/30) ≈ 0.948比特/条目。引入P是必要的，因为与ASQ情况不同，最优解可能不是输入的子集，并且允许我们对编码施加实际约束（例如，将P定义为BF16或FP32中可表示的所有值）。注意ECASQ推广了ASQ，后者是b=log₂s且P=X的特例。虽然可以通过设置s=|P|移除单独的基数参数，但这样做忽略了表示码本所需的空间。解码器必须同时知道选定的字典Q和熵编码元数据；因此对s设限可提供与实现无关的直接开销上界。若存储的重构值及其熵模型描述符分别使用w和c比特，则其元数据成本为(w+c)|Q|/d比特/条目。由于Q有序，电平的位置隐式提供了其符号到值的关联，因此无需额外排列。我们的实验使用w=c=16。  

为求解ECASQ，我们借鉴先前熵感知量化工作（例如[12 (https://arxiv.org/html/2608.18147#bib.bib8)]）并引入拉格朗日乘子λ'>0以写出修正代价函数MSE(Q,X) + λ'·H(X̂)。实践中，通常使用vNMSE(Q,X) + λ·H(X̂)作为代价，当λ=λ'/‖X‖²时与上式等价。注意这仍推广了ASQ（当λ=0时），并且通过对λ进行二分搜索可以满足熵约束。  

我们注意到标准ASQ动态规划(1)不适用于拉格朗日参数化。这是因为新代价函数不可分离，且量化值x_k的期望频率（因此熵）不仅取决于其右侧值x_j，还取决于其相邻量化值。因此，我们设计最优和近似动态规划以克服此挑战。  

下文推导ECASQ的最优和近似算法。第4节呈现该问题的最优算法，运行时间为O(d²·s)，需要O(d²)空间。随后，第5节介绍近似算法，仅需O(d)空间，且对于b>1，其解MSE不超过熵约束为b-1比特的最优解。推广到P≠X的情况见附录D。  

## 4 最优动态规划  
我们最优动态规划背后的直觉是：由于量化向量X̂中量化值的期望频率仅取决于Q中的两个相邻值，因此可以通过跟踪最后两个量化值来考虑其对熵的贡献。为便于表述，首先假设P=X，即算法只能在输入条目上放置量化值，与ASQ方法相同。我们提出一个新的动态规划，使用两个辅助数组，对任何1≤ℓ≤j≤d定义：  
U(ℓ,j) = ∑_{k=ℓ}^{j-1} ∑_{i=k+1}^{j} (x_j - x_i)(x_i - x_k)  
D(ℓ,j) = ∑_{i=ℓ}^{j} (x_i - x_ℓ)(x_j - x_i)  
以及辅助函数h(y)= { -y log₂y  若y>0  
{ 0         其他情况  

假设已对所有i∈[s], ℓ,j∈[d]求解dp_OPT(i,ℓ,j)，则最优代价为：  
min_{ℓ∈[d]} (dp_OPT(i,ℓ,d) + λ × h(U(ℓ,d)/d))  

dp(i,ℓ,j)通过递推求解，对i≥3：  
dp_OPT(i,ℓ,j) = C[ℓ,j] + min_{k∈[ℓ]} T_{i,ℓ,j}(k)  
其中 T_{i,ℓ,j}(k) ≜ dp_OPT(i-1,k,ℓ) + λ × h((U(k,ℓ)+D(ℓ,j))/d)  
对i=3的停止条件为：  
dp_OPT(3,ℓ,j) = C[1,ℓ] + C[ℓ,j] + λ [h(D(1,ℓ)/d) + h((U(1,ℓ)+D(ℓ,j))/d)]  

直观地，计算每个(i,ℓ,j)条目的值需要O(d)时间（检查所有k选项），而动态规划有s·d²个单元格，因此运行时间和空间边界分别为O(d³·s)和O(d²·s)。下文使用通用动态规划技术将其改进为O(d²·s)和O(d²)。  

注意递推考虑Q中的重复值作为ℓ∈[ℓ]；可以去除这些重复以获得无重复的Q，满足|Q|≤s。  

### 4.1 使用SMAWK优化运行时间  
SMAWK算法[1 (https://arxiv.org/html/2608.18147#bib.bib41)]是加速动态规划的通用技术。该算法的输入是一个满足四边不等式（也称为Monge性质）的矩阵M∈ℝ^{d×d}：∀a≤b≤c≤d: M_{a,c} + M_{b,d} ≤ M_{a,d} + M_{b,c}。算法随后在O(d)时间内找到行最小值。即，它输出k₁,...,k_d使得k_y是第y行的最小值。四边不等式的一个重要性质在以下引理中陈述（证明见附录A）。  

{可重述的} 引理四边形  
设g为凹函数，v和m为单调递减函数，则G(k,j)≜g(v(k)+m(j))满足四边不等式。  

固定i和ℓ。设允许的前驱索引为K_ℓ={k:1≤k<ℓ}。