论文(CVPR 2019):https://arxiv.org/abs/1911.09464
源码(Pytorch ):https://github.com/aliyun/alibabacloud-quantization-networks
现有的低比特量化方法可以分为两类:以近似(Approximation)为基础的方法和以优化(Optimization)为基础的方法。
针对以上思考,作者认为将量化操作近似表征为一个简单的非线性(全程可微)函数,如表征为 Sigmoid、ReLU 和 Maxout 函数,那么就不需要对梯度进行近似设计了。
基于此,作者给出了自己的方法,如 Fig 1(d) 所示。具有可学习的偏移和尺度的多个 Sigmoid 函数的线性组合表征了量化函数。
进一步地,如 Fig 2 所示,作者在训练过程中会不断地松弛 Sigmoid 函数的陡度,以减轻 Quantization Gap 问题。
接下来作者首先阐述了这一创新的激发点,然后从非线性函数的角度对 Quantization 做了新解释,最后说明了网络需要学习什么。
定义非线性激活函数为 g ( x ) g(x) g(x)。在感知机(Perceptron)中常常会使用阶跃函数来产生二进制输出:
这种形式就很像量化操作。其零点处不可微、其他地方导数为 0 。于是又诞生了 Sigmoid 函数用于早期前馈神经网络:
此时处都平滑可微了。使得参数可以利用反向传播来更新。到目前 ReLU 函数成为了频繁使用的激活函数之一,而 ReLU 的变体改进是 Maxout:
其中 { a j } \{a_j\} {aj} 和 { c j } \{c_j\} {cj} 是可学习参数。Maxout 的形式表明,复杂的凸函数可以由 k k k 个简单的线性函数的组合来近似。这种思想也启发了作者对量化函数的设计。
量化操作会将连续输入映射为离散整数,而二进制量化操作可视为单位阶跃函数。受公式 (3) 设计的启发,量化操作可以表征为如 Fig 2(e) 所示的具有指定偏置和尺度的多个单位阶跃函数的组合:
其中 x x x 为全精度 Weight/Activation; y y y 为量化后整数,有 y ∈ Y y \in \mathcal{Y} y∈Y ;量化间隔数为 n n n 个,而 Y \mathcal{Y} Y 的元素个数为 n+1 个; A \mathcal{A} A 为标准单位阶跃函数; β \beta β 为对输入的全局缩放因子; s i s_i si 和 b i b_i bi 分别为阶跃函数的缩放因子和偏移因子。需要注意的是, β \beta β 和 b i b_i bi 是可学习参数,而 s i = y i + 1 − y i s_i=y_{i+1}-y_{i} si=yi+1−yi。全局的补偿 o = 1 2 ∑ i = 1 n s i o=\frac12\sum^n_{i=1}s_i o=21∑i=1nsi 是用来保持量化后的输出是关于零点对称的(说明本质上是对称量化)。当整数集 Y \mathcal{Y} Y 给定后, n n n、 s i s_i si 和 o o o 就可以立即得到。
作者将公式 (4) 中的每个单位阶跃函数替换为 Sigmoid 函数,就使得整个量化过程变得可微起来。但是因为最终在推理阶段,必须使用理想量化函数(即用回公式 (4)),所以需要在训练阶段逐渐缩小理想量化函数和软量化函数的差距。受蒸馏(Distillation)思想的启发,作者在 Sigmoid 函数中引入了温度因子 T:
当 T T T 值变大时, Sigmoid 函数就越像阶梯函数,但同时梯度开始饱和。在训练阶段,作者从一个小的 T T T 开始以保证稳定有效的学习,然后 T T T 随着训练轮数的增加而增加,直到最后接近理想量化函数。
设需要被量化的全精度 Weight 或 Activation 为集合 X = { x d , d = 1 , . . . , D } \mathcal{X}=\{x_d,d=1,...,D\} X={xd,d=1,...,D} ,而量化函数会独立地作用于 x d x_d xd :
其中 β \beta β 和 α \alpha α 分别为输入和输出的缩放因子, b i b_i bi 为第 i i i 个量化间隔的起始点,注意 b 0 = − ∞ b_0=-\infin b0=−∞。
在训练阶段,设 Loss 为 ℓ \ell ℓ ,则会计算如下梯度:
其中 g d i = σ ( T ( β x d − b i ) ) g_d^i=\sigma(T(\beta x_d-b_i)) gdi=σ(T(βxd−bi)) ,而 n n n, s i s_i si 和 o o o 的梯度可以直接由 Y \mathcal{Y} Y 得到。
在实际训练时,作者会将 公式 (6) 的量化函数插入到合适的位置。若一个层的输入原来为 x x x,那么经过量化后就会变成 Q ( x ) \mathcal{Q}(x) Q(x);参与计算的 W W W 就会变成 Q ( W ) \mathcal{Q}(W) Q(W)。
在实际推理时,公式 (6) 就由下面公式替代:
Algorithm 1 总结了训练的流程。设有一个全精度的网络为 N N N ,其有 M M M 个模块,每个模块可为全连接层或卷积层。设第 m m m 个模块的所有需要量化的 Activation(或叫做 Input) 为 X ( m ) \mathcal{X}^{(m)} X(m) ,所有需要量化的 Weight 为 Θ ( m ) \Theta^{(m)} Θ(m) 。所有的 X ( m ) \mathcal{X}^{(m)} X(m) 共享同一套量化参数 { α X ( m ) , β X ( m ) , b X ( m ) } \{\alpha_{\mathcal{X}}^{(m)}, \beta_{\mathcal{X}}^{(m)}, \bm{b}_{\mathcal{X}}^{(m)}\} {αX(m),βX(m),bX(m)}。所有的 Θ ( m ) \Theta^{(m)} Θ(m) 共享同一套量化参数 { α Θ ( m ) , β Θ ( m ) , b Θ ( m ) } \{\alpha_{\Theta}^{(m)}, \beta_{\Theta}^{(m)}, \bm{b}_{\Theta}^{(m)}\} {αΘ(m),βΘ(m),bΘ(m)}。
Table 2 为 ResNet-18 的比较结果。
Table 3 为ResNet-50 的比较结果。
Table 5 为针对量化间隔 { b i } \{b_i\} {bi} 的消融实验。Linear 表示均匀量化,Non-uniform 表示非均匀量化。
Fig 3 为针对 Temperature 的消融实验,固定 Temperature 然后训练得出实验结果。
Table 6 为针对是否使用预训练模型的消融实验。可以发现用预训练模型来做 QAT 还是有好处的。
Table 7 为量化后模型在 VU9P FPGA 设备上允运行的开销情况。