G IT R E -B ASIN : M ERGING M ODELS M ODULO P ERMUTATION S YMMETRIES
Anonymous authors
Paper under double-blind review
A BSTRACT
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.
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 and CIFAR-100.
Additionally, we investigate 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.
I NTRODUCTION
We investigate the unreasonable effectiveness of stochastic gradient descent (SGD) algorithms on the high-dimensional non-convex optimization problems of deep learning.
In particular,
Why does SGD thrive in optimizing high-dimensional non-convex deep learning loss landscapes despite being noticeably less robust in other non-convex optimization settings, like policy learning (Ainsworth et al., 2021), trajectory optimization (Kelly, 2017)
and recommender systems (Kang et al., 2016)?
What are all the local minima?
When linearly interpolating between initialization and final trained weights, why does the loss smoothly and monotonically decrease (Goodfellow & Vinyals, 2015; Frankle, 2020; Lucas et al., 2021; Vlaar & Frankle, 2021)?
How can two independently trained models with different random initializations and data batch orders inevitably achieve nearly identical performance?
Furthermore, why do their training loss curves often look identical?
We posit that the final phenomenon referenced in item 3 points to the existence of some yet uncharacterized invariance(s) in the training dynamics that cause independent training runs to exhibit almost identical characteristics.
Hecht-Nielsen (1990) noted the permutation symmetries of hidden units in neural networks; briefly, one can swap any two units of a hidden layer in a network and – assuming weights are adjusted accordingly – network functionality will not change.
Recently, Entezari et al.
(2021) conjectured that these permutation symmetries may let us linearly connect points in weight space with no detriment to the loss.
Conjecture 1 (Permutation invariance, informal (Entezari et al., 2021)) .
Most SGD solutions belong to a set whose elements can be permuted so that no barrier (as in Definition 2.2) exists on the linear interpolation between any two permuted elements.
[Table]
We refer to such solutions as being linearly mode connected (LMC)
(Frankle et al.
If true,
Conjecture 1 will both materially expand our understanding of how SGD works in the context of deep learning and offer a credible explanation for the preceding phenomena, in particular.
Contributions.
In this paper, we attempt to uncover what invariances may be responsible for the phenomena cited above and the unreasonable effectiveness of SGD in deep learning.
We make the following contributions:
Matching methods.
We propose three new algorithms, grounded in concepts and techniques from combinatorial optimization, to align the weights of two independently trained models.
Where appropriate, we prove hardness results for these problems and propose approximation algorithms.
Our fastest method identifies permutations in mere seconds on current hardware.
Relationship to optimization algorithms.
We demonstrate by means of counterexample that linear mode connectivity is an emergent property of training procedures, not of model architectures.
We connect this result to prior work on the implicit biases of SGD.
Experiments, including zero-barrier LMC for ResNets.
Empirically, we explore the existence of linear mode connectivity modulo permutation symmetries in experiments across MLPs, CNNs, and ResNets trained on MNIST, CIFAR-10, and CIFAR-100.
We contribute the first-ever demonstration of zero-barrier LMC between two independently trained ResNets.
We explore the relationship between LMC and model width as well as training time.
Finally, we show evidence of our methods’ ability to combine models trained on independent datasets into a merged model that outperforms both input models in terms of test loss (but not accuracy) and is no more expensive in compute
or memory than either input model.
B ACKGROUND
Although our methods can be applied to arbitrary model architectures, we proceed with the multilayer perceptron (MLP) for its ease of presentation (Bishop, 2007).
Consider an L -layer MLP,
[Equation]
where σ denotes an element-wise nonlinear activation function.
Furthermore, consider a loss, L (Θ) , that measures the suitability of a particular set of weights Θ towards some goal, e.g., fitting to a training dataset.
Central to our investigation is the phenomenon of permutation symmetries of weight space.
Given Θ , we can apply some permutation to the output features of any intermediate layer, ℓ , of the model, denoted by a permutation matrix P ∈ S d ,
[Equation]
for σ , an element-wise operator.
It follows that as long as we reorder the input weights of layer ℓ + according to P ⊤ , we will have a functionally equivalent model.
To be precise, if we define Θ ′ to be identical to Θ with the exception of
[Equation]
then the two models are functionally equivalent: f ( x ; Θ) =
f ( x ; Θ ′ ) for all inputs x .
This implies that for any trained weights Θ , there is an entire equivalence class of functionally equivalent weight assignments, not just one such Θ , and convergence to any one specific element of this equivalence class, as opposed to any others, is determined only by random seed.
We denote a functionalitypreserving permutation of weights as π (Θ) .
[Figure]
Consider the task of reconciling the weights of two, independently trained models, A and B , with weights Θ A and Θ B , respectively, such that we can linearly interpolate between them.
We assume that models A and B were trained with equivalent architectures but different random initializations, data orders, and potentially different hyperparameters or datasets, as well.
Our central question is: Given Θ A and Θ B , can we identify some π such that when linearly interpolating between Θ A and π (Θ B ) , all intermediate models enjoy performance similar to Θ A and Θ B ?
We base any claims of loss landscape convexity on the usual definition of multi-dimensional convexity in terms of one-dimensional convexity per
Definition 2.1 (Convexity) .
A function f : R D → R is convex if every one-dimensional slice is convex, i.e., for all x, y ∈ R D , the function g ( λ )
= f
((1 − λ ) x + λy ) is convex in λ .
Due to Definition 2.1, it suffices to show that arbitrary one-dimensional slices of a function are convex in order to reason about the convexity of complex, high-dimensional functions.
In practice, we rarely observe perfect convexity but instead hope to approximate it as closely as possible.
Following Draxler et al.
(2018), Entezari et al.
(2021), Frankle et al.
(2020) and others, we measure approximations to convexity via “barriers.”
Definition 2.2 (Loss barrier (Frankle et al., 2020)) .
Given two points Θ A , Θ B such that L (Θ A )
≈
L (Θ B ) , the loss barrier is defined as max λ ∈ [0 , 1] L ((1 − λ )
Θ A + λ Θ B ) −
( L (Θ A ) + L (Θ B )) .
Loss barriers are non-negative, with zero indicating an interpolation of flat or positive curvature.
P ERMUTATION S ELECTION M ETHODS
We introduce three methods of matching units between model A and model B .
Further, we present an extension to simultaneously merging multiple models in Appendix A.10 and an appealing but failed method in Appendix A.11.
3. M ATCHING A CTIVATIONS
Following the classic Hebbian mantra, “[neural network units] that fire together, wire together” (Hebb, 2005), we propose associating units across two models by performing regression between their activations.
Matching activations between models is compelling since it captures the intuitive notion that two models must learn similar features to accomplish the same task (Li et al., 2016).
Provided activations for each model, we aim to associate each unit in A with a unit in B .
It stands to reason that a linear relationship may exist between the activations of the two models.
We fit this into the regression framework by constraining ordinary least squares (OLS) to solutions in the set of permutation matrices, S d .
For activations of the ℓ ’th layer, let Z ( A ) , Z ( B ) ∈ R d × n denote the d -dim.
activations for all n training data points in models A and B , respectively.
Then,
[Equation]
where ⟨ A , B ⟩ F = (cid:80) i,j A
i,j B i,j denotes the Frobenius inner product between real-valued matrices
A and B .
Conveniently, eq.
(1) constitutes a “linear assignment problem” (LAP) (Bertsekas, 1998) for which efficient, practical algorithms are known.
Having solved this assignment problem on each layer, we can then permute the weights of model B to match model A as closely as possible
[Equation]
for each layer ℓ , producing weights Θ ′ with activations that align as closely possible with Θ A .
Computationally, this entire process is relatively lightweight: the Z ( A ) and Z ( B ) matrices can be computed in a single pass over the training dataset, and, in practice, a full run through the training dataset may be unnecessary.
Solving eq.
(1) is possible due to well-established, polynomial-time algorithms for solving the linear assignment problem (Kuhn, 2010; Jonker & Volgenant, 1987; Crouse, 2016).
Also, conveniently, the activation matching at each layer is independent of the matching at every other layer, resulting in a separable and straightforward optimization problem; this advantage will not be enjoyed by the following methods.
Dispensing with regression, one could similarly associate units by matching against a matrix of cross-correlation coefficients in place of Z ( A ) ( Z ( B ) )
⊤ .
We observed correlation matching to work equally well but found OLS regression matching to be more principled and easier to implement.
Activation matching has been independently studied in Tatro et al.
; Singh & Jaggi (2020) albeit not from the perspective of OLS regression.
3. M ATCHING W EIGHTS
Instead of associating units by their activations, we could alternatively inspect the weights of the model itself.
Consider the first layer weights, W 1 ; each row of W 1 corresponds to a single feature.
If two such rows were equal, they would compute exactly the same feature (ignoring bias terms for the time being).
And, if [ W ( A )] i,
[Equation]
≈ [ W
[Equation]
B )
[Equation]
] j, : , it stands to reason that units i and j should be associated.
Extending this idea to every layer, we are inspired to pursue the optimization
[Equation]
We can re-express this in terms of the full weights,
[Equation]
resulting in another matching problem.
We term this formulation the “sum of bilinear assignments problem” (SOBLAP).
Unfortunately, this matching problem is thornier than the classic linear assignment matching problem presented in eq.
Unlike LAP, we are interested in permuting both the rows and columns of W ( B ) ℓ to match W ( A ) ℓ , which fundamentally differs from permuting only rows or only columns.
We formalize this difficulty as follows.
Lemma 1.
The sum of a bilinear assignments problem (SOBLAP) is NP-hard and admits no polynomial-time constant-factor approximation scheme for L > 2 .
Lemma 1 contrasts starkly with classical LAP, for which polynomial-time algorithms are known.
Undeterred, we propose a approximation algorithm for SOBLAP.
Looking at a single P ℓ while holding the others fixed, we observe that the problem can be reduced to a classic LAP,
[Equation]
This leads to a convenient coordinate descent algorithm: go through each layer and greedily select its best P ℓ .
Repeat until convergence.
We present this in Algorithm 1.
Although we present Algorithm 1 in terms of an MLP without bias terms, in practice our implementation can handle the weights of models of nearly arbitrary architectures, including bias terms, residual connections, convolutional layers, attention mechanisms, and so forth.
We propose an extension of Algorithm 1 to merging more than two models at a time in Appendix A.10.
[Figure]
Algorithm 1: P ERMUTATION C OORDINATE D ESCENT
Given: Model weights Θ A
= (cid:110) W ( A ), . . .
, W ( A ) L (cid:111) and Θ B = (cid:110) W ( B ), . . .
, W ( B ) L (cid:111)
Result:
A permutation π = { P 1 , . .
, P L − 1 } of Θ B such that vec (Θ A ) · vec ( π (Θ B )) is approximately maximized.
Initialize: P 1 ← I , . .
, P L − 1 ← I repeatfor ℓ ∈ R ANDOM P ERMUTATION (1 , . . .
, L − 1) do
P ℓ ← S OLVE LAP (cid:16) W ( A ) ℓ
P ℓ − 1 ( W ( B ) ℓ ) ⊤ + ( W ( A ) ℓ +1 )
⊤ P ℓ +1 W ( B ) ℓ +1 (cid:17) end until convergence
Lemma 2.
Algorithm 1 terminates.
Our experiments showed this algorithm to be fast in terms of both iterations necessary for convergence and wall-clock time, generally on the order of seconds to a few minutes.
Unlike the activation matching method presented in Section 3.1, weight matching ignores the data distribution entirely.
Ignoring the input data distribution and therefore the loss landscape handicaps weight matching but allows it to be much faster.
We therefore anticipate its potential application in fields such as finetuning (Devlin et al., 2019; Wortsman et al., 2022b;a), federated learning (McMahan et al., 2017;
Kone ˇ cn ´ y et al., 2016a;b), and model patching (Matena & Raffel, 2021; Sung et al., 2021; Raffel, 2021).
In practice, we found weight matching to be surprisingly competitive with data-aware methods.
Section 5 studies this trade-off.
3. L EARNING P ERMUTATIONS WITH A S TRAIGHT - THROUGH E STIMATOR
Inspired by the success of straight-through estimators (STEs) in other discrete optimization problems (Bengio et al., 2013; Kusupati et al., 2021; Rastegari et al., 2016; Courbariaux & Bengio, 2016), we attempt here to “learn” the ideal permutation of weights π (Θ B ) .
Specifically, our goal is to optimize
[Equation]
where ˜Θ B denotes an approximation of π (Θ B ) , allowing us to implicitly optimize π .
However, eq.
(3) involves inconvenient non-differentiable projection operations, proj ( · ) , complicating the optimization.
We overcome this via a “straight-through” estimator: we parameterize the problem in terms of a set of weights ˜Θ B
≈ π (Θ B ) .
In the forward pass, we project ˜Θ B to the closest realizable π (Θ B ) .
In the backwards pass, we then switch back to the unrestricted weights ˜ Θ B .
In this way, we
[Figure]
are guaranteed to stay true to the projection constraints in evaluating the loss but can still compute usable gradients at our current parameters, ˜Θ B .
Conveniently, we can re-purpose Algorithm 1 to solve proj (˜Θ B ) .
Furthermore, we found that initializing ˜Θ B =
Θ A performed better than random initialization.
This is to be expected immediately at initialization since the initial matching will be equivalent to the weight matching method of Section 3.1.
However, it is not immediately clear why these solutions continue to outperform a random initialization asymptotically.
Unlike the aforementioned methods, Algorithm 2 attempts to explicitly “learn” the best permutation π using a conventional training loop.
By initializing to the weight matching solution of Section 3.
and leveraging the data distribution as in Section 3.1, it seeks to offer a best-of-both-worlds solution.
However, this comes at a very steep computational cost relative to the other two methods.
I NTERLUDE : A C OUNTEREXAMPLE
Before arguing for the presence of linear mode connectivity, we step back to consider its limitations.
In Appendix A.6 we present a counterexample proving that there exist models that do not enjoy LMC under any permutation of weights.
Since our two networks have small width, we can simply inspect all four possible permutations of the intermediate units.
We visualize this in Figure 6.
Critically, no permutation results in LMC.
We present this counterexample in part to establish some basic intuition for Conjecture 1 but more importantly to highlight that any success interpolating between permuted networks is due to inherent bias in the optimization algorithms used and the types of basins they are attracted to, not the network architectures themselves.
To the extent we achieve LMC in practice, it follows that it is an artifact of our optimization methods, not our architectures.
We also note that there are invariances beyond permutation symmetries: It is possible to move features between layers, re-scale layers, and so forth.
Prior works noted the feature/layer association (Nguyen et al., 2021) and re-scaling invariances (Ainsworth et al., 2018).
The importance of these other symmetries and their interplay with optimization algorithms remains unclear.
E XPERIMENTS
Our base methodology is to separately train two models, A and B , starting from different random initializations and with different random batch orders, resulting in trained weights Θ A and Θ B ,
[Figure]
respectively.
We then evaluate slices through the loss landscape, L ((1 − λ )
Θ A + λπ (Θ B )) for λ ∈ [0 , 1] , where π is selected according to the methods presented in Section 3.
Ideally, we seek a completely flat or even convex one-dimensional slice.
As discussed in Section 2, the ability to exhibit this behavior for arbitrary Θ A , Θ B empirically suggests that the loss landscape contains only a single basin modulo permutation symmetries.
We remark that a failure to find a π such that linear mode connectivity holds cannot rule out the existence of a satisfactory permutation.
Given the astronomical number of permutation symmetries,
Conjecture 1 is essentially impossible to disprove for any realistically wide model architecture.
5. L OSS L ANDSCAPES B EFORE AND A FTER M ATCHING
We present results for models trained on MNIST (LeCun et al., 1998), CIFAR-10 (Krizhevsky, 2009), and ImageNet (Deng et al., 2009) in Figure 2.
Na ¨ ıve interpolation ( π (Θ B ) =
Θ B ) substantially degrades performance when interpolating.
On the other hand, the methods introduced in
Section 3 can achieve much better barriers.
We achieve zero-barrier linear mode connectivity on
MNIST with all three methods, although activation matching performs just slightly less favorably than weight matching and straight-through estimator (STE) matching.
We especially note that the test loss landscape becomes convex after applying our weight matching and STE permutations!
In other words, our interpolation actually yields a merged model that outperforms both models A and
B .
We elaborate on this phenomenon in Section 5.4 and Appendix A.10.
On ImageNet we fall short of zero-barrier connections, although we do see a 67% decrease in barrier relative to na ¨ ıve interpolation.
As we demonstrate in Section 5.3, we can achieve zero-barrier LMC on CIFAR-10 with large ResNet models.
Therefore, we hypothesize that the presence of LMC depends on the model having sufficient capacity (esp.
width) to capture the complexity of the input data distribution, and that ImageNet results could be improved by expanding model width.
STE matching, the most expensive method, produces the best solutions.
Somewhat surprising, however, is that the gap between STE and the other two methods is relatively small.
In particular, it is remarkable how well Algorithm 1 performs without access to the input data at all.
We found that weight matching offered a compelling balance between computational cost and performance: It runs in mere seconds (on modern hardware) and produces high-quality solutions.
5. O NSET OF M ODE C ONNECTIVITY
Given the results of Section 5.1, it may be tempting to conclude that the entirety of weight space contains only a single basin modulo permutation symmetries.
However, we found that linear mode connectivity is an emergent property of training, and we were unable to uncover it early in training.
We explore the emergence of LMC in Figure 3.
Concurrent to our work, Benzing et al.
(2022) showed that LMC at initialization is possible using a permutation found at the end of training.
Note that the final inlaid interpolation plot in Figure 3(right) demonstrates an important shortcoming of the loss barrier metric, i.e., the interpolation includes points with lower loss than either of the two models.
However, the loss barrier is still positive due to non-negativity, as mentioned in Section 2.
E FFECT OF M ODEL W IDTH
Conventional wisdom maintains that wider architectures are easier to optimize (Jacot et al., 2018;
Lee et al., 2019).
We now investigate whether they are also easier to linearly mode connect.
We train
VGG-16 (Simonyan & Zisserman, 2015) and ResNet20 (He et al., 2016) architectures of varying widths on the CIFAR-10 dataset.
Results are presented in Figure 4.
A clear relationship emerges between model width and linear mode connectivity, as measured by the loss barrier between solutions.
Although 1 × -sized models did not seem to exhibit linear mode connectivity, we found that larger width models decreased loss barriers all the way to zero.
In
Figure 4(right), we show what is to our knowledge the premiere demonstration of zero-barrier linear mode connectivity between two large ResNet models trained on a non-trivial dataset.
We highlight that relatively thin models do not seem to obey linear mode connectivity yet still exhibit similarities in training dynamics.
This suggests that either our permutation selection methods are failing to find satisfactory permutations on thinner models or that some form of invariance other than permutation symmetries must be at play in the thin model regime.
5. M ODEL P ATCHING , S PLIT D ATA T RAINING , AND I MPROVED C ALIBRATION
[Figure]
can be merged constructively.
Algorithmmakes it possible for two ResNet models trained on disjoint, biased subsets of CIFAR-100 to be merged in weight space such that their combination outperforms both input models in terms of test loss on the combined dataset.
Inspired by work on finetuning (Wortsman et al., 2022a), model patching (Singh &
Jaggi, 2020; Raffel, 2021), and federated learning (McMahan et al., 2017; Kone ˇ cn ´ y et al., 2016a;b), we study whether it is possible to synergistically merge the weights of two models trained on disjoint datasets.
Consider, for example, an organization with multiple (possibly biased) datasets separated for regulatory (e.g., GDPR) or privacy (e.g., on-device data) considerations.
Models can be trained on each dataset individually, but training in aggregate is not feasible.
Can we combine separately trained models so that the merged model performs well on the entirety of the data?
To address this question, we split the CIFAR100 dataset (Krizhevsky, 2009) into two disjoint subsets: dataset A , containing 20% examples labelled 0-49 and 80% labelled 50-99, and dataset B , vice versa.
ResNet20 models A
and
B were trained on their corresponding datasets.
Privacy requirements mandate that we utilize a data-agnostic algorithm like Algorithm 1.
Figure 5 shows the result of merging the two models with weight matching.
For comparison, we benchmark na¨ıve weight interpolation, ensembling of the model logits, and full-data training.
As expected, merging separately trained models did not match the performance of an omniscient model trained on the full dataset or an ensemble of the two models with twice the number of effective weights.
On the other hand, we did manage to merge the two models in weight space, achieving an interpolated model that outperforms both input models in terms of test loss while using half the memory and compute required for ensembling.
Furthermore, the merged model’s probability estimates are better calibrated than either of the input models as demonstrated in Figure 11.
Accuracy results are presented in Figure 10.
Algorithm 1 also vastly outperformed
R ELATED W ORK
(Linear) mode connectivity.
Freeman & Bruna (2017) introduced the concept of mode connectivity, i.e., that SGD solutions in the loss landscape are connected by constant-loss curves in weight space.
Further explorations were undertaken in Garipov et al.
; Draxler et al.
(2018), among others.
Tatro et al.
(2020) explored non-linear mode connectivity modulo permutation symmetries.
Frankle et al.
(2020) demonstrated a connection between linear mode connectivity and the lottery ticket hypothesis.
Juneja et al.
(2022) demonstrated that LMC does not always hold, even when fine-tuning.
Hecht-Nielsen (1990); Chen et al.
(1993) noted the existence of permutation symmetries, and Brea et al.
(2019) implicated them as a source of saddle points in the loss landscape.
Recently, the prescient work of Entezari et al.
(2021) conjectured that SGD solutions could be linear mode connected modulo permutation symmetries and offered a battery of experiments buttressing this claim.
Unlike previous works on LMC we accomplish zero-barrier paths between two independently-trained models with an algorithm that runs on the order of seconds.
Loss landscapes and training dynamics.
Li et al.
; Yosinski et al.
(2014) investigated whether independently trained networks learn similar features, and to what extent they transfer.
Jiang et al.
(2021) argued that independently trained networks meaningfully differ in the features they learn in certain scenarios.
Zhang et al.
(2019) studied the relative importance of layers.
Benton et al.
(2021) argued that SGD solutions form a connected volume of low loss.
Pittorino et al.
(2022) proposed a toroidal topology of solutions and a set of algorithms for symmetry removal.
On the theoretical front, Kawaguchi (2016) proved that deep linear networks contain no local minima.
Boursier et al.
; Chizat & Bach (2018); Mei et al.
characterized the training dynamics of one-hidden layer networks, proving that they converge to zero loss.
Godfrey et al.
; Simsek et al.
(2021) investigated the algebraic structure of symmetries in neural networks and how this structure manifests in loss landscape geometry.
Federated learning and model merging.
McMahan et al.
; Kone ˇ cn ´ y et al.
(2016a;b) introduced the concept of “federated learning,” i.e., learning split across across multiple devices and datasets.
Wang et al.
(2020) proposed an exciting federated learning method in which model averaging is done after permuting units.
Unlike this work, they merged smaller “child” models into a larger “main” model, and did so with a layer-wise algorithm that does not support residual connections or normalization layers.
Raffel (2021); Matena & Raffel (2021); Sung et al.
(2021) conceptualized the study of “model patching,” i.e., the idea that models should be easy to modify and submit changes to.
Ilharco et al.
(2022) investigated model patching for the fine-tuning of open-vocabulary vision models.
Ashmore & Gashler (2015) first explored the use of matching algorithms for the alignment of network units.
Singh & Jaggi (2020) proposed merging models by soft-aligning associations weights, inspired by optimal transport.
Liu et al.
(2022a)
; Uriot & Izzo (2020) further explored merging models taking possible permutations into account.
Wortsman et al.
(2022a) demonstrated state-of-the-art ImageNet performance by averaging the weights of many fine-tuned models.
D ISCUSSION AND F UTURE W ORK
We explore the role of permutation symmetries in the linear mode connectivity of SGD solutions.
We introduce three novel algorithms to canonicalize independent neural network weights in order to make the loss landscape between them as flat as possible.
In contrast to prior work, we linearly mode connect large ResNet models with no barrier in seconds to minutes.
Despite presenting successes across multiple architectures and datasets, linear mode connectivity between thin models remains elusive.
Therefore, we conjecture that permutation symmetries are a necessary piece, though not a complete picture, of the fundamental invariances at play in neural network training dynamics.
In particular, we hypothesize that linear, possibly non-permutation, relationships connect the layerwise activations between models trained by SGD.
In the infinite width limit, there exist satisfactory linear relationships that are also permutations.
An expanded theory and empirical exploration of other invariances – such as cross-layer scaling or general linear relationships between activations – presents an intriguing avenue for future work.
Ultimately, we anticipate that a lucid understanding of loss landscape geometry will not only advance the theory of deep learning but will also promote the development of better optimization, federated learning, and ensembling techniques.
E THICS S TATEMENT
Merging models raises interesting ethical and technical questions about the resulting models.
Do they inherit the same biases as their input models?
Are rare examples forgotten when merging?
Is it possible to gerrymander a subset of the dataset by splitting its elements across many shards?
Deployment of any form of model merging ought to be paired with thorough auditing of the resulting model, investigating in particular whether the merged model is representative of the entirety of the data distribution.
R EPRODUCIBILITY S TATEMENT
Our code is open sourced at https://github.com/REDACTED/REDACTED .
Our experimental logs and downloadable model checkpoints are fully open source at https://wandb.ai/
REDACTED/REDACTED .
A. E XTENDED R ELATED W ORK
Non-linear mode connectivity.
A flourishing set of literature exists studying non-linear mode connectivity, including but not limited to Garipov et al.
; Draxler et al.
; Kuditipudi et al.
This insightful line of work is inspirational to our own, however we take a strictly linear approach to mode connectivity as in Frankle et al.
(2020); Juneja et al.
Restricting ourselves to linear trajectories comes with the advantage of having direct implications for a single-basin theory.
However, it comes at the cost of a more challenging, discrete optimization problem.
In particular, we found that – in contrast to non-linear mode connectivity – linear mode connectivity becomes drastically harder with smaller width models.
Note additionally that most pre-existing mode connectivity work does not account for permutation symmetries of weight space, a linchpin element of our work.
A notable exception to this trend can be found in Tatro et al.
To summarize: Freeman & Bruna (2017) introduced the notion of mode connectivity and proved that loss landscapes for single-hidden layer ReLU models contain only a single basin in the infinite width limit.
Garipov et al.
(2018) and Draxler et al.
(2018) concurrently demonstrated that simple zero-barrier curves can be learned to connect the optima in weight space, thus reshaping our understanding of practical loss landscape geometries.
Kuditipudi et al.
(2019) proposes a theoretical explanation for the mode connectivity phenomenon.
Benton et al.
(2021) extends mode connectivity from one-dimensional paths to entire manifolds of low-loss, and show that these manifolds can be leveraged for state-of-the-art Bayesian ensembling of models.
Relationship with Tatro et al.
The impact of permutation symmetries on the non-linear mode connectivity of models is considered in Tatro et al.
In particular, they independently propose an algorithm more-or-less equivalent to Section 3.1 but use it in conjunction with learned non-linear mode connecting curves.
In contrast, we show that linear mode connectivity can be achieved without the need for learning non-linear paths between the aligned weights.
Our derivation of Section 3.1 from the principle of least-squares regression is novel, to the best of our knowledge.
Relationship with Singh & Jaggi (2020).
Singh & Jaggi (2020) studies model merging with “soft matchings” between units from the perspective of optimal transport.
We emphasize the following commonalities/differences with their work
: • We focus on linear mode connectivity modulo permutation symmetries and its implications for a single-basin theory.
On the other hand, Singh & Jaggi (2020) emphasizes “soft” (ie., non-permutation) matching of units via optimal transport.
• The activation matching method of Singh & Jaggi (2020) reduces to that of Section 3.
when the optimal transport regularization term is set to zero and the unit “importance” values are set to uniform across all units on all layers.
• Our weight matching and straight-through estimator methods solve for an alignment across all layers jointly.
In contrast, Singh & Jaggi (2020) executes greedy, single-pass matching looking only at weight information from the immediately previous layer when selecting permutations.
Singh & Jaggi (2020) mentions jointly solving for alignments as an avenue for future work.
• The “wts” method of Singh & Jaggi (2020) is not run on models including bias terms, skip connections, or normalization layers.
In contrast, our Algorithm 1 works with models of nearly arbitrary architecture.
• Our weight matching method (Algorithm 1) outperforms the “wts” method of Singh & Jaggi (2020).
See Appendix A.7 for more information.
• Singh & Jaggi (2020) introduces a method for merging multiple models simultaneously, but only demonstrates results on at most 8 models at a time and performs continued training after merging.
In contrast, our Algorithm 3 has been shown to work with as many asmodels at a time and does not require continued training after merging.
Our analysis of the calibration of the resulting merged models has no parallel in Singh & Jaggi (2020).
Relationship with Entezari et al.
Entezari et al.
(2021) introduces the single-basin conjecture and provides the following evidence towards it:
• Entezari et al.
(2021) provides a statistical test which fails to detect a difference in barrier statistics between independently trained models and random permutations of the same model (Fig. 5 of Entezari et al.
Our work provides stronger support for the conjecture in that we give methods that can directly “unscramble” these permutations, proving that LMC can be found (Figures 4 and 5).
Entezari et al.
(2021)’s
experimental protocol does not provide evidence for linear mode connectivity.
Rather, their experimental results suggest that barriers resulting from independent training look like the barriers resulting from random permutations.
But this result is consistent with a world in which all solutions have barriers between them – both between members of the same permutation equivalence class and between solutions in separate equivalence classes!
In other words, there may still exist multiple equivalence classes of solutions.
In contrast, we provide concrete evidence for a single-basin theory by developing algorithms that directly place independent solutions into the same basin (Figure 1).
• Although Entezari et al.
’s conjecture is an important intellectual ancestor to our work, their demonstration of linear mode connectivity is limited to a single hidden-layer MLP on MNIST (Fig. 2 of Entezari et al.
However, this result for single hiddenlayer MLP models is preceded by Freeman & Bruna (2017);
Uriot & Izzo (2020).
On the other hand, we focus on larger models and datasets that are more closely aligned with models used in practice at the time of writing.
• Entezari et al.
(2021) proposes a simulated annealing algorithm that yields modest reductions in barrier between independently trained models, yet requires multiple days to run.
On the other hand, our weight matching algorithm (Algorithm 1) completely removes barriers between models for more challenging models and datasets (Figures 2 and 5), and runs in seconds (Appendix A.5).
Moreover, our weight matching method does not require access to the training data, enabling its potential application in domains like federated learning and distributed training.
In short, the work of Entezari et al.
(2021) first proposed the “single-basin” conjecture.
Our work is the first (to the best of our knowledge) to demonstrate that linear mode connectivity can be achieved between large models independently trained on challenging datasets.
Differentiating through permutations.
Akin to differentiable permutation learning, many prior works have studied differentiable sorting (Grover et al., 2019; Prillo & Eisenschlos, 2020; Cuturi et al., 2019; Petersen et al., 2022; 2021;
Mena et al., 2018).
Blondel et al.
(2020) studied differentiable sorting and ranking with asymptotics that correspond to their non-differentiable versions.
Fogel et al.
(2015) explored recovering the linear orderings of items based on pairwise information, another form of permutation optimization.
Bengio et al.
(2013) introduced the straight-through estimator for differentiating through discrete projections that we utilize in Section 3.3.
A. E XPERIMENTAL D ETAILS A.3. M ULTI - LAYER P ERCEPTRON MODELS
In these experiments we utilized networks with 3 hidden layers of 512 units each.
ReLU activations were used between layers and no normalization was performed.
Optimization was done with Adam and a learning rate of 1 e − 3 .
A.3. VGG-16 AND R ES N ET MODELS ON CIFAR DATASETS
We utilized the VGG-16 architecture of Simonyan & Zisserman (2015) with the exception that we used LayerNorm normalization in place of BatchNorm.
Similarly we used the ResNet20 architecture of He et al.
but with LayerNorms in place of BatchNorms.
The following data augmentation was performed during training
• Random resizes of the image between 0.8 × -1.2 × • Random 32 × 32 pixel crops • Random horizontal flips • Random rotations between ± 30 ◦
Optimization was done with SGD with momentum (momentum set to 0.9).
A weight decay regularization term of 5 e − 4 was applied.
A single cosine decay schedule with linear warm-up was used.
Learning rates were initialized at 1 e − 6 and linearly increased to 1 e − 1 over the span of an epoch.
After that point a single cosine decay schedule (Loshchilov & Hutter, 2017) was used for the remainder of training.
A.3. R ES N ET 50 MODELS ON I MAGE N ET -1K
In this experiment we utilized pre-trained ResNet50 model available for download online, and one trained ourselves.
These were standard ResNet50 models, including the use of BatchNorm.
In line with prior work (Izmailov et al., 2018; Wortsman et al., 2021; Maddox et al., 2019;
Wang et al., 2021;
Anonymous, 2023), we recalculate BatchNorm statistics after performing weight interpolation.
A note of chronology.
Our work was concurrent to that of Anonymous (2023).
However after our initial publication, the recalculation of BatchNorm statistics was suggested to us by the authors of Anonymous (
Therefore, the ImageNet experiments presented in Section 5.1 ought to be viewed as following the work of Anonymous (2023)
A. T HE R ELATIONSHIP BETWEEN P ERMUTATION M ATCHING AND N ORMALIZATION L AYERS
In this section we discuss the impact that different types of common normalization layers can have on the feasibility of model merging.
• BatchNorm (Ioffe & Szegedy, 2015) generally breaks after interpolating between weights due to the so-called “variance collapse” problem (Anonymous, 2023).
Therefore, following Anonymous (2023), we recommend the recalculation of batch statistics after merging models (Izmailov et al., 2018;
Wortsman et al., 2021; Maddox et al., 2019; Wang et al., 2021).
• LayerNorm
(Ba et al., 2016) is invariant to permutations of units
and we found that architectures with LayerNorm can be merged without issue.
• InstanceNorm (Ulyanov et al., 2016) also places no restrictions on unit order, and in principle does not present any issues, although we have not run any experiments with it.
• GroupNorm (Wu & He, 2020) relies on unit indexes to organize units into groups, and therefore is not invariant to permutations of units.
In principle, permutation alignment methods would not work on architectures with GroupNorm, though we have not tested this.
[Figure]
d e n t
[Figure]
t
[Figure]
a p l a y e r
S w a p l a y e r
S w a p b o t h
[Figure]
A. A DDITIONAL I NFORMATION ON A LGORITHM
On currently available hardware (p3.2xlarge AWS instance with an NVIDIA V100 GPU)
[Figure]
we observed the following timing results with Algorithm 1,
MLP (3 layers, 512 units each): 3 seconds 2.
ResNet50 ( 1 × width): 33 seconds 3.
ResNet20 ( 32 × width): 194 seconds
In addition
[Figure]
we tested the ability of Algorithm 1 to recover a known, randomly selected permutation.
In a handful of experiments we found that Algorithm 1 was able to exactly recover the known, random permutation in just 3-4 of passes over the layers.
A.
C OUNTEREXAMPLE D ETAILS
Consider a simple 2-dimensional classification task.
Our data points are drawn x ∼
Uniform ([ − 1 , 1] 2 ) and y = 1 x 1 < 0 and x 2 > 0 .
Figure 7 provides a visualization of a sample of such data.
[Table]
We utilize an MLP architecture consisting of two hidden layers, with two units each, and ReLU nonlinearities.
Consider two weight assignments that both achieve a perfect fit to the data:
[Equation]
We predict a positive label when f ( x ) ≥ 0 and a negative label otherwise.
Intuitively, these networks are organized such that each layer makes a classification whether x 1 < 0 or x 2 > 0 .
In model A , the first layer tests whether x 2 > 0 , and the second layer tests whether x 1 < 0 , whereas in model B the order is reversed.
With a bit of algebra, it is possible to see that both f A and f B achieve perfect performance.
A. O N THE F AILURES OF G REEDY U NI - DIRECTIONAL M ATCHING
In contrast to prior work (Singh & Jaggi, 2020; Wang et al., 2020), we eschew greedy uni-directional, single-pass matching between models.
Instead we derive a weight matching algorithm from a principled, yet computationally infeasible optimization problem.
In contrast to prior work, our method can be viewed as “bi-directional”: it selects unit associations based on weights in all relevant layers, not just in the immediately previous layer.
In this section, we describe benefits of our holistic approach, including an example problem showing the failure modes of greedy uni-directional matching.
We find that matching across all layers simultaneously allows our weight matching algorithm (Algorithm 1) to exploit units’ relationships with downstream weights in a way that greedy uni-directional matching cannot.
Concretely, greedy uni-directional matching begins at the first layer and computes a matching, P 1 , considering only W ( A ), W ( B ).
After matching, P 1 is applied throughout the remainder of the network.
Then, we proceed through the layers in order repeating this process.
A.7. E XPERIMENTAL C OMPARISON
In order to further evaluate or performance relative to prior work, and Singh & Jaggi (2020) in particular, we explored experimental comparisons with VGG11 models trained on CIFAR-10 and on ResNet50 models trained on ImageNet.
VGG11 models trained on CIFAR-10.
Following an exact reproduction of experiment from Table 1 of Singh & Jaggi (2020), we merged the model weights released along with their paper.
We present results in Table 2.
We found that our weight matching method outperforms the “wts” method of
Singh & Jaggi (2020) in both implementation speed and model performance when reproducing one of their experiments.
ResNet50 models trained on ImageNet.
We applied OT-Fusion to the ImageNet experiment that we consider in Section 5.1.
We present results in Figure 8.
We found the OT-Fusion method resulted in models with 1.38% top-1 accuracy on ImageNet, only marginally improving over na¨ıve averaging.
On the other hand, we achieve 51.01%.
[Figure] [Figure]
[Figure] [Figure]
A N E XAMPLE F AILURE C ASE
Consider two networks, A and B , with the objective that they capture the identity function f ( x ) =
x ,
[Equation] [Equation]
(cid:21) x (7) where ϵ > 0 is some negligible constant.
It can be seen that these reduce to f Θ A ( x ) = x and f Θ B ( x ) =
(1 + ϵ ) x .
When aligning these models there are two possible opportunities for permutation, π = { P 1 , P 2 } .
The permuted model B then has the form
[Equation]
(cid:18)
[Equation]
(cid:21)
[Equation]
(cid:18)
P
[Equation]
(cid:20)1 + ϵ (cid:21)(cid:19) x (8)
Now, aligning with greedy uni-directional matching will result in the alignment π gud =
{ P 1 =
I , P 2 = I } .
On the other hand, our weight matching method (Algorithm 1) results in π wm = (cid:26)
P 1 =
[Equation]
, P
[Equation]
(cid:21)(cid:27) , regardless of the algorithm’s execution order.
Interpolating these matched models at λ = 0 .
5 , we have
[Equation]
Here we can see that the greedy uni-directional matching (Equation (9)) results in a merged model that fails to represent the input, identity-function models.
On the other hand, our weight matching algorithm (Algorithm 1, Equation (10)) produces a merged model that accurately reflects both of the input models, and even improves performance over the B model.
[Figure]
[Figure] [Figure]
Although our merged model is not competitive in terms of top-1 accuracy in the CIFAR-100 split data experiment, we find that it has far better calibrated probability estimates than either of the input models
Model A
Model B
T o pA cc u r a c y
ImageNet, ResNet50 (1× width)
Model A
Model B
T o pA cc u r a c y
ImageNet, ResNet50 (1× width)
[Figure]
A UXILIARY
[Figure]
A. S TRAIGHT - THROUGH E STIMATOR D ETAILS
See Algorithm 2 for a complete description of the straight-through estimator algorithm.
Given: Model weights Θ A , Θ B , and a learning rate η .
Result:
A permutation π of Θ B such that L ( 12 (Θ A + π (Θ B ))) is approximately minimized.
Initialize: ˜Θ B ← Θ A repeat π (Θ B ) ← proj (˜Θ B ) using Algorithm 1.
Evaluate the loss of the midpoint, L ( 12 (Θ A + π (Θ B ))) .
Evaluate the gradient, ∇L , using ˜Θ B in place of π (Θ B ) in the backwards pass.
Update parameters, ˜Θ B ← ˜Θ B − η ∇L .
until convergence
A. M ERGING M ANY M ODELS
We propose Algorithm 3 to merge the weights of more than two models at a time.
Following an argument similar to Lemma 2, it can be seen that Algorithm 3 terminates.
In our limited testing, we found that this algorithm converges quickly to solutions that extrapolate better than individual models and results in a merged model with better probability estimate calibration than any of the input models.
For example, we present the results of this algorithm on MLPs trained on MNIST in Table A.10.
In addition, we found that merging multiple models helps to calibrate the resulting model predictions.
We present this effect in Figure 13.
A. F AILED I DEA : A M ETHOD FOR S TEEPEST D ESCENT
Imagine standing in weight space at Θ A and trying to decide in which immediate direction to move in order to approach a Θ B -equivalent point.
There are many, many possible permutations of Θ B –
Given: Model weights Θ 1 , . . .
, Θ N
Result:
A merged set of parameters ˜Θ .
repeatfor i ∈ R ANDOM P
ERMUTATION (1 , . . .
, N ) do Θ ′ ← 1 N − 1 (cid:80) j ∈{ 1 ,...,N }\{ i } Θ j π ← P ERMUTATION C OORDINATE D ESCENT (Θ ′ , Θ i )
Θ i ← π (Θ i ) end until convergence return 1 N (cid:80) Nj =
1 Θ j
Train Loss
Train Acc.
Test Loss
Test Acc.
Seed0. Seed0.
Seed0. Seed0.
Seed0.
M ERGE M
ANY 0.
[Figure]
call them π (1) (Θ B ) , π (2) (Θ B ) , . . .
– to aim for in the distance.
Assuming that the loss landscape is in fact (quasi-)convex modulo these permutation symmetries, a natural choice would be to pick the π ( i ) (Θ B ) that corresponds to the direction of steepest descent starting from Θ A since we expect π ( i ) (Θ B ) to lie in the same basin as Θ A .
In other words,
[Equation]
Now, we are tenuously in a favorable situation: ∇L (Θ A ) is straightforward to compute, and picking the best π reduces to a matching problem.
In particular it is a SOBLAP matching problem of the same form as in Section 3.2.
In addition, there is a fast, exact solution for the single intermediate layer case ( L = 2 ).
In practice, we found that this method can certainly find directions of steepest descent, but that they are accompanied by high barriers in between the initial dip and π (Θ B ) .
A. P ROOF OF L EMMA
To lighten notation we use ⟨· , ·⟩ = ⟨· , ·⟩ F in this section.
Lemma.
Given A , B ∈ R d × d ,
[Equation]
⟩ is strongly NP-hard and has no PTAS.
Proof.
We proceed by reduction from the quadratic assignment problem (QAP) (Koopmans & Beckmann, 1957; Cela, 2013).
Consider a QAP,
[Equation]
for C , D ∈ R d × d .
Now, pick A = C
+ λ I , B
= D − λ I .
The we have,
[Equation]
P
[Equation] [Equation]
For sufficiently large λ , the tr( P Q ⊤ ) term will dominate.
Letting α = max (max i,j | C i,j | , max i,j | D i,j | ) , we can bound the other terms,
[Equation]
Now there are two classes of solutions: those where P = Q and those where P ̸ = Q .
We seek to make the best (lowest) possible P ̸ = Q solution to have worse (higher) objective value than the worst (highest) P = Q solution.
When P = Q , the highest possible objective value is
[Equation]
and similarly, the lowest possible objective value when P ̸
=
Q is
[Equation]
where the final term is due to the fact that at least one entry of P Q ⊤ must be 0.
With some algebra, it can be seen that λ > 5 dα is sufficient to guarantee that all P = Q solutions are superior to all
P ̸ = Q solutions.
Now when P = Q , all frivolous terms reduce to constants and we are left with the QAP objective:
[Equation]
completing the reduction.
QAP is known to be strongly NP-hard (Koopmans & Beckmann, 1957;
Sahni & Gonzalez, 1976) and MaxQAP is known to not admit any PTAS (Makarychev et al., 2014), thus completing the proof.
A. P ROOF OF L EMMA
Lemma.
Algorithm 1 terminates.
Proof.
We proceed by contradiction.
Consider a graph with each possible permutation π i
= { P 1 , . . .
, P L − 1 } as a vertex and directed edges π i → π j if π j can be reached from π i with a single P ℓ update, as in Algorithm 1.
(Ignore those updates that result in no change to P ℓ in order to avoid π i → π i cycles.)
Let ρ ( π ) =
vec (Θ A ) · vec ( π (Θ B )) denote the utility of a particular π .
Note that π i → π j implies ρ ( π i ) <
ρ ( π j ) .
There exist finitely many possible permutations π i , meaning that a failure to terminate must involve a cycle in the graph π 1 → · · · → π n → π 1 .
However ρ forms a total order on the vertices and therefore we have a contradiction.