nucleation

LIT-680: Proving Linear Mode Connectivity of Neural Networks via Optimal Transport

3 min read

Properties6
StatusActive — read in full 2026-10-03 (NOTE-543); the first proof of linear mode connectivity modulo permutation beyond a single hidden layer. Aligning neurons is an optimal-transport problem between the empirical distributions of their weights, so convergence rates of empirical measures in Wasserstein distance bound how wide a network must be for the permuted straight line to stay low-loss. It proves the result for wide two-layer nets trained by SGD in the mean-field regime, and for deep nets whenever neuron weights within a layer are i.i.d., which holds at initialization. The width needed grows like (1/ε) to the power of the previous layer's width, a rate shown to be tight, and falls if weights concentrate near a low-dimensional subspace. The proofs rest on independence of neurons, which the authors call the central and least realistic assumption.
Tagsloss-landscapes, mathematics, anthology-candidate
ExtendsLIT-652 — The Role of Permutation Invariance in Linear Mode Connectivity of Neural Networks
Compared againstLIT-661 — Git Re-Basin: Merging Models modulo Permutation Symmetries
Cited as “paper” byNOTE-543 — Proving Linear Mode Connectivity of Neural Networks via Optimal Transport
Cited as “source” byTHEORY-112 — Aligning independently trained solutions over the architecture's symmetries removed the linear barrier between them for MNIST MLPs and 32×-wide ResNet-20s on CIFAR-10, two-thirds of it for an ImageNet ResNet-50, and not all of it for VGG-16 or ResNet-20 at 1× width, the barrier after weight matching falling with width; for 6- to 8-layer ViTs and GPT-2s, permutations alone removed 63–90% of it, and adding an orthogonal map on the residual stream, learned against the loss at the path's midpoint, took it to zero or near zero except on BookCorpus (0.42), THEORY-114 — In image classification, linear mode connectivity appeared during training and not at initialization: copies of ResNet-20 and VGG-16 on CIFAR-10 and of ResNet-50 and Inception-v3 on ImageNet, trained from a shared state, became stable to SGD noise in test error only after 1.5–20% of training (LeNet on MNIST was stable from the start), and between independently initialized MLPs on MNIST and CIFAR-10 the barrier after weight-matching alignment was high at initialization and fell gradually through training, THEORY-1172 — Linear mode connectivity emerges during training, not at initialization: copies trained from a shared state become stable to SGD noise after 1.5–20% of training, and independently trained networks become linearly connected after alignment only gradually, THEORY-1188 — Once the architecture's symmetries are factored out, most of the linear barrier between independently trained solutions disappears; the group that matters is the full one, permutations for MLPs and CNNs but for transformers also an orthogonal map on the residual stream, without which a barrier remains; and the barrier left after alignment falls with width

Graph View

Damien Ferbach, Baptiste Goujaud, Gauthier Gidel and Aymeric Dieuleveut (2024), Proceedings of the 27th International Conference on Artificial Intelligence and Statistics (AISTATS 2024), PMLR 238 — ARXIV-2310.19103

Key takeaways

  • Alignment is optimal transport. Permuting network B’s neurons at layer ℓ to minimize Σ_i ‖[W_A]_i − [W̃_B]_i‖² is the Monge problem between the empirical distributions of the two layers’ weight vectors. By Birkhoff’s theorem its value is their Wasserstein distance (App. B, Lemma B.3). If each layer’s weights are i.i.d. draws from one distribution, that distance shrinks with width at the known rates for empirical measures, the activations of A, permuted B and every network on the line between them stay close, and with a convex loss the line stays low-loss.
  • Two-layer networks after training (Theorem 3.1). Two wide two-layer nets trained independently by SGD in the mean-field regime (small step size, bounded Lipschitz non-linearity, which excludes ReLU) are, with high probability, linearly connected modulo a permutation of the hidden layer, with an output-level and a squared-loss version.
  • Deep networks at initialization (Lemma 5.1, Theorem 5.2). With Gaussian i.i.d. weights and ReLU allowed, LMC modulo permutation holds once the width of layer ℓ exceeds m̃_ℓ = Õ((T_ℓ/ε)^{m̃_{ℓ−1}}). For two-layer networks that is ε^{−m₀}, against Entezari et al.’s ε^{−(2m₀+4)}. A matching lower bound (Theorem 5.3) shows the recursive exponential growth is tight. An “effective width” device lets every layer have the same width.
  • Lower-dimensional weights need less width (Theorem 5.4). If weight distributions are approximately supported on a subspace of dimension e·m̃_{ℓ−1}, the required width drops until the approximation breaks down. The result extends to independent sub-Gaussian weights.
  • Empirically (§6), across learning rates for a 3×512 MLP on MNIST, a lower approximate dimension of the matrices being matched goes with lower barriers. A weight matching weighted by the activation covariance beats plain weight matching.

Standing in the record

Filed with the batch ADR-027 opened loss-landscapes for. It extends Entezari et al. (LIT-652): it proves their conjecture beyond their single case, from one hidden layer at random initialization to trained two-layer nets in the mean-field regime and to deep nets with independent weights. Its bound is better than theirs in the case they share: ε^{−m₀} against ε^{−(2m₀+4)} for the hidden width.

It is compared against Git Re-Basin (LIT-661) in §6. It takes that paper’s weight-matching and activation-matching objectives as baselines and adds a third, weight matching under the activation covariance, which “constantly and substantially outperforms naive Weight Matching” across learning rates (Fig. 2). It also gives an account of a tension between the two earlier papers. Git Re-Basin found no linear connectivity at initialization, where Entezari et al.’s theorem proves it. The width the bound requires at initialization is “excessive as compared to the typical width of neural networks used in practice”, so both can hold.

Appendix B.14 relates the result to Kuditipudi et al.’s dropout stability (LIT-671) in a simplified one-hidden-layer setting: networks with independent sub-Gaussian neuron weights become dropout stable at the same rate in width as they become linearly connected. It calls this “a first step” towards unifying the two views, so no relation is declared.

It carries mathematics because its method is the theory of optimal transport and the convergence of empirical measures, applied to a question about networks.

Backlinks