量化训练之可微量化参数—LSQ


提示:阅读这篇文章需要了解量化训练的基本过程,详见:

LEARNED STEP SIZE QUANTIZATION

摘要

在推理时以低精度操作运行的深度网络相比高精度的替代方案具有功耗和空间优势,但需要克服精度降低时维持高准确性的挑战。在这里,我们提出了一种用于训练这种网络的方法,即学到的步长量化。该方法在使用2、3或4位精度进行权重和激活量化的模型,以及在训练能达到完全精度基准准确性的3位模型时,在ImageNet数据集上实现了迄今为止最高的准确性。

我们的方法建立在现有的学习量化网络权重方法的基础上,通过改进量化器本身的配置方式来实现。具体而言,我们引入了一种新颖的方法来估计和缩放每个权重和激活层的量化器步长处的任务损失梯度,使其可以与其他网络参数一起学习。这种方法在给定系统所需的不同精度水平上都可以工作,并且只需要对现有的训练代码进行简单的修改。

可以看出,文章是把量化参数 step size (也叫 scale)也当作参数训练,而这种量化参数也进行求导训练的技巧也叫作可微量化参数

注:在这之后,高通也发表了增强版的 LSQ+,把另一个量化参数 zero point 也进行训练,从而把 LSQ 推广到非对称量化中。

普通量化训练

在普通量化训练中需要加入伪量化节点 (Fake Quantize),这些节点做的事情就是把输入的 float 数据量化一遍后,再反量化回 float,以此来模拟量化误差,同时在反向传播的时候,发挥 STE 的功能,把导数回传到前面的层。

img

Fake Quantize 的过程可以总结成以下公式 (为了方便讲解 LSQ,这里采用 LSQ 中的对称量化的方式):
$$
\begin{align}
\overline v &= \text{round}(clip(v/s, -Q_N, Q_P)) \newline
\hat v &= \overline v \times s \tag{2}
\end{align}
$$
参数解释:

  • $v$ 是 float 的输入
  • $\bar{v}$ 是量化后的数据 (仍然使用 float 来存储,但数值由于做了 round 操作,因此是整数)
  • $\hat{v}$ 是反量化的结果
  • $−Q_N$ 和 $Q_P$ 分别是量化数值的最小值和最大值 (在对称量化中,$Q_N$ 、$Q_P$ 通常是相等的)
  • $s$ 是量化参数。

由于 round 操作会带来误差,因此 $\hat{v}$ 和 $v$ 之间存在量化误差,这些误差反应到 loss 上会产生梯度,这样就可以反向传播进行学习。每次更新 weight 后,我们会得到新的 float 的数值范围,然后重新估计量化参数 $s$:
$$
s=\frac{|v|_{max}}{Q_P} \tag{3}
$$
之后,开始新一次迭代训练。

LSQ

训练量化步长

可以看到,上面这个过程的量化参数都是根据每一轮的权重计算出来的,而整个网络在训练的过程中只会更新权重的数值。

LSQ 想做的,就是把这里的 $s$ 也放到网络的训练当中,而不是通过权重来计算。

也就是说,每次反向传播的时候,需要对 $s$ 求导进行更新。

这个导数可以这样计算:把 (1)(2) 式统一一下得到:
$$
\begin{align} \hat v&=round(clip(v/s, -Q_N, Q_P))\times s \newline \ &=\begin{cases}-Q_N \times s & v/s <= -Q_N \newline round(v/s)\times s & -Q_N < v/s < Q_P \newline \ Q_P \times s & v/s >= Q_P \end{cases} \end{align}
$$
然后对 $s$ 求导得到:
$$
\frac{\partial \hat v}{\partial s}= \begin{cases} -Q_N & v/s <= -Q_N \newline round(v/s)+\frac{\partial round(v/s)}{\partial s}\times s & -Q_N< v/s<Q_P \tag{6} \newline Q_P & v/s >= Q_P \newline \end{cases}
$$
$round(v/s)$ 这一步的导数可以通过 STE 得到:
$$
\begin{align} \frac{\partial round(v/s)}{\partial s}&=\frac{\partial (v/s)}{\partial s} \tag{7} \newline &=-\frac{v}{s^2} \notag \end{align}
$$
最终得到论文中的求导公式:
$$
\frac{\partial \hat v}{\partial s}= \begin{cases} -Q_N & v/s <= -Q_N \newline -\frac{v}{s}+round(v/s) & -Q_N< v/s <Q_P \tag{8} \newline Q_P & v/s >= Q_P \newline \end{cases}
$$
作者在实验中发现,这种简单粗暴的训练方式有一个好处。

