Source-linked AI summary
Rigging the Lottery: Making All Tickets Winners
Utku Evci, Trevor Gale, Jacob Menick, Pablo Samuel Castro, Erich Elsen
TL;DR
Sparse training methods that first train dense models limit the size and efficiency of trainable sparse networks. RigL dynamically updates sparse connectivity using weight magnitudes and gradient information while maintaining fixed parameter and computational budgets. Across tasks, it achieves higher accuracy or quality for a given computational budget than existing sparse-training approaches, though sparse hardware and software support remains limited.
Problem
Dense-to-sparse training limits sparse-model size to the largest trainable dense model and performs computation for parameters that are zero or will be pruned.
Method
RigL starts from a random sparse network and periodically removes low-magnitude connections while activating new connections using instantaneous gradient information.
Results
For a given computational budget, RigL achieves higher accuracy than existing dense-to-sparse and sparse-to-sparse training algorithms.
Takeaways & Limitations
RigL can improve deployed sparse models, improve large sparse models under limited training iterations, and support extremely large sparse models when combined with sparse primitives.
Takeaways & Limitations
The extremely large sparse-model scenario remains unexplored because hardware and software support for sparsity is lacking.
Abstract
from arXiv · showhide
Many applications require sparse neural networks due to space or inference time restrictions. There is a large body of work on training dense networks to yield sparse networks for inference, but this limits the size of the largest trainable sparse model to that of the largest trainable dense model. In this paper we introduce a method to train sparse neural networks with a fixed parameter count and a fixed computational cost throughout training, without sacrificing accuracy relative to existing dense-to-sparse training methods. Our method updates the topology of the sparse network during training by using parameter magnitudes and infrequent gradient calculations. We show that this approach requires fewer floating-point operations (FLOPs) to achieve a given level of accuracy compared to prior techniques. We demonstrate state-of-the-art sparse training results on a variety of networks and datasets, including ResNet-50, MobileNets on Imagenet-2012, and RNNs on WikiText-103. Finally, we provide some insights into why allowing the topology to change during the optimization can overcome local minima encountered when the topology remains static. Code used in our work can be found in github.com/google-research/rigl.
1. Introduction
RigL addresses the limits of dense-to-sparse training by dynamically changing sparse-network connectivity while maintaining computational efficiency. The method achieves higher quality than prior techniques at a given computational cost and can outperform dense-to-sparse training algorithms.
- Method: RigL combines weight magnitudes and gradient information to jointly optimize parameters and connectivity.The topology is updated during training rather than held fixed.
- Motivation: Dense-to-sparse methods require at least the memory and FLOPs of dense-model training, limiting the largest trainable sparse model.They may also perform substantially more computation than dense training.
- Motivation: Static sparse training can underperform because a fixed topology may trap optimization in unfavorable regions of the loss landscape.The paper studies why changing nonzero-weight connectivity during training aids optimization.
- Method: RigL trains sparse networks by maintaining memory and computational cost proportional to network density.It avoids requiring a lucky initialization and is introduced as the Rigged Lottery method.
- Results: RigL achieves higher quality than previous techniques for a given computational cost and finds more accurate models than current dense-to-sparse methods.The evaluation covers computer vision and natural-language tasks.
2. Related Work
Prior sparse-training work spans magnitude pruning, structured pruning, and dynamic topology methods, while lottery-ticket studies question whether suitable sparse initializations can be generated from scratch.
- Pruning: Magnitude pruning evolved from single-shot removal to repeated pruning and retraining, but reaching 90% sparsity can require ten retraining rounds.These methods target sparse networks after dense training.
- Pruning: Alternative pruning approaches, including L0 Regularization and Variational Dropout, achieved about the same accuracy-versus-sparsity trade-off on ResNet-50 and Transformer architectures.The cited comparison concerns several non-magnitude methods.
- Structured pruning: Structured pruning removes channels or neurons to produce dense networks that are easier to accelerate, whereas RigL finds smaller networks requiring fewer FLOPs to run.The paper compares RigL with state-of-the-art structured pruning methods.
- Dynamic sparse training: Dynamic sparse methods keep sparsity throughout training, using strategies such as magnitude-based removal with random regrowth or parameter redistribution across layers.Dynamic Sparse Reparameterization can increase FLOPs when parameters shift toward early convolutional layers.
- Lottery tickets: Lottery-ticket research found that large networks may require weights after thousands of optimization steps rather than the original initialization.Whether masks and parameters can be generated de novo remains unclear.
3. Rigging The Lottery
RigL trains a randomly sparse network while preserving its sparse connectivity budget, periodically replacing low-magnitude connections with gradient-selected ones. Its design fixes layerwise sparsity and uses scheduled topology updates to maintain computational efficiency.
- RigL algorithm: RigL starts with a random sparse network and periodically updates connectivity during training.The method removes connections by magnitude and activates new ones using instantaneous gradient information.
- Sparsity distribution: Layerwise sparsity is fixed during training to make the final FLOP budget controllable.The method avoids reallocating parameters between layers while considering multiple sparsity-distribution strategies.
- Sparsity distribution: Uniform distribution assigns each layer the total sparsity, while keeping the first layer dense because sparsifying it disproportionately harms performance.The first layer has little effect on total model size but a disproportionate effect on performance when sparsified.
- Sparsity distribution: Erdős–Rényi and ERK distribute sparsity using layer dimensions, with ERK additionally incorporating convolutional kernel dimensions.ERK scales sparse convolutional-layer parameters using input and output channels together with kernel width and height.
- Update schedule: Connectivity updates occur every ΔT iterations until Tend, with the updated fraction controlled by α and fdecay.The paper uses cosine annealing for fdecay and amortizes dense-gradient computation when updates are sufficiently infrequent.
- RigL algorithm: RigL grows new connections with the highest-magnitude gradients among inactive connections.Newly activated connections are initialized to zero, so they initially do not affect network outputs while receiving subsequent gradients.
4. Empirical Evaluation
RigL is evaluated across vision and language tasks, showing strong sparse-model performance under fixed resource budgets. Its dynamic connectivity, gradient-based growth, and ERK allocation improve accuracy or efficiency relative to static, dense-to-sparse, and alternative dynamic methods.
- ResNet-50: RigL achieves the highest accuracy among evaluated sparse ResNet-50 methods while consistently requiring fewer FLOPs than the other methods.The comparison uses uniform sparsity and extended training runs up to five times the original training steps.
- ResNet-50: 72.75% Top-1 Accuracy is achieved by a 96.5% sparse ResNet-50 with ERK sparsity, around 3.5% higher than extended magnitude pruning.Smaller dense models and static sparse models do not reach a comparable level.
- MobileNet: RigL trains 75% sparse MobileNets with no performance loss and consistently outperforms the baselines as sparsity increases.The experiments use five times the original number of training steps.
- MobileNet: A 75% sparse MobileNet-v1 with width multiplier 1.98 improves Top-1 Accuracy by 4.3% over a dense baseline with the same FLOPs and parameter count.The result demonstrates higher accuracy for a sparse model matched to the dense baseline’s resource counts.
- Analysis: ERK allocation improves final accuracy over uniform allocation but requires approximately twice as many FLOPs.ERK reallocates parameters toward layers with few parameters, with a larger advantage at high sparsity.
- Analysis: RigL can escape a local minimum reached by static sparse training, whereas retraining with static sparsity cannot.The method removes low-magnitude connections and activates high-gradient connections.
5. Discussion & Conclusion
RigL trains sparse neural networks efficiently by dynamically changing connectivity while maintaining a fixed computational budget. Its topology updates remove low-magnitude connections and add connections with high expected gradients, supporting stronger optimization and accuracy within a given budget.
- RigL achieves higher accuracies than existing dense-to-sparse and sparse-to-sparse algorithms for a given computational budget.
- RigL is presented for deployment accuracy, limited-iteration training of large sparse models, and extremely large sparse models combined with sparse primitives.
- RigL dynamically drops low-magnitude connections and grows new connections using instantaneous gradient information during training.Connections are updated at regularly spaced intervals while the network remains sparse.
- Magnitude-based dropping limits loss damage by removing parameters with low magnitude and a small Taylor-approximation remainder.
- RigL’s dynamic topology can improve the sparse-network optimization landscape by replacing flat directions with newly activated, high-gradient connections.
B. Comparison with Bayesian Structured Pruning Algorithms
Compared with structured pruning methods on MNIST multilayer perceptrons, RigL searches sparse connectivity directly and produces compact, FLOP-efficient networks. Repeated dynamic training further improves compactness, while experiments also examine input connectivity, sparsity distributions, momentum, and lottery-ticket initialization.
- RigL finds compact networks with fewer parameters, fewer inference FLOPs, and fewer training resources than structured pruning approaches.
- Repeating RigL from the 408-100-69 architecture produces a network half the size and requiring one third the FLOPs of Dai et al. (2018)’s best architecture.
- RigL reallocates input connections toward informative central MNIST pixels and discards uninformative pixels during training.
- ERK sparsity distribution performs better than other distributions across four sparse training methods.
- For SNFS, momentum coefficient 0.99 performs best, while coefficients 0 and 0.9 perform similarly.The passage suggests this may relate to the batch size of 4096.
- Fixed-topology training performs significantly worse than RigL, and restarting with the final topology and original initialization does not improve RigL.The experiment concludes that there are no special lottery tickets in this setting.
F. Effect of Update Schedules on Other Dynamic Sparse Methods
Alternative update schedules for SET and SNFS show that cosine scheduling with particular update intervals and drop fractions performs well across dynamic sparse methods. The preferred schedule depends on the interaction between update interval and drop fraction.
- Cosine scheduling with ΔT = 50 and α = 0.1 performs best across SET and SNFS in the reported sweep.
- Higher drop fractions tend to work better with longer update intervals, as illustrated by SET with ΔT = 1000 and α = 0.5.
- Constant decay performs well at low initial drop fractions such as α = 0.1 but worsens as α increases.
- Inverse-power schedules with k = 1 and k = 3 perform similarly at low α, but k = 3 degrades noticeably as the update interval increases.
- Linear inverse-power decay with k = 1 provides results similar to cosine scheduling in the reported comparison.
H. Calculating FLOPs of models and methods
The paper estimates training and inference cost by counting layer-wise multiplications and additions, distinguishing dense, static sparse, pruning, SNIP, SET, and RigL methods. RigL’s average cost combines sparse iterations with infrequent dense-gradient calculations.
- A sparse model’s forward-pass FLOPs are computed by summing layer-wise multiplications and additions for the chosen layer sparsities.Different sparsity distributions can therefore produce different prediction costs.
- Static sparse and dense training scale with 3∗fS and 3∗fD FLOPs, respectively.
- Pruning requires Et[3 ∗fD ∗(1 −st)] FLOPs, where st is the model sparsity at iteration t.
- SET and SNIP are treated as scaling with 3∗fS FLOPs under the stated accounting assumptions.
- SNFS averages 2∗fS + fD FLOPs because sparse forward and backward computation is paired with dense-gradient calculation every iteration.
- RigL averages (3∗fS∗ΔT +2∗fS+fD) FLOPs by calculating dense gradients only every ΔT iterations.
I. Hyper-parameters used in Charachter Level Language Modeling Experiments
The experiments specify the language-model architecture, optimization settings, sparsity, and mask-update schedule, while comparing alternative update schedules for RigL.
- Hyper-parameters: The model uses a 128-dimensional shared embedding, vocabulary size 256, GRU state size 512, and two readout layers of widths 256 and 128.Training uses next-step prediction with cross-entropy loss and Adam.
- Hyper-parameters: Training uses learning rate 7e-4, L2 regularization 5e-4, sequence length 512, and batch size 32.The supplied passage also states that gradients are clipped.
- Update schedules: The alternative RigL update schedules evaluated are Constant, Exponential with k=3, and Linear.These schedules are shown as separate left, middle, and right panels.
- Training schedule: Models use 75% sparsity and train for 200,000 iterations, with magnitude pruning scheduled from iterations 50,000 to 150,000 every 1,000 iterations.Sparse networks are initialized with a uniform sparsity distribution.
- Training schedule: RigL uses a cosine update schedule with α = 0.1 and ΔT = 100, continuing mask updates until training ends.The passage reports that continuing updates performed slightly better than stopping at iteration 150,000.
J. Additional Plots and Experiments for CIFAR-10
Additional CIFAR-10 experiments examine training loss, test accuracy across mask-update intervals, and the relationship between RigL optimization and generalization.
- Training loss and generalization: RigL achieves similar final training loss to pruning, although its accuracy is around half a percent lower in the reported experiments.Training longer with RigL further decreases final loss, and test accuracies start matching pruning.
- Training loss and generalization: Static sparse training has poor performance reflected in higher training loss across all sparsity levels.The comparison is made using final training loss.
- Training loss and generalization: RigL improves optimization but generalizes slightly worse than pruning.This conclusion follows the reported training-loss and test-accuracy comparison.
- Mask-update intervals: The best final test accuracies occur at mask-update interval ΔT = 100 for both Uniform and ERK sparsity distributions.Both curves follow a pattern similar to the ImageNet-2012 sweeps, with initial drop fraction α fixed at 0.3.
- Layer-wise sparsity: Figure 12 plots the layer-wise sparsities of ResNet-50 under the Erdős-Rényi-Kernel distribution.The passage identifies the distribution and the plotted quantity but does not state a comparative outcome.
L. Performance of Algortihms at Training 95 and 96.5% Sparse ResNet-50
The additional high-sparsity ResNet-50 experiments report RigL performance alongside final sparse-model loss and the experiment table for 95% and 96.5% sparsity.
- Performance at increased sparsity: RigL achieves higher performance than pruning on ResNet-50 even without extending the training length.The result is reported for experiments with increased sparsity.
- Additional plots: Figure 11 includes final training loss for sparse models and RigL performance at different mask-update intervals.The caption identifies the two plotted quantities without reporting numerical values.
M. Bugs Discovered During Experiments
The experiment audit documents implementation bugs, their corrections, resulting performance changes, and additional methodological clarifications affecting sparse-training comparisons.
- Audit scope: The audit states that the initial bugs lowered accuracy for all sparse-training techniques without changing the general conclusion that RigL was more effective.This limitation applies to the initial implementations.
- Implementation bugs: Replica-specific random operations caused SET, SNFS, and RigL replicas to choose different connections to drop and grow.The bug was fixed with stateless random operations, improving SET by 0.1–0.3% on Figure 2-left.
- Implementation bugs: Masked-parameter gradients were not automatically aggregated across replicas for RigL and SNFS.Explicit ALL-REDUCE fixed the synchronization bug and improved performance around 0.5–1% at default training lengths.
- SNIP experiments: The corrected SNIP metric is saliency, defined as gradient times parameter magnitude, rather than gradient magnitude alone.Using gradient magnitudes alone produced much worse-than-random performance, while the correction improved SNIP substantially.
- Revisions: The revision reorders SET and NeST, replaces Idrop with Iactive, fixes a pruning-FLOPs typo, and emphasizes zero initialization at growth.These are reported as notation, calculation, and methodological clarifications.