DoReFa-Net：使用低位宽梯度训练低位宽卷积神经网络

论文地址：https://arxiv.org/abs/1606.06160
官方Tensorflow代码：https://github.com/tensorpack/tensorpack/tree/master/examples/DoReFa-Net
PyTorch代码链接：https://github.com/XJTUWYD/DoReFa_Cifar10

摘要

我们提出了DoReFa-Net——它是一种使用低位宽参数梯度来训练低位宽权重和激活值的卷积神经网络的方法。尤其，在反向传播阶段，参数梯度可以在传递到下一层卷积层之前被随机量化到低位宽。由于前向/反向阶段卷积都是在对低位宽权重和激活值/梯度上操作的，这样的话DoReFa-Net可以使用低位宽卷积核来加速训练和推理。而且，低位卷积可以在CPU、FPGA、ASIC和GPU上高效实现，因此DoReFa-Net打开了一扇在这些硬件上加速训练低位宽神经网络的大门。我们在SVHN和ImageNet数据集上的实验证明了DoReFa-Net可以达到与32位参数相当的精度。举个例子，从AlexNet派生出来的 1位权重、2位激活值、6位梯度的 DoReFa-Net AlexNet模型从头开始训练可以在ImageNet上达到 46.1%的top-1精度。

1. Introduction

据我们所知，无论是xnornet还是bnn，都没有做反向传播阶段梯度的量化，之前没有任何工作可以在反向传播阶段将梯度量化到8位以下而仍旧保持相当的预测精度。在BNN和XNOR-Net中，虽然权重是二值化的，但是梯度仍旧是全精度浮点数，因此在反向传播时反卷积依旧是 1位数和32位数之间的运算，这导致BNN和XNOR-Net的训练时间主要花在反向阶段。

这篇论文的贡献主要在于：

1. 我们泛化了二值化神经网络的方法创造了DoReFa-Net，它可以拥有任意位宽的权重、激活之和梯度。由于前向传播和方向传播都是对低精度的数进行操作，所以DoReFa-Net使用位卷积核来加速训练过程中的前向和反向；
2. 由于位卷积可以在CPU、FPGA、ASIC和GPU上高效实现，DoReFa-Net打开了一扇在这些硬件上加速训练低位宽神经网络的大门。尤其是FPGA和ASIC的功效很高，可以减少低位宽网络训练的功耗；
3. 我们探索了权重、激活值和梯度的配置空间，例如使用1位权重、1位激活值和2位梯度可以在SVHN数据集上达到93%的精度。在我们的实验中，梯度通常需要比激活值更大的位宽，而激活值通常需要比权重更大的位宽，这样能保证与32位浮点参数相比精度不会掉很多。我们将我们的方法命名为"DoReFa-Net"；

2. DoReFa-Net

在本节我们将DoReFa-Net公式化，提出了一种方法来训练低位宽权重、激活值和梯度的方法，我们注意到权重和激活值都是确定性量化的，而梯度需要随机量化。

我们首先强调怎样在DoReFa-Net中开发位操作卷积核，然后阐述了具体的量化权重、激活值和梯度的方法。

2.1 在低位宽神经网络中使用位卷积核

1位的点积核计算为：

上面的式子同样也可以用在计算低位宽定点整数之间的乘加运算。假定 x是一个M位定点整数序列集合，$x=\sum_{m=0}^{M-1}c_{m}(x)$，y是一个K位定点整数序列集合，$y=\sum_{k=0}^{K-1}c_{k}(y)$，这里的$(c_{m}(x))_{m=0}^{M-1}和(c_{k}(y))_{k=0}^{K-1}$都是位向量，x与y的点积可以由位操作计算：

