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$个特征图的输出 { O p } p = 1 P \left\{\mathbf{O}_{p}\right\}_{p=1}^{P} {Op​}p=1P​,，假定 { B 1 , … , B P } \left\{B_{1}, \ldots, B_{P}\right\} {B1​,…,BP​}表示生成feature map的block，对于$H$中的第$p$个输入，作者使用一个可训练的adaptor ${\varphi }_p(\cdot )$ 对得到的特征图${\mathbf{O}}_p$进行加工，并生成新的特征图${\varphi }_p\left({\mathbf{O}}_p\right)$。作者这么做的动机为compensate the distribution discrepancy between the low-precision model and full-precision model。It ensures the quantized activations { O p } p = 1 P \left\{\mathbf{O}_{p}\right\}_{p=1}^{P} {Op​}p=1P​ to be compatible to the full-precision calculation in $H$。在具体实现的时候，作者使用一个$1\times 1$卷积和一个BN层进行实现。
在$H$中，从adaptor输入的特征图${\varphi }_p\left({\mathbf{O}}_p\right)$和$H$自身的特征图${\mathbf{g}}_p$相加，然后送入$\mathrm{ReLU}(\cdot )$当中，得到$${\mathbf{g}}_p=\mathrm{ReLU}\left({\varphi }_p\left({\mathbf{O}}_p\right)+{\mathbf{g}}_{p-1}\right)$$关于该网络的优化，对于训练样本 { x i , y i } i = 1 N \left\{\boldsymbol{x}_{i}, \boldsymbol{y}_{i}\right\}_{i=1}^{N} {xi​,yi​}i=1N​，主干网络$F$和混合精度网络$F\circ H$，训练的损失函数为 min ⁡ { θ F , θ H } ∑ i = 1 N L ( F ( x i ; θ F ) , y i ) + L a u x ( ( F ∘ H ) ( x i ; θ H , θ F ) , y i ) \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} {θF,θH}min​i=1∑N​L(F(xi​;θF),yi​)+Laux​((F∘H)(xi​;θH,θF),yi​)​式中，${\mathbf{\theta }}^F$和${\mathbf{\theta }}^H$分别表示$F$和$H$网络参数，$\mathcal{L}$表示目标损失函数，${\mathcal{L}}_{aux}$表示辅助网络的loss，在分类任务中，这两个loss都设置为交叉熵损失函数。从上式可以看出，${\mathbf{\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量化基本不掉点，总来的说，应该还是有一些效果的，对于在其他量化方法的性能提升，就需要具体的实验验证，不过作者提供的这种训练的思路还是可取的，缺点就是会让训练的时候的网络变大，训练变慢。