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