论文阅读笔记：N2UQ: Towards Accurate Quantization via Generalized Straight-Through Estimation

论文：https://arxiv.org/pdf/2111.14826

代码：https://github.com/liuzechun/Nonuniform-to-Uniform-Quantization

1 背景

量化模型与全精度模型仍然有不可忽视的性能差距，特别是在极低比特量化的情况下。例如，2-bit经典均匀量化的ResNet50在ImageNet上取得了67.1的top1准确率，相比于量化前的ResNet下降了9.9%。这种性能差距主要是由于用有限的量化等级去表示真实权重和激活导致的量化误差，以及均匀量化器不能灵活地适应不同的输入值分布。

为了更好地拟合底层分布和减少量化误差，之前的一些方法提出非均匀量化，即根据真实值分布的密度来调整量化分辨率。然而，非均匀量化的精度提高通常以牺牲硬件效率为代价。由于非均匀量化的输出是浮点权重和激活，他们的乘法运算不能够通过二进制地按位运算直接加速，解决这一问题的常用方法是建立查找表，将浮点值映射到二进制数字，如图1（a）所示，这种非均匀量化引起的后处理需要花费更大的硬件面积和额外的消耗。

本研究的目标是开发一种新的量化方法，在保持均匀量化的硬件投影简单性的同时，提供灵活性，以实现非均匀量化的性能优势。尽管希望量化器的输出(即量化的权重/激活)具有统一的量化等级，以免引起额外的后处理任务，但每个输出等级并不一定需要代表实值输入的相等范围。如图1 ( b )所示，在学习输入值的阈值时，强制输出量化等级是等距的，以便在拟合量化的底层实值分布时具有更多的灵活性。这种量化器设计称为非均匀到均匀量化器。

然而，优化这种非均匀到均匀量化器是困难的，由于阈值参数的梯度难以计算，无法通过现有的梯度估计量化方法（即直通估计器STE）来解决。STE简单地将输入梯度估计为一个与输出梯度相等的阈值操作，其定义无法在训练中纳入阈值差异，也无法用梯度更新新阈值。

为了规避这一问题，作者重新研究了随机二值化中最早推导出的STE，并推导出一种新的更灵活的后向量化近似方法，命名为广义直通估计器G-STE。当所有输入区间大小相等时，它退化为STE，而对于需要非均匀输入阈值的场景，它通过梯度学习自动适应阈值，并对量化函数提供更细粒度的近似。具体来说，G-STE将随机量化的期望编码为向前确定性量化函数的后向近似，自然地将关于输入阈值参数的难解梯度计算转化为关于斜率的难解梯度计算，并在向后梯度计算中编码输入阈值差异对剩余网络的影响。

2 创新点

（1）提出了非均匀到均匀量化器( N2UQ )，通过学习输入阈值来提高量化精度，同时在实现中保持与均匀量化类似的硬件友好性。

（2）提出了G-STE 来解决N2UQ中输入阈值参数难以确定的梯度计算问题。G - STE计算随机量化的期望作为前向确定性量化的后向近似。

（3）基于熵分析，提出了一种新的考虑整体权重分布的权重正则化方法，以进一步保留权重量化中的信息。

（4）即使在将量化权重和激活固定为均匀且仅学习输入阈值的严格约束下，N2UQ在ImageNet上仍以0.5 - 1.7 %的精度超过了目前最先进的非均匀量化方法。具体来说，2-bit ResNet - 50模型在ImageNet上取得了76.4 %的top - 1准确率，与真实值的差距缩小到0.6 %，证明了N2UQ设计的有效性。

3 方法

3.1 预备知识

矩阵乘法是深度神经网络中计算代价最高的运算。通过将卷积层和全连接层中的权重和激活量化为固定比特，矩阵乘法可以通过按位运算加速：

其中，$a$ 和 $w$ 分别为量化激活 $a^q$ 和 量化权重 $w^q$ 的二进制向量。 $M$ 和 $K$ 分别用于表示 $a^q$ 和 $w^q$ 的比特数。而式（1）对于均匀量化一般成立， 当考虑非均匀量化时，情况就变复杂了。

3.2 非均匀到均匀量化

为了梳理量化的内在机制，作者总结了量化器的两个性质：

性质1（量化器的输出）：用式（1）加速量化网络的前提是量化器的输出（即量化激活 $a^q$ 和 量化权重 $w^q$）可以映射到固定比特内的二进制数字：$a^q={\sum }_{i=0}^{M-1}{a}_i{2}^i,w^q={\sum }_{j=0}^{K-1}{w}_j{2}^j$。

这个性质很直观，因为式（1）发生在二元向量之间。虽然均匀量化的 $a^q$ 和 $w^q$ 可以通过线性映射很容易地转换为二进制，但由于 n 比特非均匀量化器的输出实际上是 $2^n$ 个等间隔的浮点值，因此通常需要额外的操作和查找表将 n 比特非均匀量化的 $a^q$ 和 $w^q$ 编码为 n 比特二进制数，如图1所示。这个后处理步骤需要更多额外的消耗。

