PaperlayerView exact revision on arXiv

Source-linked AI summary

Git Re-Basin: Merging Models modulo Permutation Symmetries

Samuel K. Ainsworth, Jonathan Hayase, Siddhartha Srinivasa

arXiv:2209.04836v6cs.LGcs.AI

TL;DR

Deep learning relies on SGD to solve massive non-convex optimization problems, motivating questions about why independently trained solutions share favorable structure. The paper uses three permutation-based alignment algorithms to merge models in weight space, finding strong linear mode connectivity across settings while identifying width-related limits and a counterexample.

  • Problem

    The paper asks why SGD succeeds on high-dimensional non-convex deep-learning problems and whether independently trained solutions can be connected through low-loss interpolation.

  • Method

    The paper introduces three algorithms that permute hidden units and align independently trained neural-network weights before linear interpolation and merging.

  • Results

    Zero-barrier linear mode connectivity is achieved for MNIST with all three methods and for large CIFAR-10 ResNets, while ImageNet barriers decrease 67% relative to naïve interpolation.

  • Takeaways & Limitations

    Permutation symmetries are a necessary piece but not a complete explanation of the invariances underlying linear mode connectivity in neural-network training dynamics.

  • Takeaways & Limitations

    Thin models remain difficult to linearly connect, and adversarial non-SGD solutions can lack linear mode connectivity under every unit permutation.

Abstract

from arXiv · show

The success of deep learning is due in large part to our ability to solve certain massive non-convex optimization problems with relative ease. Though non-convex optimization is NP-hard, simple algorithms -- often variants of stochastic gradient descent -- exhibit surprising effectiveness in fitting large neural networks in practice. We argue that neural network loss landscapes often contain (nearly) a single basin after accounting for all possible permutation symmetries of hidden units a la Entezari et al. 2021. We introduce three algorithms to permute the units of one model to bring them into alignment with a reference model in order to merge the two models in weight space. This transformation produces a functionally equivalent set of weights that lie in an approximately convex basin near the reference model. Experimentally, we demonstrate the single basin phenomenon across a variety of model architectures and datasets, including the first (to our knowledge) demonstration of zero-barrier linear mode connectivity between independently trained ResNet models on CIFAR-10. Additionally, we identify intriguing phenomena relating model width and training time to mode connectivity. Finally, we discuss shortcomings of the linear mode connectivity hypothesis, including a counterexample to the single basin theory.

1 INTRODUCTION

The paper investigates why SGD succeeds on difficult deep-learning optimization problems and whether permutation symmetries explain common structure among independently trained solutions. It proposes methods to align and merge models while examining the limits of linear mode connectivity.

  • Research questions: The paper asks why SGD performs well on high-dimensional non-convex deep-learning landscapes and how independently trained models can be linearly connected.These questions concern optimization behavior, local minima, and relationships between trained solutions.
  • Permutation symmetries: Permutation symmetries may provide an uncharacterized invariance underlying similarities between independent training runs.Hidden units can be exchanged with corresponding weight adjustments without changing network functionality.
  • Geometric structure: The paper frames deep-learning loss landscapes as containing extensive geometric repetition from permutation symmetries.Table 1 compares the number of model permutation symmetries with an upper estimate for atoms in the observable universe.
  • Contributions: The paper proposes three algorithms to align independently trained model weights and investigates their relationship to optimization procedures and linear mode connectivity.The methods are grounded in combinatorial optimization, with approximation algorithms proposed where matching problems are hard.
  • Experiments: The study evaluates permutation-based connectivity across multiple architectures and datasets, including independently trained ResNets.Experiments cover MLPs, CNNs, and ResNets on MNIST, CIFAR-10, and CIFAR-100.

2 BACKGROUND

This section defines permutation symmetries and linear interpolation between equivalent-architecture models, then motivates loss barriers as a practical measure of approximate convexity. Permuting hidden-unit representations can preserve function while potentially bringing two trained solutions into a shared basin.

  • Setup: An L-layer MLP is used to present the methods, with weights, biases, activations, and a loss function defining the model and optimization setting.The activation σ is element-wise, and L(Θ) measures how suitable weights Θ are for a goal such as fitting a dataset.
  • Permutation symmetry: Permutation symmetry reorders intermediate-layer output features while corresponding input weights in the next layer are reordered by P^T.The transformation preserves the network's computed function when the adjacent weights are adjusted consistently.
  • Permutation symmetry: Every trained weight assignment belongs to an equivalence class of functionally identical assignments, with the selected representative depending on training randomness.The paper denotes a functionality-preserving permutation of weights as π(Θ).
  • Linear mode connectivity: The central task is to find a permutation π such that interpolating between ΘA and π(ΘB) maintains performance comparable to both independently trained models.Models may differ in random initialization, data order, hyperparameters, or datasets while sharing an equivalent architecture.
  • Geometric intuition: Figure 1 depicts model B being permuted into a functionally equivalent π(ΘB) positioned in the same basin as ΘA before merging.The figure presents permutation as a teleportation between equivalent representations in weight space.
  • Loss barriers: Loss barriers quantify deviations from convex interpolation, with zero indicating flat or positively curved interpolation between similarly performing endpoints.The barrier is the maximum interpolated loss above the endpoint loss level.

