EWSG 基于元素级梯度缩放_LoBob

EWGS：基于 (element-wise) 元素级梯度缩放的网络量化

论文出处：《NetworkQuantizationwithElement-wiseGradientScaling》时间：2021CVPR 单位：YonseiUniversity

0. 前言 ¶

做量化训练的时候，要加入量化节点，即 dequantize(quantize(x,b,s),b,s)，那么在后向传播的时候发现这个 quantize 的操作导数处处为 0，那么直接不管的话，梯度在 quantize 的 round 这个 op 之后就为 0 了，就没法训练下去了。2013 年的时候，Bengio 用 straight-throughestimator(STE，直通估计) 来解决 quantize 导数为 0 的问题，其实就是 bypass。直接就忽略 quantize 这个操作的误差了，得到的梯度也是 mismatch 的，必然是次优的解决办法。STE 存在问题有大佬已经写过了，可以看看这个连接，我就不 copy-pasty 了，量化训练之补偿 STE：DSQ 和 QuantNoise - 知乎 (zhihu.com)。那怎么修改 STE 的问题呢？一个最简单的想法就是，找一个可导的 quantize，这个 idea 在 2017-2019 有部分论文在做块工作，比如 HWGQ、Lq-Net、DSQ 等。2020 后做可导 quantize 的工作没怎么关注了，因为不好落地，对 speedup 有影响。这篇工作 EWGS 的思路和方法很简单，很符合直觉，个人觉得是值得阅读的文章。

1、EWGS 公式 ¶

一句话说 EWGS：给出离散值 (也就是量化值) 的梯度，EWGS 会根据量化误差来自适应缩放梯度，让做梯度更新的时候方向和模值更加准确。

总结 EWSG 的工作：

1、考虑了 quantize 输入和输出的误差，自适应的放大或者缩小梯度

2、将比例因子与任务损失二阶导数联系起来，使用 Hessian 矩阵的迹来估计因子

3、不需要广泛的超参数搜索。

l 和 u 是一组 fp32 数的最大值和最小值，下面的公式就是把输入 x 映射到 0，1 的浮点数值

$x_{n}=\operatorname{clip}\left(\frac{x-l}{u-l},0,1\right)(1)$

把上一步归一化后的数值乘上 int 数值可以表征的大小，就做完了从 fp32 到 in8 的量化，round 的操作是取整，分子是做 fp32->int8，分布是把 int8 转回去 fp32，因为这是做量化训练。这 x_n 和 x_q 是有误差的，来源于 round 这个操作。

$x_{q}=\frac{\operatorname{round}\left(\left(2^{b}-1\right)x_{n}\right)}{2^{b}-1}(2)$

之后就可以输入量化后的输出了 Q_w 和 Q_a,Q_a 因为经过了 Relu 后是非负数，那么就直接用 x_q 表示；而 Q_w 是对称量化，有负数的，那么先 - 0.5 就把 x_q 的移到了 [-0.50.5]，乘以 2 就表示正确了。

$Q_{W}(x)=2\left(x_{q}-0.5\right),Q_{A}(x)=x_{q}(3)$

这里有一个很重要的细节，就是对量化后卷积层 / 全连接层的输出加了一个α缩放因子，这一点 trick。

这个公式就是 EWSG 的公式了

$g_{x_{n}}=g_{x_{q}}\left(1+\delta\operatorname{sign}\left(g_{x_{q}}\right)\left(x_{n}-x_{q}\right)\right)(4)$

STE 是这样的， $\mathcal{G}_{\mathbf{x}_{n}}=\mathcal{G}_{\mathbf{x}_{q}}$ , 直接将导数为 0 或者不可导的变成了 1，直接直通。因为这些下标太难打了，我直接截图了。一个是公式 (1) 的梯度，一个是公式 (2) 的梯度。δ是大于 0 的数值，当δ等于 0 的时候，EWSG 就是 STE 了。

STE 是次优的原因：

（1）多个 $x_n$ 可以产生相同的 $x_q$ ，大白话就是多个 fp32 的数值变成同一个 int 值；

