ReActNet：通过优化激活函数得到更高精度的二值神经网络

作者‖ Cocoon

编辑‖ 3D视觉开发者社区

3D开发者社区链接：https://mp.weixin.qq.com/s/ObAT1PJl6ZcXBhEQFOipKA

概述

论文链接： https://arxiv.org/pdf/2003.03488.pdf

代码地址： https://github.com/liuzechun/ReActNet

通常认为的二值神经网络，指模型权重以及激活值均为二进制数值的模型，被认为是最有前途的神经网络压缩方法之一。但显然，二值的表达与浮点数的表达存在着差距，而本篇所介绍的$ReActNet$，其目的就是希望在不引入额外计算开销的前提下，能够缩小这之间的差距。

论文首先构建了一个$Baseline$网络，具体的构建方式为使用不引入参数的$shortcut$修改某一浮点数网络，并将其二值化。该$Baseline$网络在精度和效率上有较好的$trade$-$off$，比大多数已有的二值网络效果都要好。

通过大量的实验，文章发现，二值网络的性能对激活函数的分布非常敏感。基于该发现，文章提出将传统的$Sign$以及$PReLU$函数进行“推广”，并将“推广”后的函数分别命名为$RSign$以及$RPReLU$，这样做的目的是希望在不引入额外计算开销的前提下，显式地学习分布的“$reshape$”和“$shift$”。

最后，文章提出了所谓的“$DistributionalLoss$”来进一步地对二值网络进行强化，目的是使其与浮点数网络学到近似的输出分布。

合并了上述想法后所攒出来的网络即命名为$ReActNet$，其性能远超出当时已有的二值网络。

方法

一比特卷积

对于权重已经一比特化的卷积层来说，其乘法运算可以被按位运算$XNOR$以及$popcoun$t操作替换，具体为：
$${\mathcal{X}}_b\ast {\mathcal{W}}_b=\mathrm{popcount}\left(\mathrm{XNOR}\left({\mathcal{X}}_b,{\mathcal{W}}_b\right)\right)$$

