模型量化论文阅读#4----EWGS:Network Quantization with Element-wise Gradient Scaling

在量化中，因为量化函数是不可微分的，所以一旦涉及到量化函数的反向传播时，就需要对量化函数的梯度进行近似，目前常用的近似是STE，从而避免量化函数的零梯度问题。所谓STE就是一个直通器，它只是传播相同的梯度而不考虑量化函数的输入和输出之间的离散误差。这篇论文就针对STE的缺点提出了一种元素级的梯度放缩算法，叫做EWGS。给定量化函数量化输出的梯度，EWGS自适应地放大或缩小每个梯度元素，并使用缩放后的梯度作为量化函数输入的梯度，通过反向传播来训练量化的网络。缩放是根据每个梯度元素的符号和量化函数的连续输入和离散输出之间的误差进行的。

EWGS量化的前向过程：

Step1: 对连续的输入值进行归一化和截断处理；

xn=clip⁡(x−lu−l,0,1)x_{n}=\operatorname{clip}\left(\frac{x-l}{u-l}, 0,1\right)xn=clip(u−lx−l,0,1)
其中，l,ul,ul,u 分别表示量化的上下边界，这两个值是可以被学习的。

Step2: 量化操作；
xq=round⁡((2b−1)xn)2b−1x_{q}=\frac{\operatorname{round}\left(\left(2^{b}-1\right) x_{n}\right)}{2^{b}-1}xq=2b−1round((2b−1)xn)
其中，bbb表示量化的bit数。

Step3: 针对量化后的激活值和权重单独进行处理；
QW(x)=2(xq−0.5),QA(x)=xqQ_{W}(x)=2\left(x_{q}-0.5\right), Q_{A}(x)=x_{q}QW(x)=2(xq−0.5),QA(x)=xq

另外，为了调整每一层的输出的尺度，在额外为每一层训练一个参数α\alphaα ，将其乘以每一层的激活值。
前向过程与其余的量化方法基本是一致的，该论文的核心在反向过程中的量化函数的梯度近似上。

EWGS量化的反向过程：

下式即为EWGS算法中梯度的近似公式：

gxn=gxq(1+δsign⁡(gxq)(xn−xq))g_{x_{n}}=g_{x_{q}}\left(1+\delta \operatorname{sign}\left(g_{x_{q}}\right)\left(x_{n}-x_{q}\right)\right)gxn=gxq(1+δsign(gxq)(xn−xq))

在STE中，gxn=gxqg_{x_{n}}=g_{x_{q}}gxn=gxq，所以STE可以视为上述的一种特例。

上式中，xqx_{q}xq即为量化后的值，xnx_{n}xn为量化前的连续值，gxqg_{x_{q}}gxq即为损失函数关于xqx_{q}xq的梯度，gxng_{x_{n}}gxn即为损失函数关于xnx_{n}xn的梯度.signsignsign是取符号的函数，δ\deltaδ是一个尺缩因子，后面会专门介绍。
可以看出，梯度gxng_{x_{n}}gxn与量化后的误差以及sign(gxq)sign(g_{x_{q}})sign(gxq)有关。

为了有助于理解，可以参考下图。举个例子，对于0.51，1和1.49三个数经过量化函数之后，得到的值都是1.但是这三个值处于不同的位置，因此在反向更新时，他们不应该是拥有相同的更新梯度（STE中梯度就是一样的）。假设−gxq- g_{x_{q}}−gxq大于0，即xqx_{q}xq要向着增大的方向更新，可以认为其向着2的方向更新，对应图1的(a)图，那么为了使得xqx_{q}xq增大，xnx_{n}xn自然也要向着增大的方向更新，在STE中，0.51，1和1.49三个值更新的梯度都是相同的。但是我们能看出，这三个值处于不同的位置，其更新的梯度理应是不一样的，对于1.49来说，其只需要增大到1.5就可以使得xqx_{q}xq增大到2，而对于0.51来说，为了使得xqx_{q}xq增大到2，其需要增加的值更大，因此其对应的梯度也要大一点。而上式恰好可以实现这种想法。这在图1中也可以体现出来。

如何获得尺缩因子δ\deltaδ

