Wasserstein Barycenter-based Model Fusion and Linear Mode Connectivity of Neural Networks
Authors: Aditya Kumar Akash
Sixu Li
Nicolas Garcia Trillos
Venue: Submitted to ICLR 2023 (2023)
Keywords: Linear mode connectivity
Neural network model fusion
Wasserstein barycenter
Federated learning

Abstract: 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.

Under review as a conference paper at ICLR 2023
WASSERSTEIN BARYCENTER-BASED MODEL FUSION AND LINEAR MODE CONNECTIVITY OF NEURAL NET- WORKS
Anonymous authors Paper under double-blind review
ABSTRACT
Based on the concepts of Wasserstein barycenter (WB) and Gromov- Wasserstein barycenter (GWB), we propose a unified mathematical framework for neural net- work (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 interpretation of node in network 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 ex- ploiting 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 proper permutation of weights is applied to one of the models.
function of the
INTRODUCTION
The increasing use of edge devices like mobile phones, tablets, and vehicles, along with the sophisti- cation in sensors present in them (e.g. cameras, GPS, and accelerometers), has led to the generation of enormous amount data. However, data privacy concerns, communication costs, bandwidth limits, and time sensitivity prevent the gathering of local data from edge devices into one single centralized location. These obstacles have motivated the design and development of federated learn- ing strategies which are aimed at pooling information from locally trained neural networks (NNs) with the objective of building strong centralized models without relying on the collection of local data McMahan et al. (2017); Kairouz et al. (2019). Due to these considerations, the problem of NN fusion-i.e, combining multiple models which were trained differently into a single model-is a fundamental task in federated learning.
A standard fusion method for aggregating models with the same architecture is FedAvg McMahan et al. (2017), which involves element- wise averaging of the parameters of local models. This is also known as vanilla averaging Singh & Jaggi (2019). Although easily implementable, vanilla averaging performs poorly when fusing models whose weights do not have one-to-one correspondence. This happens because even when models are trained on the same dataset it is possible to obtain models that differ only by a permutation of weights Wang al. (2020); Yurochkin et al. (2019); this feature is known as permutation invariance property of neural networks. Moreover, vanilla averaging is not naturally designed to work when using local models with different architectures (e.g.. differ- ent widths). In order to address these challenges, Singh & Jaggi (2019) proposed to first find the best alignment between the neurons (weights) of different networks by using optimal transport (OT) Villani (2008); Santambrogio (2015); Peyré & Cuturi (2018) and then carrying out vanilla aver- aging step. Other approaches, like those proposed Wang et al. (2020); Yurochkin et al. (2019), interpret nodes of local models as random permutations of latent "global nodes" modeled according to a Beta- Bertaull process prior Thibaux & Jordan (2007). By using "global nodes", nodes from different input NNs can be embedded into common space where comparisons and aggregation are meaningful. Most works in the literature discussing the fusion problem have mainly focused on the aggregation of fully connected (FC) neural networks and CNNs, but have not, for the most

---

Under review as a conference paper at ICLR 2023
part, explored other kinds of architectures like RNNs and LSTMs. One exception to this general state of the art is the work Wang et al. (2020), which considers the fusion of RNNs by ignoring hidden-to-hidke weights during the neurons' matching, thus discarding some useful information in the pre-trained RNNs. For more references on the fusion problem see in the Appendix.
A different line of research that has attracted considerable attention in the past few years is the quest for a comprehensive understanding of the loss landscape of deep neural networks, a fundamental component in studying the optimization and generalization properties of NNs Li et al. (2018); Mei et al. (2018), 8); Neyshabur et al. (2017); Nguyen et al. (2018); Izmailov et al. (2018). Due to over- parameterization, scale, and permutation invariance properties of neural networks, the loss land- scapes of DNNs have many local minima Keskar et al. (2016); Zhang et al. (2021). Different works have asked and answered affirmatively the question of whether there exist paths of amall-increasing loss connecting different local minima found by SGD Garipov et al. (2018); Draxler et al. (2018), 8). This phenomenon is often referred to as mode connectivity Garipov et al. (2018) and the loss in- crease along paths between two models is often referred to as (energy) barrier Draxler et al. (2018). has been observed that low-barrier paths are non-linear, i.e., linear interpolation of two different models will not usually produce neural network with small loss. These observations suggest that, from the perspective of local structure properties of loss landscapes, different SGD solutions belong to different (well-separated) basins Neyshabur et al. (2020). However, recent work Entezari et al. (2021) has conjectured that local minima found by SGD do end up lying on the same basin of the loss landscape after a proper permutation of weights is applied to one of the models. The question of how to find these desired permutations remains in general elusive.
The purpose of this paper is twofold. On one hand, we present large family of barycenter-based fusion algorithms that can be used to aggregate models within the families of fully connected NNs, CNNs. ResNets, RNNs and LSTMs. The most general family of fusion algorithms that we intro- duce relies on the concept of Gromov- Wassersttsnn barycenter (GWB), which allows us to use the information in hidden-to-hidee
layers in RNNs and LSTMs in contrast to previous approaches in the literature like that proposed in Wang et al. (2020). In order to motivate the GWB based fusion algorithm for RNNs and LSTMs, we first discuss a Wasserstein barycenter (WB) based fusion algo- rithm for fully connected, CNN, and ResNet models which follows closely the OT fusion algorithm from Singh & Jaggi (2019). By creating a link between the NN model fusion problem and the problem of computing Wasserstein (or Gromov- -assersttt)))) barycenters, our aim is exploit the many tools that have been developed in the last decade for the computation of WB (or GWB) -see see the Appendix for references- and to leverage the mathematical structure of OT problems. Using our framework, we are able to fuse models with different architectures and build target models with arbitrary specified dimensions (at least in terms of width). On the other hand, through several nu- merical experiments in a variety of settings (architectures and datasets), we provide new evidence backing certain aspects of the conjecture put forward in Entezari et al. (2021) about the local struc- ture of NNs' loss landscapes. Indeed, we find out that there exist sparse couplings between different models that can map different local minima found by SGD into basins that are only separated by low energy barriers. These sparse couplings, which can be thought of as approximations to actual permutations, are obtained using our fusion algorithms, which, surprisingly, only use training data to set the values of some hyperparameters. We explore this conjecture in imaging and natural language processing (NLP) tasks and provide visualizations of our findings. Consider, for example, Figure (left), which is the visualization of fusing two FC NNs independently trained on the MNIST dataset. We can observe that the basins where model and permuted model 2 (i.e. model 2 after multiplying its weights by the coupling obtained by our fusion algorithm) land are close to each other and are only separated by low energy barriers.
Our main contributions can then be summarized follows: (a) we formulate the network model fusion problem as series of Wasserstein GGromov-Wsseerstrrr barycenter problems, bridging in this way the NN fusion problem with computational OT: (b) we empirically demonstrate that our framework is highly effective at fusing different types of networks, including RNNs and LSTMs. (c) we visualize the result of our fusion algorithm when aggregating two neural networks a 2D-plane. By doing this we not only provide some illustrations on how our fusion algorithms perform, but also present empirical evidence for the conjecture made in Entezari et al. (2021), casting light over the loss landscape of a variety of neural networks.
At the time of completing this work, we became aware of two very recent preprints which also explore the conjecture made in Entezari et al. (2021) empirically. In particular, Ainsworth et al.
2

