Paper Reading||Learned Step Size Quantization

论文(ICLR 2020)：https://arxiv.org/abs/1902.08153
非官方源码(Pytorch)：https://github.com/zhutmost/lsq-net

Motivation

1. 早期量化器的参数都是固定的。固定映射方案虽然简单、具有吸引力，但它不能保证优化网络性能。
2. 量化误差最小化方案可能会完美地最小化量化误差，但如果不同的量化映射实际上没有最小化任务误差，则仍然不是最佳的。
3. QAT量化如何以最佳方式最大化任务性能，仍然是一个悬而未决的问题。

Method

Definition：QAT Quantization Procedure

 　作者总结道目前的真量化方案均需要先定义量化器，对推理过程中的Weight 和Activation进行真量化得到 $\hat{w}$ 和 $\hat{x}$ 。而 $\hat{w}$ 和 $\hat{x}$ 直接进行乘积 (或卷积) 运算得到离散形式的 $\hat{y}$ ，最后再恢复成全精度值 $y$ 。如 Fig 1 所示。

Proposal：Learned Step Size

 　首先作者自己定义了伪量化的有关公式：

 　公式中 $s$ 为 Step size，即量化尺度； $v$ 为量化器的输入，可以是全精度的 Weight 或 Activation ；$clip(a,n,p)$ 会将数据 a “钳制”到 [n,p] 范围内， $round(a)$ 即对数据a进行四舍五入操作； $\bar{v}$ 为真量化结果，为离散的整数；而 $\hat{v}$ 为反量化后的结果。 $v\to \bar{v}\to \hat{v}$ 整个过程可以称为伪量化。

 　于是对于之前提到的Motivations，作者提出了一个非常简洁的方法：在公式 (1) (2) 的基础上，将 Step size $s$ (也就是量化尺度)定义为可学习参数。这只需简单修改代码 (可能这就是为什么没有官方代码的原因吧) 。

 　那么，有 $\hat{v}$ 对 $s$ 梯度：

 　上式是可以推出来的 ，其中还用到了直通估计器)的思想：

 　公式 (3) 直观的感受就是 $s$ 对 $clip$ 内的数据更敏感。进一步地，作者在Fig 2 中阐明了可行性。

 　Fig 2 (a) 图展示的是均匀量化的基本过程，将连续的 $v$ 转为离散整数 $\hat{v}$ ，其中 Transition Point 称为过渡点，比如1和2之间的1.5就是过渡点。而 (b) 图是梯度示意图，与QIL和PACT算法进行对比。可以看见在 $clip$ 区间内：

• PACT的梯度值为0，这很不利于 $s$ 更新。
• QIL虽然有梯度值，但是太单调，且根据钟形分布，最终加和的梯度值基本总是正值，那么根据梯度更新原理， $s$ 会不断减小，直到最终加和的梯度值为0才收敛，这导致局部最优。
• 而对于LSQ，既表现有梯度值，又具有分段性，对 Transition Point 附近的值更加敏感，最终加和的梯度值可正可负。也就是说，当加和梯度值为负时 (过渡点左侧的数据较多)，则根据梯度更新原理会使得 $s$ 变大，可以理解为过渡点也变大了 (拉远过渡点和过渡点左侧数据的距离) ， $clip$ 范围变大，更多的点进入 $clip$ 范围；当加和梯度值为正时 (过渡点右侧的数据以及 $clip$ 外右侧数据较多)，则根据梯度更新原理会使得 $s$ 变小，可以理解为过渡点也变小了 (拉远过渡点和过渡点右侧数据的距离) ， $clip$ 范围缩小，更多的点逃离 $clip$ 范围。

Proposal：Scale for Step Size Gradient

 　作者首先定义了一个指标 $R$ ：

 　其中 ${\nabla }_{s}L$ 表示对 $L$ 求 $s$ 的梯度； $L$ 即为 Loss Function ； $||w||$ 即表示对 $w$ 求 $l_2$-norm 。作者认为 $R$ 越接近1越好，即表明网络中的平均权重更新幅度与平均 $s$ 更新幅度大致相同 。这样才能实现稳定的收敛。但作者在实验中发现大部分模型的 $R$ 都很大，也就是说 $s$ 更新的步伐太大了，这导致局部极小值的重复超调，延长收敛时间。

 　为了纠正这一点，作者进一步提出 Gradient Scale $g$ ，将 $s$ 的损失乘以 $g$ ，以减小平均 $s$ 更新幅度。