3 PERMUTATION SELECTION METHODS

The paper introduces three methods for aligning independently trained models by permuting hidden units, enabling weight-space merging and improved linear mode connectivity. These methods differ in their use of activations, weights, or learned straight-through optimization, with distinct computational trade-offs.

  • Overview: Three algorithms align independently trained neural networks by selecting unit permutations before merging their weights.The methods include activation matching, weight matching, and straight-through estimator matching.
  • 3.1 MATCHING ACTIVATIONS: Activation matching uses constrained regression between layer activations and solves each layer’s permutation as a linear assignment problem.The activation matrices are computed over training data, and efficient polynomial-time assignment algorithms provide the matching.
  • 3.2 MATCHING WEIGHTS: The full weight-alignment objective is NP-hard for more than two layers, so the paper uses coordinate descent that repeatedly solves layerwise assignment problems.The underlying sum of bilinear assignments problem permutes both rows and columns, unlike classical linear assignment.
  • Results: Across MNIST, CIFAR-10, and ImageNet, permutation-based methods substantially improve over naïve interpolation; MNIST achieves zero-barrier connectivity with all three methods.Weight and activation matching perform similarly, while straight-through matching performs best at much higher computational cost.
  • 3.2 MATCHING WEIGHTS: Weight matching avoids the input data distribution and is generally much faster, but it can be disadvantaged by ignoring the loss landscape.Experiments found it surprisingly competitive with data-aware methods and typically convergent within seconds to a few minutes.
  • 3.3 LEARNING PERMUTATIONS WITH A STRAIGHT-THROUGH ESTIMATOR: Straight-through estimator matching learns an approximate permutation through a conventional training loop, enforcing realizable projections forward while using unrestricted gradients backward.It produces the best solutions but has a very steep computational cost relative to the other methods.
  • Results: Linear mode connectivity emerges gradually during training and is difficult to obtain at initialization, while the CIFAR-10 MLP exhibits greater variance than the MNIST MLP.The figure reports loss barriers over training time and uses different y-axis scales in its inlaid plots.

4 A COUNTEREXAMPLE TO UNIVERSAL LINEAR MODE CONNECTIVITY

The paper gives a counterexample showing that linear mode connectivity is not universal across all solutions. Its findings restrict the single-basin claim to solutions produced by SGD-like optimization and leave other invariances unresolved.

  • Counterexample: Adversarial, non-SGD solutions can exist for which no permutation of hidden units produces linear mode connectivity.This counterexample supports the view that common optimization procedures are biased toward LMC-compatible solutions.
  • Scope of the conjecture: The counterexample does not contradict the conjecture because that conjecture is restricted to SGD solutions.The precise mechanism producing this optimization bias remains an open question.
  • Open limitations: Permutation symmetry is not the only possible invariance: features can move between layers and layers can be rescaled, but their interaction with optimization remains unclear.These additional symmetries define an unresolved boundary for explanations based only on permutations.

5 EXPERIMENTS

