量化训练之可微量化参数——LSQ
作者:互联网
(本文首发于公众号,没事来逛逛)
有读者让我讲一下 LSQ (Learned Step Size Quantization) 这篇论文,刚好我自己在实践中有用到,是一个挺实用的算法,因此这篇文章简单介绍一下。阅读这篇文章需要了解量化训练的基本过程,可以参考我之前的系列教程。
LSQ 是 IBM 在 2020 年发表的一篇文章,从题目意思也可以看出,文章是把量化参数 step size (也叫 scale) 也当作参数进行训练。这种把量化参数也进行求导训练的技巧也叫作可微量化参数。在这之后,高通也发表了增强版的 LSQ+,把另一个量化参数 zero point 也进行训练,从而把 LSQ 推广到非对称量化中。
这篇文章就把 LSQ 和 LSQ+ 放在一起介绍了。由于两篇文章的公式符号不统一,为了防止符号错乱,统一使用 LSQ 论文中的符号进行表述。
普通量化训练
在量化训练中需要加入伪量化节点 (Fake Quantize),这些节点做的事情就是把输入的 float 数据量化一遍后,再反量化回 float,以此来模拟量化误差,同时在反向传播的时候,发挥 STE 的功能,把导数回传到前面的层。
Fake Quantize 的过程可以总结成以下公式 (为了方便讲解 LSQ,这里采用 LSQ 中的对称量化的方式):
\[\begin{align} \overline v&=round(clip(v/s, -Q_N,Q_P)) \tag{1} \\ \hat v&=\overline v \times s \tag{2} \end{align} \]其中,\(v\) 是 float 的输入,\(\overline 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 \tag{4} \\ &=\begin{cases}-Q_N \times s & v/s <= -Q_N \\ round(v/s)\times s & -Q_N < v/s < Q_P \tag{5} \\ 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 \\ round(v/s)+\frac{\partial round(v/s)}{\partial s}\times s & -Q_N< v/s<Q_P \tag{6} \\ Q_P & v/s >= Q_P \\ \end{cases} \]\(round(v/s)\) 这一步的导数可以通过 STE 得到:
\[\begin{align} \frac{\partial round(v/s)}{\partial s}&=\frac{\partial (v/s)}{\partial s} \tag{7} \\ &=-\frac{v}{s^2} \notag \end{align} \]最终得到论文中的求导公式:
\[\frac{\partial \hat v}{\partial s}= \begin{cases} -Q_N & v/s <= -Q_N \\ -\frac{v}{s}+round(v/s) & -Q_N< v/s <Q_P \tag{8} \\ Q_P & v/s >= Q_P \\ \end{cases} \](上面这堆公式敲得非常辛苦,给个赞不过分吧o)
作者在实验中发现,这种简单粗暴的训练方式有一个好处。
假设把量化范围固定在 [0, 3] 区间,(即 \(Q_N=0\),\(Q_P=3\))。下面 A 图表示量化前的 \(v\) 和反量化后的 \(\hat{v}\) 之间的映射关系(假设 \(s=1\)),这里面 round 采用四舍五入的原则,也就是说,在 0.5 这个地方 (图中第一道虚线),\(\hat{v}\) 是会从 0 突变到 1 的,从而带来巨大的量化误差。
因此,从 0.5 的左侧走到右侧,梯度应该是要陡然增大的。
在 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} \\ \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\) 的方式,但还是那句话,视具体任务、具体网络结构而定,可以自己调整 (比如我通常就按照 \(v\) 取 90% 左右的区间来估计 \(s\) 和 \(\beta\) 的初始值),甚至你可以用 weight equalize 先预处理一遍网络的权重再来跑 LSQ+ 的算法。
实验
这两篇文章都只给出了分类任务的实验,我觉得应该增加一点别的任务来体现算法的通用性。这里就不列举实验结果了,感兴趣的同学可以看看论文。值得注意的一点是在低比特 (4bit 以下) 的情况下,精度也可以保持得比较好。
一点思考
我自己在一个 GAN 类型的网络上尝试过 LSQ+ 算法,当时被它的效果惊艳到。
这个问题的背景是这样的:最开始的时候,我用普通的量化训练 (8bit) 加上一些蒸馏的技巧来量化这个网络,结果和全精度模型差不多。后来,团队的小伙伴对这个 GAN 网络做了巨量的压缩,同时用了一些技巧大大增强了这个网络的生成能力。然后,我的量化算法在这个网络上就失效了,精度损失非常明显。期间尝试了很多种方案,但都没法拯救。
我自己在分析这个网络权重的时候,发现一个现象,随着网络被压缩得越来越小,权重的数值范围是在逐渐增大的,换句话说,这个网络本身的信息量在逐渐增大。对量化来说,这是件很可怕的事情,因为留给我量化的信息容量是固定的,就只有 8 比特。随着网络信息量增大,每次做量化训练时,round 带来的误差也会更大,这可能使得网络的梯度变得非常不稳定。甚至我会想,是不是 8 比特的信息量就不可能承载得了新网络的容量?
后来,在万念俱灰之下,尝试了 LSQ+ 算法,结果一下子把精度提高了一个档次,我感觉我又活过来了!事后分析的时候,我觉得一个很重要的原因就是:LSQ+ 在前向传播的时候,\(s\) 本身也在控制调整权重的数值分布,而且这种调整是可微的,可以用损失函数进行学习,是一种动态的调整。相比仅仅更新 weight 来调整数值分布的做法,LSQ 多了一条路径来学习。
最后,给需要做量化部署的同学提个醒,在导出量化模型进行部署时,需要根据训练好的 \(s\) 来确定权重的 minmax 大小,因为在 LSQ 的前向传播中,模型权重的数值范围是受 \(s\) 影响的,最终也是根据 \(s\) 反应到损失函数上的。
参考
- Learned Step Size Quantization
- LSQ+: Improving low-bit quantization through learnable offsets and better initialization
- https://www.yuque.com/yahei/hey-yahei/quantization-retrain_differentiable
欢迎关注我的公众号:大白话AI,立志用大白话讲懂AI。
标签:partial,训练,微量,frac,LSQ,量化,hat 来源: https://www.cnblogs.com/jermmyhsu/p/16188288.html