Training Quantized Neural Networks with a Full-precision Auxiliary Module

Training Quantized Neural Networks with a Full-precision Auxiliary Module

文章鏈接

Introduction

作者認爲，量化模型性能下降的主要原因在於在訓練過程中，quantizer是不可導的，無法直接使用SGD來優化網絡。目前的解決方案可分爲兩類，一類是直接使用使用STE，或這是使用一個“鬆弛”的quantizer，另一類則是用全精度網絡來指導量化網絡訓練，如knowledge distillation。作者借鑑第二類方法，提出了一種全精度網絡+量化網絡的訓練方法，在前向推理階段，去掉全精度網絡，可得到一個精度較高的量化網絡。

Method

作者提出的網絡訓練框架如上圖所示。藍色圖部分表示爲$F$，爲量化網絡。粉紅色框記作$H$，表示全精度網絡，即爲auxiliary module。在訓練階段，兩塊的loss都會計算，全精度網絡和量化量化網絡同時進行訓練，在測試階段，$H$將會被捨棄，只保留$F$。
在上圖中，$H$是由一系列的adaptor和aggregators組成，這個輔助性的$H$從$F$中接收$P$個特徵圖的輸出$\left\{\mathbf{O}_{p}\right\}_{p=1}^{P}$,，假定$\left\{B_{1}, \ldots, B_{P}\right\}$表示生成feature map的block，對於$H$中的第$p$個輸入，作者使用一個可訓練的adaptor $\phi_{p}(\cdot)$ 對得到的特徵圖$\mathbf{O}_{p}$進行加工，並生成新的特徵圖$\phi_{p}\left(\mathbf{O}_{p}\right)$。作者這麼做的動機爲compensate the distribution discrepancy between the low-precision model and full-precision model。It ensures the quantized activations $\left\{\mathbf{O}_{p}\right\}_{p=1}^{P}$ to be compatible to the full-precision calculation in $H$。在具體實現的時候，作者使用一個$1\times 1$卷積和一個BN層進行實現。
在$H$中，從adaptor輸入的特徵圖$\phi_{p}\left(\mathbf{O}_{p}\right)$和$H$自身的特徵圖$\boldsymbol{g}_{p}$相加，然後送入$\operatorname{ReLU}(\cdot)$當中，得到$\boldsymbol{g}_{p}=\operatorname{ReLU}\left(\phi_{p}\left(\mathbf{O}_{p}\right)+\boldsymbol{g}_{p-1}\right)$關於該網絡的優化，對於訓練樣本$\left\{\boldsymbol{x}_{i}, \boldsymbol{y}_{i}\right\}_{i=1}^{N}$，主幹網絡$F$和混合精度網絡$F \circ H$，訓練的損失函數爲$\begin{aligned} \min _{\left\{\boldsymbol{\theta}^{F}, \boldsymbol{\theta}^{H}\right\}} \sum_{i=1}^{N} \mathcal{L}\left(F\left(\boldsymbol{x}_{i} ; \boldsymbol{\theta}^{F}\right), \boldsymbol{y}_{i}\right) +\mathcal{L}_{a u x}\left((F \circ H)\left(\boldsymbol{x}_{i} ; \boldsymbol{\theta}^{H}, \boldsymbol{\theta}^{F}\right), \boldsymbol{y}_{i}\right) \end{aligned}$式中，$\boldsymbol{\theta}^{F}$和$\boldsymbol{\theta}^{H}$分別表示$F$和$H$網絡參數，$\mathcal{L}$表示目標損失函數，$\mathcal{L}_{a u x}$表示輔助網絡的loss，在分類任務中，這兩個loss都設置爲交叉熵損失函數。從上式可以看出，$\boldsymbol{\theta}^{F}$是由$F$和$F \circ H$所共享，在求導的時候，作者說將兩路的梯度取一個平均(不太理解爲什麼不直接求導？)，作者給的理由是梯度方向更爲準確，論文後面的一句話也沒看懂，In other words, the full-precision module $H$ provides direct gradient for $F$ using weight sharing during back-propagation。整個訓練算法的流程如下如下所示，如下圖來看，作者確實是對梯度取了平均，即對loss取了平均。

作者還討論了這種方法和其他方法的優劣，如和knowledge distillation或者是用中間特徵圖來計算分類loss的區別，具體可參見原論文。

Experiment

在分類實驗的時候，作者將第一個卷積層和最後的FC層量化到8-bit，其他層量化到超低bit。結果如下圖


從這個結果來看，如果單看絕對精度，2-bit量化沒有特別高，而且還用DOReFa量化，效果不是很好，但是和自身相比的話，確實也說明了輔助網絡的存在能夠提升量化性能，同時，作者還在檢測任務上進行了實驗，4-bit量化基本不掉點，總來的說，應該還是有一些效果的，對於在其他量化方法的性能提升，就需要具體的實驗驗證，不過作者提供的這種訓練的思路還是可取的，缺點就是會讓訓練的時候的網絡變大，訓練變慢。