Linear Mode Connectivity and the Lottery Ticket Hypothesis

TL;DR

The paper introduces instability analysis to ask whether SGD noise steers neural network training toward the same minimum. It demonstrates that standard vision models become stable to SGD noise early in training, yielding linear mode connectivity between networks trained with different data orders or augmentations. Applying this framework to Iterative Magnitude Pruning (IMP) clarifies when sparse subnetworks (lottery tickets) can train to full accuracy: matching subnetworks exist only when they are stable to SGD noise, which can occur at initialization for small models or early in training for large ones. Extending IMP with rewinding enables discovering matching subnetworks in large-scale settings, linking lottery-ticket phenomena to optimization dynamics and suggesting pruning opportunities earlier in training.

Abstract

We study whether a neural network optimizes to the same, linearly connected minimum under different samples of SGD noise (e.g., random data order and augmentation). We find that standard vision models become stable to SGD noise in this way early in training. From then on, the outcome of optimization is determined to a linearly connected region. We use this technique to study iterative magnitude pruning (IMP), the procedure used by work on the lottery ticket hypothesis to identify subnetworks that could have trained in isolation to full accuracy. We find that these subnetworks only reach full accuracy when they are stable to SGD noise, which either occurs at initialization for small-scale settings (MNIST) or early in training for large-scale settings (ResNet-50 and Inception-v3 on ImageNet).

Figures (32)

• Figure 1: A diagram of instability analysis from step 0 (left) and step $k$ (right) when comparing networks using linear interpolation.
• Figure 2: Error when linearly interpolating between networks trained from the same initialization with different SGD noise. Lines are means and standard deviations over three initializations and three data orders (nine samples total). Trained networks are at 0.0 and 1.0.
• Figure 3: Linear interpolation instability when starting from step $k$. Each line is the mean and standard deviation across three initializations and three data orders (nine samples in total).
• Figure 4: Linear interpolation instability on the test set when making two copies of the state of the network at step $k$ and either (1) training for the remaining $T-k$ steps (blue) or (2) training for $T$ steps with the learning rate schedule reset to step 0 (orange).
• Figure 5: Test error when linearly interpolating between subnetworks trained from the same initialization with different SGD noise. Lines are means and standard deviations over three initializations and three data orders (nine in total). Percents are weights remaining.