[论文解读] Efficient Neural Network Training via Forward and Backward Propagation Sparsification

Xiao Zhou, Weizhong Zhang+3 more|arXiv (Cornell University)|Nov 10, 2021

Domain Adaptation and Few-Shot Learning参考文献 51被引用 15

一句话总结

本文提出了一种针对深度神经网络的完全稀疏训练方法，通过解耦权重与结构参数的更新，实现了高效的前向传播与反向传播。该方法引入了一种仅需两次前向传播的方差缩减策略梯度估计器，实现了完全稀疏的反向传播，并相比先前方法将训练速度提升了整整一个数量级。

ABSTRACT

Sparse training is a natural idea to accelerate the training speed of deep neural networks and save the memory usage, especially since large modern neural networks are significantly over-parameterized. However, most of the existing methods cannot achieve this goal in practice because the chain rule based gradient (w.r.t. structure parameters) estimators adopted by previous methods require dense computation at least in the backward propagation step. This paper solves this problem by proposing an efficient sparse training method with completely sparse forward and backward passes. We first formulate the training process as a continuous minimization problem under global sparsity constraint. We then separate the optimization process into two steps, corresponding to weight update and structure parameter update. For the former step, we use the conventional chain rule, which can be sparse via exploiting the sparse structure. For the latter step, instead of using the chain rule based gradient estimators as in existing methods, we propose a variance reduced policy gradient estimator, which only requires two forward passes without backward propagation, thus achieving completely sparse training. We prove that the variance of our gradient estimator is bounded. Extensive experimental results on real-world datasets demonstrate that compared to previous methods, our algorithm is much more effective in accelerating the training process, up to an order of magnitude faster.

研究动机与目标

• 解决现有稀疏训练方法在稀疏前向传播的同时仍需密集反向传播所导致的效率低下问题。
• 通过重新思考结构参数的梯度估计方式，实现完全稀疏的前向与反向传播。
• 开发一种可扩展、内存高效的训练方法，支持更大、过参数化的网络训练，且不增加内存开销。
• 在PyTorch和TensorFlow等标准深度学习框架中，实现显著的计算加速。

提出的方法

• 将稀疏神经网络训练建模为全局稀疏性约束下的约束优化问题，联合优化权重与结构参数（伯努利掩码概率）。
• 将训练过程解耦为两个步骤：通过标准链式法则反向传播进行权重更新（通过掩码结构实现稀疏化），并通过一种新型方差缩减策略梯度估计器（VR-PGE）进行结构参数更新。
• 采用VR-PGE，仅通过两次前向传播即可估计结构参数的梯度，从而消除对反向传播的依赖，实现完全稀疏计算。
• 通过约束伯努利参数之和来控制全局稀疏性，确保训练过程中网络规模保持一致。
• 在前向与反向传播中均使用稀疏子网络，将FLOPs降低至密集基线的约ρ²倍，其中ρ为通道保留率。
• 引入一种策略，减少子网络采样的频率（例如每50次迭代一次），实现模型在CPU上完整存储，子网络在GPU上运行，从而降低内存压力。

实验结果

研究问题

• RQ1 我们能否在不依赖密集反向传播的前提下，实现在深度神经网络中完全稀疏的训练（包括前向与反向传播）？
• RQ2 方差缩减策略梯度估计器能否替代基于链式法则的梯度估计方法用于结构参数，同时保持训练的稳定性和收敛性？
• RQ3 完全稀疏训练是否能在实践中带来显著的计算加速，尤其是在标准深度学习框架中？
• RQ4 所提出的方法能否在不带来高昂内存或计算成本的前提下，扩展至更大规模的网络？

主要发现

• 在VGG-19上，该方法在仅保留0.8%参数（ρ = 0.8）的情况下，训练计算时间最高可加速7.41倍。
• 在CIFAR-10上使用ResNet32时，该方法在保留3.06%参数的情况下实现4.50倍加速，同时保持90.82%的验证准确率。
• 所提出的VR-PGE估计器的方差被严格证明有界，确保了梯度更新的稳定与可靠。
• 该方法在所有稀疏度水平下均保持高准确率，即使在极低的通道保留率下也仅有可忽略的性能下降（例如VGG-19上从93.84%降至90.82%）。
• 实验结果表明，该方法显著快于先前的稀疏训练基线，加速倍数最高可达一个数量级。
• 将子网络采样频率降低至每50次迭代一次，可保持最终准确率，从而实现高效的内存管理，并支持更大模型的可扩展性。