Lottery Tickets on a Data Diet: Finding Initializations with Sparse Trainable Networks
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.