ICML 2022 | 基于特殊激活函数与额外跳跃连接的稀疏网络训练改进算法

2023-01-30 12:01

©PaperWeekly 原创 · 作者 | GlobalTrack

已有很多研究关注于使用网络剪枝算法寻找可以部署于一般硬件的稀疏神经网络。一般的方法获得稀疏掩码后使用于一般密集网络训练相同的策略训练剪枝后的稀疏网络。现在还不清楚似乎是否针对密集神经网络设计的训练策略对训练稀疏神经网络是最优的。

与之前大部分剪枝算法关注于寻找最优的稀疏掩码不同（主要观点是好的稀疏掩码训练的更好），本文关注于改善稀疏神经网络的训练过程。之前研究发现稀疏学习可能会遇到较差的梯度流动和混乱的优化轨迹，为了解决这两个问题，本文指出设计新的稀疏学习策略是必要的。

论文标题：

Training Your Sparse Neural Network Better with Any Mask

论文链接：

https://proceedings.mlr.press/v162/jaiswal22a/jaiswal22a.pdf

代码链接：

https://github.com/VITA-Group/ToST

方法

本文给出两个策略：幽灵软神经元（Ghost Soft Neuron，GSw）与幽灵跳跃连接（Ghost Skip Connection，GSk）。

1.1 软神经元策略

稀疏度较高的稀疏学习由于移除了大部分连接，容易遇到层崩坏（Layer Collapse）问题。这使得稀疏网络由于梯度不能有效后传，网络几乎不能训练。

本文指出一个可能的原因是常用的 ReLU 激活层是针对密集网络设计的。ReLU 激活函数的梯度会突然降到 0。ReLU 激活函数的这种不平滑特性是稀疏学习的一种阻碍。

针对 ReLU 不适合稀疏学习的问题，本文提出将 ReLU 暂时替换为 Swish 和 Mish 激活函数的策略。Swish 和 Mish 是光滑非单调的激活函数。非单调属性允许了小的负输入的梯度，可以得到更稳定的梯度流。

1.2 跳跃连接策略

层崩坏也可能是梯度流动的阻碍。跳跃连接常用于避免梯度消失的问题。受此启发，本文提出在稀疏训练时通过添加临时跳跃连接。

新加的跳跃连接输入在 3*3 卷积前，输出位于激活函数前。

另外一个问题是是否需要将改变激活函数策略和添加跳跃连接策略保留到训练完成。实验中发现可以略微提升精度。但这回改变原始骨干网络的结构并会增加额外的硬件延迟。本文的实验中坚持原先模型的结构。

1.3 其他调整

重缩放初始化：本文指出之前提出的初始化策略不适合稀疏学习的要求。本文的初始化策略受 Gradinit: Learning to initialize neural networks for stable and efficient training 启发采用保留原始稀疏掩码的初始化，仅仅通过一个学习得到的尺度系数重缩放。本文指出受益于 BatchNorm 层，本文的重缩放初始化策略不会损害原始的初始化过程。

标签平滑：标签平滑（Label Smoothing, LS）改变标签为均匀分布硬标签的混合：，是平滑系数。本文方法将标签平滑嵌入稀疏掩码训练过程中。

实验

表 2 给出了各种剪枝算法与本文 ToST 算法组合后在 ResNet-32、ResNet-56 模型与 CIFAR 数据集上的性能比较。这里选择了初始化剪枝经典算法 SNIP, Grasp, SynFlow，基于 IMP 的彩票假说算法 LTH 与随机剪枝算法。稀疏度范围 {90%，95%，98%}。本文提出的 ToST 算法在所有稀疏度和数据集条件下都超过了所有比较的剪枝算法。