Paper Reading||Overcoming Oscillations in Quantization-Aware Training

论文(arxiv 2022)：https://arxiv.org/abs/2203.11086
暂无源码

1 Motivation

1. 作者发现，在 QAT 过程中使用直通估计器(STE)时，权值似乎会在相邻的量化水平之间随机振荡，这就会为优化过程引入有害的噪声。
2. 这种权重振荡，会破坏训练期间收集的批归一化层的参数统计结果，导致较差的推理准确性。这种影响在具有深度可分离层的高效网络(如 Mobilenets 和 EfficientNets )的低比特(≤4比特)量化中表现得特别明显。
3. 大多数的 QAT 算法还无法克服这种权重振荡。

2 Observation

 　作者首先探讨了为什么量化感知训练中会出现权重振荡现象，以及这种现象如何影响实践中的神经网络训练。

2.1 Quantization-aware training

 　本文中，作者使用的量化映射方式为均匀对称量化，权重和激活均会被量化。定义量化函数为 $q(\cdot )$ ，输入向量为 $\text{w}$ ，量化的输出为 $\hat{\text{w}}$ ，那么有：

 　其中 $\lfloor \cdot \rceil$ 为四舍五入操作， $clip(\cdot ,\alpha ,\beta )$ 是下界和上界分别为 $\alpha$ 和 $\beta$ 的截断函数， $s$ 为量化尺度， $n$ 和 $p$ 为量化阈值的上界和下界。
 　原始浮点权重 $\text{w}$ 仅作为优化的代理，通常被称为潜在权值或阴影权值。
 　实际中为了解决四舍五入函数不可导问题，往往会采用直通估计器来将四舍五入函数的梯度估计为1，那么可以定义关于 $\text{w}$ 的 损失函数$\mathcal{L}$ 为：

2.2 Oscillations

 　Défossez 等人(2021年)观察到，STE 有个反直觉、非常有趣的副作用—— STE 会使得潜在权值在相邻量化电平之间的决策边界附近震荡。
 　为了更好阐述这一现象，首先在一个 1D 小型实例模型(Toy Model)上进行相关实验。设原浮点域中最优浮点权重为 $w_{\ast }$ ，即原浮点域中最后收敛的权值；设输入数据为 $\mathrm{x}$ ，其从分布 $X$ 中采样出来，而 $\mathbb{E}[{\mathrm{x}}^2]={\sigma }^2$ ， $0<{\sigma }^2<\infty$ 。那么该 Toy模型的优化目标为最小二乘问题(
)：

 　其中 $q(\cdot )$ 为式(1)中的量化器， $w$ 为潜在权重。作者使用 STE 对梯度进行优化，于是就观察到了 Fig 1 左图的现象。图中，当潜在权重 $w$ 接近最优值 $w_{\ast }$ 时， $w$ 开始在最优值以上的量化电平 $w_{\uparrow }$ 和以下的量化电平 $w_{\downarrow }$ 之间的过渡点附近振荡，而不是收敛到最优值 $q(w_{\ast })$ 附近的区域。

 　权重围绕决策阈值振荡，这是因为式(2)中的梯度，在过渡点以上时会将潜在权重 $w$ 向下推向 $w_{\downarrow }$ ，而在过渡点以下时会将潜在权重向上推向 $w_{\uparrow }$ 。
 　但是，作者发现，该振荡的频率与学习速率无关。这将在 Appendix A.3 中讨论。作者还发现，振荡的频率由最优值与其最近的量化电平的距离决定，即 $d=|w_{\ast }-q(w_{\ast })|$ 。具体来说，假设 $q(w_{\ast })=w_{\uparrow }$ ，那么过渡点下方的梯度就是过渡点上方梯度的 $k=\frac{s}{d}-1$ 倍，其中 $s$ 为量化尺度(也叫量化间隔)，有 $s>d$ 。也就是说， $w$ 需要行进 $k$ 次迭代才能穿过过渡点，然后梯度才能反转。
 　当然，从概率的角度来看， $w$ 最终量化到最近的量化电平 $q(w_{\ast })$ 的概率还是较大的，这种概率行为是与随机四舍五入(Stochastic Rounding)类似的，但该行为的随机性是源于梯度的有偏、离散性，并不是 Sampling 导致的。
 　最后作者还发现，振荡不是普通 STE 所独有的。在 STE 的几个变体中也存在这样的振荡现象，作者在 Fig 1中、右部分展示了这个情况。