性质2（量化器的输入）：量化函数 $x^q=F_Q(x^r)$ 的输入和输出步长的均匀性可以通过一个适当的量化器设置来分离，即输出 $x^q$ 中的量化是均匀的，输入 $x^r$ 的量化不一定需要均匀。

这一性质源于观察到量化器表示具有多个量化级别 $x^q$ 的实值变量 $x^r$，当量化等级固定时，实值分布是多样的。因此，在将 $x^r$ 划分为固定量化等级 $x^q$ 时，允许可学习的阈值理论上可以结合更多的表示灵活性，以实现更高的精度。

3.2.1 向前传播：阈值量化学习

基于这两个性质，作者提出了用于激活量化的非均匀到均匀量化器，其前向通式为：

其中 $n$ 为比特数，$T$ 表示可学习阈值， $x^r,x^q$ 分别表示输入实值变量和输出量化变量。N2UQ的目标是在量化器中编码阈值学习，以允许更多的灵活性，同时输出统一量化的权重和激活，以适应快速的比特操作，而不需要量化和矩阵乘法之间的后处理步骤。

3.2.2 向后传播：广义直通估计器G-STE

然而，通过式（2）反向传播是困难的，因为：（I）式（2）关于输入 $x_r$ 的导数几乎处处为零，（II）对于阈值参数的梯度计算很难处理。

对于问题（I），以前的量化工作采用直通估计器STE来近似量化函数中的后向梯度。STE只是将传入梯度分配给阈值操作，作为传出梯度：

$x_{i,l}^q$ 和 $x_{i,l}^r$ 表示第 $l$ 层的第 $i$ 个量化/实值变量。这种简单的近似函数对均匀量化器的效果很好，然而STE在量化器的输入和输出间隔中隐式地强制等轴纵横比，因为它将量化函数视为向后传播中的恒等函数，这阻碍了量化器在固定输出量化值的同时允许可学习的输入阈值。此外，STE无法处理问题（II），即关于可学习阈值的梯度计算，因为STE基本上只是绕靠梯度，就好像量化函数不存在一样，因此，为了提出一种更细颗粒度和更灵活的后向近似，可以处理输入阈值的梯度计算，作者在啊随机二值化中重新审视STE，并对到处广义直通估计器G-STE。

引理1。在二值化中，对向前确定性二值化函数进行梯度逼近的直通估计器STE可以从随机二值化函数的期望得到。

在随机二值化中，实值变量根据与 -1/1 的距离随机二值化得到 -1/1：

其中 ${\tilde{x}}_{i,l}^b$ 表示随机二值化变量。为了更新 $W_{ij,l}$ （第 $l$ 层的权重连接到第 $l-1$ 层的神经元 $j$ 和第 $l$ 层的神经元 $i$），通过随机二值化函数计算期望梯度：

其中 $E,E_{{\tilde{x}}_{i,l}^b},E_{/{\tilde{x}}_{i,l}^b}$ 分别是整个网络上的期望，只有随机二值化变量的期望和除随机二值化变量外其他部分的期望。特别地，

这就得到了二值化函数的通用直通估计器，他在阈值附近传输相同的梯度，而当实值输入离阈值太远时（即 $x^r>1$ 或 $x^r<-1$），则不考虑梯度。同时，通过对等式（4）中的概率设置硬阈值（即 $p=0.5$ ），可以得到前向传播中广泛采用的确定性二值化函数：

到此，作者展示了STE在前向确定性二值化函数的后向近似中编码了随机二值化的期望。

引理2。量化函数 $x^q$ 可以看做是具有不同阈值的二值化函数 $x^b$ 的和：$x^q=x_{thr{e}_1}^b+x_{thr{e}_2}^b+...+x_{thr{e}_n}^b$。

这个引理是不言而喻的，如图2（b）所示，由此，作者将随机二值化的概念推广到量化中，得到G-STE来解决非均匀到均匀量化器中的阈值学习问题。

从第一个量化段开始，将初始点记为 $s$，其长度记为 $a_1$。为了简化几件，这里省略下标，循序随机二值化的概念，在 $[s,s+a_1]$ 范围内，实值变量可以随机量化为 0/1，其概率与它们到 $s/s+a_1$ 的距离成正比：

x ~ q { 0 , 1 } \widetilde{x}^{q_{\{0,1\}}} x q{0,1}​ 是该段内的随机量化输出。与式（6）相似，这个量化段的推导可以从等式（8）的期望中计算出来。

这样，通过改变斜率，阈值参数 $a_1$ 对网络的影响被合理的编码在后向近似函数式（9）中。此外，在前向传递中，不使用需要随机种子生成过程的随机量化，而是通过将概率阈值设置为p=0.5来采用确定性量化。

