Source-linked AI summary
Picking Winning Tickets Before Training by Preserving Gradient Flow
Chaoqi Wang, Guodong Zhang, Roger Grosse
TL;DR
Training large neural networks is expensive, and pruning usually occurs only after training. GraSP prunes at initialization by preserving gradient flow, outperforming baselines at high sparsity, including a 6.2% top-1 advantage over SNIP for ImageNet ResNet-50 at 90% pruning.
Problem
Pruning typically targets trained networks, leaving the expensive training process unresolved and making sparse subnetworks difficult to train from scratch with comparable performance.
Method
GraSP prunes at initialization by removing weights whose removal least decreases the network’s post-pruning gradient norm.
Results
Across datasets and architectures, GraSP performs especially well at high pruning ratios, beating SNIP by 6.2% in ImageNet ResNet-50 top-1 accuracy at 90% pruning.
Takeaways & Limitations
GraSP shows that networks can be pruned before training while remaining competitive with traditional pruning algorithms that require prior training.
Takeaways & Limitations
Connection-sensitivity pruning can block information flow because it evaluates each weight in isolation, and GraSP still trails traditional pruning algorithms.
Abstract
from arXiv · showhide
Overparameterization has been shown to benefit both the optimization and generalization of neural networks, but large networks are resource hungry at both training and test time. Network pruning can reduce test-time resource requirements, but is typically applied to trained networks and therefore cannot avoid the expensive training process. We aim to prune networks at initialization, thereby saving resources at training time as well. Specifically, we argue that efficient training requires preserving the gradient flow through the network. This leads to a simple but effective pruning criterion we term Gradient Signal Preservation (GraSP). We empirically investigate the effectiveness of the proposed method with extensive experiments on CIFAR-10, CIFAR-100, Tiny-ImageNet and ImageNet, using VGGNet and ResNet architectures. Our method can prune 80% of the weights of a VGG-16 network on ImageNet at initialization, with only a 1.6% drop in top-1 accuracy. Moreover, our method achieves significantly better performance than the baseline at extreme sparsity levels.
1 INTRODUCTION
Large networks offer strong optimization and generalization but are costly to train and deploy, motivating pruning at initialization to preserve efficient gradient flow. GraSP addresses this by pruning weights whose removal least decreases the post-pruning gradient norm.
- Motivation: Overparameterized networks perform well, but their training and inference are computationally expensive, while pruning mainly reduces test-time resource requirements.
- Limitation of prior work: SNIP’s isolated per-weight gradient criterion can remove connections important to information flow because pruning changes gradients through weight interactions.
- GraSP: GraSP prunes weights whose removal causes the least decrease in the gradient norm after pruning, thereby accounting for each connection’s role in network gradient flow.The method is described as easy to implement and conceptually simple.
2 RELATED WORK AND BACKGROUND
The section situates initialization-time pruning within prior work on pruning before, during, and after training, dynamic sparse training, signal propagation, dynamical isometry, mean-field theory, and the Neural Tangent Kernel (NTK). It emphasizes that pruning at initialization must account for training dynamics, while NTK theory provides a foundation for analyzing those dynamics.
- Pruning after training: Most pruning algorithms operate after training by removing redundant weights, while magnitude-based methods may incorrectly measure individual weight importance.These approaches prune pre-trained networks to reduce performance degradation, typically using weight thresholds.
- Pruning during training: Training-time pruning methods seek sparse networks through adaptive dropout rates or L0 regularization but require roughly the same computational cost as training dense networks.Generalized dropout tunes individual dropout rates, while L0 methods address discontinuities during optimization.
- Dynamic Sparse Training: Dynamic Sparse Training methods change sparsity during training through prune-redistribute-regrowth cycles, including sparse momentum approaches that dynamically determine sparse masks.Representative methods include work by Bellec et al., Mocanu et al., Mostafa and Wang, and Dettmers and Zettlemoyer.
- Pruning before training: Pruning before training is more challenging because removing weights affects training dynamics; prior approaches include the Lottery Ticket Hypothesis and SNIP’s connection-sensitivity criterion.The Lottery Ticket Hypothesis links a pruned network structure with its corresponding initialization, while SNIP identifies trainable sparse subnetworks at initialization.
- Signal propagation: Dynamical-isometry and mean-field studies analyze signal propagation and input-output Jacobian singular values to derive initialization schemes, including one that trained 10,000-layer vanilla ConvNets.These works share the spirit of the paper’s pruning criteria.
- Neural Tangent Kernel: For infinitely wide networks with proper initialization, the Neural Tangent Kernel exactly captures output-space training dynamics and remains constant throughout training, with close empirical agreement reported for wide practical networks.The section presents NTK analysis as the foundation for justifying the proposed method.
3 REVISITING SINGLE-SHOT NETWORK PRUNING (SNIP)
This section formalizes foresight pruning as selecting an initialization mask before training and revisits SNIP’s loss-preservation criterion, arguing that preserving loss at initialization is less justified than preserving training dynamics.
- Problem Formulation: Foresight pruning selects a binary mask before training to minimize the trained network’s expected loss under a fixed sparsity constraint.The formulation removes ⌈p·d⌉ weights, while the training algorithm A returns the trained masked network.
- Problem Formulation: Because globally minimizing the foresight-pruning objective is intractable, the method instead relies on heuristics with good practical performance.The formulation treats the training algorithm, such as SGD, as known and returns trained weights from the mask.
- Revisiting SNIP: SNIP was the first foresight-pruning algorithm and removes connections according to their individual sensitivity to the loss.Its sensitivity is defined by considering the effect of removing one initialized weight in isolation.
- Revisiting SNIP: SNIP preserves the loss of the randomly initialized network, extending a criterion motivated by pruning trained networks.The section notes that loss preservation underlies methods such as optimal brain damage and optimal brain surgery.
- Revisiting SNIP: At initialization, however, the loss is no better than chance, making loss preservation a less clear criterion for foresight pruning than preserving early training dynamics.The section argues that training dynamics are more important at the beginning of training.
4 GRADIENT SIGNAL PRESERVATION
GraSP prunes at initialization by preserving or increasing gradient flow, using Hessian-aware scores to identify weights whose removal least harms optimization. Its scores account for weight coupling, can be computed efficiently without explicitly constructing the Hessian, and are linked to faster training through the NTK.
- Criterion: The Hessian captures dependencies between weights, making GraSP more informative than an identity-Hessian approximation that recovers SNIP up to absolute value.The method accounts for the observation that different weights are highly coupled, so the Hessian is far from the identity.
- Criterion: GraSP scores each weight by its predicted change in gradient flow after removal, prioritizing weights whose removal does not reduce gradient flow.A negative score indicates reduced gradient flow after removal, whereas a positive score indicates increased gradient flow; larger scores therefore indicate lower importance.
- Implementation: For pruning ratio p, GraSP removes the top p fraction of weights by score and computes the Hessian-gradient product through higher-order automatic differentiation.This avoids explicitly constructing the Hessian while keeping the method efficient and easy to implement.
- Training dynamics: GraSP implicitly encourages large NTK eigenvalues in output-gradient directions, which should produce efficient training because larger eigenvalue directions train faster.In practice, increasing gradient norm may increase the loss, motivating a temperature term on the logits to smooth predictions.
5 EXPERIMENTS
Experiments evaluate GraSP across image-classification datasets, architectures, pruning ratios, and baselines. GraSP outperforms random pruning, remains competitive with SNIP at moderate sparsity, and better preserves gradient norms and performance at high sparsity.
- Experimental setup: GraSP is evaluated on CIFAR-10/100, Tiny-ImageNet, and ImageNet using VGGNet and ResNet architectures.Experiments include comparisons with random pruning, SNIP, dynamic sparse training methods, and traditional pruning on pretrained networks.
- Comparison with baselines: At pruning ratios of 85% and 90%, GraSP and SNIP achieve results close to the baselines, while remaining below pruning methods applied to trained networks.The paper also evaluates the methods at 95% and 98% pruning ratios on CIFAR-10/100 and Tiny-ImageNet.
- ImageNet evaluation: At 60% pruning on ImageNet, GraSP and SNIP perform almost identically and remain close to the original network’s performance.ImageNet experiments use ResNet-50 and VGG16 at pruning ratios of 60%, 80%, and 90%.
- Gradient-flow analysis: At high sparsity, GraSP better preserves gradient norms than SNIP and achieves better results, whereas random pruning has a much lower gradient norm and performs worst.Both SNIP and GraSP produce lower gradient norms at high sparsity, such as 98%, and pruned networks generally underfit the training data.
- Gradient-flow analysis: GraSP avoids pruning too aggressively in individual layers because its pruning strategy aims to preserve gradient flow after pruning.This produces a different layer-wise allocation of remaining weights from SNIP.
6 DISCUSSION AND CONCLUSION
GraSP prunes network weights at initialization by preserving gradient flow, while remaining competitive with traditional pruning algorithms that require prior training.
- GraSP prunes network weights at initialization while preserving gradient flow through the network.The criterion is motivated by preserving gradient flow after pruning.
- GraSP remains competitive with traditional pruning algorithms that require first training the network.
- GraSP can also be interpreted as aligning the large eigenvalues of the Neural Tangent Kernel with the targets.