# Wasserstein Barycenter-based Model Fusion and Linear Mode Connectivity of Neural Networks

arXiv: https://arxiv.org/abs/2210.06671 · HTML: https://mlsanity.com/paper/2210.06671/

## Key facts

- Cited by 15 papers in the map (data through 14 Oct 2026).
- Builds on 39 of its 63 references that are in the map.
- First posted on arXiv 13 Oct 2022.

## Builds on (most cited first, 30 of 39)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Advances and Open Problems in Federated Learning](https://mlsanity.com/paper/1912.04977/) (2019)
- [Character-level Convolutional Networks for Text Classification](https://mlsanity.com/paper/1509.01626/) (2015)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [Averaging Weights Leads to Wider Optima and Better Generalization](https://mlsanity.com/paper/1803.05407/) (2018)
- [Visualizing the Loss Landscape of Neural Nets](https://mlsanity.com/paper/1712.09913/) (2017)
- [Computational Optimal Transport](https://mlsanity.com/paper/1803.00567/) (2018)
- [Exploring Generalization in Deep Learning](https://mlsanity.com/paper/1706.08947/) (2017)
- [Loss Surfaces, Mode Connectivity, and Fast Ensembling of DNNs](https://mlsanity.com/paper/1802.10026/) (2018)
- [Federated Learning with Matched Averaging](https://mlsanity.com/paper/2002.06440/) (2020)
- [Git Re-Basin: Merging Models modulo Permutation Symmetries](https://mlsanity.com/paper/2209.04836/) (2022)
- [A Mean Field View of the Landscape of Two-Layers Neural Networks](https://mlsanity.com/paper/1804.06561/) (2018)
- [Optimal Transport for Domain Adaptation](https://mlsanity.com/paper/1507.00504/) (2015)
- [Bayesian Nonparametric Federated Learning of Neural Networks](https://mlsanity.com/paper/1905.12022/) (2019)
- [What is being transferred in transfer learning?](https://mlsanity.com/paper/2008.11687/) (2020)
- [Essentially No Barriers in Neural Network Energy Landscape](https://mlsanity.com/paper/1803.00885/) (2018)
- [The Role of Permutation Invariance in Linear Mode Connectivity of Neural Networks](https://mlsanity.com/paper/2110.06296/) (2021)
- [Model Fusion via Optimal Transport](https://mlsanity.com/paper/1910.05653/) (2019)
- [Fast Computation of Wasserstein Barycenters](https://mlsanity.com/paper/1310.4375/) (2013)
- [Algorithmic Regularization in Learning Deep Homogeneous Models: Layers are Automatically Balanced](https://mlsanity.com/paper/1806.00900/) (2018)
- [FedBE: Making Bayesian Model Ensemble Applicable to Federated Learning](https://mlsanity.com/paper/2009.01974/) (2020)
- [Gromov-Wasserstein Learning for Graph Matching and Node Embedding](https://mlsanity.com/paper/1901.06003/) (2019)
- [Scalable Gromov-Wasserstein Learning for Graph Partitioning and Matching](https://mlsanity.com/paper/1905.07645/) (2019)
- [Optimal Transport for Multi-source Domain Adaptation under Target Shift](https://mlsanity.com/paper/1803.04899/) (2018)
- [Multilevel Clustering via Wasserstein Means](https://mlsanity.com/paper/1706.03883/) (2017)
- [Scalable Bayes via Barycenter in Wasserstein Space](https://mlsanity.com/paper/1508.05880/) (2015)

## Built on by (most cited first, 15 of 15)

- [Deep Model Fusion: A Survey](https://mlsanity.com/paper/2309.15698/) (2023)
- [Re-basin via implicit Sinkhorn differentiation](https://mlsanity.com/paper/2212.12042/) (2022)
- [Transformer Fusion with Optimal Transport](https://mlsanity.com/paper/2310.05719/) (2023)
- [Equivariant Deep Weight Space Alignment](https://mlsanity.com/paper/2310.13397/) (2023)
- [Proving Linear Mode Connectivity of Neural Networks via Optimal Transport](https://mlsanity.com/paper/2310.19103/) (2023)
- [FedFisher: Leveraging Fisher Information for One-Shot Federated Learning](https://mlsanity.com/paper/2403.12329/) (2024)
- [Simultaneous linear connectivity of neural networks modulo permutation](https://mlsanity.com/paper/2404.06498/) (2024)
- [On Linear Mode Connectivity of Mixture-of-Experts Architectures](https://mlsanity.com/paper/2509.11348/) (2025)
- [FedPIA -- Permuting and Integrating Adapters leveraging Wasserstein Barycenters for Finetuning Foundation Models in Multi-Modal Federated Learning](https://mlsanity.com/paper/2412.14424/) (2024)
- [Foldable SuperNets: Scalable Merging of Transformers with Different Initializations and Tasks](https://mlsanity.com/paper/2410.01483/) (2024)
- [Heterogeneous Federated Reinforcement Learning Using Wasserstein Barycenters](https://mlsanity.com/paper/2506.15825/) (2025)
- [Functional Equivalence in Attention: A Comprehensive Study with Applications to Linear Mode Connectivity](https://mlsanity.com/paper/2606.17830/) (2026)
- [A Lipschitz spaces view of infinitely wide shallow neural networks](https://mlsanity.com/paper/2410.14591/) (2024)
- [Robustness and Regularization in Hierarchical Re-Basin](https://mlsanity.com/paper/2510.09174/) (2025)
- [Partial Fusion of Neural Networks: Efficient Tradeoffs Between Ensembles and Weight Aggregation](https://mlsanity.com/paper/2605.22350/) (2026)

## Abstract (arXiv)

Based on the concepts of Wasserstein barycenter (WB) and Gromov-Wasserstein barycenter (GWB), we propose a unified mathematical framework for neural network (NN) model fusion and utilize it to reveal new insights about the linear mode connectivity of SGD solutions. In our framework, the fusion occurs in a layer-wise manner and builds on an interpretation of a node in a network as a function of the layer preceding it. The versatility of our mathematical framework allows us to talk about model fusion and linear mode connectivity for a broad class of NNs, including fully connected NN, CNN, ResNet, RNN, and LSTM, in each case exploiting the specific structure of the network architecture. We present extensive numerical experiments to: 1) illustrate the strengths of our approach in relation to other model fusion methodologies and 2) from a certain perspective, provide new empirical evidence for recent conjectures which say that two local minima found by gradient-based methods end up lying on the same basin of the loss landscape after a proper permutation of weights is applied to one of the models.