• 对于权重的 $s$ 的 Gradient Scale ，有 $g=\frac{1}{\sqrt{N_W{Q}_P}}$ 。 $N_W$ 是该权重矩阵的元素总个数。
• 对于激活的 $s$ 的Gradient Scale ，有 $g=\frac{1}{\sqrt{N_F{Q}_P}}$ 。 $N_F$ 是该激活矩阵的元素总个数。

 　至于为什么 $g$ 选取为如上所示，作者在Fig 4 中进行直观展示。

 　Fig 4 以权重量化器的 Gradient Scale 为例，若 Gradient Scale 取1，则 $R$ 还是其本身，在各种 bit 量化中均很大； Gradient Scale 取 $\frac{1}{\sqrt{N_W}}$ ，则确实有效地降低了 $R$ ，但在不同 bit 量化下出现了差异，2-bit量化和8-bit量化下的 $R$ 差了一个数量级。于是作者最后取 Gradient Scale 为 $\frac{1}{\sqrt{N_W{Q}_P}}$ (图中作者应该笔误了，是 $Q$ 不是 $L$ )，使得各个 bit 量化下差异不大，均有很好地抑制作用。

 　进一步地，作者在 Table 3 中探索了 $g$ 的其他取值的效果。可以看见作者的取值还是比较合理的，尽管这是认为设置的超参数，存在局部最优性。

Experiment

 　在所有的实验中，模型的第一层和最后一层始终使用8-bit量化，因为第一层和最后一层具有高精度已成为量化网络的标准做法，并已被其他人证明可提高性能。需要注意的是作者还在训练中使用 Cosine 学习率技巧。所有实验均在ImageNet数据集上进行。

 　接着作者讨论了权重衰减超参数对于训练的影响。如 Table 2 所示。

 　Table 1 展示了与其他方法的对比。需要注意的是LSQ的所有模型首尾层不量化。

 　Fig 3 是作者与其他方法的可视化对比。

 　如 Table 4 所示，最后作者还做了个 LSQ+Distillation 的实验，以探索更高的性能优化。

Appendices

Step Size Gradient Scale Derivation

 　在附录中，作者进一步论证了 $s$ 的 Gradient Scale $g$ 为何取为 $\frac{1}{\sqrt{N_W{Q}_P}}$。

 　首先作者观察到权重矩阵的 $l_2$-norm 的期望与其元素个数的平方根成正比，即跟 $\sqrt{N_W}$ 成正比。于是作者认为：

 　接着作者观察到 $||{\nabla }_{w}L||$ 和 ${\nabla }_{s}L$ 可以近似看成同一阶，而有：

 　接着作者假设 $\frac{\partial {\hat{w}}_i}{\partial s}$ 接近于1 (根据Fig 2)，并把所有的 $\frac{\partial L}{\partial {\hat{w}}_i}$ 看作是以0为中心分布的非相关性随机变量，于是有：

 　同样地假设 $\frac{\partial \hat{w}}{\partial w}=1$，则有

 　于是根据公式 (4) (6) (7) (8) (9) 可得：

 　于是自然而然地，把 $s$ 的 Gradient Scale $g$ 何取为 $\frac{1}{\sqrt{N_W{Q}_P}}$了，因为这样就可以有 $R\cdot g\approx 1$。

Implementation

 　最后作者附上了代码实现细节：

个人思考

该文章的创新其实很简洁，就是在自定义的量化公式的基础上让 Step Size成为可学习参数。其他的论证过程都很充分，个人困惑的是对 Step Size Gradient Scale 的取值为 $\frac{1}{\sqrt{N_W{Q}_P}}$ 的论证。也就是附录里的论证。有点牵强，看了半天没看懂，写的也很匆忙。