Across datasets and architectures, permutation matching substantially improves interpolation between independently trained models. Connectivity improves with width and training, while model merging can outperform both inputs but remains subject to metric and search limitations.

  • Experimental setup: Models are separately trained from different initializations and batch orders, then evaluated along interpolation paths after permuting one model’s weights.The goal is a flat or convex loss slice between the aligned models.
  • Loss landscapes: Zero-barrier linear mode connectivity is achieved on MNIST with all three matching methods, and matched interpolation can outperform both input models.Naive interpolation substantially degrades performance, whereas weight matching and STE matching make the test-loss landscape convex on MNIST.
  • Loss landscapes: On ImageNet, matching decreases the barrier by 67% relative to naive interpolation but does not achieve zero-barrier connectivity.The authors hypothesize that larger width may improve ImageNet connectivity by providing sufficient capacity.
  • Matching methods: Weight matching offers a fast balance of cost and quality, while STE matching performs best and activation matching is slightly less favorable on MNIST.Weight matching runs in seconds on current hardware and does not require access to the input data distribution.
  • Onset and limitations: Linear mode connectivity emerges during training rather than holding throughout weight space, and failure to find a permutation cannot disprove a satisfactory permutation exists.The astronomical number of possible permutations makes the conjecture effectively impossible to disprove for realistically wide models.
  • Onset and limitations: The loss-barrier metric can remain positive even when interpolation contains points with lower loss than either endpoint.Because barriers are non-negative, this metric does not by itself capture every favorable interpolation shape.
  • Effect of model width: Wider CIFAR-10 models exhibit lower loss barriers, with larger widths reaching zero; thin 1× models do not appear linearly mode connected.The study includes VGG-16 and ResNet20, culminating in zero-barrier connectivity between large ResNet models.
  • Model merging: Weight matching merges ResNet models trained on disjoint biased CIFAR-100 subsets into a model with lower test loss than either input, using half the ensemble’s memory and compute.The merged model does not match full-data training or ensembling, but it is better calibrated than either input model.

6 RELATED WORK

Related work situates Git Re-Basin among studies of mode connectivity, loss landscapes and training dynamics, and federated learning or model merging. These lines of work motivate permutation-aware analysis and distributed model combination.

  • Mode connectivity: Prior studies established nonlinear and linear mode connectivity, including connections to the lottery ticket hypothesis and permutation symmetries.The paper extends this literature by studying connectivity after aligning hidden-unit permutations.
  • Loss landscapes and training dynamics: Research on loss landscapes examined shared features, feature differences, layer importance, and connected low-loss volumes across independently trained networks.These studies provide context for asking whether independently trained solutions occupy related regions of weight space.
  • Federated learning and model merging: Federated-learning work introduced learning across distributed devices and datasets, including permutation-aware averaging and model-combination methods.The paper differs from earlier permutation-based federated learning by targeting independently trained models and supporting broader architectures, including residual connections.

7 DISCUSSION AND FUTURE WORK

The paper argues that permutation symmetries help explain linear mode connectivity but are not a complete account, and it highlights unresolved technical, ethical, and research questions around model merging.

  • Future work: Permutation symmetries may be necessary but are not sufficient to explain linear mode connectivity, motivating study of broader linear relationships between activations.The paper points to cross-layer scaling and general linear relationships as future directions.
  • Risks of model merging: Model merging raises ethical questions about inherited biases, forgotten rare examples, and possible dataset gerrymandering.These concerns apply to the behavior and representativeness of merged models.
  • Risks of model merging: Deployment of merged models should include thorough auditing for representativeness across the entire data distribution.The paper specifically recommends checking whether the merged model represents the entirety of the data distribution.
  • Limitations and future work: The authors present their techniques as non-universal and anticipate additional failure modes beyond those currently identified.They frame understanding when and why permutation-based merging is feasible as an open research problem.

A.2 EXTENDED RELATED WORK

The paper positions permutation-aware linear connectivity as distinct from nonlinear paths and soft matching, while claiming broader, faster, and more direct evidence than prior work.

  • Mode connectivity: Unlike nonlinear mode-connectivity methods, this work seeks linear connectivity after aligning model weights through permutation symmetries.The linear restriction is presented as having direct implications for a single-basin theory.
  • Model matching: The paper’s weight matching jointly aligns units across all layers, using downstream and upstream weight relationships rather than greedy single-pass matching.This holistic approach is contrasted with methods that use only immediately preceding-layer information.
  • Model matching: Algorithm 1 works with models of nearly arbitrary architecture, including settings with bias terms, skip connections, or normalization layers.The comparison states that Singh and Jaggi’s “wts” method was not run on such models.
  • Single-basin conjecture: The paper strengthens the single-basin conjecture by directly finding permutations that place independently trained solutions into the same basin.It contrasts this with earlier statistical evidence that did not establish linear mode connectivity.
  • Single-basin conjecture: The authors claim the first demonstration, to their knowledge, of linear mode connectivity between large models independently trained on challenging datasets.Their weight matching removes barriers on larger models and datasets than the single-hidden-layer MNIST demonstrations discussed for prior work.

A.3.2 VGG-16 AND RESNET MODELS ON CIFAR DATASETS

