[Paper Review] Stabilizing the Lottery Ticket Hypothesis

Jonathan Frankle, Gintare Karolina Dziugaite+2 more|arXiv (Cornell University)|Mar 5, 2019

Advanced Neural Network Applications34 references146 citations

TL;DR

The paper shows that pruning early in training, via rewinding to a few percent into training rather than initialization, can yield highly sparse subnetworks that match or exceed the original network’s accuracy on CIFAR-10 and ImageNet, and introduces stability as a key explanation.

ABSTRACT

Pruning is a well-established technique for removing unnecessary structure from neural networks after training to improve the performance of inference. Several recent results have explored the possibility of pruning at initialization time to provide similar benefits during training. In particular, the "lottery ticket hypothesis" conjectures that typical neural networks contain small subnetworks that can train to similar accuracy in a commensurate number of steps. The evidence for this claim is that a procedure based on iterative magnitude pruning (IMP) reliably finds such subnetworks retroactively on small vision tasks. However, IMP fails on deeper networks, and proposed methods to prune before training or train pruned networks encounter similar scaling limitations. In this paper, we argue that these efforts have struggled on deeper networks because they have focused on pruning precisely at initialization. We modify IMP to search for subnetworks that could have been obtained by pruning early in training (0.1% to 7% through) rather than at iteration 0. With this change, it finds small subnetworks of deeper networks (e.g., 80% sparsity on Resnet-50) that can complete the training process to match the accuracy of the original network on more challenging tasks (e.g., ImageNet). In situations where IMP fails at iteration 0, the accuracy benefits of delaying pruning accrue rapidly over the earliest iterations of training. To explain these behaviors, we study subnetwork "stability," finding that - as accuracy improves in this fashion - IMP subnetworks train to parameters closer to those of the full network and do so with improved consistency in the face of gradient noise. These results offer new insights into the opportunity to prune large-scale networks early in training and the behaviors underlying the lottery ticket hypothesis

Motivation & Objective

• Investigate why pruning-at-initialization fails for deeper networks and whether pruning during early training can yield trainable subnetworks.
• Assess the impact of rewinding to early training iterations on subnetwork performance and stability.
• Introduce and analyze stability to pruning and data order as mechanisms influencing lottery tickets.

Proposed method

• Modify iterative magnitude pruning (IMP) to rewind to weights from early in training (k% through) instead of iteration 0.
• Evaluate IMP with and without rewinding on CIFAR-10 using Lenet, Resnet-18, and VGG-19, comparing against random pruning.
• Measure two forms of stability: stability to pruning and stability to data order, using L2 distance between masked weights after training.
• Extend experiments to large-scale ImageNet models (ResNet-50, Inception-v3, SqueezeNet) with rewinding.
• Analyze how later rewinding improves stability and accuracy, and how this relates to the lottery ticket hypothesis.

Experimental results

Research questions

• RQ1 Can subnetworks identified by IMP at initialization be found to train to similar accuracy in deeper networks?
• RQ2 Does pruning later in early training (rewinding) yield smaller, trainable subnetworks that match or exceed the original network’s performance?
• RQ3 Is subnetwork stability (to pruning and to data order) a predictor of finding winning tickets?
• RQ4 How does rewinding affect high-sparsity subnetworks on large-scale tasks like ImageNet?

Key findings

• IMP fails to find winning tickets at initialization for deeper networks like Resnet-18 and VGG-19 without learning-rate adjustments.
• Rewinding to early training iterations (0.1%–7%) enables subnetworks with 50%–99% sparsity that can match the full network’s accuracy on CIFAR-10.
• On ImageNet, rewinding to 4.4%, 3.5%, and 6.6% into training yields subnetworks that are 70%, 70%, and 50% smaller that can match original accuracy for Resnet-50, Inception-v3, and SqueezeNet, respectively.
• Subnetworks found by IMP are substantially more stable to pruning and to data order than randomly pruned subnetworks, and this stability correlates with higher accuracy.
• Later rewinding iterations consistently improve stability and accuracy for subnetworks that do not initially yield a winning ticket.
• The results suggest a revised Lottery Ticket Hypothesis with Rewinding, indicating opportunities to prune early in training, not strictly at initialization.