同样，可以从这种概率意义上获得剩余片段的前向确定性量化函数和相应的后向近似函数。 然后，通过对 $2^n-1$ 段求和，可以自然地推导出具有 $n$ 位量化后向近似的非均匀到均匀量化器，如图2所示。 将这种提出的反向梯度近似方法命名为广义直通估计器（G-STE），并获得以下定理。

定理1。广义直通估计器：

n 位非均匀到均匀量化器的后向梯度逼近函数：
使用Nonuniform-to-Uniform量化器的随机版本的期望作为其反向近似来解决阈值学习问题，这里 d 0 = s , d i = s + ∑ j = 1 i a j , i ∈ { 1 , … , 2 n − 1 } d_0=s, d_i=s+\sum_{j=1}^ia_j, i∈\{1,…,2^n-1\} d0​=s,di​=s+∑j=1i​aj​,i∈{1,…,2n−1}。

具有G-STE反向逼近的非均匀到均匀量化器的本质是在反向逼近函数中编码随机量化的平均值，同时在前向传递中具有硬阈值的确定性量化函数。 这有助于将在前向传递中强制输出为特定值的阈值参数的影响同化为在后向传递中逼近压缩或扩展的量化概率，其中阈值参数ai对剩余神经网络的影响反映在 后向函数的斜率。 反过来, 对阈值参数梯度可以很容易地计算为，

可以看到，当所有区间大小相等时，G-STE 退化为 STE，而在训练期间，它学会了使用从网络计算的梯度自动调整输入阈值。 具体来说，输出级别被缩放到 { 0 , 2 / ( 2 n − 1 ) , … , 2 } \{0, 2/(2^n-1),…,2\} {0,2/(2n−1),…,2} 。 另外应用两个可学习的缩放参数 ${\beta }_1,{\beta }_2$ 来缩放量化前的输入 $x^r:=x^r\times {\beta }_1$ 和量化后的输出 $x^q:=x^q\times {\beta }_2$ 。${\beta }_1,{\beta }_2$ 初始化为 1，$a_i$ 初始化为 $2/(2^n-1)$并强制为正。 与经典的均匀量化器相比，所提出的 nonuniform-to-uniform 量化器每层仅引入 $2^n+2$ 个额外参数，考虑到大量网络权重，这可以忽略不计。

3.2.3. 熵保权重正则化

此外，作者提出了权重正则化，以鼓励量化权重具有更多的信息承载能力。一个重要的观察结果是，实值网络中的权重通常很小，例如， 对于预训练的ResNet-18网络中的权重约为$10^{-2}$，但量化权重通常在以下范围内扩展：[−1, 1]。 这种幅度失配将导致量化权重塌陷到接近零的几个量化级别。经典的解决方案是使用 tanh 函数并除以最大绝对权重值，将权重重新缩放为[−1, 1]。 然而，这些重新缩放的权重可能由几个极值主导，并且仍然没有完全占据[−1，1]，如图3所示，这会导致巨大的信息损失。

从信息论的角度来看，当量化权重包含更高的熵时，会保留更多的信息。 因此，作者在量化器之前对实值权重进行正则化，以获得量化权重的最大熵： $maxH=-p_{i}log(p_i)$， 满足 ${\sum }_{i=1}^N{p}_i=1$。 这里 $p_i$ 是量化到第 $i$ 个量化级别的实值权重的比例，$N$ 是量化级别的总数。 基于拉格朗日乘子，很容易得到最优的 KaTeX parse error: Undefined control sequence: \* at position 3: p^\̲*̲_i=\frac{1}{N},…，表示当实值权重被量化到多个量化级别的比例相等时，量化后的权重中的信息熵达到最大值。

给定量化函数 $F_Q=round((clip(w^{r’},-1,1)+1)\times \frac{2^n-1}{2})\times \frac{2}{2^n-1}-1$，作者凭经验解决了当实值权重归一化为 $W^{r^{\prime }}=\frac{2^{(n-1)}}{2^n-1}\frac{|W^r|}{||W^r|{|}_{l1}}{W}^r$，对应的量化权重 $w^q=F_Q(w^{r’},n)$ 近似均匀分布在所有量化级别。这里，$W_r$ 是实值权重过滤器，|$W_r$| 表示 $W_r$ 中的条目数。与从图像生成并随批次变化的激活不同，权重是静态参数，可以高度灵活地调整每个值。因此，不是通过学习阈值参数来确定所有量化的权重值，而是采用等距阈值进行权重量化，让每个权重单独学习其值并根据整体统计数据缩放权重。

4 效果

4.1 与SOTA方法对比

量化后在ImageNet上的效果如表1。

MobileNet上各种方法的效果对比如表2。

4.2 消融实验

各组件的消融实验。

不同权重正则化的消融实验。