本文介绍了ICLR2020论文中的LSQ方法，这是一种针对权重和激活的量化技术。LSQ通过引入可学习的量化步长，解决了固定量化映射可能带来的性能损失问题。此外，它还提出了一种梯度缩放策略，以确保步长更新的稳定性和收敛性。实验结果显示，LSQ在ImageNet数据集上与其他量化方法相比表现出优越的性能。
论文(ICLR 2020)：https://arxiv.org/abs/1902.08153
非官方源码(Pytorch)：https://github.com/zhutmost/lsq-net
LSQ 概要
Motivation
- 早期量化器的参数都是固定的。固定映射方案虽然简单、具有吸引力，但它不能保证优化网络性能。
- 量化误差最小化方案可能会完美地最小化量化误差，但如果不同的量化映射实际上没有最小化任务误差，则仍然不是最佳的。
- QAT量化如何以最佳方式最大化任务性能，仍然是一个悬而未决的问题。
Method
Definition：QAT Quantization Procedure
作者总结道目前的真量化方案均需要先定义量化器，对推理过程中的Weight 和Activation进行真量化得到 w ^ \hat{w} w^ 和 x ^ \hat{x} x^ 。而 w ^ \hat{w} w^ 和 x ^ \hat{x} x^ 直接进行乘积 (或卷积) 运算得到离散形式的 y ^ \hat{y} y^ ，最后再恢复成全精度值 y y y 。如 Fig 1 所示。
Proposal：Learned Step Size
首先作者自己定义了伪量化的有关公式：
公式中
s
s
s 为 Step size，即量化尺度；
v
v
v 为量化器的输入，可以是全精度的 Weight 或 Activation ；
c
l
i
p
(
a
,
n
,
p
)
clip(a,n,p)
clip(a,n,p) 会将数据 a “钳制”到 [n,p] 范围内，
r
o
u
n
d
(
a
)
round(a)
round(a) 即对数据a进行四舍五入操作；
v
ˉ
\bar{v}
vˉ 为真量化结果，为离散的整数；而
v
^
\hat{v}
v^ 为反量化后的结果。
v
→
v
ˉ
→
v
^
v \rightarrow \bar{v} \rightarrow \hat{v}
v→vˉ→v^ 整个过程可以称为伪量化。
于是对于之前提到的Motivations，作者提出了一个非常简洁的方法：在公式 (1) (2) 的基础上，将 Step size s s s (也就是量化尺度)定义为可学习参数。这只需简单修改代码 (可能这就是为什么没有官方代码的原因吧) 。
那么，有 v ^ \hat{v} v^ 对 s s s 梯度：
上式是可以推出来的 ，其中还用到了直通估计器)的思想：
公式 (3) 直观的感受就是
s
s
s 对
c
l
i
p
clip
clip 内的数据更敏感。进一步地，作者在Fig 2 中阐明了可行性。
Fig 2 (a) 图展示的是均匀量化的基本过程，将连续的 v v v 转为离散整数 v ^ \hat{v} v^ ，其中 Transition Point 称为过渡点，比如1和2之间的1.5就是过渡点。而 (b) 图是梯度示意图，与QIL和PACT算法进行对比。可以看见在 c l i p clip clip 区间内：
- PACT的梯度值为0，这很不利于 s s s 更新。
- QIL虽然有梯度值，但是太单调，且根据钟形分布，最终加和的梯度值基本总是正值，那么根据梯度更新原理， s s s 会不断减小，直到最终加和的梯度值为0才收敛，这导致局部最优。
- 而对于LSQ，既表现有梯度值，又具有分段性，对 Transition Point 附近的值更加敏感，最终加和的梯度值可正可负。也就是说，当加和梯度值为负时 (过渡点左侧的数据较多)，则根据梯度更新原理会使得 s s s 变大，可以理解为过渡点也变大了 (拉远过渡点和过渡点左侧数据的距离) ， c l i p clip clip 范围变大，更多的点进入 c l i p clip clip 范围；当加和梯度值为正时 (过渡点右侧的数据以及 c l i p clip clip 外右侧数据较多)，则根据梯度更新原理会使得 s s s 变小，可以理解为过渡点也变小了 (拉远过渡点和过渡点右侧数据的距离) ， c l i p clip clip 范围缩小，更多的点逃离 c l i p clip clip 范围。
Proposal：Scale for Step Size Gradient
作者首先定义了一个指标 R R R ：
其中 ∇ s L \nabla_{s}L ∇sL 表示对 L L L 求 s s s 的梯度； L L L 即为 Loss Function ； ∣ ∣ w ∣ ∣ ||w|| ∣∣w∣∣ 即表示对 w w w 求 l 2 l_{2} l2-norm 。作者认为 R R R 越接近1越好，即表明网络中的平均权重更新幅度与平均 s s s 更新幅度大致相同 。这样才能实现稳定的收敛。但作者在实验中发现大部分模型的 R R R 都很大，也就是说 s s s 更新的步伐太大了，这导致局部极小值的重复超调，延长收敛时间。
为了纠正这一点，作者进一步提出 Gradient Scale g g g ，将 s s s 的损失乘以 g g g ，以减小平均 s s s 更新幅度。
- 对于权重的 s s s 的 Gradient Scale ，有 g = 1 N W Q P g=\frac{1}{\sqrt{N_{W}Q_{P}}} g=NWQP1 。 N W N_{W} NW 是该权重矩阵的元素总个数。
- 对于激活的 s s s 的Gradient Scale ，有 g = 1 N F Q P g=\frac{1}{\sqrt{N_{F}Q_{P}}} g=NFQP1 。 N F N_{F} NF 是该激活矩阵的元素总个数。
至于为什么 g g g 选取为如上所示，作者在Fig 4 中进行直观展示。
Fig 4 以权重量化器的 Gradient Scale 为例，若 Gradient Scale 取1，则 R R R 还是其本身，在各种 bit 量化中均很大； Gradient Scale 取 1 N W \frac{1}{\sqrt{N_{W}}} NW1 ，则确实有效地降低了 R R R ，但在不同 bit 量化下出现了差异，2-bit量化和8-bit量化下的 R R R 差了一个数量级。于是作者最后取 Gradient Scale 为 1 N W Q P \frac{1}{\sqrt{N_{W}Q_{P}}} NWQP1 (图中作者应该笔误了，是 Q Q Q 不是 L L L )，使得各个 bit 量化下差异不大，均有很好地抑制作用。
进一步地，作者在 Table 3 中探索了 g g g 的其他取值的效果。可以看见作者的取值还是比较合理的，尽管这是认为设置的超参数，存在局部最优性。
Experiment
在所有的实验中，模型的第一层和最后一层始终使用8-bit量化，因为第一层和最后一层具有高精度已成为量化网络的标准做法，并已被其他人证明可提高性能。需要注意的是作者还在训练中使用 Cosine 学习率技巧。所有实验均在ImageNet数据集上进行。
接着作者讨论了权重衰减超参数对于训练的影响。如 Table 2 所示。
Table 1 展示了与其他方法的对比。需要注意的是LSQ的所有模型首尾层不量化。
Fig 3 是作者与其他方法的可视化对比。
如 Table 4 所示，最后作者还做了个 LSQ+Distillation 的实验，以探索更高的性能优化。
Appendices
Step Size Gradient Scale Derivation
在附录中，作者进一步论证了 s s s 的 Gradient Scale g g g 为何取为 1 N W Q P \frac{1}{\sqrt{N_{W}Q_{P}}} NWQP1。
首先作者观察到权重矩阵的
l
2
l_2
l2-norm 的期望与其元素个数的平方根成正比，即跟
N
W
\sqrt{N_{W}}
NW 成正比。于是作者认为：
接着作者观察到 ∣ ∣ ∇ w L ∣ ∣ ||\nabla_{w}L|| ∣∣∇wL∣∣ 和 ∇ s L \nabla_{s}L ∇sL 可以近似看成同一阶，而有：
接着作者假设 ∂ w ^ i ∂ s \frac{\partial\hat{w}_{i}}{\partial s} ∂s∂w^i 接近于1 (根据Fig 2)，并把所有的 ∂ L ∂ w ^ i \frac{\partial L}{\partial\hat{w}_{i}} ∂w^i∂L 看作是以0为中心分布的非相关性随机变量，于是有：
同样地假设 ∂ w ^ ∂ w = 1 \frac{\partial\hat{w}}{\partial w}=1 ∂w∂w^=1，则有
于是根据公式 (4) (6) (7) (8) (9) 可得：
于是自然而然地，把 s s s 的 Gradient Scale g g g 何取为 1 N W Q P \frac{1}{\sqrt{N_{W}Q_{P}}} NWQP1了，因为这样就可以有 R ⋅ g ≈ 1 R \cdot g \approx 1 R⋅g≈1。
Implementation
最后作者附上了代码实现细节：
个人思考
该文章的创新其实很简洁，就是在自定义的量化公式的基础上让 Step Size成为可学习参数。其他的论证过程都很充分，个人困惑的是对 Step Size Gradient Scale 的取值为 1 N W Q P \frac{1}{\sqrt{N_{W}Q_{P}}} NWQP1 的论证。也就是附录里的论证。有点牵强，看了半天没看懂，写的也很匆忙。
扫一扫
4530
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
