量化训练、推理神经网络-WAGE

1. WAGE背景介绍

WAGE（Training and Inference with Integers in Deep Neural Networks）是ICLR2018的Best Paper，这篇论文提出了一种量化神经网络的框架，可以将神经网络中的参数（Weight）、激活值（Activation）、梯度（Gradient）以及误差（Error）进行可选比特长度的整型量化。在本文中使用$W^i$、$a^i$、$g^i$、$e^i$表示第$i$层的参数、激活值、梯度以及误差，用下标${\cdot }_q$来表示已进行量化的值。在WAGE框架中还需要定义好对4种值的量化函数$Q_W(\cdot )$、$Q_a(\cdot )$、$Q_g(\cdot )$和$Q_e(\cdot )$，用$[\cdot ]$表示量化之后的整型值在计算机中存储所需要的比特长度。

WAGE的前向传播（推理）过程如下式（1.1）：
$$a_q^i=Q_a(Q_W(W_i)\cdot a_q^{i-1})\text{\hspace{1em}(1.1)}$$

WAGE的反向传播（训练）过程如下式（1.2）：
$$e_q^i=Q_e(\frac{\partial \mathcal{L}}{\partial a^i}),g_q^i=Q_g(\frac{\partial \mathcal{L}}{\partial W^i})\text{\hspace{1em}(1.2)}$$

2. 移位线性映射和随机舍入

在WAGE中使用线性映射将浮点数映射到$k$比特有符号整型值，那么整型值之间的最小不可分距离$\sigma$可以从式（2.1）中计算出：
$$\sigma (k)=2^{1-k},k\in {\mathbb{N}}_{+}\text{\hspace{1em}(2.1)}$$
那么WAGE中的将浮点数$x$量化为$k$比特的有符号整型值的截断量化函数被定义为：
Q ( x , k ) = C l i p { σ ( k ) ⋅ r o u n d ( [ x σ ( k ) ] ) , − 1 + σ ( k ) , 1 − σ ( k ) } (2.2) Q(x, k) = Clip \{ \sigma(k) \cdot round([\frac{x}{\sigma(k)}]), -1+\sigma(k), 1-\sigma(k) \} \tag{2.2} Q(x,k)=Clip{σ(k)⋅round([σ(k)x​]),−1+σ(k),1−σ(k)}(2.2)
$Clip$函数将整型值限制在区间$[-1+\sigma (k),1-\sigma (k)]$中，去掉$-1$以保持对称性。在对浮点值进行量化之前会进行一个缩放（移位）运算，否则经过式（2.2）的处理，浮点值大部分会被饱和截断。缩放函数在本文被定义为$Shift$操作：
$$Shift(x)=2^{round(lo{g}_{2}x)}\text{\hspace{1em}(2.3)}$$
最终WAGE会对反向传播中的梯度值进行随机舍入运算，以确保梯度能够有效地传播来进行训练。$Q_g(\cdot )$算子会通过一个16比特的随机数产生器来约束量化后的$k_g$比特整型值的随机性，在接下来的章节中会详细介绍四种量化算子的实现。

3. 参数初始化

Batch Normalization能够预防神经网络训练中的梯度爆炸和梯度消失问题，并缓解网络对参数初始化的需求，但是在整型运算中实现Batch Normalization十分复杂。由于种种原因，论文作者改进了何恺明提出的参数初始化方法：
$$W\sim U(-L,+L),其中L=max(\sqrt{6/n_{in}},L_{min}),L_{min}=\beta \sigma \text{\hspace{1em}(3.1)}$$
$n_{in}$是输入层的神经元数量，$L_{min}$约束均匀分布能够覆盖的最小区间（至少要大于$\sigma$，因此$\beta >1$）。

4. 量化函数

4.1 参数量化函数$Q_W(\cdot )$