The appendix examines merging and counterexamples across CIFAR experiments, finding that alignment can succeed while thin-model connectivity and universal single-basin claims remain limited.

  • Normalization: LayerNorm architectures can be merged without issue, whereas BatchNorm generally breaks after weight interpolation because of variance collapse.The paper recommends recalculating batch statistics after merging BatchNorm models.
  • Algorithm 1: Algorithm 1 exactly recovered known random permutations in 3-4 passes over the layers in a handful of experiments.This result concerns recovery of deliberately selected permutations.
  • Counterexample: A two-hidden-layer ReLU MLP counterexample has no permutation that yields linear mode connectivity between two perfectly fitting models.The models implement the same classification task while testing the two input conditions in reversed layer order.
  • Counterexample: The counterexample suggests that linear mode connectivity can depend on training procedures rather than model architecture alone.The precise connection between optimization algorithms and linear mode connectivity remains unclear.
  • Counterexample: The counterexample does not contradict Entezari et al.’s conjecture because that conjecture restricts its claim to a high-probability subset containing SGD solutions.The restriction to SGD solutions is described as a load-bearing part of the conjecture.
  • Matching strategy: Weight matching jointly considers all relevant layers, unlike greedy matching that proceeds layer by layer using limited local information.The paper presents this as enabling use of downstream-weight relationships.

A.7.1 EXPERIMENTAL COMPARISON

Experimental comparisons show that the paper’s weight matching improves over prior merging methods in speed and performance, including a large ImageNet accuracy difference.

  • ResNet50/ImageNet: 51.01% top-1 accuracy is achieved on ImageNet, compared with 1.38% for OT-Fusion, which only marginally improves over naïve averaging.The comparison uses ResNet50 models and recalculates BatchNorm statistics after interpolation for all methods.
  • VGG11/CIFAR-10: 4.5× faster implementation and better model performance are reported for weight matching versus Singh and Jaggi’s “wts” method on VGG11/CIFAR-10.The experiment reproduces a prior study using its published model weights.
  • Alignment passes: Weight matching can continue optimization passes until convergence, whereas OT-Fusion is limited to a single pass over the network layers.The paper identifies the number of alignment passes as substantially affecting performance.
  • Failure case: In an identity-function example, greedy matching produces a merged model that fails to represent the inputs, while weight matching accurately reflects both models and improves over model B.The example compares the two alignment choices under linear interpolation.

A.10 MERGING MANY MODELS

Algorithm 3 extends model merging to more than two models, and experiments report improved test loss, extrapolation, and probability calibration. The approach is evaluated on MNIST MLPs and CIFAR-100 split-data models, with calibration comparable to ensembling at lower resource cost.

  • Algorithm 3 merges the weights of more than two models simultaneously.
  • The multi-model merging algorithm converges quickly to solutions that extrapolate better than individual models and improve probability-estimate calibration.
  • CIFAR-100 split-data merging produced better-calibrated probabilities than either input model, though top-1 accuracy was not competitive.Calibration matched model ensembling while requiring 2× less memory and compute.
  • 43% lower test loss was reported when five separately trained MNIST MLPs were merged.The merged model also appeared to have better out-of-distribution performance than any input model.

A.11 FAILED IDEA: A METHOD FOR STEEPEST DESCENT

A proposed steepest-descent matching method selects permutations using the gradient at one model and reduces the choice to a matching problem. In practice, it found steepest-descent directions but encountered high barriers before reaching the target permutation.

  • The method chooses among permuted versions of ΘB according to the steepest-descent direction from ΘA.This choice assumes the loss landscape is convex modulo permutation symmetries.
  • Selecting the best permutation becomes a SOBLAP matching problem, with a fast exact solution for a single intermediate layer.
  • In practice, steepest-descent directions were accompanied by high barriers between the initial loss decrease and the target permutation.
  • Algorithm 3 applied to 32 MNIST MLPs produced calibration competitive with model ensembling.The merged model showed improved calibration on both training and test datasets.

A.12 PROOF OF LEMMA 1

The proof reduces a permutation optimization problem to quadratic assignment, then uses a sufficiently large penalty to separate equal- and unequal-permutation cases. It also establishes termination of Algorithm 1 through finite-state progress.

  • The hardness proof proceeds by reduction from the quadratic assignment problem (QAP).
  • A sufficiently large λ makes all P = Q solutions better than all P ≠ Q solutions.The supplied bound is λ > 5dα.
  • When P = Q, the objective reduces to the QAP objective after the remaining terms become constants.
  • QAP is strongly NP-hard, and MaxQAP admits no PTAS.
  • Algorithm 1 terminates because each update strictly increases a utility over finitely many possible permutations, preventing cycles.
Loading 2209.04836v6…