2.3 Oscillations in practice

 　本小节表明，这种震荡确实存在于神经网络中，对他们的优化具有重要意义。
 　Fig 2 展示了在ImageNet上的接近收敛的 MobileNetV2 中3 bits 量化权重的优化过程，训练方法为 LSQ 。

 　可以观察到，许多权值似乎在两个相邻的量子化水平之间随机振荡。在 Fig 3中，还可以看到在网络假设收敛(迭代完毕)之后，很大一部分潜在权值正好位于量化电平之间的过渡点附近(决策边界)。这进一步强化了一个观察结果——即相当大比例的权重会出现波动，无法收敛。

 　基于上述现象，作者总结出两个需要待解决的问题：一个是振荡会污染批归一化层的统计参数；一个是振荡会阻碍网络的优化稳定性。

2.3.1 THE EFFECT ON BATCH-NORM STATISTICS

 　权重的振荡现象，意味着每一层的输出统计数据可以在每次梯度更新后会显著变化。由于振荡引起的分布偏移，这些 BN 的统计参数可能会被破坏，导致精度显著下降。事实上，有两个因素影响 BN 的破坏程度：权值位宽和每个输出通道的权值数量。简单来说，当前通道的权重数越少，位宽 $b$ 越低，那么振荡就越剧烈，对 BN 参数的破坏越深。

 　在 Table 1 中，作者使用 KL 散度来定量分析实际每 Iter 数据统计分布和 BN 估计分布之间的差异性。可以观察到，BN估计分布总体都不那么准确，且在深度可分离层(DW)中分布的差异性表现得更明显。
 　解决这个问题的一种简单易行的方法是，在训练完毕后用一小部分数据重新估计BN层的统计信息。这种方法，作者称为批处理归一化(BN)重估计。作者认为，由于权重的振荡，BN 重估计在 QAT中是必不可少的。

 　在表2中，可以看到 BN 重估计不仅提高了MobileNetV2 的最终量化精度，而且降低了不同初始化种子之间的方差。我们还观察到，随着MobileNetV2位宽的下降，准确度的差距变大，

2.3.2 THE EFFECT ON TRAINING

 　除了对BN统计数据造成伤害外，振荡还会对训练过程本身产生负面影响。
 　为了说明这一点，作者在两个振荡状态之间随机采样振荡权值后，MobileNetV2可以实现更低的训练损失(和更高的验证精度)，如 Table 3 所示，同时多次随机抽样反复实验。

 　综上，作者在第4节提出了两个解决方法，来缓解这种振荡现象。

3 Overcoming oscillations in QAT

 　作者在这一小节首先引入了一个量化振荡的度量，然后再提出了两种技术，阐述了客服振荡的方法。

3.1 Quantifying Oscillations

 　为了测量振荡，定量化分析振荡，作者利用指数滑动平均(EMA)来计算振荡的频率随时间的变化，接着定义了一个最小频率作为判断权值振荡的阈值。
在迭代 $t$ 中出现一个振荡需要满足两个条件：

1. 权重的整型值发生改变。即 ${\text{w}}_{\mathrm{i}\mathrm{n}\mathrm{t}}^t\ne {\text{w}}_{\mathrm{i}\mathrm{n}\mathrm{t}}^{\text{t}-1}$ ，其中 ${\text{w}}_{\mathrm{i}\mathrm{n}\mathrm{t}}^t=\mathrm{c}\mathrm{l}\mathrm{i}\mathrm{p}(\lfloor \frac{\text{w}}{s},n,p\rceil )$ 。
2. 整数域的变化方向需要与之前整数域的变化方向相反，所以要有 $o^t=\mathrm{s}\mathrm{i}\mathrm{g}\mathrm{n}({\Delta }_{\mathrm{i}\mathrm{n}\mathrm{t}}^t)\ne \mathrm{s}\mathrm{i}\mathrm{g}\mathrm{n}({\Delta }_{\mathrm{i}\mathrm{n}\mathrm{t}}^{\tau })$ ，其中 $\tau$ 为迭代的上一个整数域变化情况，而 ${\Delta }_{\mathrm{i}\mathrm{n}\mathrm{t}}^t={\text{w}}_{\mathrm{i}\mathrm{n}\mathrm{t}}^t-{\text{w}}_{\mathrm{i}\mathrm{n}\mathrm{t}}^{\text{t}-1}$ 为改变的方向。

 　为了跟踪振荡的情况，作者计算了指数滑动平均(EMA)：