假设把量化范围固定在 [0, 3] 区间,(即 $Q_N$=0,、$Q_P$=3)。下面 A 图表示量化前的 $v$ 和反量化后的 $\hat{v}$ 之间的映射关系(假设 $s$=1),这里面 round 采用四舍五入的原则,也就是说,在 0.5 这个地方 (图中第一道虚线),$\hat{v}$ 是会从 0 突变到 1 的,从而带来巨大的量化误差。

image-20231031224108474

因此,从 0.5 的左侧走到右侧,梯度应该是要陡然增大的。

image-20231031224211063

在 B 图中,作者就对比了 QIL、PACT 和 LSQ (前面两个是另外两种可微量化参数的方法) 在这些突变处的梯度变化,结果发现,QIL 和 PACT 在突变处的梯度没有明显变化,还是按照原来的趋势走,而 LSQ 则出现了一个明显突变 (注意每条虚线右侧)。因此,LSQ 在梯度计算方面是更加合理的。

保持训练稳定

此外,作者还认为,在计算 $s$ 梯度的时候,还需要兼顾模型权重的梯度,二者差异不能过大,因此,作者设计了一个比例系数来约束 $s$ 的梯度大小:
$$
R=\frac{\partial_s L}{s}/\frac{||\partial_w L||}{||w||} \approx 1 \tag{9}
$$

同时,为了保持训练稳定,作者在 $s$ 的梯度上还乘了一个缩放系数 g,对于 weight 来说,$g=1/\sqrt{N_W Q_P}$,对于 feature 来说,$g=1/\sqrt{N_F Q_P}$,$N_W$ 和 $N_F$ 分别表示 weight 和 feature 的大小。

而在初始化方面,作者采用 $\frac{2|v|}{\sqrt{Q_P}}$ 的方式初始化 $s$。

到这里,LSQ 的要点基本讲完了,其实,精华的部分就是把$s$作为量化参数进行训练,至于后面的梯度约束、初始化等,在不同网络结构、不同任务中可能需要灵活调整,没必要完全照论文来。

LSQ+

LSQ+ 的思路和 LSQ 基本一致,就是把零点 (zero point,也叫 offset) 也变成可微参数进行训练。

加入零点后,(1)(2) 式就变成了:
$$
\begin{align} \overline v&=round(clip((v-\beta)/s, -Q_N,Q_P)) \tag{10} \newline \hat v&=\overline v \times s + \beta \tag{11} \end{align}
$$

(高通这个零点计算方式和我之前使用的差得比较多,我自己使用的时候是遵照我之前文章的风格 $v/s+\beta$ 来计算的,因此大家也可以灵活调整)

之后就是按照 LSQ 的方式分别计算导数 $\frac{\partial \hat{v}}{\partial s}$ 和 $\frac{\partial \hat{v}}{\partial \beta}$,再做量化训练。

论文还给出了一些初始化 $s$ 和 $\beta$ 的方式,但还是那句话,视具体任务、具体网络结构而定。

参考


文章作者: QT-7274
版权声明: 本博客所有文章除特別声明外,均采用 CC BY 4.0 许可协议。转载请注明来源 QT-7274 !
评论
  目录