DoReFa-Net 低比特量化训练笔记

DoReFa-Net 这篇论文的核心思路是：对卷积神经网络的权值、激活值和梯度分别做低比特量化，从而把训练和推理中绝大多数的计算都压缩到低比特的位运算层面。下面按照论文的脉络整理一下各个模块的做法。

位运算与点乘的关系

论文的出发点是：两个整数之间的点乘可以用位运算来表示。

\mathbf{x} \cdot \mathbf{y} = \mathrm{bitcount}(\mathrm{and}(\mathbf{x}, \mathbf{y})), \quad x_i, y_i \in \{0,1\}\ \forall i.

将这个思路推广到向量之间的点乘，两个向量中整数的表示位数也可以不同，这就为权值和激活值使用不同的比特宽度奠定了基础。

\begin{aligned} \mathbf{x} \cdot \mathbf{y} &= \sum_{m=0}^{M-1} \sum_{k=0}^{K-1} 2^{m+k}\, \mathrm{bitcount}\big[\mathrm{and}(c_m(\mathbf{x}), c_k(\mathbf{y}))\big], \\ c_m(\mathbf{x})_i,\ &c_k(\mathbf{y})_i \in \{0,1\}\ \forall i, m, k. \end{aligned}

直通估计器（STE）

量化操作本身（取整/截断）不可导，反向传播需要近似处理。论文在梯度近似上使用了 STE（Straight-Through Estimator），本质上就是自定义 forward 和 backward 函数：forward 做量化，backward 把梯度直接透传。

本文中的 STE 对应的是圆整函数，也就是 quantize_k，具体算法参见代码实现。

\begin{aligned} \textbf{Forward:}\quad & r_o = \frac{1}{2^k - 1}\,\mathrm{round}\big((2^k - 1) r_i\big) \\ \textbf{Backward:}\quad & \frac{\partial c}{\partial r_i} = \frac{\partial c}{\partial r_o}. \end{aligned}

权值量化

量化权值时通常会引入缩放因子。XNOR-NET 使用的是随输入变化的缩放因子，DoReFa-Net 则使用固定的缩放因子，实现更简洁。

\begin{aligned} \textbf{Forward:}\quad & r_o = \mathrm{sign}(r_i) \times \mathbf{E}(|r_i|) \\ \textbf{Backward:}\quad & \frac{\partial c}{\partial r_i} = \frac{\partial c}{\partial r_o}. \end{aligned}

当比特数大于 1 时，权值量化采用以下方式，其中的 quantize_k 就是上面介绍的 STE 实现。当 k=1 时，具体的表达形式与多比特情形略有不同，但实验表明两者的效果差别不大。

\begin{aligned} \textbf{Forward:}\quad & r_o = f_\omega^k(r_i) = 2\,\mathrm{quantize}_k\!\left(\frac{\tanh(r_i)}{2\max(|\tanh(r_i)|)} + \frac{1}{2}\right) - 1. \\ \textbf{Backward:}\quad & \frac{\partial c}{\partial r_i} = \frac{\partial r_o}{\partial r_i}\frac{\partial c}{\partial r_o}. \end{aligned}

激活值量化

激活值的量化方式如下。

f_\alpha^k(r) = \mathrm{quantize}_k(r).

梯度量化

梯度的量化相比权值和激活值要更复杂一些，主要原因是梯度的数值范围更大，且没有明确的边界，因此论文对梯度采用随机量化。

\tilde{f}_\gamma^k(\mathrm{d}r) = 2\max_0(|\mathrm{d}r|)\left[\mathrm{quantize}_k\!\left(\frac{\mathrm{d}r}{2\max_0(|\mathrm{d}r|)} + \frac{1}{2}\right) - \frac{1}{2}\right].

所谓随机量化，就是在量化时额外叠加一个噪声项。

N(k) = \frac{\sigma}{2^k - 1}\ \text{where}\ \sigma \sim Uniform(-0.5, 0.5).

实验表明这个噪声项比较关键，去掉后效果会明显下降。

f_\gamma^k(\mathrm{d}r) = 2\max_0(|\mathrm{d}r|)\left[\mathrm{quantize}_k\!\left(\frac{\mathrm{d}r}{2\max_0(|\mathrm{d}r|)} + \frac{1}{2} + N(k)\right) - \frac{1}{2}\right].

需要注意的是，梯度的量化只在反向传播过程中进行。

\begin{aligned} \textbf{Forward:}\quad & r_o = r_i \\ \textbf{Backward:}\quad & \frac{\partial c}{\partial r_i} = f_\gamma^k\!\left(\frac{\partial c}{\partial r_o}\right). \end{aligned}

整体算法流程

将上述三部分合在一起，整体算法如下。

Algorithm 1  Training an L-layer DoReFa-Net with W-bit weights and A-bit
activations using G-bit gradients. Weights, activations and gradients are
quantized according to Eqn 9, Eqn 11, Eqn 12, respectively.

Require: a minibatch of inputs and targets (a_0, a*), previous weights W,
         learning rate eta
Ensure:  updated weights W^{t+1}

  {1. Computing the parameter gradients:}
  {1.1 Forward propagation:}
   1: for k = 1 to L do
   2:     W_k^b  <- f_omega^W(W_k)
   3:     a-bar_k <- forward(a_{k-1}^b, W_k^b)
   4:     a_k    <- h(a-bar_k)
   5:     if k < L then
   6:         a_k^b <- f_alpha^A(a_k)
   7:     end if
   8:     Optionally apply pooling
   9: end for

  {1.2 Backward propagation:}
  Compute g_{a_L} = dC/da_L knowing a_L and a*.
  10: for k = L to 1 do
  11:     Back-propagate g_{a_k} through activation function h
  12:     g_{a_k}^b   <- f_gamma^G(g_{a_k})
  13:     g_{a_{k-1}} <- backward_input(g_{a_k}^b, W_k^b)
  14:     g_{W_k^b}   <- backward_weight(g_{a_k}^b, a_{k-1}^b)
  15:     Back-propagate gradients through pooling layer if there is one
  16: end for

  {2. Accumulating the parameter gradients:}
  17: for k = 1 to L do
  18:     g_{W_k} = g_{W_k^b} * (dW_k^b / dW_k)
  19:     W_k^{t+1} <- Update(W_k, g_{W_k}, eta)
  20: end for

其中所有的复杂操作——包括 forward、backward_input 和 backward_weight——都是低比特运算，这是论文声称能显著降低训练开销的核心所在。

关于首层和末层的特殊处理：

第一层卷积因为直接与网络输入相连，权值不做量化，但卷积后的激活值仍然需要量化。
最后一层全连接层在类别数较小的情况下也不做量化，但从 FC 层传出的梯度是要量化的。

推理阶段的融合优化

在具体实现上，可以将算法中的第 3、4、6 步骤进行融合，从而减少推理时的内存和时间开销（相当于省掉了 h() 这一中间步骤）。融合后可以通过对比的方式直接得到结果。

a_k^b = f_\alpha(h(a_k))

论文笔记
量化
模型压缩

2019 · 05 · 06