[Paper Review] Lottery Tickets on a Data Diet: Finding Initializations with Sparse Trainable Networks

Mansheej Paul, Brett W. Larsen+3 more|arXiv (Cornell University)|Jun 2, 2022

Anomaly Detection Techniques and Applications4 citations

TL;DR

This paper investigates the critical role of early pre-training in iterative magnitude pruning (IMP) for finding sparse, trainable subnetworks. It shows that training on just a small fraction of easy examples—rather than the full dataset—can yield equally good initializations, and identifies loss landscape properties like linear mode connectivity as strong predictors of IMP success.

ABSTRACT

A striking observation about iterative magnitude pruning (IMP; Frankle et al. 2020) is that $\unicode{x2014}$ after just a few hundred steps of dense training $\unicode{x2014}$ the method can find a sparse sub-network that can be trained to the same accuracy as the dense network. However, the same does not hold at step 0, i.e. random initialization. In this work, we seek to understand how this early phase of pre-training leads to a good initialization for IMP both through the lens of the data distribution and the loss landscape geometry. Empirically we observe that, holding the number of pre-training iterations constant, training on a small fraction of (randomly chosen) data suffices to obtain an equally good initialization for IMP. We additionally observe that by pre-training only on "easy" training data, we can decrease the number of steps necessary to find a good initialization for IMP compared to training on the full dataset or a randomly chosen subset. Finally, we identify novel properties of the loss landscape of dense networks that are predictive of IMP performance, showing in particular that more examples being linearly mode connected in the dense network correlates well with good initializations for IMP. Combined, these results provide new insight into the role played by the early phase training in IMP.

Motivation & Objective

• To understand why the first few hundred steps of dense pre-training are essential for successful iterative magnitude pruning (IMP).
• To investigate whether a small subset of training data, especially 'easy' examples, can replace full-dataset pre-training for IMP.
• To identify loss landscape properties of the dense network initialization that predict successful IMP outcomes.
• To compare the role of data in IMP pre-training versus learning rate warmup in dense training.
• To guide more efficient training strategies by identifying optimal data and pre-training regimes for sparse network discovery.

Proposed method

• Pre-train dense networks on pruned datasets: randomly selected subsets and subsets filtered by EL2N scores (example difficulty).
• Measure IMP success by evaluating test accuracy of sparse subnetworks found after pruning and fine-tuning from each pre-training point.
• Compute the train loss barrier between two independently trained networks from the same initialization to assess loss landscape stability.
• Use linear mode connectivity between dense network weights trained on different data orders as a proxy for landscape geometry.
• Ablate learning rate warmup duration and data composition (easy, random, hard) to evaluate stability in dense training.
• Correlate pre-training performance (test accuracy) and loss landscape properties (train loss barrier, mode connectivity) with final IMP success.

Experimental results

Research questions

• RQ1 Can a small subset of training data, particularly 'easy' examples, replace full-dataset pre-training to achieve a matching initialization for IMP?
• RQ2 How do loss landscape properties—specifically train loss barrier and linear mode connectivity—predict the success of IMP?
• RQ3 Does pre-training on easy data reduce the number of steps needed to find a good initialization for IMP?
• RQ4 How does the role of data in IMP pre-training compare to its role in learning rate warmup for dense training?
• RQ5 Are there geometric properties of the dense network's loss landscape at initialization that correlate strongly with subsequent sparse network performance?

Key findings

• Training on just 10% of randomly selected data for a few hundred steps yields a pre-training initialization that enables IMP to find sparse subnetworks with test accuracy comparable to full-dataset pre-training.
• Pre-training on the easiest examples (by EL2N scores) reduces the number of pre-training steps required to find a good initialization for IMP compared to random or full data subsets.
• The train loss barrier between two independently trained networks from the same initialization correlates more strongly with final IMP performance than the pre-training test accuracy.
• Linear mode connectivity of the dense network's loss landscape on individual examples is predictive of IMP success: higher connectivity correlates with better sparse subnetwork performance.
• In contrast to IMP pre-training, using easy data during learning rate warmup leads to worse training stability and performance, especially at longer warmup durations.
• The loss landscape geometry of the dense network at initialization—particularly the train loss barrier and example-level mode connectivity—provides stronger predictive power for IMP success than the pre-training accuracy alone.