论文(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 的论证。也就是附录里的论证。有点牵强,看了半天没看懂,写的也很匆忙。