在上面的等式中，计算复杂度为$MK$，与位宽成正比。这里应该如何理解$c_{m}(x)与c_{k}(y)$呢？假定现在输入和权重都是8位整数，output feature map上每一个点都是27个输入和27个权重的乘加和（卷积核为3x3，通道数也为3），那么这里的$c_{m}(x)$应该是27个8位输入整数中第m位的序列集合，里面包含有27个元素，$c_{m}(x)_{i}$应该为27个8位输入整数中第i个整数的第m位。

2.2 直通估计器

使用直通估计器（STE）的理由可以通过一个简单实例进行说明。设有一个简单的阈值函数——ReLU 函数，即 f(x) = max(0,x)。此外，设网络一开始就有某套初始权重。这些 ReLU 的输入（乘上了权重的信号）可以是负数，这会导致 f(x) 的输出为 0。对于这些权重，f(x) 的导数将会在反向传播过程中为 0，这意味着该网络无法从这些导数学习到任何东西，权重也无法得到更新。STE 的概念也由此而来。STE 会将输入的梯度设置为一个等于其输出梯度的阈值函数，而不管该阈值函数本身的实际导数如何。

一个简单的例子是在伯努利分布采样中定义的STE为：

这里的c是目标函数，由于从伯努利分布中采样是一个不可微分的过程，$\frac{\partial q}{\partial p}$没定义，因此反向传播中不能由链式法则直接计算出$\frac{\partial c}{\partial p}$，然而由于q和p的期望相同，我们可以使用定义好的梯度$\frac{\partial c}{\partial q}$对$\frac{\partial c}{\partial p}$做近似，并且构建了一个如上所示的STE，换句话来讲，STE给出了一个对$\frac{\partial q}{\partial p}$的定义。

在我们的工作中广泛使用的STE是量化器——将一个真实的浮点输入$r_{i}\in [0, 1]$量化为k位输出$r_{0}\in [0, 1]$，定义的STE如下：
很明显量化器STE的输出q是一个由k位表示的真实数，由于$r_{o}$是一个k位定点整数，卷积计算可以由等式（3）高效执行，后面跟着正确的缩放即可。

2.3 权重的低位宽量化

在本节中，我们详细描述了权重如何做到k位量化的。
在之前的工作中，STE被用来做二值化权重，比如在BNN中，权重被下面的STE二值化：

在XNOR-Net中，权重按照下面的STE二值化，不同之处在于权重在二值化之后进行了缩放：

在XNOR-Net中，缩放因子$E_{F}(|r_{i}|)$是对应卷积核的权重绝对值均值。理由是引入这个缩放因子将会增加的权重的表达范围，同时仍然可以在前向传播时卷积做位运算。然而，通道级的缩放一因子同样带来一个问题——反向传播时无法在梯度和权重之间做位卷积操作。因此，在我们的实验中，我们使用一个常量缩放因子来替代通道级缩放。在本篇论文中，我们对于所有的二值化权重使用同一个缩放因子：

当$k>1$时，我们使用k位表达的权重，我们将STE $f_{w}^{k}$应用到权重如下：

这里在量化到k位之前，我们先使用 tanh 将权重限制在 $[-1,1]$ 之间。通过$\frac{tanh(r_{i})}{2max(|tanh(r_{i})|)}+\frac{1}{2}$将数值约束在$[0,1]$之间，最大值是相对于全层权重而言的。然后通过 $quantize_{k}=\frac{1}{2^k-1}round((2^k-1)r_{i})$将浮点数转换为k位定点数，范围为$[0,1]$，最后通过映射变换将$r_{0}$约束到$[-1,1]$。

需要注意的是，当k=1时，等式9不同于等式7，它提供了一个不同的二值化权重的方法，然而，我们发现在实验中这种区别不重要。

2.4 激活值的低位宽量化

接下来我们将详细描述我们如何获得低位宽的激活值，由于它是下一层卷积的输入，所以对能否使用位操作卷积替代浮点卷积至关重要。

