Source-linked AI summary
Sparse Networks from Scratch: Faster Training without Losing Performance
Tim Dettmers, Luke Zettlemoyer
TL;DR
Training sparse neural networks without repeated dense prune-and-retrain cycles remains challenging. The paper introduces sparse momentum, which maintains sparsity from random initialization by using momentum to prune, redistribute, and regrow weights. Across MNIST, CIFAR-10, and ImageNet, it reaches dense-level performance, achieves state-of-the-art sparse results, and provides speedups up to 5.61x.
Problem
Prior sparse configurations require starting from dense networks and repeatedly pruning and retraining, making sparse training expensive.
Method
Sparse momentum maintains sparse weights from random initialization and uses exponentially smoothed gradients to guide pruning, layer redistribution, and regrowth.
Results
Sparse momentum outperforms other sparse algorithms across MNIST, CIFAR-10, and ImageNet, rivals dense performance, and provides speedups up to 5.61x.
Takeaways & Limitations
Sparse momentum provides a single-run approach to sparse training that can match dense performance while reducing training time.
Abstract
from arXiv · showhide
We demonstrate the possibility of what we call sparse learning: accelerated training of deep neural networks that maintain sparse weights throughout training while achieving dense performance levels. We accomplish this by developing sparse momentum, an algorithm which uses exponentially smoothed gradients (momentum) to identify layers and weights which reduce the error efficiently. Sparse momentum redistributes pruned weights across layers according to the mean momentum magnitude of each layer. Within a layer, sparse momentum grows weights according to the momentum magnitude of zero-valued weights. We demonstrate state-of-the-art sparse performance on MNIST, CIFAR-10, and ImageNet, decreasing the mean error by a relative 8%, 15%, and 6% compared to other sparse algorithms. Furthermore, we show that sparse momentum reliably reproduces dense performance levels while providing up to 5.61x faster training. In our analysis, ablations show that the benefits of momentum redistribution and growth increase with the depth and size of the network. Additionally, we find that sparse momentum is insensitive to the choice of its hyperparameters suggesting that sparse momentum is robust and easy to use.
1 Introduction
Sparse learning trains networks with sparse weights from random initialization, avoiding repeated dense training and pruning while targeting dense-level performance. Sparse momentum uses momentum to prune, redistribute, and regrow weights, with strong results across MNIST, CIFAR-10, and ImageNet-1k.
- Motivation: Sparse configurations can match dense-network errors and train faster, but prior approaches found them through expensive prune-and-retrain cycles.The paper targets sparse training without starting from a dense network or repeatedly retraining it.
- Contribution: Sparse momentum trains sparse networks from random initialization in a single run while maintaining sparsity throughout training.Its cycle prunes small weights, redistributes weights across layers, and regrows high-momentum missing connections.
- Contribution: Sparse momentum uses exponentially smoothed gradients to identify layers and missing connections that can reduce error efficiently.Layer redistribution follows mean momentum magnitude, while within-layer growth targets zero-valued weights with high momentum magnitude.
- Evaluation: MNIST, CIFAR-10, and ImageNet-1k experiments establish state-of-the-art sparse performance against compression and sparse-training methods.The comparisons include both methods that create sparsity from dense models and methods that maintain sparse weights throughout training.
- Evaluation: 2.74x–5.61x speedups are estimated for sparse convolution, while dense convolution yields 1.07x–1.36x speedups compared with dense baselines.The estimates concern training sparse convolutional networks to dense performance levels on CIFAR-10.
2 Related Work
Prior sparse-network research largely produced sparse models by pruning dense networks and retraining them, while newer work also investigated sparse configurations and sparse training throughout optimization. These approaches motivate sparse momentum’s single-run alternative.
- From Dense to Sparse Neural Networks: Earlier pruning research used derivative-based criteria or heuristics to remove weights while preserving efficient post-pruning training.The related-work lineage includes second-order pruning and training-efficiency heuristics.
- From Dense to Sparse Neural Networks: The influential train-dense, prune, retrain cycle became a foundation for later compression methods and extensions.Extensions addressed recurrent models, continuous pruning, pruning-cost optimization, layerwise pruning, growth-pruning cycles, and soft weight-sharing.
- Interpretation and Analysis of Sparse Neural Networks: Winning lottery tickets showed that sparse initializations can reach dense-like predictive performance and train equally fast, but finding them requires multiple prune-and-retrain cycles from dense networks.Follow-up work sought faster ways to find such configurations.
- Sparse Neural Networks Throughout Training: Sparse momentum differs by reaching dense performance from random initialization in one training run while accelerating training.This contrast is stated against methods that derive sparse configurations from dense networks.
- Sparse Neural Networks Throughout Training: Sparse-throughout-training methods use prune-redistribute-regrowth cycles, including Bayesian posterior sampling in DEEP-R and simpler sparse evolutionary training approaches.The passage characterizes DEEP-R as theoretically rigorous but computationally expensive and challenging for large networks and datasets.
3 Method
Sparse momentum maintains sparsity through pruning, layer redistribution, and regrowth guided by exponentially smoothed gradients. It applies these operations after epochs, supports practical sparse-training experiments, and includes safeguards for layer-capacity edge cases.
- 3.1 Sparse Learning: Sparse learning is defined as training networks that maintain sparsity throughout training while matching dense-network predictive performance.The method separates the problem into pruning, redistribution across layers, and regrowing weights.
- 3.2 Sparse Momentum: At each epoch’s end, sparse momentum normalizes layer momentum, prunes the smallest-magnitude weights, redistributes them across layers, and regrows high-momentum missing connections.The figure specifies p = 50% pruning in its illustrated cycle and decays p over training.
- 3.2 Sparse Momentum: Exponentially smoothed gradients reduce mini-batch oscillations and emphasize recent gradients when estimating consistently error-reducing weights.Momentum magnitudes also estimate how efficiently layers and zero-valued weights could reduce error.
- 3.2 Sparse Momentum: Layer regrowth is proportional to each layer’s mean momentum magnitude, while pruning removes low-magnitude weights and growth selects zero-valued weights with largest momentum.This implements the paper’s divide-and-conquer treatment of pruning, redistribution, and regrowth.
- 3.2 Sparse Momentum: The algorithm redistributes excess regrowth requests when a layer lacks capacity and avoids pruning important dense layers that cannot grow further.These are explicit edge-case rules outside the simplified algorithm presentation.
- Training Procedure: Cosine prune-rate decay reaches zero on the last epoch, while sensitivity analysis reports similar performance for cosine and linear schedules and insensitivity to starting prune rate.The commonly used settings were p = 0.2 and α = 0.9, selected on MNIST and CIFAR-10.
- Experimental Setup: Experiments include partially dense and fully sparse settings for modified AlexNet, VGG16, LeNet-5, and ResNet-50 architectures.The partially dense setup keeps selected convolutional or final layers dense, whereas the fully sparse setup begins with no dense layer.
- Experimental Setup: Masked weights enforce sparsity after each gradient update, and the implementation wraps PyTorch networks with modular pruning, redistribution, and growth components.The software adapts PyTorch networks to sparse momentum with fewer than 10 lines of code, while sparse tensors are simulated using masks.
4 Results
Across MNIST, CIFAR-10, and ImageNet, sparse momentum matches or outperforms competing sparse methods while retaining dense-level performance at substantially lower weight counts. On CIFAR-10, the resulting sparse networks offer notable estimated training speedups, especially with sparse convolution.
- Sparse momentum is the only method with consistently strong performance across both evaluated MNIST LeNet models.It matches the LeNet-5 Caffe dense baseline using 8% of the weights.
- 8% of weights matches the LeNet-5 Caffe dense baseline, while 5% matches dense performance for CIFAR-10 VGG16-D.On CIFAR-10, sparse momentum also outperforms SNIP across all evaluated models.
- Sparse momentum outperforms competing methods on MNIST and ImageNet Top-1 accuracy, while remaining competitive in fully sparse ImageNet training.On CIFAR-10, it is better than dynamic sparse in 4 of 5 cases, although confidence intervals often overlap in the partially dense setup.
- 4.1 Speedups and Weights Needed for Dense Performance Levels: Sparse-convolution speedups are estimated from saved FLOPS, whereas dense-convolution speedups use removable all-zero convolutional channels.The dense-convolution estimate highlights the gap between currently available dense algorithms and efficient sparse training.
- 4.1 Speedups and Weights Needed for Dense Performance Levels: 5.61x is the highest estimated sparse-convolution speedup for Wide Residual Networks on CIFAR-10.AlexNet achieves about 3.0x speedup, while dense-convolution speedups are much lower; the sparse momentum procedure itself corresponds to a 0.973x±0.029x slowdown relative to a dense baseline.
5 Analysis
The analysis examines how momentum-based redistribution and growth contribute to sparse-network performance, and tests sensitivity to prune-rate and momentum settings.
- Ablation Analysis: Redistribution becomes increasingly important as network size grows, with omitted redistribution raising ImageNet Top-1 error from 3.42% to 9.71%.
- Ablation Analysis: Momentum growth substantially improves ResNet-50 performance over random growth on ImageNet, whereas smaller networks can perform well with redistribution and random growth.
- Ablation Analysis: Random growth without redistribution cannot find good weights, suggesting that larger networks require smarter growth strategies for good performance.
- Sensitivity Analysis: Sparse momentum is robust to prune rates from 0.2 to 0.4, with only slight improvement at smaller rates; cosine and linear schedules perform similarly.
- Sensitivity Analysis: Momentum values from 0.7 to 0.9 have overlapping confidence intervals, and a Wilcoxon test found no evidence of greater sensitivity than a dense control, W(16) = 22.0, p = 0.58.
6 Conclusion and Future Work
The paper concludes that sparse momentum can outperform other sparse algorithms, rival dense-network performance, and accelerate training while remaining robust to hyperparameter choices.
- Sparse momentum uses mean momentum magnitude to grow and redistribute weights, outperforming other sparse algorithms on MNIST, CIFAR-10, and ImageNet.
- Sparse momentum can rival dense neural-network performance while yielding training speedups and remains robust to hyperparameter choices.
- Specialized sparse convolution and sparse matrix-multiplication algorithms are identified as an important future research goal for realizing sparse-network benefits.
Appendices
The appendices record major revisions, including added dense-performance results, sensitivity analyses, ImageNet settings, fuller training details, and relocated feature analysis.
- The 2019-08-23 overhaul added results showing the weight levels at which sparse networks reach dense performance.
- The revision added sensitivity analyses for momentum, prune rate, and the prune-rate schedule.
- ImageNet experiments were added for both fully sparse and partially dense settings, alongside corrected reporting for the baseline ResNet-50 model.
- Algorithm 1 gained full training-procedure details, while the more detailed sparse-momentum algorithm moved to the appendix.
- Sparse-versus-dense feature analysis gained statistical tests and moved to the appendix, where it was no longer considered a main result.
B.1 Dense vs Sparse Features
The appendix analyzes class specialization in dense and sparse convolutional features, finding that sparse networks learn broader class-useful features, especially in deeper layers.
- The channel-activation analysis aggregates activation by class, normalizes class contributions, and uses each channel’s maximum class proportion as its specialization measure.
- A specialization value of 1/Nc represents equal activity across classes, while larger deviations indicate greater specialization toward a class.
- The study compares AlexNet-s, VGG16-D, and WRN 28-2 on CIFAR-10 using enough weights to reach dense performance levels.
- Sparse convolutional channels have lower class specialization than dense channels, indicating features useful for a broader range of classes; this trend intensifies with depth.
- The authors conclude that sparse networks might rival dense networks by learning more general features with lower class specialization.
C.1 Tuned ResNet-50 on ImageNet
The tuned fully sparse ResNet-50 on ImageNet uses a cosine learning-rate schedule, label smoothing, and learning-rate warmup; results are reported in Table 5. The sparse momentum procedure redistributes and regrows weights using layerwise and weightwise momentum information.
- The tuned fully sparse ResNet-50 uses cosine learning-rate scheduling, label smoothing of 0.9, and learning-rate warmup.
- The fully sparse ImageNet results for the tuned ResNet-50 are reported in Table 5.
- Sparse momentum computes mean momentum contributions across layers before redistributing weights.The algorithm accumulates total momentum and counts nonzero weights for the layers under consideration.
- Pruning removes weights below a threshold determined from the smallest nonzero weights.The algorithm stops gradient flow for weights below the pruning threshold.
- Regrowth considers momentum only for zero-valued weights and enables gradient flow for those exceeding the regrowth threshold.The mask is updated using the boolean OR operator for weights with the largest momentum magnitude.