---

Under review as a conference paper at ICLR 2023
Rew modd! Doe moddd Fasd wold
(1. w)
2.
Lager
Layer
Layeth
Laper%
Figure 1: Left: The test error surface of FC NNs trained on MNIST. The permuted model 2 is model 2 after multiplying its weights by the coupling obtained by fusion algorithm. Right: The illustration of our interpretations of FC NNs. Following our definitions, node v = (72. w), where 72 probability measure on layer N2 and w N2 + R is the weight function corresponding to node v. For example, the scalar w(22) is the weight between nodes v and 22, and we use W2 as the shorthand notation of (22)).
(2022) demonstrates that there is zero-barrier LMC (after permutation) between two independently trained NNs (including ResNet) provided the width of layers is large enough. In Benzing et al. (2022), the conjecture is explored for FC NNs, finding that the average of two randomly initialized models using the permutation revealed through training gives a non-trivial NN. Compared to our work, none of these two works explored this conjecture for recurrent NNs; GWB fusion method of particular relevance for this aim. To the best of our knowledge, we thus provide the first-ever exploration of the conjecture posited in Entezari et al. (2021) for NLP tasks.
highlight that
1.1 NOTATION
We first introduce some basic notation and briefly review a few relevant concepts from OT. simplex of histograms with n bins denoted by Σn {a IR E; a; 1}. The set of couplings between histograms a E En, and b Eng denoted by [(a, b) = {IT E R=1x== ITI ng = a, IT 1, b). where 1, == (1,...,1)² E IR.. For any 4-way tensor = [Lnj i.j.K. and matrix II [ip i.3" we define the censor-matric multiplication of L and as the matrix & Ex., Ligke isj"
1.2 OPTIMAL TRANSPORT AND WASSERSTEIN BARYCENTERS
Let X' be an arbitrary topological space and let c: X' x X' - [0, oc) be a cost function assumed to satisfy c(x, x)) = for every z. We denote by M (X) the space of (Borel) probability measures
on X. For 1i fuy 12 X, define discrete measures in M) where a En,, b Eng, and 8g denotes the Dirac delta measure at x E X. The Wasserstein "distance" between l and v, relative the cost e, is defined
E a,oz, and - b
((, = inf II, NET
(1)
where C= [e(xi, Wi)],, is the "cost" matrix between (. {usls II [mo],, E T(u, v) the coupling matrix between a and v, and (A, B) == tr(AA the Frobenius inner product.
Let (T') / E MM (X) be collection of discrete probability measures. The Wasserstein barycenter problem (WBP) Agueh & Carlier (2011) associated with these measures reads
Σ W(nr))
min (EMT(X)
(2)
A minimizer of this problem is called Wasserstein barycenter (WB) of the measures [ and can be understood as an average of the input measures. In the sequel we will use the concept of WB to define fusion algorithms for FC NN, CNN, and ResNet. For RNN and LSTM the fusion reduces to solving a series Gromov- Wasserstein barycenter-like problems (see the reviews of GWBP in the Appendix).
3

---

Under review as a conference paper at ICLR 2023
2 WASSERSTEIN BARYCENTER BASED FUSION
In this section, we discuss our layer-wise fusion algorithm based on the concept of WB. First we introduce the necessary interpretations of nodes and layers of NNs in Section 2.1. Next in Section 2.2, we describe how to compare layers and nodes across different NNs so as to make sense of aggregating models through WB. Finally we present our fusion algorithm in Section 2.3.
2.1 NESTED DEFINITION OF FULLY CONNECTED NN
For a fully connected network N, we use v to index the nodes in its 1-th layer N₁. Let y denote a probability measure on the I-th layer defined as the weighted sum of Dirac delta measure over the nodes in that layer, i.e.,
Σ on E Mf(N). VEN,
=
(3)
N
We interpret a node from the 1-th layer as an element in N₁ that couples a function on the domain Ni-1 (previous layer) with probability measure. In particular, the node is interpreted as v == (N-11, w), where `\ 1 a measure on the previous layer Ni-1 and w represents the weights between the node v and the nodes in previous layer - These weights can be interpreted as function - + and we use the notation W/ to denote the value of function w evaluated at the q-th node in the previous layer Ni-1- For the first layer i... / = 1, the nodes simply represent placeholders for the imput features. The above interpretation is illustrated in Figure (right). This interpretation of associating nodes with function of previous layer allows to later define "distance" between nodes in different NNs (see Section 2.2) and motivated from TI" spaces and distance Gareía Trillos & Slepčev (2015); Thorpe et al. (2017) which is designed for comparing signals with different domains.
2.2 COST FUNCTIONS FOR COMPARING LAYERS AND NODES
Having introduced our interpretations of NNs, we now define the cost functions for comparing layers and nodes which will be used to aggregate models through WB. Consider the I-th layers N₁ and N of two NNs N and N' respectively. We use Wasserstein distance between the measures "/t and 7 over Nj and N2 respectively to define distance between the layers:
dn(n, M) = W(n, 21) = inf (CI, IT.)
(4)
where matrix 11, = [ is coupling between the measures "/ and 7/: and Cj is the cost matrix give by Ci [c, e²) n, where qt is a cost function between nodes the 1-th layers.
Following our inductive interpretation of NNs, the cost function et can also be defined inductively. Consider nodes v and o' from 1-th layer of NNs N and N' respectively. For the first layer = 1, we pick a natural candidate for cost function, namely e1 (v, e') := Luge', a reasonable choice given that all networks have the same input layer. For 2 2, recall our interpretation of nodes = (1-1, w), o' (1/11: w'), where ` and 7/_1 denotes the respective measures associated with previous layer ( - 1 and we w' denotes the respective weight functionn for nodes v and v'. Since the domains of the weight functions w and w' are layers in different NNs, it is not clear how to compare them directly. However in TLP interpretation, after finding a suitable coupling between the support measures 7 and 2/11- one can couple the functions w and w' and use a direct L2-comparison. Motivated by computational and methodological considerations, we use a slight modification of the TIP distance and decouple the problem for the measures from the weights. Specifically, we define (c,,,') := 4u(7-1.7% + dw (w, w'); where dy is the Wasserstein distance (as defined in equation 4) between the measures 7 and 7/-1 from layers - 1. And dw is defined using the optimal coupling of weight functions' support measures, i.c.
dw (w, w') = Σ (wy - w')) (mI) =: (L(W, (II+)))),
(5)
9,9
where I(w, w') : [(we - -w()²] and (II, = [ is the optimal coupling between
n- and 7/-1- Note that d 9.a x/1) 1 is fixed constant when comparing any nodes on the
4

---

Under review as a conference paper at ICLR 2023
I-th layers Ni and N/. For simplicity, we let cı(u, 1') (W, W in what follows, and the information of support measures N and M-11 is implicitly included in their optimal coupling (14--)). Here we have omitted bias terms to case the exposition of our framework, but a natural implementation that accounts for bias terms can be obtained by simply concatenating them with the weight matrix.
We set (II1)* equal to the identity matrix normalized by the size of input layer given that this is a solution to equation 4 when the cost 01 is defined as 01 (v, 8) == Lu Other choices of cost function ej are possible. e.g. the activation-based cost function proposed in Singh & Jaggi (2019).
2.3 FUSION ALGORITHM
In the following we consider n. input FC NNs , We use Nj to denote the I-th layer of the i-th network N and kij to denote the number of nodes in that layer, i.e. kj = |ill- Let of to be the probability measure on layer Ni similar to definition in equation with the support points being nodes in that layer. We denote the target model (i.e. the desired fusion output) by Neww and use .2.....ku km to denote the sizes of its layers N2 networks, including the target model, have the same input layer and the same number of layers m.
Amww respectively. We assume that all
Based on the discussion in Sections 2.1 and 2.2, we now describe an inductive construction of the layers of the target network Noew by fusing all n input NNs. First, NTwww is set to be equal to N this is the base case of the inductive construction and simply means that we set the input layer of new to be the same as that of the other models; we also set 7/1 := 7/. Next, assuming that the fusion has been completed for the layers 1 to ( - 1 (1 > 2), we consider the fusion of the (-th layer. For the simplicity of notations, we drop the index / while referring to nodes and their corresponding weights in this layer. In particular, we use 2t and w' to denote the nodes in layer Ní and their corresponding weights. To carry out fusion of the 1-th layer of the input models, we aggregate their corresponding measures through finding WB which provides us with sensible "average" I-th layer for the target model. Hence, we consider the following WBP over 7 ""."
k
min (T.
Σ ((a, f = E CC. II)
s.t.
N1-
ki Σ
deg Dj = (1,,1, (6)
Here the measure y is the candidate [-th layer "average" of the input models and is forced to take a specific form (notice that we have fixed the size of its support and the masses assigned to its support points). Nodes ny in the support of 7 are set to take the form dy = (7--,, t0y). i.e. the measure 7-1 obtained when fusing the (1 - 1) 3-th layers is the first coordinate in all the Mi- This plugs the current layer of the target model with its previous layer. As done for the input models, wy is interpreted as a function from the 1)-th layer into the reals, and represents the actual weight vector from the (I - 1) layer to the j-th node in the I-th layer of the new model. Cj- [c((u,, (e(l,,,') jgg the cost matrices corresponding to WBP equation 6, where e a cost function between nodes on the 1-th layers (see in Section 2.2). Let Wi and Wi: to be the weight function matrices of the I-th layer of target models Neww and input model N' respectively (e.g. Wi- := (wi...., ",.....,,,y77 and define -wxx)
L(W,, Wf) = [(wie- - ) igg..." where Wjg denotes the function w, evaluated at the q-th node in layer - 1 and similarly for w.s. The cost matrices Cj can now be rewritten
C = [c(e,, u0)], [dw, wg]] j.g = L(W, w/) @ (II)1)".
where (m is the optimal coupling between measures 7-1 and 7 Combining equation 7 with equation gives us the following optimization problem which we solve to obtain the fused layer:
1 " min BWW; (1))) W
n Σ (CWW, wi) @ (IIL II).
(8)
In order to solve the minimization problem 8, we can follow a strategy discussed in Cuturi & Doucet (2014); Anderes et al. (2016); Claici et al. (2018), i.e., alternatingly update weights Wi and couplings { (remember that the (II; are computed once and for all and are fixed in equation 8. In particular, after initializing weight matrices, we alternate between two steps until some stopping criterion is reached:
5

---

Under review as a conference paper at ICLR 2023
Step 1: For fixed Wi, we update the couplings {. Note that the minimization of B(W; (Ii h:) over the couplings {TT)): splits into OT problems, each of which can be solved using any of the algorithms used in computational OT (e.g. Sinkhorn's 's algorithm Cuturi (2013)).
For fixed couplings { we update the weights Wi. Note that for fixed couplings the objective B(W; {Ijjk) is quadratic in Wi and hence we obtain the following update formula:
Step 2:
Wi 4- kikht
(TH
where : is elementwise division. We refer to the above fusion algorithm as Wasserstein barycenter- based fusion (WB fusion). The pseudo-code for this algorithm and corresponding computational complexity can be found in the Appendix, where we also provide some details on how to adapt our fusion method to handle convolutional layers and skip-connections.
GROMOV- WASSERSTEIN
SECCENTER-BASHE FUSION
In this section we discuss extension of our fusion framework to cover RNNs and LSMMs. Compared to FC networks, RNNs contain "self-loops" in each layer thidden-to-hidden recurrent connections) which allows information to our interpretation of neurons in the FC case, node vy the (-th layer will be represented as o'g == [(7i_1 wy);; (7É, k2)], where wy is the weight function between inputs of the preceding layer and hidden states, and hig is the weight function between hidden states; 7/-1 and )i are the probability measures corresponding layer I --1 1 and layer I respectively. This definition comes weight functions hy are supported on the 1-th layer itself, whereas wyJ is supported on the (1 - 1)-th layer.
be passed from one step of the neural network to the next. Similar
from the observation that hidden-to-hidde
Having carried out the fusion of the first 1 layers we consider the following problem to fuse the 1-th layers:
min B(W, HG (T)))) == - n - E (CW,, wi) s (mi-1) + ann H() @ IT, m (10) (.))
where our a hyperparameter that balances the importance input-to-hidden weights and hidden- to-hidden weights during the fusion; we'll use (IT)'* to denote an optimal IT.. We use Hi and Hi to denote the hidden-to-hidden weight function matrices of layer N and Ni respectively, and we let L(H,, H)) = [(hja h/ -.9.q.= L(W,, Wf) defined the same as in the fully connected case. Notice that this is a GW-like barycenter problem.
We provide more detailed explanation on how derive optimization problem 10 and adapt the GWB fusion for RNNs discussed in this section to the LSTM case in the Appendix. In Section 4.3 we show that the models obtained when setting ah 0 in equation 10 greatly outperform the models obtained when setting ag = O, justifying in this way the use of GWBs.
4 EXPERIMENTS
Overview: We present an empirical study of our proposed WB and GWB based fusion algorithms assess its performance in comparison to other state of the fusion methodologies and reveal new insights about the loss landscapes for different types of network architectures and datasets. We first consider the fusion of models trained on heterogeneous data distributions. Next we present results for WB fusion of FC NNs and deep CNNs, and draw connections between workings of WB fusion and LMC of SGD solutions. Finally, we consider GWB fusion and present results on RNNs, LSTMs and extend the conjecture made in Entezari et al. (2021) for recurrent NNs.
For baselines, we consider vanilla averaging and the state-of-the-ar fusion methodolo- gies like OT fusion Singh & Jaggi (2019) and FedMA Wang al. (2020). For fair comparison under the experimental settings of one-shot fusion we consider FedMA without the model retraining step and restrict its global model not outgrow the base models. We refer to this "one-shot
Baselines:

---

Under review as a conference paper at ICLR 2023
FedMA". For RNNs and LSTMs, our baselines additionally include slightly modified versions of WB based fusion and OT fusion where we ignore the hidden-to-hiddee connections. Other methods which require extensive training not applicable in one-shot model aggregation settings.
Base models & General-setup: For our experiments on FC NNs, we use MLPNET introduced
in Singh & Jaggi (2019), which consists of 3 hidden layers of sizes (400, 200, 100). Addition- ally, we introduce MLPLARGE and MLPSMALL with hidden layers of size {800, 400, 200} and {200, 100, 50) respectively. For deep CNNs, we use VGGI1 Simonyan & Zisserman (2014) and RESNETI 88 He et al. (2016). For recurrent NNs, we work with RNNs and LSTMs with one hidden layer of size 256 and 256 respectively. Hyperparameters are chosen using a validation set and final results are reported on a held out test set. More training details are provided in the Appendix.
Visualization methodology: We visualize the result of fusing two pre-trained models on two- dimensional subspace of NNs' loss landscape by using the method proposed in Garipov et al. (2018). In particular, each plane is formed by all affine combinations of three weight vectors corresponding to the parameters of base model 1, base model 2 and permuted model 2 (i.e. base model 2 after multiplying its weights by the coupling obtained by our fusion algorithm) respectively.
4.1 WB FUSION UNDER HETEROGEENOUSS DATA DISTRIBUTIONS
Setup: We first apply WB fusion in aggregating models trained heterogeneous data distributions which is a setting often found in federated learning where the clients have local data generated from different distributions and privacy concerns prevent data sharing among them. Here we follow the setup described in Singh & Jaggi (2019). To simulate heterogeneous data-split on MNIST digit classification one of the models (named A) is trained with special skill to recognize one of the digits (eg. digit 4) that is not known to the other model, named B. Model B is trained on 90% of the training data for remaining digits while model A uses the other 10% data. Under this data split, we consider two settings. For the first setting, the base models are fused into target model of the same architecture (MLPNET). For the second setting, we consider the fusion of two small base models (MLPSMALL) into large target model (MLPNET). This simulates the setting where clients in federated learning are constrained by memory resources to train smaller models. In both cases we use model fusion to aggregate knowledge learned from the base models into a single model, more memory-efficien strategy than its ensemble-basee counterparts.
Figure 2 shows the results of single shot fusion when different proportions of the base models are considered. We find that (a) WB fusion consistently outperforms the baselines, (b) for certain combinations WB produces fused models with accuracy even better than the base model and demonstrates successful one shot knowledge aggregation. Note that for each proportion of model aggregation (x-axis), the results are reported over multiple runs where one of the base models is randomly chosen to initialize the target model in the fusion algorithm. We find that WB fusion is more stable against initialization as indicated by the lower variance in Figure 2. For fusion into different architectures vanilla averaging is not applicable, and we do not include "one-shot FedMA" for comparison here since it is not clear how to assign different proportions base models in FedMA, or to specify a target architecture different from the base models.
Quantitative results:
medde
R8
Tt
Ap
RicSer
7
70
86
WSR saior OT fucion
VS fusion
Varilo
0T tudoon
0.0
0.2
0.4
0.6
0.0
0.0
0.4
0.6
0.8
towards model
Weight fowards model
Figure 2: Left/ / Right: Test accuracy % for fused models when base models are trained on het- erogeneous data distributions and combined with various proportions into target model of same / different architecture. Some models obtained by WB fusion outperform even the base models.
7

---

Under review as a conference paper at ICLR 2023
4.2 WB FUSION UNDER HOMOGENEOUS DATA DISTRIBUTIONS AND CONNECTIONS TO LMC
Setup: In this section we perform WB fusion for various models and architectures, and provide loss landscape visualizations which reveal workings of the fusion algorithm and shed light on linear model connectivity (LMC) of SGD solutions after applying appropriate permutations. We first con- sider fusion of FC NNs on the MNIST dataset Deng (2012) and train MLPNET following Singh & Jaggi (2019). For this we consider two different settings. In the first setting, the target model has the same architecture as the base models. For the second one, we fuse the base models into larger model MLPLARGE. As noted before, the latter scenario is relevant for federated learning, given the limitations of memory and computational resources on edge devices. Next, we consider fusion of deep CNNs like 1, RESNET18 18 trained on CIFARIO dataset Krizhevsky et al. (2009). For all these cases, we fuse 2 trained models initialized differently. For the skip-connection and fusion into different architectures, FedMA is not directly applicable and hence not considered for comparisons.
Quantitative results: Table 1 contains the results of fusion for FC NNs and deep CNNs. We find that (a) WB fusion produces models at par outperforms other fusion methods for all considered model types and datasets, (b) for fusion into different architectures and ResNets, we find that WB fusion is more effective and robust.
Table I: Performance comparison (Test accuracy + standard deviation %) of different fusion algo- rithms under various network architectures and datasets. "BASE" means initializing target model with one of the base models. For each case, the target model obtained by WB fusion gets the highest test accuracy and smallest standard deviation.
MNIST
CIFARIO
MLPNET/BASE
MLPLARGE
VGG11/BASE
RESNET18/BASE
MODEL AVG
96.31
90.11
91.56 0.31
VANILLA AVO
86.50 :: 4.60
30.82 =: 4.49
20.56 = 3,90
ONE-SHOT FEDMA
97.89 0.10
1.01
OT
97.81 0.12
91.53
2.61
85.39 0.93
71.37 = 6.53
97.92
94.83 1.18
85.39 0.98
73.75 +: 4.39
Visualizations: Figure I (left) contains the visualization of fusing two MLPNET trained on MNIST dataset under WB framework and Figure 3 (left) contains the fusion result of WB fusion of two VGG11 models trained CIFAR1O. We find that (a) the couplings obtained in WB fusion (refer to equation 8) between the layers of target model and base models are sparse, i.e. they are almost permutations; (b) the basins of the permuted model 2 (obtained by multiplying the weights of base model 2 by the found couplings) and base model 1 lic close to each other and are separated by low energy barrier. These visualizations thus provide new empirical evidence in support of the conjecture made in Entezari et al. (2021). They also shed light on the workings on WB fusion algorithm. In particular, equation can be interpreted as coordinate- wise averaging of the permuted models. Since permuted models land in basins that are separated by low energy barriers, their linear interpolation gives a good fused model.
4.3
GWB FUSION FOR RECURRENT NEURAL NETWORKS
Sctup: In this section, we consider the fusion of NNs like RNNs and LSTMs on sequence based tasks. We use 4 different datasets for this setting: i) MNIST Deng (2012): Images of 28 x 28 dimen- sions are interpreted as 28 length sequences of vectors R ii) SST-2 Socher al. (2013): Binary classification task of predicting positive and negative phrases; iii) AGNEWS Zhang et al. (2015): Corpus of news articles from 4 classes; and iv) DBpedia Zhang al. (2015): Ontology classifica- tion dataset containing 14 non-overlapping classes. For the NLP tasks, use pre-trained Glo embeddings Pennington et al. (2014) of dimensions 100 and 50 for RNNs and LSTMs respectively. The embedding layer is not updated during the model training. We set the target model to have the same architecture as the base models.
Quantitative results: Table contains the result of fusion for various datasets and model architec- tures. We find that (a) our GWB framework outperforms other fusion algorithms for each combi- nation of model type and dataset, which highlights the importance of using hidden-to-hidde con-
8

---

Under review as a conference paper at ICLR 2023
340
Figure 3: Visualizations of the fusion results on the test error surface, which is function of network weights in a two-dimensional subspace, for different models and datasets. Left: Fusion of two VGGI1 models trained on CIFARIO dataset using WB framework. Middle: Fusion of two LSTM models trained on MNIST dataset. Right: Fusion of two LSTM models trained on DBpedia dataset. We can observe that in all these cases the basins of permuted model 2 (obtained by multiplying the weights of base model 2 by the found coupling) and base model 1 lie close to each other and are separated by low energy barrier.
nections for the fusion of recurrent NNs; (b) the accuracy gains for GWB over WB is different for different tasks, which indicates that relative importance of hidden-to-hideee connections is task de- pendent; (c) the accuracy of fused model is higher for LSTMs in comparison to RNNs, which we attribute to the fact that LSTMs have four hidden states and thus four input-to-hiddes and hidden- to-hidden weight matrices. More information for each hidden node allows the algorithm to uncover better couplings. Our results in (a) and (b) show the usefulness of hyperparameter a (set between 20) from equation 10 in balancing the relative importance of -odden-to-hidden weights.
Table 2: Performance comparison (Test accuracy + standard deviation %) of different fusion algo- rithms under various network architectures and datasets. For each case, target model obtained by GWB fusion reaches the highest test accuracy and small standard deviation.
MNIST
AGNEWS
DRRSDDIA
SST-3
RNN
RNN
MODEL
96.88
98 0.00
0.12
92.35 0.17
97.12
02.62:00.11
1.03
90.31 0.27
VANILA Ave
26.54 10.70
31.92 1.86
10.77 ::
76.01 : 3.80
50.56 4: 1.65
2.11
73.91 2:5
76255= 1.02
OT
36,78 d: 14.30
68.33 7.07
50.00 2: 4.30
86.19 = 2.14
37.91 4.86
17.00
78.92 2.07
82:18:0000
ONE-sHoT FunMA
34.10
06.08
86.30
12.10
70.17 2.27
6:26
828T
81.39 2.87 2.87
03.27
1.86
61.01 3.887
87.06
56.15 5.07
87.50
2:589
82.60 1.05
0.77
Figure (middle, right) contains visualization of fusing LSTM models under the GWB framework. As noted for the FC NNs and deep CNNs visualizations, find that (a) the couplings found by GWB fusion algorithm are sparse, and (b) these couplings map different local minima into neighboring basins that are separated by low energy barriers. This empirical evidence suggests that the original conjecture in Entezari et al. (2021) can be extended to richer network architectures and tasks (RNNs and LSTMs on NLP datasets).
Visualizations:
CONCLUSION
In this paper we have proposed neural network fusion algorithms that based the concept of Wasserstein/Grome Gronooo-aaassssteiii Wasserstein barycenter. Our fusion algorithms allow us to aggregate models within a variety of NN architectures, including RNN and LSTM. Through extensive experimentation we: 1) illustrated the strengths of our algorithms 2) provided new empirical evidence backing recent conjectures about the linear mode connectivity of different neural networks with architectures such as RNN or LSTM and for different imaging and NLP datasets. Limitations and future work: NNs with ReLU activation are also scale-invariant across the layers which is currently not handled in our cost functions. Although the empirical evidence in Entezari et al. (2021); Du et al. (2018) suggests that the models trained on same datasets using SGD converges to solutions with more balanced weights, it might be the case that for certain heterogeneous settings the weights across models become less balanced. For future work we would like to explore fusion using scale-invarian cost functions and apply WB/GWB fusion algorithms to federated learning.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Martial Aguch and Guillaume Carlier. Barycenters in the wasserstein space. SIAM Journal on Mathematical Analysis, 43(2):904-224, 2011.
Samuel Ainsworth, Jonathan Hayase, and Siddhartha Srinivasa. Git re-basin: Merging models 0000:2230904
modulo permutation symmetries. arXiv preprint arXiv:2209. 04836, 2022.
Ethan Anderes, Steffen Borgwardt, and Jacob Miller. Discrete wasserstein barycenters: Optimal transport for discrete data. Mathematical Methods of Operations Research, 84(2)3899409, 2016.
Frederik Benzing, Simon Schug, Robert Meier, Johannes von Oswald, Yassir Akram, Nicolas Zuc- chet, Laurence Aitchison, and Angelika Steger. Random initialisations performing above chance and how to find them. arXiv preprint arXiv:2209.07509 2022.
Sebastian Claici, Edward Chien, and Justin Solomon. Stochastic wasserstein barycenters. arXiv
preprint arXiv: 057557 2018.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in neural information processing systems, pp. 2292-2300, 2013.
Marco Cuturi and Arnaud Doucet. Fast computation of wasserstein barycenters. 2014.
Li Deng. The mnist database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine, 29(6):141-)) 2012.
Felix Draxler, Kambis Veschgini, Manfred Salmhofer, and Fred Hamprecht. Essentially no barriers in neural network energy landscape. In International conference on machine learning, pp. 1309- 1318. PMLR, 2018.
Simon Du, Wei Hu, and Jason Lee. Algorithmic regularization in learning deep homogeneous models: Layers automatically balanced. Advances in Neural Information Processing Systems, 31, 2018.
Rahim Entezari, Hanie Sedghi, Olga Saukh, and Behnam Neyshabur. The role of permutation invariance in linear mode connectivity of neural networks. arXiv preprint 10.0022 2021.
Nicolás García Trillos and Dejan Slepčev. Continuum limit of total variation point clouds. Archive for Rational Mechanics and Analysis, pp. 1-49, 2015. ISSN 1432-0673, doi: 10.1007/ 011-0929---910000000001100900090
500205-015-0929-2. URL http:///p//// /xxx doi oeri.orgg100 1007
Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry P Vetrov, and Andrew G Wilson. Loss surfaces, mode connectivity, and fast ensembling of dnns. Advances in neural information pro- cessing systems, 31, 2018
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp- 770-778, 2016.
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wil- son. Averaging weights leads to wider optima and better generalization. arXiv preprint arXiv:1803.05407, 2018.
Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Ad- vances and open problems in federated learning. arXiv preprint 17X1v:1912.04977, 2019.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv: (60.006600 2016.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
10

---

Under review as a conference paper at ICLR 2023
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss land- scape of neural nets. Advances neural information processing systems, 31, 2018.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-effciennt learning of deep networks from decentralized data. In Artificial Intelli- gence and Statistics, pp. 1273-1282. PMLR, 2017.
Song Mei, Andrea Montanari, and Phan-Minn Minh Nguyen. A mean field view of the landscape of two- layer neural networks. Proceedings of the National Academy of Sciences, 115(3)):77655 2018.
Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nati Srebro. Exploring general- ization deep learning. Advances in neural information processing systems, 30, 2017.
Behnam Neyshabur, Hanie Sedghi, and Chiyuan Zhang. What is being transferred in transfer learn- ing? Advances in neural information processing systems, 33:512-523, 2020.
Quynh Nguyen, Mahesh Chandra Mukkamala, and Matthias Hein. On the loss landscape of class of deep neural networks with no bad local valleys. arXiv preprint arXiv: 10749, 2018.
Jeffrey Pennington, Richard Socher, and Christopher D Manning. Glove: Global vectors for word representation. In Proceedings of the 2014 conference on empirical methods in natural language processing (EMNLP), pp. 1532-1543, 2014.
Gabriel Peyré and Marco Cuturi. Computational optimal transport. 2018
Filippo Santambrogio. Optimal transport for applied mathematicians. Birküuser, NY, 55(58-63):94, 2015.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409.1556 2014.
Sidak Pal Singh and Martin Jaggi. Model fusion via optimal transport. arXiv preprint
2019.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher Manning, Andrew Ng, and Christopher Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of the 2013 conference on empiricul methods in natural language pro- cessing, pp. 1631-1642, 20133
Romain Thibaux and Michael I. Jordan. Hierarchical beta processes and the indian buffet process. 2007.
Matthew Thorpe, Serim Park, Soheil Kolouri, Gustavo K Rohde, and Dejan Slepčev. A transporta- tion (P distance for signal analysis. Journal of mathematical imaging and vision, 1972)1 2017.
Cédric Villani. Optimal transport: old and new, volume 338. Springer Science & Business Media, 2008.
Hongyi Wang, Mikhail Yurochkin. Yuckai Sun. Dimitris Papailiopoulos, and Yasaman Khazaeni. Federated learning with matched averaging. arXiv preprint arXiv:2002.06444 2020.
Mikhail
Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, Trong Nghia Hoang. and Yasaman Khazaeni. Bayesian nonparametric federated learning of neural networks. arXiv preprint arXiv:1905. 12022, 2019.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64(3):107- 115, 2021.
Xiang Zhang, Junbo Zhao, and Yann LeCun. Character-level convolutional networks for text clas- sification. Advances in neural information processing systems, 28, 2015.
11


Reviews:
[{"review_id": "yvMQY_R4hEI", "reviewer": "Reviewer_CYR8", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper considers the fusion algorithm to aggregate different neural network models trained locally. The Wasserstein/Gromov-Wasserstein barycenter is the main technique used in the construction. Extensive numerical experiments demonstrate the good performance of the proposed fusion algorithm, and provide empirical evidence for the conjecture made in Entezari et al. (2021).", "summary_of_the_review": "Overall I like the idea of applying Wasserstein/Gromov-Wasserstein barycenter for NN fusion, which is actually a natural idea in federated learning, such as [1][2] (which might be related to this work). And extensive numerical results are provided to justify the good performance. However, it is hard to judge the novelty of this work as barycenters are not new techniques and their usage in federated learning also exists.  \n\n[1] Nguyen, T. A., Nguyen, T. D., Le, L. T., Dinh, C. T., & Tran, N. H. (2022). On the Generalization of Wasserstein Robust Federated Learning. arXiv preprint arXiv:2206.01432.\n[2] Farnia, F., Reisizadeh, A., Pedarsani, R., & Jadbabaie, A. (2022). An Optimal Transport Approach to Personalized Federated Learning. arXiv preprint arXiv:2206.02468.\n\n", "strength_and_weaknesses": "Strength:\n- the proposed fusion algorithms allow aggregating models within a variety of NN architectures \n- by formulating the fusion problem as a series of Wasserstein/Gromov-Wasserstein barycenter problems, bridge the NN fusion problem with computational OT\n- empirically demonstrate the effectiveness of fusing different types of networks, including RNNs and LSTMs\n\nWeaknesses:\n- Although the usage of Wasserstein/Gromov-Wasserstein barycenter in NN fusion seems to be new, the Wasserstein barycenter problem itself is well studied, and thus the contribution of this work is rather limited\n- The proposed fusion algorithm is performed layer by layer, in which the architecture of the whole network seems to be unused.\n\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Uj57HzGKmHA", "reviewer": "Reviewer_obb2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes a mathematical formalism for model fusion motivated by the TLP distance which is interesting. They illustrate the performance of their method by fusing various networks on both homogenous and heterogenous task settings, where they seem to perform slightly better than the OTFusion method upon which the given paper builds. Also, they show additional results by extending OTFusion to handle RNNs and LSTMs by using the Gromov-Wasserstein distances. Further, they also provide interesting visualizations and touch upon the Linear Mode Connectivity (LMC) conjecture in NLP tasks.  Despite all of this,  it must be said that nevertheless, it is but a formal rewrapping of the existing model fusion approach OTFusion --- which is obviously also considering Wasserstein barycenters. Furthermore, their empirical comparisons seem to be off and it seems that they have not been performed fairly, which casts a doubt on the seeming advantage of their method.", "summary_of_the_review": "In general, I liked the paper at first glance. However, when one digs deep, it becomes rather clear that presently both the discussion and the experiments are misleading and not properly carried out. Nevertheless, I am willing to increase my score towards acceptance, when all of the above-stated points have been adequately addressed.\n\n---- POST-REBUTTAL ----\nI thank the authors for their rebuttal. It becomes quite clear from the rebuttal, as the authors themselves admit, that there should not be visible differences between the baselines on FCNs and VGG networks. Likewise, the authors also make it clear that when using the EMD solver, the default choice in the prior work OTFusion, their method performs similarly. However, these aspects will hardly be clear to most readers. It is completely alright not to focus on activation-based costs for the reasons mentioned, but that does not mean it is fair to provide an incomplete/inaccurate picture. Then about the novelty of formalism: it is again totally fine to have the TLP formalism, but all you need to emphasize is the differences/similarities to prior work. In particular, these cannot be relegated to supplementary or openreview discussions. Also, regarding 'WB', as elaborated before I disagree with the authors and surely the argument cannot rely on merely looking at the abstract instead of the entire main paper. Further, specifying 'WB\" for their algorithm is just plain misrepresentation, to be frank. \n\nPresently, as the authors admit, there are many significant and extensive changes yet to be made. So, in this highly-inaccurate shape, the paper cannot be accepted. Moving forward, I would recommend the authors truly depict and represent their own work wrt prior works, and then focus on consolidating their own novel contributions (like GW, theoretical analysis etc.) in a future submission. Thus, I maintain my score. ", "strength_and_weaknesses": "### Strengths:\n- Their formalism of defining a measure for each neuron and the TLP interpretation is interesting and could be useful for future extensions. \n- Support RNNs and LSTMs via Gromov-Wasserstein barycenters extends the model fusion direction further. \n- Some evidence is provided for LMC conjecture on RNNs and LSTMs. \n\n\n### Weaknesses:\n- **Novelty:** The reality of the introduced formalism is that it basically devolves into the model fusion known as OTFusion (Singh & Jaggi, 2019). I happened to read the given paper's Appendix D, which made me compare the two more in detail and which makes this aspect quite clear. The reason is that many of the considered elements in the formalism when instantiated reduce their whole abstract framework to something pretty obvious, for e.g., the function w_j defined from the previous layer to the reals basically turns out to be the weight vector (as the domain is discrete). Apart from the row-wise scaling factors, I don't see much difference between the resulting cost functions in eqn 26 and 27; and this should also not change the resulting transport map much. The paper mentions that OTFusion introduces this without a derivation but that is false; see section S9 of their paper where they derive the barycentric projection. All of this, i.e., the similarities as well as the differences, should be discussed in detail and within the main text so that reader can better realize these aspects.\n\n- **Incomplete discussion & comparison**: The question then arises of where does the empirical difference come from. I believe this can be reconciled by looking at the choice of ground cost used, which is based on comparing the weight vectors. Now, the thing is that under this cost, the fusion problem is a bilinear assignment problem (Ainsworth et al., 2022), and so the transport map at each layer depends on the previous. Thus it makes sense that the current paper has to run more iterations of the OTFusion procedure under this ground metric. However, this paper hides the fact that many of the results in OTFusion are based on what is called 'activations-based alignment'. To repeat Ainsworth et al., 2022, now the problem is a linear assignment problem since the transport map at each layer can be computed independently and thus there is indeed no need of running more iterations. But, unfortunately, this paper seems to ignore this completely. \n - **Improper empirical comparisons**: In fact, the empirical comparisons themselves seem to be rather off as well. For instance, in Figure 2 where they discuss the fusion performance for homogenous and heterogenous tasks, it seems that OTFusion implementation is done strangely. This can be seen if we take a look at the results from that paper in the given settings (which this paper follows), namely Figure 2 and 4 of the OTFusion paper. There OTFusion seems to have much smaller variance and is rather stable, with clear gain in performance. I suspect the difference comes from not using activations based alignment as specified in the caption of that paper. Additionally, for most of the results in OTFusion they seem to use exact OT instead of Sinkhorn since they get a better performance with the former. So before the resulting gains from the paper's method are to be evaluated they should at least ensure that their comparisons are proper, and the experiments fairly reproduced and implemented (such as activations, exact OT, and likewise for results in Table 1 & 2 besides Figure 2). \n\n- To properly show if their method is significantly different from OTFusion, another strategy could be to do the visualizations as in Figure 3 where they consider OTFusion (both the choices of ground cost, weights and activations), their own method, and vanilla average (or something else that is appropriate). I think this should clarify if there is indeed much difference from OTFusion.\n\n- The **discussion of the related work** should be corrected. (i) Firstly, as per the above points, the similarities and differences wrt OTFusion should be properly acknowledged in the section where they introduce their paper. Next, naming their method as 'Wasserstein Barycenters'  WB fusion is plainly incorrect and deceptive. OTFusion, although uses OT in its acronym and not WB, makes it explicitly clear that they use Wasserstein Barycenters to fuse the model. They define WB and the variational problem underneath, it is obvious from their algorithm that they are solving this variational problem, and even more they explicitly discuss the exact update as Barycentric projection. I suggest the authors necessarily rename their method to something which is more representative of what they propose, e.g., TLPFusion. (ii) Besides, OTFusion, it should also be mentioned that there have also been attempts in the literature to handle the quadratic assignment in the context of model fusion Liu et al., 2022, https://proceedings.mlr.press/v162/liu22k/liu22k.pdf. Their approach might be better able to handle the case of RNNs/LSTMs. Furthermore, it was already mentioned explicitly in Wang et al. 2020 https://arxiv.org/pdf/2002.06440.pdf that a Gromov-Wasserstein barycenter based idea could be used in this precise context --- this should also be acknowledged in this work. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "7OQXFivN5J", "reviewer": "Reviewer_fEwc", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a new unified mathematical framework for neural network (NN) model fusion, which is based on the Wasserstein/Gromov-Wasserstein barycenters (WB/GWB), i.e., formulating as a series of optimal transport (OT) problems. The proposed mathematical framework is universal and can be applied to a broad class of NN architectures. The authors experimentally demonstrate that the framework is effective at fusing different types of NNs, including RNNs and LSTMs, and interpret the fusion by visualizing the aggregation of two neural networks in a 2D plane. This may be new experimental evidence about the linear mode connectivity of the loss landscape.", "summary_of_the_review": "Overall I recommend the rejection for this current manuscript. The major reason is that (1) The theoretical and empirical soundness of the proposed WB/GWB model fusion method is somehow limited by the lack of in-depth theoretical analysis and more experimental comparisons, (2) the claim that this work brings new insight, and contribute to the understanding to the linear mode connectivity phenomenon, is a bit over-claimed. It is possible I missed some important details as I am not an expert in the model fusion line of research, and I definitely welcome the authors to correct me if I made some factual mistakes and get involved in the discussion.", "strength_and_weaknesses": "#### Strengths\n1. I like the idea of using optimal transport and Wasserstein/Gromov-Wasserstein barycenters to model the model weight fusion, and to my best knowledge, I think it is novel. The description of the corresponding fusion algorithm is also clear.\n2. I like the idea of visualizing the model fusion and loss landscape in 2D (Figure 3 and related paragraphs).\n\n#### Weaknesses\n1. One major concern is that (please correct me if I am wrong), I think the proposed WB/GWB-based model fusion lack of sufficient & in-depth theoretical analysis (or any theoretical guarantees) on how the test performance is preserved/affected by the proposed fusion method. In terms of experiments, the advantages to baselines like OT fusion and FedMA in terms of test performance are not consistent or significant. So I think a more careful theoretical and experimental comparison to the baselines is also needed.\n2. In terms of the relation to the linear mode connectivity phenomenon/conjecture. Currently, I think only from the visualizations in Figure 3; it is confusing how the experiments on WB/GWB model fusion bring new insights or contribute to our understanding of the linear mode connectivity phenomenon. Is the model fusion setup a special case considered in Entezari et al.(2021) or other related literature? I understand Figure 3 can be thought of as some empirical evidence, but what is the new insight that contributes to that line of research? If the contribution is somehow limited to empirical support, I think it is a bit overclaimed and misleading to address this contribution in the title, abstract, and introduction. Another question is: will other model fusion baselines yiled similar behavior when visualized similarly to Figure 3?\n\n#### Minor Issues\n1. The authors should use `\\citep` and `\\citet` for in-text citations of different formats. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "1HTI4oobmqO", "reviewer": "Reviewer_shzY", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a method for neural network model fusion based on the concepts of Wasserstein barycenter, and Gromov-Wasserstein barycenter. The method uses OT couplings from each previous layer to construct couplings between the subsequent layer that minimize a cost based off the previous layers' couplings. Significant empirical validation of this method is given, including applications to model fusion in data hetergenous settings, data homogenous settings, as well as diverse model architectures and some exploration of the linear mode connectivity hypothesis.", "summary_of_the_review": "This work addresses an important problem and proposes a method with extensive empirical validation. However, the empirical results are not completely convincing, and the method is only marginally novel compared to previous work. Therefore, I think it is slightly below the threshold of acceptance.", "strength_and_weaknesses": "Strengths:\n- Model fusion is an important problem\n- Applying optimal transport to this problem is of significant interest\n- Some of the empirical performance is significantly better than existing work\n\nWeaknesses:\n- The empirical performance is not convincingly improved on the whole\n- There seems to be only marginal innovation over the previous work Singh & Jaggi 2019\n- There is no theoretical development of their ideas", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject