[Paper Review] Picking Winning Tickets Before Training by Preserving Gradient Flow

Chaoqi Wang, Guodong Zhang+1 more|arXiv (Cornell University)|Feb 18, 2020

Advanced Neural Network Applications37 references148 citations

TL;DR

GraSP prunes neural networks at initialization by preserving gradient flow, enabling up to 80% weight pruning (ImageNet) with minimal accuracy loss.

ABSTRACT

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.

Motivation & Objective

• Motivate pruning before training to save training resources.
• Propose a gradient-flow–based pruning criterion that accounts for inter-weight dependencies.
• Demonstrate effectiveness across CIFAR, Tiny-ImageNet, and ImageNet with VGGNet and ResNet architectures.
• Analyze how pruning affects training dynamics and relate findings to Neural Tangent Kernel.

Proposed method

• Introduce Gradient Signal Preservation (GraSP) as a pruning criterion.
• Compute Hessian-gradient products to evaluate how pruning a weight affects gradient flow.
• Score weights via S(-θ) = -θ ⊙ (H g) and prune the top p fraction with lowest scores.
• Train the resulting sparse network from initialization and evaluate performance.
• Leverage NTK insights to relate pruning effects to optimization dynamics.

Experimental results

Research questions

• RQ1 Can networks be pruned effectively at initialization without training the full dense model?
• RQ2 Does preserving gradient flow during pruning improve trainability and final accuracy at high sparsity?
• RQ3 How does GraSP compare to SNIP and other baselines across modern architectures and datasets?
• RQ4 What is the role of initialization and batch size in GraSP performance?

Key findings

• GraSP can prune up to 80% of VGG-16 weights on ImageNet at initialization with only a 1.6% drop in top-1 accuracy.
• GraSP consistently outperforms SNIP at extreme sparsity across CIFAR-10/100, Tiny-ImageNet, and ImageNet.
• GraSP yields better maintenance of gradient flow than random pruning and often approaches or surpasses late-reset lottery tickets and some DST baselines.
• Pruned networks with GraSP exhibit faster loss decrease on training and better gradient norms at high sparsity than SNIP.
• GraSP aligns with NTK-based predictions, encouraging high-variance directions in output space gradients to be retained for efficient optimization.
• GraSP shows robustness to different initializations and batch sizes, especially with common initializations like Kaiming.