其中${\mathcal{W}}_b$以及${\mathcal{X}}_b$分别表示二值权重矩阵以及二值激活。具体地，权重和激活通过以下的方式进行二值化：
$$x_b=\mathrm{Sign}\left(x_r\right)=\{\begin{array}{l}+1,\text{ if }{x}_r>0\\ -1,\text{ if }{x}_r\le 0\end{array},w_b=\frac{{\Vert {\mathcal{W}}_r\Vert }_{l1}}{n}\mathrm{Sign}\left(w_r\right)=\{\begin{array}{l}+\frac{{\Vert {\mathcal{W}}_r\Vert }_{l1}}{\Vert },\text{ if }{w}_r>0\\ -\frac{{\Vert {\mathcal{W}}_r^n\Vert }_{l1}}{n},\text{ if }{w}_r\le 0\end{array}$$
其中，下标$b$和下标$r$分别指代二值和浮点数。$\frac{{\Vert {\mathcal{W}}_r\Vert }_{l1}}{n}$表示权重绝对值的平均值，其作用为作为尺度因子减少浮点数与二值之间的差距，在$XNOR-Net$（Rastegari 等，2016）中，也应用了类似的尺度因子。然而，随着后续所谓的$ReAct$操作的引入，这个尺度因子将变得不再重要。

Baseline网络

许多二值化工作选择的浮点数网络是$ResNet$，然而，该论文认为对更为紧实的轻量级网络进行二值化更具有实际意义，因此选择了$MobileNetV1$网络结构作为浮点数网络。受$Bi-Real$网络（Liu等，2018）的启发，论文也加入一些$shortcut$，而且将在$MobileNetV1$中的$3\ast 3$的逐深度卷积以及$1\ast 1$的逐点卷积块替换成了 $3\ast 3$ 和 $1\ast 1$ 的普通卷积，且并行地加入了$shortcut$， 示意图如下所示：

此外，还提出了一个新的结构来处理下采样的操作，文章希望能够在不增加额外开销的情况下进行处理。然而，文章需要输入与输出的通道数是一致的，这样才能够在不引入额外的浮点数矩阵乘法的前提下进行二值化，且只使用一个简单的$identity$ $shortcut$进行激活的传播。

正如上图中的左图所示，复制了输入通道，将两个具有相同输入的$block$进行$concat$进而解决掉输入输出通道数不一致的问题，且在$shortcut$过程中使用平均池化的操作来进行空间上的下采样。在$Baseline$中，除了第一层和最后输出的全连接层以外，都是一比特的，这样的架构设计是硬件友好的。

ReActNet

图像分类的神经网络本质是学习一个从输入图像到输出分布的映射关系。因此，文章认为一个性能良好的二值神经网络就应该学到和浮点数网络相似的输出分布。但是，二值网络的离散性限制了其像浮点数网络那样学习复杂分布的能力。为了解决这个问题，$XNOR-Net$提出计算尺度缩放因子，并将其与激活相乘的方案，其后续的工作进一步地提出了如何通过反向传播进行这些因子的更新。

与$XNOR-Net$不同，文章着重研究了激活分布，并且发现，激活分布的微小变化会极大地影响一比特$CNN$的语义特征表达。但是反过来讲，一比特的$CNN$学习分布的能力本来就有限，为了解决这样的问题，文章引入了广义的激活函数（$generalized$ $activation$ $functions$），所谓的广义激活有着可学习的系数，进而提高一比特$CNN$学习分布的能力。

分布对于一比特$CNN$的重要性

在浮点数网络中，似乎不那么关注分布，这是因为激活都是连续的浮点数，对它们来说，$reshape$或者$shift$分布并不费力气。然而，对于一比特的$CNN$来说，激活只能够选择$1$或$-1$两个值，如果输入的特征分布有一点偏移，很可能会导致特征图受到巨大的影响，如下图所示：

一比特的$CNN$显然不像浮点数的$CNN$那样对特征信息具备鲁棒性，输入特征的一点变化会导致特征引入过多不必要的背景信息或过少的特征信息，进而使得$BNN$的识别负担加大。

广义激活函数

基于以上的观察，我们提出了一个$ReAct$操作，目的是显式地$reshape$或$shift$激活的分布。$ReAct$操作将传统的$Sign$以及$PReLU$广义化，分别命名为$ReAct-Sign$以及$ReAct-PReLU$，以下将简称为$RSign$和$RPReLU$。

实际上，$RSign$相对于$Sign$来说，有着逐通道的可学习阈值，具体为：
$$x_i^b=h\left(x_i^r\right)=\{\begin{array}{ll}+1,& \text{ if }{x}_i^r>{\alpha }_i\\ -1,& \text{ if }{x}_i^r\le {\alpha }_i\end{array}$$
其中下标$i$表示不同的通道。

类似地，$RPReLU$的定义为：
$$f\left(x_i\right)=\{\begin{array}{ll}{x}_i-{\gamma }_i+{\zeta }_i,& \text{ if }{x}_i>{\gamma }_i\\ {\beta }_i\left(x_i-{\gamma }_i\right)+{\zeta }_i,& \text{ if }{x}_i\le {\gamma }_i\end{array}$$

其中，${\gamma }_i$以及${\zeta }_i$都是可学习的参数，用来控制原先的$PReLU$的形状变化，下标$i$也是指不同的通道数。

$RSign$与$Sign$，$ RPReLU$与$PReLU$的示意图分别为：

通过$RSign$以及$RPReLU$的引入，可以大幅减轻一比特$CNN$学习激活分布的压力，进而使得一比特$CNN$更为专注的学习更有意义的特征。

后续的实验结果也表明，这样的设计将大幅的提高$Baseline$的精度。

优化

$RSign$以及$RPReLU$中的参数可以和网络中的其他参数一同进行端到端的优化。

在$RSign$中的${\alpha }_i$可以通过链式法则进行求解，具体为：
$$\frac{\partial \mathcal{L}}{\partial {\alpha }_i}=\sum _{x_i^r}\frac{\partial \mathcal{L}}{\partial h\left(x_i^r\right)}\frac{\partial h\left(x_i^r\right)}{\partial {\alpha }_i}$$
其中$\mathcal{L}$表示损失函数，$\frac{\partial \mathcal{L}}{\partial h\left(x_i^r\right)}$表示更深网络的梯度。

类似地，在$RPReLU$中，梯度可以通过以下的公式进行计算：
∂ f ( x i ) ∂ β i = I { x i ≤ γ i } ⋅ ( x − γ i ) ∂ f ( x i ) ∂ γ i = − I { x i ≤ γ i } ⋅ β i − I { x i > γ i } ∂ f ( x i ) ∂ ζ i = 1 \begin{aligned} &\frac{\partial f\left(x_{i}\right)}{\partial \beta_{i}}=\mathbf{I}_{\left\{x_{i} \leq \gamma_{i}\right\}} \cdot\left(x-\gamma_{i}\right) \\ &\frac{\partial f\left(x_{i}\right)}{\partial \gamma_{i}}=-\mathbf{I}_{\left\{x_{i} \leq \gamma_{i}\right\}} \cdot \beta_{i}-\mathbf{I}_{\left\{x_{i}>\gamma_{i}\right\}} \\ &\frac{\partial f\left(x_{i}\right)}{\partial \zeta_{i}}=1 \end{aligned} ​∂βi​∂f(xi​)​=I{xi​≤γi​}​⋅(x−γi​)∂γi​∂f(xi​)​=−I{xi​≤γi​}​⋅βi​−I{xi​>γi​}​∂ζi​∂f(xi​)​=1​
其中，$\mathbf{I}$表示指示函数， I { ⋅ } = 1 \mathbf{I}_{\{\cdot\}}=1 I{⋅}​=1的条件是其括号中的不等式成立，反之则为0。

损失函数： Distributional Loss

因为文章认为如果二值网络能够和浮点数网络学到类似的分布，那么其性能就能够得到提升，所以文章特意设计了一个分布的损失函数，命名为$Distributional$ $Loss$，其数学表达为：
$${\mathcal{L}}_{\text{Distribution }}=-\frac{1}{n}\sum _c\sum _{i=1}^n{p}_c^{{\mathcal{R}}_{\theta }}\left(X_i\right)\log \left(\frac{p_c^{{\mathcal{B}}_{\theta }}\left(X_i\right)}{p_c^{{\mathcal{R}}_{\theta }}\left(X_i\right)}\right)$$
其物理含义为浮点数网络输出的分布与二值网络输出的分布之间的$KL$散度。下标$c$表示类别，$n$表示$batch$ $size$。

实验与结论

实验设计

• 数据集：

数据集为$ILSVRC12$ $ImageNet$分类数据集，相比起$MNIST$等数据集来说，更具挑战性。数据增强方法同原版$MobileNet$一致。

• 训练策略：

训练方法参照$Bi-real$中推荐的标准二值化方法，并采用了Brais等（2020）的两步走训练策略。具体地，在第一步，随机地使用二值的激活和浮点数的权重进行训练，在第二步，将第一步得到的参数作为初始值，然后再用权重和激活都是二值的网络进行$finetune$。在这两步中，都是用的$Adam$优化器，学习率调整策略为线性下降的方式，初始的学习率为$5e-4$。训练的$batchsize$为$256$，共训练$600k$次迭代。权重衰减在第一步中为$1e-5$，在第二步中为$0$。

• $Distributional$ $Loss$：

在两步中，都使用了文章所介绍的 $Distributional$ $Loss$ ，而不是原先的网络输出与标签之间的交叉熵损失。

• $OPs$计算

与Brais等（2020）的$OPs$计算方式一致，文章分别计算二进制操作$BOP$与浮点运算$FLOPs$，然后整体的$OPs$计算方法为：
$$\mathrm{O}\mathrm{P}\mathrm{s}=\mathrm{B}\mathrm{O}\mathrm{P}\mathrm{s}/64+\mathrm{F}\mathrm{L}\mathrm{O}\mathrm{P}\mathrm{s}$$

和SOTA的比较

将$ReActNet$与$SOTA$的量化与二值化方法进行比较，具体的比较结果如下表所示：

消融实验

在消融实验中，主要探究以下几个技巧在网络中所起的作用：

(1) $Block$复制以及$Concat$；

(2) $Distributional$ $Loss$；

(3) $ReAct$操作。

下表为消融实验的结果表：

• $Block$复制以及$Concat$

浮点数的$shortcut$对于$BNN$来说，是至关重要的，然而，降采样的通道数是输出通道数的两倍，这违反了输入输出通道数一致的限制，因此，为了能够加入$identity$ $shortcut$， 文章对其做了$Block$复制和$Concat$的操作。示意图见下图$(a)$：

这样的设计相对于没有$shortcut$的设计（上图$(b)$）而言，提高了$2.9$个点。相关的提升也可以见$Table2$中的第二行和第四行。

• $Distributional$ $Loss$

  由$Table2$中的结果可见，该损失大概提升了$1.4$个点。

• $ReAct$操作

  由$Table2$中的结果可见，$RSign$和$RPReLU$分别使得$Baseline$网络提升了$4.9$以及$3.6$个点。将这两个策略都加上的时候，$ReActNet$-$A$相比起$Baseline$提升了$6.9$个点，与浮点数网络的精度差距缩减到理论$3$个点以内。更有趣的是，从验证集精度的变化曲线看来，使用原版本的$PReLU$以及$Sign$使得曲线震荡的非常严重，文章认为这个现象是因为$PReLU$的斜率变化使得符号改变，进而使得后面所有的层都受其影响而导致的。此外，文章也认为这间接的验证了其假设，即一比特$CNN$对分布异常敏感。而文章所提出的$ReAct$操作则在鲁棒性和精确性上有着很好的表现。

可视化

为了更好的理解，文章将学习到的参数以及中间层激活的分布进行了可视化。

$ReAct$操作中学习参数的可视化

如下图所示，文章使用颜色条来表示每一层的学习参数，相比起传统的使用$PReLU$的二值网络来说，使用$RPReLU$所学习到的参数有正有负，更接近于浮点数网络中的分布，而传统的几乎只学到了正值。而且如下图中的$(c)$所示，$shift$系数的绝对值也较大，说明对于一比特的$CNN$而言，显式地进行$shift$也是有必要的。

中间层激活分布的可视化

在下图中，展示了$Baseline$与$ReActNet$内部的激活分布的直方图：

显然，有了$ReAct$操作后，整体的分布显得更为丰富，此外，从左到右的第二个子图的对比也能够发现，经由$ReAct$操作后，$Sign$的输出的正负更加均匀了，说明在表达二进制的特征时，更好地利用了黑白像素。

结论

文章分析了激活分布对二值化网络性能的影响，并针对该问题，提出了带有学习参数的$Sign$和$PReLU$，分别命名为$RSign$与$RPReLU$，目的是更好地接近浮点数网络的激活分布。在实验中，文章以$MobileNetV1$作为待二值化的浮点数网络，在其基础上加入了$Shortcut$等修改，并进行二值化，进而作为$Baseline$网络。在该$Baseline$网络上加入$RSign$及$RPReLU$等操作后，结果有较大程度的提升，进而说明了$ReAct$操作的有效性。此外，在希望二值网络尽可能输出与浮点数网络一样的分布的背景下，文章设计了名为$Distribution$ $Loss$ 的损失函数，实际上为二值网络输出分布与浮点数网络输出分布之间的$KL$散度，消融实验的结果表明，该损失函数也为精度的提升贡献了自己的作用。

参考文献

1. Rastegari M, Ordonez V, Redmon J, et al. Xnor-net: Imagenet classification using binary convolutional neural networks[C]//European conference on computer vision. Springer, Cham, 2016: 525-542.
2. Liu Z, Wu B, Luo W, et al. Bi-real net: Enhancing the performance of 1-bit cnns with improved representational capability and advanced training algorithm[C]//Proceedings of the European conference on computer vision (ECCV). 2018: 722-737.
3. Howard, A.G., Zhu, M., Chen, B., Kalenichenko, D., Wang, W., Weyand, T., Andreetto, M., Adam, H.: Mobilenets: Efficient convolutional neural networks for mobile vision applications. arXiv preprint arXiv:1704.04861 (2017)
4. Brais Martinez, Jing Yang, A.B.G.T.: Training binary neural networks with real-tobinary convolutions. International Conference on Learning Representations (2020)