3.2 Oscillation dampening

 　为了抑制振荡行为，作者定义了一个正则化项，以鼓励潜在权值接近量化电平而不是过渡点：

 　其中 $s,n,p$ 为式(1)所定义的量化参数。那么整个模型的目标函数就写为： $\mathcal{L}={\mathcal{L}}_{\mathrm{t}\mathrm{a}\mathrm{s}\mathrm{k}}+\lambda {\mathcal{L}}_{\mathrm{d}\mathrm{a}\mathrm{m}\mathrm{p}\mathrm{e}\mathrm{n}}$ 。而对于 Dampen，有梯度：

 　可以看到，由于 Clip 函数的作用，只有在量化过程中没有裁剪的权值才会收到正则化效果，这避免了对LSQ算法的量化尺度更新的影响。
 　但这种正则化的缺点是，其不仅会影响振荡的权值，而且还会阻碍不处于振荡状态的权值的移动更新。

3.3 Iterative freezing of oscillating weights

 　作者提出了另一种更有针对性的方法。通过在训练中冻结特定权重来防止这些权重的振荡。对振荡的定量化如式(4)所示。如果任何权重的振荡频率超过阈值 $f_{\mathrm{t}\mathrm{h}}$ ，该权重将被冻结，直到训练结束。
 　需要注意的是，作者是在整数域进行冻结，这样能避免量化尺度 $s$ 变化而导致潜值的不同舍入。
 　当一个权重振荡时，它不一定在振荡的两个量化电平上花费相同的时间。正如作者在2.2节的 Toy Model 中所展示的那样，权重的每个状态的可能性会线性地依赖于量化状态到最优值的距离。因此，随着时间的推移，所有量化值的期望确实会对应于真正的最优值。为了将权重冻结到其更频繁的状态，以让状态更接近最优值，作者使用指数滑动平均(EMA)来记录以前整数值的情况。在满足冻结条件时，作者通过对 EMA 结果进行四舍五入，冻结权重的整数状态就是最频繁的整数状态。
 　更加细节的描述展示在 Algorithm 1 中。

 　此算法可以与任何基于梯度的优化器结合使用，而不限于特定的量化公式或梯度估计器。

4 Experiments

 　本节，作者评估了所提出的两个方法的有效性，并在 ImageNet 上与其他 QAT方法做了比较。

4.1 Experimental setup

• 量化配置：作者所做的实验均基于 LSQ，对权重和激活均进行量化，同时更新量化尺度 $s$ 。具体来说，作者保持第一层和最后一层的权值为8位，并对除归一化层外的所有层的输入量化。整个量化均使用 Per-tensor 粒度量化。
• 优化配置：在所有情况下，作者从一个预先训练的全精度网络开始，并使用 MSE 估计并初始化权值和激活的量化尺度。使用动量为0.9的 SGD ，并使用余弦退火学习速率衰减训练。在消融研究时，作者只对重量量化进行了20个 Epoch 的训练。在5.3节中，作者对重量和激活量化进行了90个 Epoch 训练。根据网络和量化位宽，我们训练的学习率为0.01或0.0033。

4.2 Ablation studies

• 振荡抑制。在 Table 4 的前三行中，可以看到，当增加正则化强度时，振荡权重的比例下降，BN 重估计精度前后之间的差距缩小。然而，过多的抑制会损害最终的准确性，这表明过度的正则化会抑制量化级别之间的权重的有益移动。解决这一问题的方法是在训练过程中逐步增加正规化强度。这允许潜在权重在训练的第一阶段更自由地移动，同时通过应用强正则化减少接近收敛的有害振荡。而作者发现，对于系数 $\lambda$ 的余弦退火调度在实际应用中效果很好。

• 迭代权重冻结。在 Table 5 中，可以看到随着阈值的降低，剩余振荡的数量显著减少，只有一些低频振荡留在网络中。还可以看到，pre-BN 重估计精度更接近后 post-BN 重估计精度，正如人们所期望的那样，在训练结束时有小的振荡。但是当振荡判定阈值太低时，在训练的早期阶段会有太多的重量会被冻结，这对最终的准确性产生负面影响。为了解决这个问题，作者对振荡判定阈值同样采用了余弦退火算法，于是就可以使用一个更强的冻结阈值，在接近训练结束时可以冻结几乎所有的振荡。

5.3 Comparison to other QAT methods

 　在 Table 6 中，展示了 MobileNetV2 的结果；Table 7 和 Table 8 展示了 MobileNetV3-Small 和 EfficentNet-lite 的结果。可以发现作者的效果还是有竞争性的。

个人思考

 　该文章角度新颖，确实做了很有意义的事情。但是关于振荡这个话题，最早其实是另外一个2021年的文章提出的，当然作者也在文中提到了该文章。其实作者来举例子的 Toy Model 观点就是2021年文章里的内容，但引申了很多实验和观察发现。
 　但我个人认为用 Toy Model 来引证实际模型中的过渡点振荡问题还是不够说服力。关于振荡的原因还需要详细考究。