（2） $x_n$ 提供的相同梯度对每个 $x_q$ 都有不同的影响，大白话就是虽然梯度 $x_n$ 的一样的，但是对应的 $x_q$ 是不一样的，这个就是 mismatch 的问题。

2、如何确定δ数值，基于海森矩阵的方法 ¶

这边就是公式推导了。

将 EWSG 公式 (即可公式 4) 展开，凑成有导数的形式，x_n-x_q 就是量化误差了，也就是符号ℇ

$\begin{aligned} g_{x_{n}}&=g_{x_{q}}+\frac{g_{x_{n}}-g_{x_{q}}}{x_{n}-x_{q}}\left(x_{n}-x_{q}\right)\\ &=g_{x_{q}}+\frac{g_{x_{q}+\epsilon}-g_{x_{q}}}{\epsilon}\left(x_{n}-x_{q}\right)\\ \end{aligned}\\$

$g_{x_{n}}\approx g_{x_{q}}+g_{x_{q}}^{\prime}\left(x_{n}-x_{q}\right)$

$g_{x_{n}}\approx g_{x_{q}}\left(1+\frac{g_{x_{q}}^{\prime}}{\left|g_{x_{q}}\right|}\operatorname{sign}\left(g_{x_{q}}\right)\left(x_{n}-x_{q}\right)\right) g_{x_{q}}^{\prime}=\frac{\partial g_{x_{q}}}{\partial x_{q}}$

其中，这项就是导数的导数也就是二阶信息，也是常说的海森信息

$g_{x_{q}}^{\prime}=\frac{\partial g_{x_{q}}}{\partial x_{q}}$

所以，δ的数值就确定了

$\delta=\frac{g_{x_{q}}^{\prime}}{\left|g_{x_{q}}\right|}$

海森矩阵的公式推导基于了一个假设 (没怎么看懂，也不想深入探究，摆烂)，得出这么个公式，

$\mathbb{E}\left[\mathbf{v} \mathbf{v}^{T}\right]=I$

代入并且进行变换，

$\begin{aligned} \operatorname{Tr}(H) &=\operatorname{Tr}(H I)=\operatorname{Tr}\left(H \mathbb{E}\left[\mathbf{v}^{T}\right]\right) \\ &=\mathbb{E}\left[\operatorname{Tr}\left(H \mathbf{v} \mathbf{v}^{T}\right)\right]=\mathbb{E}\left[\mathbf{v}^{T} H \mathbf{v}\right] \end{aligned}$

最后δ的公式如下：N 是海森矩阵中对角线元素的个数，G 是由梯度 Gx 的分布决定的梯度表示。

但这个变换对于计算的意义我还是没看懂，因为这样还是要计算海森矩阵，估计也是用 pyHessian 的 library 算的，是用其他近似的方法求个海森矩阵，具体在 HAWQ(v1、v2、v3)（下次一定写这三篇工作）。

个人觉得这个变换很凑数，也可能自己没看看懂那个假设，有看懂的大佬麻烦指正我！。

$\delta=\frac{\operatorname{Tr}(H) / N}{G} (5)$

实验中发现梯度很多是 0，这样梯度的平均值偏向于 0，把 G 设置成偏大的数：

$3 \sigma\left(\mathcal{G}_{\mathbf{x}_{q}}\right)$

3、实验 ¶

实验我就不贴 ResNet 的效果了，因为大差不差，直接 STE+trick 也能有差不多的性能，直接看 mobilenet 的实验，是比 PROFIT(ECCV2020) 的差，打不过啊。但 EWGS 的好处在于简单，且推理的时候并没有增加其他操作，戏虐说辞 "无痛涨点"。其他的实验图意义不大，就不看了。

4、总结 ¶

EWGS 结果没有超过前人的工作，我比较欣赏的是他的方法很简单，而且不给推理增加负担，是篇好工作。CVPR2020 也有一篇做量化训练的时候修改梯度的，UnifiedINT8，通过修改梯度的方向和数值来缓解 mismatch 带来的影响。但 EWGS 从数学上个人觉得更加可解释和合理。故记录一下。

本文总阅读量88次