[论文解读] Dynamic Sparse Training: Find Efficient Sparse Network From Scratch With Trainable Masked Layers

Junjie Liu, Zhe Xu+3 more|arXiv (Cornell University)|May 14, 2020

Advanced Neural Network Applications参考文献 31被引用 34

一句话总结

本文引入了动态稀疏训练（DST），一种端到端的方法，通过可训练的裁剪阈值同时学习稀疏网络结构和权重，在训练过程中实现细粒度、逐步裁剪与恢复。

ABSTRACT

We present a novel network pruning algorithm called Dynamic Sparse Training that can jointly find the optimal network parameters and sparse network structure in a unified optimization process with trainable pruning thresholds. These thresholds can have fine-grained layer-wise adjustments dynamically via backpropagation. We demonstrate that our dynamic sparse training algorithm can easily train very sparse neural network models with little performance loss using the same number of training epochs as dense models. Dynamic Sparse Training achieves the state of the art performance compared with other sparse training algorithms on various network architectures. Additionally, we have several surprising observations that provide strong evidence for the effectiveness and efficiency of our algorithm. These observations reveal the underlying problems of traditional three-stage pruning algorithms and present the potential guidance provided by our algorithm to the design of more compact network architectures.

研究动机与目标

• 说明需要高效稀疏网络来减少推理时的内存与计算要求。
• 提出一个端到端的稀疏训练框架，能够同时学习权重和逐层裁剪掩码。
• 实现每个训练步骤中的细粒度、逐步裁剪与恢复，而非仅在各个训练轮之间进行。
• 通过反向传播和可训练阈值机制自动调整逐层裁剪率。
• 在 MNIST、CIFAR-10 与 ImageNet 上的多种架构中展示最先进的性能。

提出的方法

• 将裁剪表示为每一层的可训练阈值，适用于神经元/过滤器的层级裁剪。
• 使用从单元阶跃函数 S 作用于 |W| - t 获得的二进制掩码 M，以得到稀疏的 W ∘ M。
• 引入基于直通估计器的导数用于训练阈值向量 t。
• 加入稀疏正则化项 Ls = sum exp(-ti) 以鼓励更高的稀疏性。
• 用可训练的掩码层替换密集层，使得性能和结构梯度都能进行反向传播。
• 表明阈值和掩码可以在每个训练步骤更新，从而实现细粒度的裁剪与恢复。

实验结果

研究问题

• RQ1 在端到端框架中，是否可以通过可训练的裁剪阈值共同学习权重与稀疏结构？
• RQ2 在训练过程中进行逐步裁剪与恢复是否优于在稀疏训练中使用预设的裁剪计划？
• RQ3 逐层、可训练阈值如何影响最终的稀疏模式及跨架构的模型性能？
• RQ4 DST 如何为通过观察到的稀疏模式设计紧凑结构提供指导？

主要发现

• DST 在 MNIST 的 Lenet-300-100 上几乎可裁剪 98% 参数，几乎不损失性能（稀疏结果：剩余 2.48%）。
• 在 MNIST Lenet-5-Caffe 上，稀疏训练保留 1.64% 参数，接近全密度精度（稀疏：99.11%）。
• 用于序列化 MNIST 的 LSTM 模型在稀疏精度方面达到可比甚至更高的水平，且参数裁剪率超过 99%。
• 在 CIFAR-10 的 VGG-16 与 WideResNet 上，DST 在高稀疏度下优于 Sparse Momentum 与 Dynamic Sparse Reparameterization（如 VGG-16：DST 剩余 8.82% 且稀疏精度 93.93%，而其他方法剩余 10%）。
• 在 ImageNet（ResNet-50）上，DST 达到更高的 top-1/top-5 准确率，且稀疏度略高（例如剩余 ~9.87-19.24% 取决于设定）相较于基线。
• DST 显示在不同 α 值下的一致稀疏模式，表明存在逐层的冗余，并为架构设计提供了指引。