在BNN和XNOR-Net中，激活值和权重二值化的方式相同。然而，我们在按照XNOR-Net中二值化激活值的方式量化激活值时并没有复现他们的结果，同时BNN中二值化的方式造成AlexNet在ImageNet上精度直接掉了几个百分点。因此，我们将STE应用在每一层的激活值$r$上，在这里我们假设前层的输出已经过了一个有界函数$h$，可以确保$r\in [0,1]$。在DoReFa-Net中，激活值$r$量化到k位的方式可以简单表述为：

2.5 梯度的低位宽量化

我们已经证明了确定性量化可以产生低位宽的权重和激活值。然而，我们发现随机量化对于低位宽梯度是十分必要的，这也和Gupta等人在2015做的关于16位权重和梯度的实验一致。

为了将梯度量化到低位宽，保持梯度的无界是非常重要的，同时梯度应该比激活值的范围更广。回顾等式（11），我们通过可微分的非线性激活函数将激活值约束到了 [0,1]，然而，这种构造不存在梯度，因此我们对梯度设计了如下的k位量化方式：

这里的$dr=\frac{\partial c}{\partial r}$是对一些层的输出r的导数，最大值是对梯度张量所有维度（除了batch size）的统计，然后在梯度上用来放缩变化将结果映射到 [0,1] 之间，然后在量化之后又放缩回去。

为了更进一步地补偿梯度量化引入的潜在偏差，我们引入了额外的函数 $N(k)=\frac{\sigma}{2^k-1}$，这里的$\sigma \sim Uniform(-0.5,0.5)$。因此噪声具有和可能的量化误差相同的幅值。我们发现人工噪声对性能的影响很大，最后，我们做k位梯度量化的表达式如下：

梯度的量化仅仅在反向传播时完成，因此每一个卷积层的输出上的STE为：

2.6 DoReFa-Net训练算法

我们给出了DoReFa-Net的训练算法——Algorithm 1，假设网络具有前馈线性拓扑，像BN层、池化层这样的细节在这里不详细展开。要注意的是，所有昂贵的操作如 forward，backward_input，backward_weight（无论是卷积层还是全连接层），都是在低位宽上做的。通过构造，在这些低位宽数字和定点整数之间总是存在仿射映射的，因此，所有昂贵贵的操作都可以通过定点整数之间的点积等式（3）来加速。

2.7 第一层和最后一层

在DCNN的所有层中，第一层和最后一层相对于其他层最为不同，因为它们是网络的输入和输出。对于第一层，输入经常是一幅图片，每个像素点都是8位，另一方面，输出层通过产生one-hot向量，这个和定义中的位向量很接近，一个有趣的问题是由于这些差异性导致了当第一层和最后一层做低位宽量化时会呈现出不同的行为。

在Han最近的工作中，当将网络的权重矩阵变为稀疏矩阵时，相同稀疏比例下，第一层卷积相比于其他卷积层对预测精度的影响更大。基于这样的直觉，同时观察到第一层的通道数很小，所以计算量占比很小，所以在接下来的实验中我们不量化第一个卷积层（除非有特别说明）。然而，第一层的输出是会被量化到低位宽的，因为它将作为下一层卷积的输入。

类似地，当网络的输出类别数目较小时，为了保证预测精度，我们通常不量化最后一个全连接层，但是最后一个全连接层的梯度是要被量化的。

2.8 通过融合非线性函数和round来减少运行时的内存占用

低位宽神经网络的一个重要动机就是节省推理时的运行内存占用，Algorithm 1的一个简单实现就是以全精度存储激活值$h_{a_{k}}$，它将会在运行时占用大量内存。尤其地，如果h涉及到浮点运算，那么非位操作的数量是不可忽视的。

比如融合Step 3、Step 4、Step 6这三个步骤，融合Step 11、Step 12和Step 13都将可以节省中间变量需要的内存。

3. 实验结果

3.1 配置空间探索

3.2 ImageNet

3.3 训练曲线

3.4 Making First And Last Layer Low Bitwidth