根据下面的恒等式：
gxn=gxq+gxn−gxqxn−xq(xn−xq)=gxq+gxq+ϵ−gxqϵ(xn−xq)\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}gxn=gxq+xn−xqgxn−gxq(xn−xq)=gxq+ϵgxq+ϵ−gxq(xn−xq)

其中，ϵ=xn−xq\epsilon = x_{n} - x_{q}ϵ=xn−xq,根据量化函数四舍五入的性质，∣ϵ∣≤0.52b−1|\epsilon| \leq \frac{0.5}{2^{b}-1}∣ϵ∣≤2b−10.5, 可以认为 ϵ\epsilonϵ足够小，从而可以得到下述近似式子：
gxn≈gxq+gxq′(xn−xq)g_{x_{n}} \approx g_{x_{q}}+g_{x_{q}}^{\prime}\left(x_{n}-x_{q}\right)gxn≈gxq+gxq′(xn−xq)

其中，gxq′=∂gxq∂xqg_{x_{q}}^{\prime}=\frac{\partial g_{x_{q}}}{\partial x_{q}}gxq′=∂xq∂gxq是损失函数关于xqx_{q}xq的二阶导数。于是我们可得到：
gxn≈gxq(1+gxq′∣gxq∣sign⁡(gxq)(xn−xq))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)gxn≈gxq(1+∣gxq∣gxq′sign(gxq)(xn−xq))
根据上式，可以得到δ\deltaδ = gxq′∣gxq∣\frac{g_{x_{q}}^{\prime}}{\left|g_{x_{q}}\right|}∣gxq∣gxq′，但是因为计算二阶导数即海森矩阵H计算量很大。论文改用Hessian矩阵H中对角线元素的平均值来近似二阶导数，假设主对角线在矩阵H中占主导地位。
有了上述假设，我们不需要获得海森矩阵H，只需要获得H的迹Tr(H)Tr(H)Tr(H)即可。而：

Tr⁡(H)=Tr⁡(HI)=Tr⁡(HE[vvT])=E[Tr⁡(HvvT)]=E[vTHv]\begin{aligned} \operatorname{Tr}(H) &=\operatorname{Tr}(H I)=\operatorname{Tr}\left(H \mathbb{E}\left[\mathbf{v} \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}Tr(H)=Tr(HI)=Tr(HE[vvT])=E[Tr(HvvT)]=E[vTHv]

其中， I是一个单位矩阵，v\mathbf{v}v是一个从Rademacher distribution采样的随机向量，满足E[vvT]=I\mathbb{E}\left[\mathbf{v} \mathbf{v}^{T}\right]=IE[vvT]=I。

我们可以根据下式快速计算得到HvHvHv,从而避免直接计算H.
∂GxqTv∂xq=∂GxqT∂xqv+GxqT∂v∂xq=∂GxqT∂xqv=Hv\frac{\partial \mathcal{G}_{\mathbf{x}_{q}}^{T} \mathbf{v}}{\partial \mathbf{x}_{q}}=\frac{\partial \mathcal{G}_{\mathbf{x}_{q}}^{T}}{\partial \mathbf{x}_{q}} \mathbf{v}+\mathcal{G}_{\mathbf{x}_{q}}^{T} \frac{\partial \mathbf{v}}{\partial \mathbf{x}_{q}}=\frac{\partial \mathcal{G}_{\mathbf{x}_{q}}^{T}}{\partial \mathbf{x}_{q}} \mathbf{v}=H \mathbf{v}∂xq∂GxqTv=∂xq∂GxqTv+GxqT∂xq∂v=∂xq∂GxqTv=Hv

最后，将δ\deltaδ定义为下式：

δ=Tr⁡(H)/NG\delta=\frac{\operatorname{Tr}(H) / N}{G}δ=GTr(H)/N

关于实验的更多细节，可以参考论文原文。

@inproceedings{lee2021network, title={Network Quantization with
Element-wise Gradient Scaling}, author={Lee, Junghyup and Kim,
Dohyung and Ham, Bumsub}, booktitle={Proceedings of the IEEE/CVF
Conference on Computer Vision and Pattern Recognition},
pages={6448–6457}, year={2021} }