Paper Reading||DoReFa-Net: Training Low Bitwidth CNNs With Low Bitwidth Gradients

论文(暂无投稿)：https://arxiv.org/abs/1606.06160
源码(TensorFlow)：https://github.com/tensorpack/tensorpack/tree/master/examples/DoReFa-Net

1 Motivation

1. 先进的DCNN通常参数多、计算复杂度高，这些缺点不利于嵌入式设备中的应用。
2. 以前的工作还没有成功地在向后传递期间将梯度量化为位宽小于8的离散值并保证性能不严重减少。

 　作者强调，现在的量化工作能很好地处理和加速卷积操作(全连接层)的前向传播过程，如公式 (1) 的点积型位运算(Bitwise Operation)：

 　其中 $\mathtt{x}$ 和 $\mathtt{y}$ 为编码进制的比特向量， $bitcount$ 用于计算一个比特向量里1的个数。当然当比特向量的每一位能表示 $-1,1$ 时，位运算也可以表示为：

 　但是在反向传播过程中，梯度的格式仍为32-bit，于是训练过程中的梯度计算会耗占时间。

2 Method

 　作者注意到，虽然权重和激活可以确定量化，但梯度只能随机量化。

2.1 Using Bit Convolution Kernels In Low Bitwidth Neural Network

 　设 $\mathtt{x}$ 为 $M$-bit 定点整形的序列， 即 $\mathtt{x}={\sum }_{m=0}^{M-1}{c}_m(\mathtt{x})2^m$ 且 { c m ( x ) } m = 0 M − 1 \{c_{m}(\bm{\mathtt{x}})\}_{m=0}^{M-1} {cm​(x)}m=0M−1​ 构成比特向量； $\mathtt{y}$ 为 $K$-bit 定点整形的序列， $\mathtt{y}={\sum }_{k=0}^{K-1}{c}_k(\mathtt{y})2^k$ 且 { c k ( y ) } k = 0 K − 1 \{c_{k}(\bm{\mathtt{y}})\}_{k=0}^{K-1} {ck​(y)}k=0K−1​ 构成比特向量，那么对于低比特的定点整形来说有位操作：

 　上式的计算复杂度为 $O(MK)$ ，即取决于位宽。

2.2 Straight-Through Estimator

 　为了解决量化过程几乎处处不可导，作者也引进了 STE 的思想。一个简单的例子是将 STE 作用于伯努利抽样 $p\in [0,1]$ ：

 　就能使得原本不可导的前向传播，在反向传播时变得可导。其中 $c$ 为目标函数。那么将STE应用于量化过程中，设实数输入 $r_i\in [0,1]$ ， $k$-bit 输出 $r_o\in [0,1]$ ，就有 ${\mathsf{q}\mathsf{u}\mathsf{a}\mathsf{n}\mathsf{t}\mathsf{i}\mathsf{z}\mathsf{e}}_{\mathbf{k}}$ ：

2.3 Low Bitwidth Quantization of Weights

 　当对 Weight 进行二值化操作时，作者定义：

 　其中 $sign(r_i)=2{\mathbb{I}}_{r_i\ge 0}-1$ 能返回 -1 或 1， ${\mathbf{E}}_F(|r_i|)$ 为整个 Weight 矩阵的绝对值的均值。而当对 Weight 进行 $k$-bit 表征量化($k>1$)时，有：

 　其中 $tanh$ 将 Weight 钳制到范围 $[-1,1]$ 。 $\frac{tanh(r_i)}{2max(|tanh(r_i)|)}+\frac{1}{2}$ 的范围为 $[0,1]$ ，需要注意此时的 $max$ 是取整个 Weight 张量里的最大值。那么最后前向传播量化后的范围变成 $[-1,1]$ 。

2.4 Low Bitwidth Quantization of Activations

 　设上一层的输出已经经过一个有界函数 $h$ 的作用来保证 Activations $r\in [0,1]$ ，那么有量化：

2.5 Low Bitwidth Quantization of Gradients

 　作者发现，随机量化对梯度来说更有效。首先需要注意的是，梯度是无边界的，可能比激活具有更大的数值范围。于是作者初步设计方案如下：

 　其中 $\mathrm{d}r=\frac{\partial c}{\partial r}$，是某一层的输出 $r$ 的后向传播梯度。而 $max$ 取得是梯度张量的各个维度，出了 Mini-batch 维度(也就是说Mini-batch里的每个实例都有自己独立的缩放因子)。
 　为了补偿梯度量化所带来的潜在，作者又引入额外的噪声函数 $N(k)=\frac{\sigma }{2^k-1}$ ， $\sigma \sim Uniform(-0.5,0.5)$ 。这个人为的随机噪声反而起到关键作用，提高模型性能。于是最终梯度量化公式如下：

 　需要注意的是，上式只在反向传播中作用，作者定义了每个卷积层输出后接的 STE 构造方法：

2.6 The Algorithm for DoReFa-Net

2.7 First and the Last Layer

 　对于第一层，其输入通常为图像，其中可能包含 8-bit 特征。对于最后一层，由于输出的是 one-hot 估计向量，对 bit 也有一定要求。
 　在绝大多数情况下，作者的实验不对第一层及其输入、最后一层及其输出进行量化，但是这两层的输出会通过梯度量化器，对梯度进行量化。

2.8 Reducing Run-time Memory Footprint by Fusing Nonlinear Function and Rounding

 　实际上作者的网络并没有被完全量化，因为其中的一些中间过程涉及全精度的储存和计算，如 Algorithm 1 里的 Step 3、Step 4、Step 6都会储存全精度值。作者将这几步融合为同一操作 $a_k^b=f_{\alpha }(h({\tilde{a}}_k))$ ，就避免了中间过程的产生。类似的，Step 11、Step 12、Step 13 也能 Fuse 到一块。
 　在有 Pooling 层的情况下，上述操作会更复杂。且若 Pooling 是 Max-pooling 时，可以进行如下 Fuse 操作：

3 Experiment

 　Table 1 中，Model A 是作者自己构造的一个模型，输入为40x40 Image，有 7 个卷积层和 1 个全连接层。而 Model B，C，D 分别减少所有层的通道数的 50%，75%，87.5%。

 　Table 2 中，Initialize 意为使用预训练模型。

个人思考

 　需要注意，作者讨论模型结构时是忽略了 BN 层的。年限较早了，尝试对梯度进行量化是本文章最主要的创新点。
 　有一个问题，作者的 Weight 伪量化，量化前和量化后的范围是不一样的？这会对模型产生影响吗？