参数量化遵循式（2.2），并使用第3节中的初始化方法：
$$W_q=Q_W(W)=Q(W,k_W)\text{\hspace{1em}(4.1)}$$
经过量化之后参数的方差被缩放了，这有可能导致训练过程中的梯度爆炸。为了缓解被放大的影响，论文作者提出在层与层之间增加一个缩放参数$\alpha$：
α = m a x { S h i f t ( L m i n / L ) , 1 } (4.2) \alpha = max \{ Shift(L_{min} / L), 1 \} \tag{4.2} α=max{Shift(Lmin​/L),1}(4.2)

4.2 激活值量化函数$Q_W(\cdot )$

作者在论文中不建议在卷积层之后使用Average Pooling，因为均值计算对精度要求更高（相比于Max Pooling）。论文中假设输出的Batch均值为0，Batch Normalization退化为一个缩放运算，在WAGE中可训练的缩放算子被式（4.2）中的$\alpha$替代了，对激活值的量化方法：
$$a_q=Q_a(a)=Q(a/\alpha ,k_a)\text{\hspace{1em}(4.3)}$$

4.3 误差量化函数$Q_e(\cdot )$

相较于激活值，误差的范围更大。论文中认为在训练时梯度的方向指导了整个网络的收敛，而梯度的大小作用不大。那么这样的话，可以直接使用整型值来表达梯度的方向来进行网络的训练。通过除以一个移位因子，WAGE将误差缩放到范围$[-\sqrt{2},+\sqrt{2}]$之内，量化函数如下：
$$e_q=Q_e(e)=Q(e/Shift(max(|e|)),k_e)\text{\hspace{1em}(4.4)}$$

4.4 梯度量化函数$Q_g(\cdot )$

$$g=\frac{\partial \mathcal{L}}{\partial W}=\frac{\partial \mathcal{L}}{\partial a}\cdot \frac{\partial a}{\partial W}=e\cdot \frac{\partial a}{\partial W}\text{\hspace{1em}(4.5)}$$

首先按照梯度方向决定训练的思路，WAGE对梯度进行了缩放，其中$\eta$（$\eta =2^j,j\in \mathbb{N}$）为学习率：
g s = η ⋅ g / S h i f t ( m a x { ∣ g ∣ } ) (4.6) g_s = \eta \cdot g / Shift(max\{ |g| \}) \tag{4.6} gs​=η⋅g/Shift(max{∣g∣})(4.6)
$g_s$表征了训练的最小步长，以及参数更新的方向。如果参数占用了$k_g$比特的空间，那么最小步长所对应的整型数和浮点数分别是$\pm 1$和$\pm \sigma (k_g)$。在训练的开始阶段，设置学习率$\eta \ge 2$使$g_s>1$来加速训练，而在训练的后半段随着学习率的衰减，梯度也会变小。在参数的更新过程中，WAGE将$g_s$分解为整型部分和小数部分，然后使用16比特随机数生成器来约束浮点数随机映射到整型数：
Δ W = Q g ( g ) = σ ( k g ) ⋅ s g n ( g s ) ⋅ { ⌊ ∣ g s ∣ ⌋ + B e r n o u l i ( ∣ g s ∣ − ⌊ ∣ g s ∣ ⌋ ) } (4.7) \Delta W = Q_g(g) = \sigma(k_g) \cdot sgn(g_s) \cdot \{ \lfloor |g_s| \rfloor + Bernouli( |g_s| - \lfloor |g_s| \rfloor) \} \tag{4.7} ΔW=Qg​(g)=σ(kg​)⋅sgn(gs​)⋅{⌊∣gs​∣⌋+Bernouli(∣gs​∣−⌊∣gs​∣⌋)}(4.7)
在更新参数之后仍然会有参数值超越范围$[-1+\sigma ,1-\sigma ]$的情况存在，因此要使用$Clip$函数来使参数值饱和，经过第$t$轮的训练，可以得到：
W t + 1 = C l i p { W t − Δ W t , − 1 + σ ( k g ) , 1 − σ ( k g ) } (4.8) W_{t+1} = Clip \{ W_t - \Delta W_t , -1 + \sigma(k_g) , 1 - \sigma(k_g) \} \tag{4.8} Wt+1​=Clip{Wt​−ΔWt​,−1+σ(kg​),1−σ(kg​)}(4.8)