[Paper Review] Linear Mode Connectivity and the Lottery Ticket Hypothesis

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

Advanced Neural Network Applications34 references100 citations

TL;DR

The paper introduces instability analysis to test SGD noise sensitivity, showing networks become stable to SGD noise early in training, leading to linear mode connectivity, and uses this to explain when iterative magnitude pruning yields matching subnetworks.

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).

Motivation & Objective

• Assess how SGD noise affects neural network optimization outcomes across runs
• Determine when different SGD noise samples lead to the same (linearly connected) minima
• Use instability analysis to evaluate iterative magnitude pruning (IMP) and lottery ticket outcomes
• Generalize IMP to rewind to intermediate training states to find matching sparse subnetworks
• Provide practical insights on pruning timing and optimization dynamics in large-scale models

Proposed method

• Define instability analysis by training two copies with different SGD noise and measuring interpolation error
• Quantify linear interpolation instability as the maximal barrier along a linear path between two trained networks
• Assess stability via a 0-2% instability threshold to declare linear mode connectivity
• Apply instability analysis across initialization and various training steps for networks on MNIST, CIFAR-10, ImageNet
• Extend IMP by rewinding to step k and pruning to obtain matching subnetworks
• Analyze sparsity regimes to categorize when IMP subnetworks are matching and stable

Experimental results

Research questions

• RQ1 Does SGD noise produce different minima or a single linearly connected minimum across runs for standard networks?
• RQ2 At what training stage do networks become stable to SGD noise, and does this affect IMP outcomes?
• RQ3 Are IMP subnetworks matching only when stable to SGD noise, and can rewinding reveal matching subnetworks at larger scales?

Key findings

• Most networks are unstable to SGD noise at initialization; stability emerges early in training for all networks studied
• After a small fraction of training (e.g., ~3% for ResNet-20, ~1.5% for VGG-16, 20% for ResNet-50, 16% for Inception-v3), networks become stable to SGD noise and have linearly connected minima
• IMP subnetworks are matching only when stable to SGD noise; stability can occur at initialization in small tasks or early in training at larger scales
• IMP subnetworks rewound to intermediate training steps can become stable and matching earlier than the full network, enabling matching at nontrivial sparsities in larger models
• Stability onset times differ by architecture and task, with large-scale models requiring more training to reach stability but achieving matching subnetworks when stable
• IMP subnetworks across sparsities show stability and matching patterns aligning with the three sparsity ranges identified (trivial, IMP-only, and non-matching)