Layer-wise linear mode connectivity
Authors: Linara Adilova
Maksym Andriushchenko
Michael Kamp
Asja Fischer
Martin Jaggi
Venue: ICLR 2024 poster (2024)
Keywords: linear mode connectivity
layer-wise
federated averaging
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: Investigation of the layer-wise structure of the barriers on the line between two parametrizations of deep models

Abstract: Averaging neural network parameters is an intuitive method for fusing the knowledge of two independent models. It is most prominently used in federated learning. If models are averaged at the end of training, this can only lead to a good performing model if the loss surface of interest is very particular, i.e., the loss in the midpoint between the two models needs to be sufficiently low. This is impossible to guarantee for the non-convex losses of state-of-the-art networks. For averaging models trained on vastly different datasets, it was proposed to average only the parameters of particular layers or combinations of layers, resulting in better performing models. To get a better understanding of the effect of layer-wise averaging, we analyse the performance of the models that result from averaging single layers, or groups of layers. Based on our empirical and theoretical investigation, we introduce a novel notion of the layer-wise linear connectivity, and show that deep networks do not have layer-wise barriers between them.

Published as a conference paper at ICLR 2024
LAYER-WISG LINEAR MODE CONNECTIVITY
Linara Adilova
Maksym Andriushchenko
Ruhr University Bochum, EPFL
EPFL
linara www.ooaddehhr-aaeeee www.huuuuu-cuuuummmm de
ma maksym. andr www.auuucschhhhhoeeeee..ccccc
Michael Kamp
Asja Fischer
IKIM UK Essen, RUB, and Monash University michael kamp&uk-essen
Ruhr University Bochum as ja. www.hirerehrr-uuueeerrre -bochum.
Martín Jaggi
EPFL
martin. martin.jaggieen jaggidepf.ie
ABSTRACT
Averaging neural network parameters is an intuitive method for fusing the knowl- edge of two independent models. It most prominently used in federated learning. If models are averaged at the end of training, this can only lead to good performing model if the loss surface of interest is very particular, i.e., the loss in the midpoint between the two models needs to be sufficiently low. This is impossible guaran- tee for the non-convex losses of state-of-ther of-tth-art networks. For averaging models trained on vastly different datasets, it proposed to average only the parameters of particular layers or combinations of layers, resulting in better performing models. To get better understanding of the effect of layer-win wise averaging, performance of the models that result from averaging single layers, or groups empirical and theoretical investigation, we introduce novel
analyse the
layers. Based notion of the layer-wise linear connectivity, and show that deep networks do not have layer-wise barriers between them."
INTRODUCTION
Understanding the optimization trajectory of neural network training, relative the structure of the loss surface, can contribute significantly the development of better performing and more reliable models. The loss surface of deep networks is far from being understood. Getting path
better picture of the loss barriers between two models is a part of this challenge.
Important examples of findings that contributed to getting better understanding of such paths the discovery of non-linear paths connecting minima without increase of the loss (Garipoy al., 2018; Draxler et al., 2018), the develop- ment of analytical approaches to perform feature matching or transforming one network into an- other (Singh & Jaggi, 2020), and the analysis of linear paths between minima or minima and origin (Frankle et al., 2020; Zhang et al., 2022; Vlaar & Frankle, 2022). One of the multiple
Figure 1: CIFAR-10 with ResNet18. Heatmap shows layer-wise averaging barriers for layers on Y-axis throughout training epochs on X-axis. First row shows the full networks averaging barrier.
Coode for the experiments is published httpt::tt////.. htttpa://gittub..oom ink wakkeer/1a wwweerrrreaaaaeeee wissseii se-1mm

---

Published as a conference paper at ICLR 2024
applications for such insights is, for example, knowledge fusion performed in more efficient way than straightforward model ensembles.
The largest obstacle the way to understanding the loss surface is the depth of modern neural networks. A good performance requires multi-layer networks, but a formal analysis of the surface has only been done for one layer networks (Safran al., 2021; Simsek et al., 2021). Interestingly, layers were empirically observed to have an emergent individual behavior. For example, shallow layers were found to converge sooner during the training than deep layers (Chen et al., 2022b), using an individual learning rate for each layer can be beneficial for final performance (Dong et al., 2022), and loss behavior on the one-dimensional cuts towards initialization values differs from layer to layer (Vlaar& & Frankle, 2022).
The research in this paper directed towards understanding the layer-wise behavior of loss barriers between models. This question is of particular interest for federated learning practitioners, because understanding the reasons for the success of averaging in non-convex problems is vital for further progress. In particular, federated training with averaging models at the end is analogous to models being trained independently like in the work of Frankle et al. (2020). If, instead, aggregation (typically averaging) is performed during training. then each of the interpolated models serve a starting point for further training. We are investigating the setup of averaging at the end, for multiple end end-points during the training process, i.c., analyzing averaging if we would stop after each epoch. Investigating the dynamies of averaging during training can give exciting insights into the appearance of barriers in federated learning, but we leave it for the future work.
Our contributions are as follows:
We propose a layer-wise linear mode connectivity property and show that a wide range of models do not have layer-wise barriers (Fig. 1). For deep linear networks we show that this might be explained by convexity of the loss surface with respect to the linear cut at individual layers. We additionally investigate connectivity of groups of layers.
We show that robustness perspective can shed light the appearance of interpolation barriers, and demonstrate in a simplified setup that particular subspaces of the optimization landscape have different robustness properties.
Finally, we apply the gained understanding to the personalization setup in federated learning with layer-wise wise aggregation, conjecturing that in non-i.i.d data separation such approach might be not suitable.
RELATED WORK DISCUSSION
general investigation of the loss surface of neural networks is important for further advancements of the optimization process, in particular for federated deep learning. So far, precise mathematical analysis was possible only for shallow models (Safran et al., 2021; Simsek et al., 2021), while deep models remain black boxes. One of the empirical approaches to this problem is analyzing connectivity properties of the parameters, i.c., if two models can be connected by path on the loss surface which does not raise the loss value compared to either end points. This is generally termed as mode connectivity, where a mode is a parameterization of a neural network which usually (but not necessarily) has low loss. Starting with the exploration of non-linear paths (Draxler et al., 2018; Izmailov et al., 2018) research continued into understanding interconnected minima (Wortsman al., 2021) and investigation of the linear mode connectivity (LMC) (Frankle et al., 2020; Entezari et al., 2022). It is interesting that it was in particular observed that deeper networks have larger barriers (Entezari et al., 2022). The example of LMC being a helpful approach for loss surface understanding is the work of Yunis et al. (2022): they employed the notion of linear connectivity for building convex hulls of models, that are supposed to model the solution basins, and investigated their properties. An interesting aspect of LMC is its relation to the functional similarity of the models. In particular, Entezari et al. (2022) hypothesizes that different basins contain functionally different networks, while feature matching moves them to the same basin and thus makes them similar. Evidence from the work of Fort et al. (2019) shows that sampling weights in the surrounding of a trained model does not give much benefit in the ensembles independent training, meaning that the models are too similar. In the work Lubana et al. (2023) the conjecture is that only mechanistically similar networks be linearly connected, i.e., models should have similar behavior on semantically
2

---

Published as a conference paper at ICLR 2024
similar input features. This is also confirmed by investigation of layer-wise feature connectivity (Zhou et al., 2023), where linearly connected models were shown to have similar features in every layer. Yet, Yunis et al. (2022) show that models in the convex hulls are functionally not similar. Also, Frankle et al. (2020) empirically demonstrates absence of simple correlation between LMC and functional similarity or Euclidian distance.
It is hard to identify whether barriers between models are harmful or helpful for training in case of continuous federated averaging: It is known fact that selecting gradient directions from a point of higher loss is beneficial for training (Foret et al., 2020) and there are indications that similar dynamics are at play in a distributed setup (Zhu et al.. 2023). But it also known that specifically generated bad initializations (Liu et al., 2020) as well as minimax optimization for adversarial robustness (Tsipras et al., 2018) can result in decreasing performance, thus it cannot be always beneficial to look for high loss point for the training restart. It was even proposed to match the features in the models before averaging (and according to Entezari et al. (2022) bring them to one basin), with empirical demonstration of improved federated training (Wang et al., 2020). Different training regimes can expose interesting properties of the loss surface as well: Even if two models trained from scratch cannot be successfully averaged (Frankle et al., 2020), different tuninng of pretrained model allow for fruitful averaging of any amount of models (Wortsman et al., 2022). Analogously, starting from a pretrained model in the federated learning setting can achieve better results than training from scratch, specifically in the relevant case of non-i.i.d. data (Chen et al., 2022a).
Most of the proposed methods for combination fusion of models build layer-wise alignment of activations or weights to achieve linear connectivity (Singh & Jaggi, 2020; Ainsworth et al., 2022;
Jordan et al., 2022). Nevertheless, one layer is sometimes enough to achieve successful fusion: Rebuffi et al. (2023) combined adversarially robust and non-robust models; Bansal et al. (2021)
combined two completely different models in one; Ilharco et al. (2022), Ortiz-Jimenez et al. (2023)
performed task arithmetic, i.e., added knowledge about new task to model without causing it to forget the original task. Recently, an attempt to confirm a layer-restricted memorization was made by Maini et al. (2023), but the results indicate that memorization is happening throughout the network and not only in one layer.
The observation that deep neural networks during training converge bottom-up (i.e.. first shallow layers and last deep layers) attracted a lot of empirical investigations (Chen et al., 2022b; Li et al., 2019; Raghu et al., 2017). This can be looked at from multiple perspectives: that deeper layers are moving further away from initialization, while shallow layers stay close (Zhang et al., 2022; Andriushchenko et al., 2023b); that the loss with respect to the shallow layers is more smooth and allows for fast convergence (Chen et al., 2022b); and that the gradients for the shallow layers are vanishing with the loss becoming smaller. Some of these perspectives contradict each other, for example, it is unclear if training of shallow layers stops early because they indeed converge to the most optimal state or just because gradients propagation is not possible anymore, thus pointing to little common understanding of the aforementioned phenomenon. At the same time it points to the layer-wise difference of the training process, which consequentially means that the loss surface of the optimization task has particular layer-wise structure. The lack of understanding of the layer-wise structure leads to surprising results, e.g., that sharpness aware minimization is sufficient for improving generalization when applied only to BatchNorm layers (Mueller et al., 2023). Investigation of the layer-wise structure of the nie-dimentional cuts of the loss surface was performed for understanding optimization process: Connecting the initialization and trained model gives insights on how successful is the training (Zhang et al., 2022; Chatterji et al., 2020; Vlaar & Frankle, 2022). These works also demonstrate very different behavior of the individual layers when interpolating to the initialization. Moreover, there seem to be only a subset of layers affecting the performance of the model when reinitialized, and its size can be used complexity measure of the model (Chatterji et al., 2020).
3 EMPIRICAL LAYER- LINEAR MODE CONNECTIVITY (LLMC)
We consider a network architecture A parametrized by W that is trained on task represented by a training set Suain and a test set Swst. both sampled from a data distribution D. At each point during training, one measure both loss c(W, S) and error E(W,S) (i.e., one minus classification accuracy). They can be measured both on the training and test set. Note that in the literature on LMC, both training and test losses are used; in the context of federated learning the training loss and
3

---

Published as a conference paper at ICLR 2024
error are most insightful, since they directly influence local optimization in our experiments find, though, that both training and test losses show similar trends. In the following, we consider training loss and error and write e (TW)E((W) for e(W, Smain), E(W), Smann) Assume that we have fixed two different weight parametrizations W1 and W2. Let 6a (M,, W2) = e e(aWi + (1 ()))) and F., (M1, W2) F(aW] + (1 - a)) for [0, be the loss and error, respectively, of the network created by linearly interpolating between Wi and W2. Then Frankle et al. (2020) define the following notion of instability.
The difference between the supremum of the loss for any interpolation aup Sa (W1, W2) and the average loss of the endpoints [(((W))++ +c(Wa) is called the linear interpolation instability for the given architeclure A.
Definition 1.
Note that one can use the error instead of the loss to define a corresponding measure of instability. Note that, since is continuous value, a granularity over which the supremum is found needs to be selected in practice. This is a decisive factor that allows to look at the loss surface in less or more details. Existing abundant evidence indicates that such interpolations are smooth, nevertheless is not formally proven.
Two parametrizations Wi and W2 have a linear barrier between them if the linear interpolation instability is sufficiently high. When models are very different in performance or when both of them are not performing good, an absence of a linear barrier does not mean that the performance of the interpolation model is good, though. It is assumed in the literature that two models are in a convex valley and thus can be successfully averaged once they are sufficiently well trained (Entezari et al., 2022; Frankle et al., 2020; Wortsman et al., 2022). We consider barriers between models at various stages of training, since in federated learning averaging is happening throughout the training process.
In the following, we analogously define a layer-wise notion of instability. Let A be structured in L layers (W()), ..... W J. In our experiments, we consider both weights and bias as one set of parameters describing a layer. Let us fix a layer W Consider a parametrization that is defined by o, Wi and W22 as ( (W,"), w
w/ where j can be selected to be 1 or 2. Such parameterization essentially lies on the line between Wi and W2, projected onto the subspace of layer W Denote as Ca.i and Ea, loss and error measured in that point.
w(()) (1 a:) w))
Definition 2. (Layer-wise linear interpolation instability) The difference between supremum of the loss on the line sup, (VV,, W2) corresponding to layer W and average loss of the original models (le(W))+ee)) +e(2)) is the layer-wise linear interpolation instability for the given architecture A and selected layer.
Note that, since we use the initial weights Wi or W2 as an origin, the selection of model j E {1,2] defines the parametrization around which we consider layer-wise interpolations. Obviously, if one model is more performant than another or more robust to weight changes the loss will be different for the same a but different j.
In the following we say that averaging is successful if the resulting model performs on par with the original ones, i.e., there is no barrier between the two models the average. Difference around 2% can be assigned to the randomness of the training process, thus only if the loss value is larger is a barrier (Frankle et al., 2020). Since federated learning averages models, we consider not the supremum over a E 10, 1). but instead only the middle point a = 0.5. To make this clear, we use the term averaging barrier (avg. barrier).
Convolutional models. In the following we show empirically that there are no layer-wise avg. barriers (Def.2) for ResNet18 trained on CIFAR-10. We replace BatchNorm layers with identity, because they are known to affect the averaging (Li et al., 2021). We consider the following training setups: (i) parallel training on the full training set with different data shuffling, (ii) with same data shuffling but different initialization. (iii federated setup (without aggregation) for two clients with i..dd local training data, and (iv) with non-i.id. local training data (split by Dirichlet distribution on labels with parameter 0.1 and 12). Fig. and Appx. Fig. 6 demonstrate that for every setup there no layer- wise avg. barriers, while a linear barrier is present.
Large language models. We test how layer-wise connectivity (Def.2) behaves in large language models. We trained a small GPT-like model with 12 layers on Wikitext. The results are shown in Appx. Fig. 19. Here, compute barriers using test set, demonstrating that also in this setup there
4

---

Published as a conference paper at ICLR 2024
mostly no layer-wise avg. barriers between models. We note, that different initialization and small learning rates result in barriers in some of the shallow layers. It is interesting, that weight sharing between the first layer and the last layer seem to affect the barrier a lot: when it is used, the barrier on the last layer is as large as the full networks barrier, while when there is no weight sharing the barrier is not so pronounced. We further checked Pythia² model pairs, which are trained on different datasets, but have the same architecture. We compute the barriers on the test set of Wikitext data (Appx. Fig. 2,,21,22). Pythia models do not use weight sharing while training, but smaller models still show rather significant barriers when averaging the last layer, different from the larger model.
Cumulative layer-wise structure. The natural question arising from the demonstrated results is whether several layers combined can lead to a barrier and how many layers are needed then. For this we introduce one more notion of instability. Let us fix a subset of layers W W(++) (++)))
Consider parameterization that is defined by a, W1 and W2 as (W("), w/() a + (1
a) () ( (1 aa) (1+ "..."............
................................ a W (6+c) + a1-a) ((ite)) ..... w( where j can be selected to be 1 or 2. Denote as f,,, I...itt,....ite and Fa.... loss and error measured in such point. Note, that fix layers going one after another for the ease of mathematical notation; the definition is the same if the layers do not follow one another.
Definition 3. (Cumulative layer-wise linear averaging instability) The difference between the middle point of the loss the line 60.5,i.4 (o.,sit W2) corresponding to layers W w(++),,... ..... w(++e) .((t)) and average loss of the original models (e((W) + <(W2)) is the cumulative layer-wise linear averaging instability for the given architecture A and selected layers.
We investigate the barrier value (Def. 3) when larger sub- sets of layers are cumulated. We consider two directions of cumulation: from shallow to deep layers and from deep to shallow, i.e., in the first case, starting with the most shal- low layer and replacing it by the average we move to more layers, till the full networks are averaged. Fig. 2 shows a curious structure revealing itself in the buildup of the barriers: neither shallowest nor deepest layers cause the barrier, but the middle ones do. We demonstrate that the position of layers causing barriers is very well defined and does not depend on the federated setup, i.i.d. and non-i.i.d (Appx. Fig. 11, 12, 13, 14). We verified this via checking randomly selected layers for cumulation, sliding window cumulation, and observing larger amount of layers to cu- mulate without barrier when only shallow or deep layers are considered (Appx. Fig. 7). The effect of the learning rate is pronounced in this set of experiments: high learn- ing rate allows to see the same structure, independent of the difference in the initializations, while low learning rate results in not linearly connected shallow layers when initialization is different (Appx., Fig. 8, 9, 10). We observe similar phenomena with VGG11 (Appx., Fig. 18).
(a) Shallow cumulation barriers
(b) Deep cumulation barriers
While the cumulative structure phenomenon might have curious implications, we left it's investigation for future work. It is possibly connected to the work of Jacot (2023) which shows that a deep neural network learns simple 1- dimensional function in its middle layers. The experiment with different learning rates hints at properties of LLMC being dependent on the optimizer parameters. For exam- ple. high learning rate promotes sparser (Andriushchenko et al., 2023b; Chen et al., 2024), and potentially more similar, features in shallow layers.
Figure 2: CIFAR-10 with ResNetl 8. Full data training setup, from same ini- tialization. Heatmap visualizes cumula- tive averaging, each layer added to the group of averaged layers one by one, starting from bottom or top.
from Biderman et (2023)
5

---

Published as a conference paper at ICLR 2024
4 MINIMALISTIC EXAMPLE OF LLMC
In order to better understand the reasons behind the absence of layer-wise barriers for models with no linear connectivity, we analyze a minimalistic example of linear networks. We choose a one- dimensional diagonal linear network l(w1, W2) = (1 W1 w2)² as one of the simplest non-convex models. We observe the LLMC phenomenon in Fig. 3: full interpolation between two minima (wi, W2) and w' (w', w2) leads to barrier, while interpolating only the second layer- which results in the point (wh, 2 102 +
to a much lower loss. However, interpolating only the first layer leads high loss which is consistent with some of our experiments deep non-linear networks.
Layer-wise convexity. Fig. 3 also illustrates that the loss convex on a line in wi and w2 separately (i.e., along any coordinate-aligned slice of the loss surface) but not in (w), w2) jointly This result can be formally generalized for linear networks of arbitrary depth.
Theorem 4.1 (Layer-wise convexity). Let the squared loss of deep linear network interpolated between sets of parameters (W |= and (W"()) interpolation coefficient a be
J=1 at any layer k (1....,T) L.} with
[(a) = Y XW (a + - a) wi()) w )..tk))....r
(1)
then L(a) is convex and there are no barriers in layer-wise interpolation.
Proof. We can rewrite L(a) as
L(a) = Y X (.) (..........)))) axW(). (w)) ... + axW() wi(h) w
Y XW ..... w(t) ta XW (Wilk)) W(l))... 1/== I + a
which is convex since the second derivative is non-negative:
d²
d2
daz L(a
d (1 (Y, W) + la = da a Waw 2 2| > 0.
Convexity of L(a) implies that there are no barriers in layer-wise interpolation: for any [0, 1]. we have [(a) I aT(()) + (1 (()) consequence of convexity. In particular, if both L(0)
and (1) have low loss, then the whole line segment between them also has low loss.
This shows an interesting layer-wise structure of the loss surface of deep linear networks: while layer-
1(W1, w2) = (1 W1W2)2
5
the overall loss landscape non-convex, wise convex on the linear cut. In particular, it means that although there can exist barriers under full net- work interpolation, there are no barriers under layer- wise interpolation. Moreover, due to convexity, the layer-wise interpolation loss is expected to grow not too fast since L(a) S aL( + (1 1 ((), i.e. in the worst case the increase of L(a) will be linear in a. For shallow non-linear networks was theoretically shown that there are also convex in- terpolations in most directions with respect to the first layer (Safran et al., 2021). We will see that it is also the case even for deep non-linear networks in the next section. Of course, does not have to hold in general for non-linear networks. But it suggests that the layer-wise structure the loss surface can be much simpler than the global struc- ture which supports the empirical observation that LLMC often holds when LMC does not.
10
We)
(w).
10
Gis + mm fura
10
10
Jwn. wS)
[wc] mg)
W1
Figure 3: Minimalistic example of the LLMC phenomenon with a ID diagonal linear network: joint interpolation between w and w' leads to a barrier, while interpolating only the second layer leads to a much lower loss.

---

Published as a conference paper at ICLR 2024
5 TOWARDS UNDERSTANDING THE LLMC PHENOMENON
In the following we investigate the properties of deep neural networks that can cause the observed phenomenon of LLMC. We present a robustness view on it and explore how the perturbations in different directions change the loss value.
5.1 LLMC AND ROBUSTNESS IN THE PARAMETER SPACE
Figure 4: Layer-wise interpolations (left: model 1 - model and model 1) and robustness to random perturbations of the same norm (right: model 1 - model 2 and model 2 + /) for vision transformers trained on CIFAR-10 10 with different learning rates and data augmentations. Here axis is the different interpolation points a.
We evaluate the models from a public repository which contains vision transformers (ViTs) trained using the same initialization and randomness but different hyperparameters (such as p of SAM and learning rate). We select three pairs of models: (1) trained with small learning rate (LR) without augmentations vs. large LR with augmentations (Fig. 4), (2) small vs. large LR, both trained without augmentations (Appx. Fig. 23). (3) trained with SAM with = 0 vs. 0.1, both trained without augmentations (Appx. Fig. 24). We compute loss of layer-wise interpolations (Fig. 4, left) and random direction perturbations (Fig. 4, right) of the same norm as the perturbation induced by the layer-wise interpolation. We sample random perturbations several times and average the obtained losses. We note that it can be
layer-wise flatness in random direction.
Fig. 4 suggests that we get barrier- free interpolation at a = 0.5 for almost all layers. Interestingly, we observe no significant growth in the linear head interpolation in contrast to the LLM experiments. Instead, the most sensitive layers are the early attention (qkv) and fully-com fully-connected (net) weights. We also observe that the success of interpolations is highly asymmetric for pair of models, and for flatter models (due to larger LR or larger of SAM), the loss grows slower over the interpolation coefficient a (first row of the heatmaps). These results confirm that (i) the robustness of the model indeed affects the barrier development and (ii) loss grows monotonically with a convex trend, at least locally for not too large values of a, which is coherent with Theorem 4. 1. Moreover, the networks are much more robust to layer-wise random perturbations compared to the layer-wise direction of interpolation between models. This suggests that layer-wise averaging directions are special in the sense of having much higher curvature than random ones. We discuss this in the next section.
We also perform robustness analysis for the setups with ResNet18 and CIFAR- 10. this group of experiments we compare the loss of the models on averaging point for each of the layers with the random directions loss taken at the same distance (Appx. Sec. A.2.2). We check the robustness when
Sttppa //gl http://tttttht////hh.e -epf wharppeas-seeee 95 wessggnneeee iat from Andriushchenko al. (2023a)
7

---

Published as a conference paper at ICLR 2024
= and observe that while random directions still do not cause the growth of loss, layer- interpolation does. The most curious is that the layers that are sensitive to averaging direction perturbation coincide with the layers that are shown to be critical for ResNet18 architecture in Zhang et al. (2022). For the case of ViTis LayerNorm layers were demonstrated to be most critical along with the layers that our analysis indicates to be sensitive to perturbations. A natural question is if the layer-wise perturbation distance is just too small and therefore random perturbation does not change the loss. We answer negatively by the empirical results in Appx. Fig. 30.
5.2 SPECIAL DIRECTIONS ON THE LOSS SURFACE
Our experiments show that the impact of perturbations in the direction of another model is dif- ferent from perturbations in random directions. To analyze this phenomenon further, we in- vestigate how the impact on the loss differs for parameter changes (i) in the direction of an- other model, (ii) in the subspace spanned by the training trajectory of the two models (the training space), and (iii) the null space, i.e., the subspace perpendicular to the training space.
We train two fully connected networks with 3 hidden lay- on MNIST for 50 epochs and save checkpoints each epoch. We compute the average parameter vector of the final two models. The averaging direction for each net- work is unit vector from the final network's 's parameter vector to the average. We describe the training space by computing an orthonormal basis of the span of the 102 vectors using singular value decomposition. Similarly, we find an orthonormal basis for the null space. We then sam- ple random noise by first sampling a random unit vector from each subspace and multiplying it with magnitude a sampled uniformly from [0.001, 10.0]. Afterwards we check the impact of that noise the test loss of both final models. The results shown in Fig. 5a indicate that perturbations along the average direction indeed have the highest impact on the loss, and perturbations perpendicular to the training space have higher impact than perturba- tions in the training subspace. A possible interpretation is that minima are flat in training space (thus perturbations in training subspace have low impact), but the two final models are in distinct minima (so loss changes a lot in the averaging direction). Random directions in training space would have low likelihood of pointing towards the other minimum and thus perturbations have in expectation less impact. Perturbations in directions perpendicular to the training space have strong impact on the loss, which is reasonable since those directions did not improve the loss during training and are more likely be detrimental.
12-
networt:
retvoct
netwert eaerae
renvor (2a0e)
retrert (deen)
34-
02
(a) full network
(b) layer
(d) layer
(c) layer 2
(c) layer 4
Figure 5: Test loss of two networks with perturbations of magnitude in the train- ing subspace, null space, and along their averaging direction; perturbing the full network and separate layers.
To understand the connection of this phenomenon to the layer structure of the networks, we perform the same ex- periment but restrict the parameter vectors to individual layers, for which we compute the average direction, train- ing space and null space. As seen in Fig. 5b-5e, the overall picture changes when looking at layers: while noise from the null space has strong effect on the loss for all layers, the effect of noise in training space decreases with the depth-it is strong for the most shallow layer and nearly has no impact for the last layer. The most striking difference we find for noise in the averaging direction, though. Here, we a noticeable effect on loss only in the most shallow layer. The other three layers are nearly entirely robust noise in the averaging direction.
In order to exclude the possible effect of ReLU on the results, we experiment also with sigmoid and tanh, for which we observe the same behavior. Moreover, the results are the same when using more
8

---

Published as a conference paper at ICLR 2024
than two neural networks only the dimension of the training space increases (the increase was linear with the number of networks in our experiments). This all indicates that the averaging direction is indeed special in the training space and perturbation in the null space, perpendicular to the training space, consistently has a high impact on loss. Thus selecting a noise direction is crucial for such notions as robustness (Xu & Mannor, 2012) or flatness (Petzka al., 2021).
LLMC AND THE PERSONALIZATION PUZZLE IN FEDERATED LEARNING
Personalization in federated learning aims at reusing the knowledge from local models for mutually improving local models performance. A very common approach is to select for aggregation only the layers that carry the common knowledge. In the literature was proposed to average the deepest layers (Liang et al., 2020), shallowest (Arivazhagan et al., 2019), or even learning weights for each of the layers (Ma et al., 2022). It is very hard, though, to identify which layers carry the local ledge and which the common. Moreover, there seem to be indication that knowledge cannot be localized to particular layer at all (Maini et al., 2023). Using the insights described in previous sections, we consider averaging only the layers that produce cumulative barrier, well as the ones that have the most pronounced sensitivity to the averaging directions. We also consider the reversed setup, i.e., averaging only layers different from those listed above. To our surprise the results for all considered partial aggregations do not differ significantly (Appx. Tab. 1). We conjecture, that in the setup where the architecture is powerful enough to learn the global task and full averaging outperforms local training. none of the partial averaging approaches will be able to outperform full averaging, but it can be on par. At the same time, in the pathological non-i.i.d. case, when full averaging prevents local models from training and local training is significantly more successful, partial averaging performs on par with local training. We conclude, that no knowledge about the LLMC helps to find a more successful setup for partial averaging. This is in alignment with the conclusions of Pillutla et al. (2022), where the main benefit of partial averaging is shown to be less communication.
7 DISCUSSION AND CONCLUSIONS
In this work we investigate the fine-graines structure of barriers on the loss surface observed when averaging models. We propose a novel notion of layer-wise linear mode connectivity and show empirically that on the level of individual layers the averaging barrier is always insignificant compared to the full model barrier. We also discover a structure the cumulative averaging barriers, where middle layers are prone to create a barrier, which might have further connections to the existing investigations of the training process of neural networks. It is important to emphasize that the definition of barrier should be selected very carefully: When performance of the end points is very different, comparing to the mean performance might be misleading for understanding the existence of barrier. Our explanation of LLMC from the robustness perspective aligns with previously discovered layer criticality (Zhang et al., 2022) and shows that indeed more robust models are slower to reach barriers. Training space analysis indicates that considering random directions the loss surface might be mislcading for its understanding. So, for example, searching for non-convexity along the training path is usually unsuccessful (Xing al., 2018).
Our research opens an interesting question: How is the structure of barriers affected by the optimiza- tion parameters and the training dataset? We preliminary investigation we observe that easier tasks result in less layers sensitive to perturbations (Appx. Fig. 31). Understanding this connection can explain the effects of the optimization parameters on the optimization landscape. Together with the existing empirical evidences that an individual layer can be a powerful tool for lossless alignment of different models, e.g., (Bansal et al., 2021; Rebuffi et al., 2023), it can be claimed that the loss surface has a pronounced layer-wise wise structure. Our preliminary experiments on personalization and existing research on memorization (Maini et al., 2023) indicate that such layer-wise structure does not necessarily result in a concentration of particular knowledge in any individual layer, though. This also aligns with the common intuition that the best representation extracted from neural network is often the activation of the penultimate layer. Further investigation of the interconnection between information propagation through the network layers and the optimization process is exciting direction for future work. This can help understanding the connection between structural similarity and functional similarity of models, as well as relating proximity the loss surface to functional similarity.
very pronounced effect of learning rate and in

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
Linara Adilova conducted the research presented in the paper during an exchange semester at EPFL supported partially by the ELISE Mobility Program. Maksym Andriushchenko was supported by the Google Fellowship and Open Phil AI Fellowship. Michael Kamp received support from the Cancer Research Center Cologne Essen (CCCE). Asja Fischer acknowledges support by Deutsche Torschungsgemminschaat (DFG, German Research Foundation) under Germany's ' Excellence Strategy -EXC-2092 CASA 390781972.
REFERENCES
Samuel Ainsworth, Jonathan Hayase, and Siddhartha Srinivasa. Git re-basin: Merging models modulo permutation symmetries. In International Conference on Learning Representations, 2022.
Maksym Andriushchenko, Francesco Croce, Maximilian Müller, Matthias Hein, and Nicolas Flam- marion. A modern look at the relationship between sharpness and generalization. International Conference on Machine Learming. PMLR, 2023a.
Maksym Andriushchenko, Aditya Varre, Loucas Pillaud- Vivien, and Nicolas Flammarion. Sgd with large step sizes learns sparse features. In International Conference on Machine Learning. PMLR, 2023b.
Manoj Ghuhan Arivazhagan, Vinay Aggarwal, Aaditya Kumar Singh, and Sunav Choudhary. Feder- ated learning with personalization layers, arXiv preprint
2019.
Yamini Bansal, Preetum Nakkiran, and Boaz Barak. Revisiting model stitching to compare neural representations. In Advances in neural information processing systems, volume 34, pp. 225-236, 2021.
Stella Biderman, Hailey Schoelkopf, Quentin Anthony, Herbie Bradley, Kyle O' O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, Aviya Skowron, Lintang Sutawika, and Oskar van der Wal. Pythia: A suite for analyzing large language models across training and scaling, 2023.
Niladri S Chatterji, Behnam Neyshabur, and Hanie Sedghi. The intriguing role of module criticality in the generalization of deep networks. In International Conference on Learning Representations, 2020.
Feng Chen, Daniel Kunin, Atsushi Yamamura, and Surya Ganguli. Stochastic collapse: How gradient noise attracts sgd dynamics towards simpler subnetworks. Advances in Neural Information Processing Systems, 36, 2024.
Hong- You Chen, Cheng- Hao Tu, Ziwei Li, Han Wei Shen, and Wei-Lun Chao. On the importance and applicability of pre-training for federated learning. In International Conference on Learning Representations, 2022a.
Yixiong Chen, Alan Yuille, and Zongwei Zhou. Which layer is learning faster? a systematic exploration of layer-wise convergence rate for deep neural networks. In The Eleventh International Conference on Learning Representations, 2022b.
Xiaoyi Dong, Jianmin Bao, Ting Zhang, Dongdong Chen, Shuyang Gu, Weiming Zhang, Lu Yuan, Dong Chen, Fang Wen, and Nenghai Yu. Clip itself is a strong fine-tuner: Achieving 85.7% and 88.0% top-1 accuracy with vit-b and vit-I on imagenet. arXiv preprini 22..v:2212.06111 2022.
Felix Draxler, Kambis Veschgini, Manfred Salmhofer, and Fred Hamprecht. Essentially no barriers in neural network energy landscape. In International conference machine learning, pp. 1309-1318. PMLR, 2018.
Rahim Entezari, Hanie Sedghi, Olga Saukh, and Behnam Neyshabur. The role of permutation invariance linear mode connectivity neural networks. In International Conference on Learning Representations, 2022.
10

---

Published as a conference paper at ICLR 2024
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-awart minimization for efficiently improving generalization. In International Conference on Learning Representations, 2020.
Stanislav Fort, Huiyi Hu, and Balaji Lakshminarayanan. Deep ensembles: A loss landscape perspec- tive. arXiv preprint arXiv: 1912.02777, 2019.
Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin. Linear mode connectivity and the lottery ticket hypothesis. In International Conference on Machine Learning, pp. 3259-3269. PMLR. 2020.
Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry P Vetrov, and Andrew Wilson. Loss surfaces, mode connectivity, and fast ensembling of dnns. In Advances in neural information processing systems, volume 31, 2018.
Gabriel Ilharco, Marco Tulio Ribeiro, Mitchell Wortsman, Ludwig Schmidt, Hannaneh Hajishirzi, and Ali Farhadi. Editing models with task arithmetic. In International Conference on Learning Representations, 2022.
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson Averaging weights leads to wider optima and better generalization. In 34th Conference on Uncertainty in Artificial Intelligence 2018, UAI 2018, pp. 876-885. Association For Uncertainty in Artificial Intelligence (AUAI), 2018.
Arthur Jacot. Bottleneck structure in learned features: Low-dimension regularity tradeoff. arXiv preprint arXiv:2305. 1900000 2023.
Keller Jordan, Hanie Sedghi, Olga Saukh, Rahim Entezari, and Behnam Neyshabur. Repair: Renor- malizing permuted activations for interpolation repair. International Conference on Learning Representations, 2022.
Xiaoxiao Li, Meirui Jiang, Xiaofei Zhang, Michael Kamp, and Qi Dou. Fedbn: Federated learning on non-iid features via local batch normalization. International Conference on Learning Representations, 2021.
Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate in training neural networks. Advances in Neural Information Processing Systems, 32, 2019.
Paul Pu Liang, Terrance Liu, Liu Ziyin, Nicholas B Allen, Randy P Auerbach, David Brent, Ruslan Salakhutdinov, and Louis-Philippe Morency. Think locally, act globally: Federated learning with local and global representations. arXiv preprint arXiv:2001.01522 2020.
Shengchao Liu, Dimitris Papailiopoulos, and Dimitris Achlioptas. Bad global minima exist and sgd can reach them. Advances Neural Information Processing Systems, 33:8543-888 8552, 2020.
Ekdeep Singh Lubana, Eric J Bigelow, Robert P Dick, David Krueger, and Hidenori Tanaka. Mecha- nistic mode connectivity. In International Conference on Machine Learning, pp. 22965-23004. PMLR, 2023.
Xiaosong Ma, Jie Zhang, Song Guo, and Wenchao Xu. Layer-wised model aggregation for personal- ized federated learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10092-10101, 2022.
Pratyush Maini, Michael C Mozer, Hanie Sedghi, Zachary C Lipton, J Zico Kolter, and Chiyuan Zhang. Can neural network memorization be localized? In International Conference on Machine Learning, 2023.
Maximilian Mueller, Tiffany Vlaar, David Rolnick, and Matthias Hein. Normalization layers are all that sharpness-aware minimization needs. arXiv preprint arXiv:2306.042200 2023.
Jachoon Oh, SangMook Kim, and Se- Young Yun. Fedbabu: Toward enhanced representation for federated image classification. In International Conference Learning Representations, 2021.

---

Published as a conference paper at ICLR 2024
Guillermo Ortiz-Jimenez, Alessandro Favero, and Pascal Frossard. Task arithmetic in the tangent space: Improved editing of pre-trained models. arXiv preprint (xXiv:230511282/ 2023.
Xingchao Peng, Qinxun Bai, Xide Xia, Zijun Huang, Kate Saenko, and Bo Wang. Moment matching for multi-source domain adaptation. In Proceedings of the IEEE International Conference Computer Vision, pp. 1406-1415, 2019.
Henning Petzka, Michael Kamp, Linara Adilova, Cristian Sminchisescu, and Mario Boley. Relative flatness and generalization. In Advances in neural information processing systems, volume 34, pp. 18420-18432, 2021.
Krishna Pillutla, Kshitiz Malik, Abdel-Rahman Mohamed, Mike Rabbat, Maziar Sanjabi, and Lin Xiao. Federated learning with partial model personalization. In International Conference on Machine Learning, pp. 77716-17758 PMLR, 2022.
Maithra Raghu, Justin Gilmer, Jason Yosinski, and Jascha Sohl-Dickstein. Sveca: Singular vector canonical correlation analysis for deep learning dynamics and interpretability. In Advances in neural information processing systems, volume 30, 2017.
Sylvestre-Alvis Rebuffi, Francesco Croce, and Sven Gowal. Revisiting adapters with adversarial training. In International Conference on Learning Representations, 2023.
Itay M Safran, Gilad Yehudai, and Ohad Shamir. The effects of mild over weer-araaeetrizzttoon on the optimization landscape of shallow relu neural networks. In Conference Learning Theory, pp. 3889-3934. PMLR, 2021.
Berfin Simsek, François Ged, Arthur Jacot, Francesco Spadaro, Clément Hongler, Wulfram Gerstner, and Johanni Brea. Geometry of the loss landscape overparameterized neural networks: Symme- tries and invariances. In International Conference on Machine Learning, pp. 9722-9732. PMLR, 2021.
Sidak Pal Singh and Martin Jaggi. Model fusion via optimal transport. Advances in Neural Information Processing Systems, 33:22045- 2020.
Dimitris Tsipras, Shibani Santurkar, Logan Engstrom, Alexander Tumer, and Aleksander Madry. Ro- bustness may be odds with accuracy. In International Conference on Learning Representations, 2018.
Tiffany J Vlaar and Jonathan Frankle. What can linear interpolation of neural network loss landscapes tell us? In International Conference on Machine Learning, pp. 22325-22344 PMLR, 2022.
Hongyi Wang, Mikhail Yurochkin, Yuekai Sun, Dimitris Papailiopoulos, and Yasaman Khazaeni. Fed- erated learning with matched averaging. In International Conference on Learning Representations, 2020.
Mitchell Wortsman, Maxwell C Horton, Carlos Guestrin, Ali Farhadi, and Mohammad Rastegari. Learning neural network subspaces. In International Conference on Machine Learning, pp. 11217-11227. PMLR, 2021.
Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rebecca Roclofs, Raphael Gontijo-Lpess, Ari Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In International Conference on Machine Learning, pp. 23965 23998. PMLR, 2022.
Chen Xing, Devansh Arpit, Christos Tsirigotis, and Yoshua Bengio. A walk with sgd. arXiv preprint arXiv:1802.08877, 2018
Huan Xu and Shie Mannor. Robustness and generalization. Machine learning, 86:391-423 2012.
David Yunis, Kumar Kshitij Patel, Pedro Henrique Pamplona Savarese, Gal Vardi, Jonathan Frankle, Matthew Walter, Karen Livescu, and Michael Maire. On convexity and linear mode connectivity in neural networks. In OPT 2022: Optimization for Machine Learning (NeurIPS 2022 Workshop), 2022.
12

---

Published as a conference paper at ICLR 2024
Chiyuan Zhang. Samy Bengio, and Yoram Singer. Are all layers created equal? Journal of Machine Learning Research, 23(67):1-28, 2022.
Zhanpeng Zhou, Yongyi Yang, Xiaojiang Yang, Junchi Yan, and Wei Hu. Going beyond linear mode connectivity: The layerwise linear feature connectivity, In Advances in neural information processing systems, 2023.
Tongtian Zhu, Fengxiang He, Kaixuan Chen, Mingli Song, and Dacheng Tao. Decentralized sgd and average-direction sam are asymptoticllly equivalent. In International Conference on Machine Learning, 2023.
13

---

Published as a conference paper at ICLR 2024
A APPENDIX
A.1 EMPIRICAL LAYER-WISH MODE CONNECTIVITY
A.1.1 CIFAR-IO, RESNET18 WITHOUT NORMALIZATION
We train ResNet18 without normalization layers using warm-up learning rate schedule: starting from 0.0001 with linear mode for 100 epochs reaching 0.05. Afterwards cosine annealing is used schedule for lcarning rate decay. Batchsize is 64, training is happening for 200 epochs with SGD optimizer, momentum 0.9 and weight decay 5F 4. We use this training setup for all experiments with ResNet18. Heatmaps display the barrier size between the models when only some layers are averaged. Barrier is computed on every 20th epoch (along the X-axis).
(a) Full data parallel training, diff. initialization
(b) Li.d. separation for federated training
Rart
(d) Pathological non-i.i.d. separation for federated training; barriers are computed on errors
(c) Non-i.i.d. separation for federated training
Figure 6: Layer-wise barriers. First row shows the full linear barrier.
Figure 7: In the left plot we select layers randomly, in the right we first average all the shallowest and all the deepest layers.
14

---

Published as a conference paper at ICLR 2024
Figure 8: In the left plot we start from deep layers (so on the most shallow layer level full models are averaged), in the right from the shallow. Here the initialization is different for the two models.
Figure 9: In the left plot we start from deep layers. in the second from the shallow. The initialization same for the two models, but the data shuffling seed different and the learning rate is low (0.001 compared to 0.05).
Figure 10: In the left plot we start from deep layers, in the right from the shallow. Here the
initialization different for the two models and the learning rate low (0.001 compared to 0.05).
Figure 11: Li.d. federated data separation. (a) deep cumulation (b) shallow cumulation
15

---

Published as a conference paper at ICLR 2024
Figure 12: Non-i.i.. federated data separation with mild discrepancy. (a) deep cumulation (b)
shallow cumulation
Figure 13: Non-ii.d federated data separation with pathological discrepancy. Loss value barriers. (a)
deep cumulation (b) shallow cumulation
Figure 14: Non-i.i.d federated data separation with pathological discrepancy. Error value barriers. (a) deep cumulation (b) shallow cumulation
16

---

Published as a conference paper at ICLR 2024
A.1.2 SLIDING WINDOW GROUP AVERAGING
CIFAR-10, ResNet18 without normalization, batch size 64, learning rate 0.05, same initialization and different shuffling of the data. Averaging groups of layers with sliding window of particular size. This experiment confirms the particular structure in the group averaging demonstrated in previous experiments, because just averaging larger number of layers does not demonstrate any structure.
Interpolated into modell
Intorpol ated into model2
fdi
3.0
1.9
2.0
21
2.2
2.2
2.3
23
full
0.0
L9
2.3
2
211
22
2
2.3
1-5
3.000.0
0.2
0.3
0.4
0.4
0.4
0.2
0.1
1.5
0.0
0.1
0.2
0.5
3.7
0.5
0.8
0.5
4
0.3
2.0
37
0.0
D.7
6.9
0.9
1.0
9.5
0.2
0.3
3.5
9.4
0.6
0.4
4.
-0.0
0.8
1.0
1.2
1.4
1.5
1.6
1.
1.2
8
0.0
0.5
0.7
3.7
0.6
6.8
6.7
5.9
5.9
0.5
0.7
1.5
3.0
0
3.3
0.4
7.11
0.0
0.2
0.3
04
0.3
0.3
3
8:12
0.0
0
0.3
0
3.3
4
8:12-
0.0
0.3
0.3
0
0.5
9.3
9-13
9-13
8
3
0.8
1.2
3.0
11-15
F
R
1.1
5
11-15
0.0
0.8
0.9
122
1.2
3.0
12:16
0.0
0.G
0.6
12-18
0.0
0.4
0.4
0.7
3.7
0.5
4
13:17
0.0
8
0.7
3
0.5
15:19
3000
0.8
2.0
1.0
6.
6
15:19
0.0
0.8
0.9
RC
1.0
0.8
16 36-20-1
0.7
16 20
0.7
666
0.0
0.1
0.1
0.4
0.4
0.2
0.1
022
0.1
0.3
100
120
160
100
120
180
Figure 15: Grouping by 4 layers
Interpolated into modell
Interpolated into model2
full
0.00
2.1
2.0
2.3
2.3
full
-0.0
2.0
2.1
2.2
1-6
0.0
0.1
6.2
0.3
0.5
0.5
0.6
0.4
0.3
0.2
0.2
1-8
0.0
0.1
0.2
0.5
0.6
0.5
0.7
0.5
0.4
0.3
0.3
2-7
0.0
3.3
0.4
0.6
0.8
0.5
0.7
0.4
0.3
0.3
2.0
2-7
0.0
0.3
0.4
0.8
0.9
0.5
1.1
9.9
0.0
0.5
3.8
0.9
0.8
6.9
L2
1.4
LS
LG
1.6
12
1.0
0.9
3-9
4.0
0.5
0.8
L.0
1.1
1.0
1.3
3.4
1.0
1.0
49
0.0
0.8
1.0
L2
1.5
L.6
1.4
111
0.8
0.7
49
0.0
0.6
0.7
0.9
0.8
0.7
0.8
0.6
0.5
0.4
5-10
0.0
0.8
6.9
1.0
1.2
1.
1.0
0.5
0.7
0.7
5-10
0.0
0.6
0.8
0.9
0.9
0.7
0.8
0.6
0.5
6-11
0.00
0.8
6.7
6.7
0.9
D.5
0.5
0.6
1.5
6-11
4.0
0.6
3.8
6
7:12
6.6
6.6
7.12
8-13
6.7
B-13
9-14
6.9
9-14
16-15
18-15
11-16
11-16
6.8
0.9
0.7
6
0.5
14.19
1.0
6
14.19
0.8
9.5
5
6
6
15
40.0
Q.
1.0
6.9
1-1
36
15:20
8.8
R-
1.0
R
3.0
6
16
0.0
3.8
6.9
8.8
0.8
6.7
L.6
1.2
11
3.8
0.7
16:21
0.0
0.7
0.7
0.8
0.8
D
1.
3.1
1.
0.7
17-22
0.0
0.7
6.7
6.6
0.6
0.6
0.8
1.3
1.0
5.7
0.6
17:22
0.0
0.6
0.6
0.7
0.7
8.5
3.8
1.0
6.9
0.6
0.5
20
40
60
80
1:0
120
160
160
180
200
20
60
80
100
120
140
160
Figure 16: Grouping by 5 layers
17

---

Published as a conference paper at ICLR 2024
Interpolated into modell
interpalated madel2
0.0
1.
211
2.0
2.1
2:2
2.3
2.3
233
2.3
fu1
0.
1.9
2.
2.3
2
2.2
2
2.3
2.
40.0
1.3
1.7
1.8
1.4
1.3
1.7
1.9
166
1.5
2.0
2.0
299
1.0
1.2
1.5
1.5
1-3
0.8
29
1.0
3.5
1.7
3-5
1.3
1.1
40-0
9.9
1.2
1.5
2.6
LG
1.5
1.2
1.0
0.7
0.9
1.2
1.1
1.3
1.1
1.0
1.0
1.5
0.9
5112
0.0
1.
1.5
LG
1/8
1
1.8
1.7
1.8
5:12
-0.0
1.3
1.5
144
3.6
1.4
1.7
1.6
1-4
1.5
1.4
JAiuUO (DA250))
Pyime 10027000
7-14
0.0
1.
2.5
1.5
1.6
2.7
1.6
3.5
1.4
2.5
7414
0.6
1.4
1.5
1.6
2.7
1-6
1.7
2.6
2.4
1.3
6.9
3.0
1.0
1.0
3.3
1.2
916
1.0
1.1
9-16
1.3
10477
0.0
2.5
1.6
2.6
LG
1.5
2.6
1.3
10417
0.6
2.4
1.5
16
2.7
1.7
2.5
2.4
1.3
11:18
11-18
12:19
40.0
1-5
3.6
1.5
2-7
1.2
1.5
3/4
1.1
1.0
12419
-0.9
1.4
1.5
3.7
3/6
1.6
2.5
1.3
1.0
0.9
1.1
1.0
14-1
4.0
0.8
1.0
1.0
1.0
3.0
1.2
1.3
3.3
1.0
1.0
1421
0.6
6.8
6.9
3.6
1.0
1.2
1.2
1.2
1.0
0.9
40.0
1.4
1.4
166
1.7
1.5
15:2
0.0
1.3
1.4
1.6
1.4
40
80
100
120
140
160
180
200
20
40
60
80
100
120
140
160
180
200
Figure 17: Grouping by 7 layers
A.1.3 CIFAR-10, VGGI1
For VGG11 the training setup is the following: batch size 128, learning rate 0.05, with step wise learning rate scheduler multiplying learning rate by 0.5 every 30 steps. The training is performed for 200 epochs with SGD with momentum 0.9 and weight decay 5E 4.
Figure 18: L..d. federated data separation. (a) deep cumulation (b) shallow cumulation
18

---

Published as a conference paper at ICLR 2024
A.1.4 WIKITEXT, LARGE LANGUAGE MODELS
Training setup for GPT-like model taken from https: /ttp::///ithbb con/epfml/ 11m base 1 innes for small network with 12 layers and 256 sequence length. Training is done on Wikitext dataset.
Rrargeardd redea
Happitadio
Happlinnd
(a) Weight sharing between first layer and last layer
(b) No weight sharing applied
Figure 19: Wikitext, small GPT, full parallel data training from different initializations; Layer-wise barriers.
19

---

Published as a conference paper at ICLR 2024
We experiment with three sizes of Pythia models: 70m, 160m, and 410m.
Interpaiated into modell
Interpolated into model2
full
embed
0.0
02
02
0.3
0.3
0.3
0.3
0.3
embed
0.0
0.2
0.3
0.3
0.3
0.3
C.4
6.4
byemom
Oirput layomarm
6-pest attention layemor
000
000
0.0
4.0
0.1
0.0
0.1
0.0
Oppost attention layemorm
0.0
0.0
0.0
0.0
0.0
-0.0
-00
-00
attention rotary_cmo
attentior ratary,
D D-sttention query_eey_ valre
0.0
0.4
0.5
0.5
0.4
0.6
0.5
C-attention- seeerreeeeeeeeeeeeeee vee__
0.0
0.2
0.4
0.5
0.4
0.4
attention dense
derse
mmp-ddrree_h.Be h_to_
000
0.4
0.4
0.6
0.5
1.1 1.1
O mi d_tte_ 4h
0.0
0.2
0.3
0.4
0.4
6.5
6.5
mip dorse 40
hip dense
1-mp of_ jyyemommm
0.0
100
0.0
2.0
0.1
0.0
D.1 0.1
1-irput_ayere Jayemarm
4.0
0.0
0.0
4.0
0.0
0.0
-00
0.0
post atterrttn baymomm
post pest.attontion jayommmronnammmmmmmmmmmmmm
1-attenton-rotee rotary_cmb
40
300
0.1
0.0
0.0
0.0
1-attentior ratary_
0.0
0.0
D.0
4.0
0.
.00
J-ttention-queyyeee Jey, valu
L-atbemton-quory. boy
1 Lattentoo--onne dense
00
0. 0.1
1 Latterbon-deree
4.0
0.1
H-mp-derse_ht h_
1-nip-dense b_to
1-pp-derse__ee 20
0
D
1-mip dense_4tntt_h
0.0
0
2-ipuf_ Jayennom
2-rput_ lnypuu_lyymmrm
postattertonnnyyye Jyymomm
00
0.1 0.0
0.0 0.0
2-post pptentattettion Jyymorr
0.0
0.0
2-attenton-mtary.r
2-attentior-etaryy
2-attortinn Settqurenr key value
0.3 0.3
2 D-ttemiion-quary buy value
0.0
2-mip-darse lo 4h
00
0 0.3
2-nip- ho 4h
4.0
2-mp-derse,
2-nlp-dense 4h
60
Sinput, layemom
0.
0.1
34rput layemorm
0.0
3-post atterttttttnion Jayennmm
3.0
Spost aattaaattonnaaee antemton_llyermarm
4.0
3-attenion-roiaryy emb
00
0.0
0.0
--atannior-rotaa emb
0
attortion query
attention key
8.
2 2-attemtian-kmmse
00
D
0.1
0.2
buttterioonderse
O
derse
00
hlp
3-mip-derse, 4h
3-ndp-dence to
laemmom
00
4irput
4-post attertion, layamom
4-post atteniion layemorm
attention
4-attontion rotaty
ottertion query key vale
D.S
4 attertion query volue
aattertion derse
mip derse h_to
4-mip detoe_h_to_____
4h
0.0
D.7
Siipp ut_ bymmomm
5 Sirput byymmmm
attortion 5 atttenttion rotary_cmb
0.0
0.0
0.0
attention
D.0
5 s-attention rotary_amb
S-attertion-query queriinn--uueeeeee _er___ke__________
0.0
5 s-attemoon-quer weyue
5 atteettton dence
5 attertion derse
s-mpp-erre__h. h_
0.0
D.4
5-mip-dense__. hu
4.0
mip derse totte_
mip ttto
wyer
0.0
0.0
fina Jayer cmbed_out
0.0
embed_out out
200000
80000
1000100
1200100
L00000
200000
40000
Lommo
1200000
140000
Figure 20: Layer-wise barriers; Wikitext, Pythia models: 70m.
20

---

Published as a conference paper at ICLR 2024
interpolated Into modell
Interpolated Into model2
emeec_in
04nppt embetjin
0.0
470
attondon emb
0.0
CAA
attention C-attertion-eelay D-utteen iouuuerry kay, value
100
0.0
bestoentionnrrary.
50
0.0
0.0
0.0
000
020
150
G-attemcionggury velee valu
11
11
10.00
0 atttention dence
11
11
o attortion denic
10.00
emp-donne, 20
C-mip dende. to_ 4h
O-mp-dense 4h to
Gerip-denee 4h to
TUN
larput ayeroor
Linput layemorm
1 post, attondoon ayamoom
post attertion layemorm
bodtentton-emmryy. emb
L-ttbeniionquurry Jey valu
Lettlent Luttertt ion-wwryy tay. valoe
atteriion
1-mp-dente_h, h,30
24npu:
D
2-pest attentttnnnnot wyeraam
D.0
2-past attention Jayemom
000
UU
OAD
battenliin-nrary.. amb
2-attertion-ettn amb
0.0
UU
OLO
2-attenion-q.ery value
2 2 kuy value
2 Z-attentonn dence
2 stterttoon dense
DI
2-mp-dense_ 20 4h
2-mip dense. to 4h
2-mp-donse 4h
2-mip-dense 4h to
3-rpptt ayemorm
inppu layemorm
3-post attentinnn myernam
3-post attention Jayemom
Soatllntion 3-atteniion-quryy key value
Seateeniion-enaa anb
ii emb
Buuttendion-cuury bey value
atteriion
mmmp-deree 2/mpp
4-pest attention lyyenamm
D.O
D.O
4-past attention Jayemor
OAO
4ngttention-nnaa emb
4-attertion-elle amb 4 atttention euery sey, value
4 attention qu vauu
4 ditteettion dense 4-pp-dense_) 10 4h
4 s-tteetoon--re der
4-mip-derse to 4h_4h
4-pp-donse 4h to
Ammp--anne 4h to
irput. aycmorm
input layemorm
S-post attenton lyyernrm
S-post attention: Jayemor
Siattentioe-enaa anb
Stättertion serertiiin-tarry amb
atttention query vauu
(relathe)) Bartier
d-attention euery jey_ value
S-attertion
(retttec)) barror
Simp--enne
ayemom
input Isyemorm
DO UA
S-pest attuntion uyernamm
6-past attention Jayemor
Aidtnntion-ooaary emb
e-attertien-etan emb
D
6 -tttentton q, vauu
6 atttentionn every value value
6 oeneee
t- s-ateetion-derr
cmmp-donse 4h
é-ipp dense to 4h
6.m cense 4h to.
mipp danse 4h to
7-rpst, lyyernrm
74nput leyemorm
T-post, attention lyyenorm
7-past attention layemar
7-attemtonn 7-attention vauue
Saattention--naay emb
tittertionnrttary
Fattention-cuery teyeeety_yyyeeeeeeeeeeeeeeeeee veluv
7-m
7.mip dense
mi
E-rpst lyernnrm
8anput Jayemorm
8-pest, attuntion lyarnnmm
8-past attention Jayemom
B-attention-onie cootaoyrooorr...ooo. emb
8-attertion cnttartiio-retaryy emb
8- B-attention-p try Jey value
vatuu B B-attention query say value
pattertion-dene
B-ahertion-der
Bimpp donse. h 4h
mip dense to 4h
mip 4cnte_4h_to__
mip dense 4h_to to
purpst, lyyernom
94npui layemorm
P-pest, attuntion layanorm
9-post attention layemor
S-attentoon- rom
9-attertion romary_mmb
9-attention- perention-..ry.eeccc
B-attentior-queyy valuv
donse
D
4h.to_h
10-rpst lyernnmm
10-nput layemar
U C
10spest, attundon ayaroom
10spost attention Tayemom
U
10 attenbor rotary_emm cmmb
I0-attertion--co rotary emb
U
10- I0-wthentio-- rry Jey, value
10-attent any value
10 attertion dense
10 attoriion derse
10 mb conse h 4h
10 mip denst. to
10 mp ccnte_4h_to___
10 mlp 2o
1i-rptt lyernnem
11nppu layemorm
11-pest attunion ayarnom
11.post attention layemorm
11-attentoon rom
11-attertion-coe rmmb 11-attent
bey_um
attertinn
attordon der
11
11 donse
11 dette____ttt_____
40
11 dense 4h fiaal yor nom embed out
final layer norm
embed out
20000
400000
60030
20000
100000
120000
1400000
40000
100000
12000g
1400000
Figure 21: Layer-wise barriers; Wikitext, Pythia models: 160m.
21

---

Published as a conference paper at ICLR 2024
Figure 22: Layer-wise barrier;; Wikitext, Pythia models: 410m.
22

---

Published as a conference paper at ICLR 2024
A.2 ROBUSTNESS PERSPECTIVE ON LLMC
A.2.1 CIFAR-10, VISION TRANSFORMERS
Figure 23: Layerwise interpolations (left) and robustness random perturbations of the same norm (right) for vision transformers trained on CIFAR-10 10 with different learning rates.
Figure 24: Layerwise interpolations (left) and robustness random perturbations of the same norm
(right) for vision transformers trained on CIFAR- 10 with different perturbation radii of SAM.
23

---

Published as a conference paper at ICLR 2024
A.2.2 CIFAR-10. 10, RESNETI8 8 WITHOUT NORMALIZATION
Robustness of the layers the perturbations in the averaging direction and random directions of the same norm. Here we show the development while training (along X-axis) for each of the layers.
(a) a 0.5
(b) = 1.0
Figure 25: Full dataset training, CIFAR-10 with ResNet18 without normalization.
(a) 0.5
(b) = 1.0
Figure 26: Federated i.i.d. training without aggregation, CIFAR-10 with ResNet18 without normal- ization.
(a) a = 0.5
(b) = 1.0
Figure 27: Federated non-i.i.d. training without aggregation, CIFAR-10 with ResNet18 without normalization.
24

---

Published as a conference paper at ICLR 2024
Fast
(a) 0.5
(b) 1.0
Figure 28: Federated pathological non-i.idd training without aggregation, CIFAR-10 10 with ResNet18 without normalization. The robustness is calculated with respect loss.
(a) a 0.5
(b) a = 1.0
Figure 29: Federated pathological non-ii.d. training without aggregation, CIFAR-10 with ResNet18 without normalization. The robustness is calculated with respect to error.
25

---

Published as a conference paper at ICLR 2024
A.3 FULL DISTANCE PERTURBATION
CIFAR- 10, ResNet18 without normalization, batch size 64, learning rate 0.05, same initialization and different shuffling of the data. Checking perturbations of the norm equal to the distance between full models, not between separate layers. It can be seen that only one layer direction exhibits high loss that will correspond the full networks barrier, thus is not the size of perturbation that defines the loss growth.
modelL robusenese
modell alpha-b.5
D
6.3
0.
D.
0J
01
36
96
mmv1
2.
1.8
15
15
2.6
0.6
0.7
0.8
0.4
0.3
0.4
86
36
kesrl 0com22
03000
bysrl 1 cony2
2.
8.8
8.
8.
kowri-k-com2
kyari-1-camd
0.6
0.6
0.6
0A
comr
byar2 6 con12
ay5r2 0002000000 0owrraaseee
kear21-comi
6.7 0.7
94 0.7
2.
0.4
944
0.4
2.6
0.5 0.4
0.3 03
0.8
0.2
15020.0022
bypr330-cmm22
26
0.5 6.4
0.3 0.3
0.2 0.2
02
byor3 0-dorsaanslee
Myer30 dww.samckco
2.6
0.5 6.4
94
8.4
0.3
6.3
0.4
34
0.5
0.7
koar3/1ccomi
lyver3 conv2
228
0.5 4.4
0.3
84
0.2
0.2
01
3/2
0.
0.0
keyer3-1-con2
2.0 0.4 0.4
0.3 0.3
0.2
0.2
0.1
36
0.0
byar440 conv.
2.6
0.5 0.4
0.3
0.3
0.2
0.2
0.
35
0.5
2.6
0.4
0.3
0.3
0.1
bstr400com2
kyer 6
bypr4
2.8
kyperf-baamm
kyar&
2.6
byort
40
00
80
10%
220
140
160
180
Figure 30: Random direction perturbations when averaging and when interpolating.
26

---

Published as a conference paper at ICLR 2024
A.3.1 CIFAR- AND CIFAR- 10, MOBILENET
MobileNet implementation and training hyperparameters were taken from https:/ /et:://gtttti com/ jioon- Cb//ed In particular we use batchsize 128, learning rate 0.1 and decay it by 0.1 on half training and 0.75 of training. Training is done for 320 epochs.
clientO interpolation into clientl
cliento interpolation into client1
conv1
full
2200068866777707664440033333
55211570212332576775888333
fill
22018181111111110100000
bnl
230011111141311114111111000
conv1
225311222.22222222222222222
kyers-0-conw2
kyers-O-cowvl
byers-O.bl1
231022820033222222233333222
#2200117211172220094427337
2300100121e111111111111111111111
kyers-O-conv1
kyers 0-con22
byyers-Obbl1
21523331333333442222222222
17775226666572222222233
7 8
kyyrr---bb22
22001000113100000000011111111111
layers-O-b22
7:0...202.022221141.111
byers-1-conw1
2:222252/72322/24457733333
byers-1-conv1
7 3.5 3.74::
byers-1-bn1
22207078770800080555500000000000
layers-i-bn1
22.00.09220222211111
kyerss1-conw2
2:377
kyers-1-com2
5:000
byers-1-bn2
220060770900001100999000000
layers-1-bn2
21126111111151
ayers-2-comw1
22118:91.222222228:28888
byers-2-convl
kyers-2-bm1
23005770.00890036424000
layers-2bbn1
2.01
1111111111110000
byers/2-com22
2 6
kyers-2-com2
byers-2-bn2
2.30.70.51.11
byyers-2-b22
2
0tf.22.32.52242
baeers3-conv1
23181552355
byers-3-conv1
5.1.935566444
kyers-5-bb1
layers:3-bn1
2.01.220.3222222222220
3
2.342
byers-3-convp
75:7 5:8 5
byers-3-bb2
030.50.80$10111
byers-3-bn2
2.2.0....2222222222
byers-4-comv1
231.1162002.7
kyers-4-cov1
2
byers-4-bn1
loyers-4-bn1
kyers-4-comw2
4yers-4-conv2
byers-4-b22
230.81
31.31.
61111118111111 81
byyer-4-bb2
2.1212222.22222222223
kyers-S-comw1
2:00181:41722222224222333
kyers-S-conv1
2.2252882:5755533000
leyers-5-bn1
2:0.1:0::12:31115:1777
layers-S-bn1
7:0202225222222222222
byers:S-comw2
byers-S-conv2
byers bn2
2.0708009121117111111.133
bayers-5-b22
012,02222222222222222273
kyers-6-conw1
2307111.41822222283000
byers
0.2.
550
byers-
20050808122312213101777
byers-6-bn1
22020222332252222022
ayeor-66comw2
kyers-6-conv2
kyers-4-bn2
2307152:3222:222:255222
layers-6-b22
2.122222229331123
byers/7-comw1
23071013255322222522
byerr 7:cov1
212
017779353983550 55
byers-7-bn1
2.00111:2::2122225200
byers-7-bn1
2.02
layers---cow2
2212082366.0000
kyers-7-conv2
byers-7-bn2
230.607051116171020111
layers-7-bn2
2.02.0
2.0.02.2242.22.2.52277.
kyers-&-comw1
20050708111777112220222000
2.02 5 9:36
byers-&-bn1
2206677071611111111220011.11
layers-S-bn1
2.02.0
2202021232232222222222
kyeor-8-coow2
2:221212285322224122222222
byers
5.
bn2
bayors-8bb22
2.01.99.0
2.12222441992 2.12..2.22222
kyyers---onwv
22066778807077770666677666680800
byers-9-conv1
byers-9-bl1
layers-9-bn1
kyers-9-comw2
221:2232122723317773333332
byers-9-conv2
775
lyyers-9-bn2
2 30000606607001888006600000000000
byers-8-bn2
2.01 01 2. 21..11111111111111
1yeos-10-comw1
kyers-10-conv1
7.0..1922242727.5111
kyers-10-bn1
2 20000800070080053300000000
byerss10-bb1
2.....82.02222111
byers-10-conw2
byers-10-conw2
775
5:15
kyers-10-bn2
2 30000886607007777444445330000300
layers-10-bn2
2.0199199
layers c1-com
300077091100111000000000000000
byers-11-conv1
7:0:011322222222233333
layers11-bnn
2 30580806600077777044400030000000
byers-11-bn1
2:010100:2:2311110000
laeer-11-comvz
21191922001121222222111112222222
byers-11-conw2
77488.755
layer-11-bb22
200000007077077777555555000002
lyyen=11-bm2
7202.022.52222222
leyers-12-coww1
30067700088088000777770000000000
byerr
1y-cov12-conv
2:021242:303122222
layers-12-bn1
byers-12-bnn1
262202223222222
myers-12-conw2
byars-12-comu
layers-12-bn2
300066170770707744550000000
byers-12-bn2
0001:111111191111
Iinear
31192022221222222222222222222222
Inear
25 50
(a) CIFAR-10
(b) CIFAR-100
Figure 31: Robustness of the layers the perturbations in the averaging direction. Same architecture
(MobileNet) shows different sensitive layers when the task is changing from CIFAR-10 -O to CIFAR- 100.
27

---

Published as a conference paper at ICLR 2024
B PERSONALIZED FEDERATED LEARNING
We select the setup with CIFAR-I0 and ResNet18 without normalization layers. A popular approach construct personalized local datasets to create label based non-i.id. distributions, either by allocating only a subset of labels to each local learner or by using a Dirichlet distribution. We construct two data separations using Dirichlet distributions with parameter and 0.01, where the second one is significantly more pathological. Note that the average performance of the local models
Table I: Average test accuracy among local models for layer-wise personalization methods on labels shift task.
Averaging
CIFAR-10
CIFAR-100
mode
non-iid
path.
path.
Full
61.82
24.24
18.43
No (local training)
50.29
91.75
53.94
Body
61.73
91.72
56.46
Classifier
50.31
93.9
55.12
Critical
51.11
94.12
55.22
Not critical
50.69
94.52
49.01
Middle
50.59
94.16
50.42
Not middle
50.38
92.37
55.87
has variance of around 8% makes the results in Tab. 1 nearly identical.
We confirm these findings in the setup considered by Oh al. (2021), i.e., MobileNet trained on CIFAR-100 with 100 clients and only several classes available to each client.
Additionally we perform experiments on DomainNet dataset (Peng et al., 2019) using ResNet18 without normalization layers. This dataset can be seen as example of feature shift task, if different domains are assigned to different local learners. Classification task is then rather complex, with 345 classes, the resulting accuracy is subpar. Nevertheless, the overall trend of indistinguishable partial aggregation stays true (Tab. 2).
Table 2: Average test accuracy and test loss among local models for layer-wise personalization methods on feature shift task.
Averaging mode
Accuracy
Loss
Full
27.36
11.25
No (local training)
23.04
9.095
Body
22.12
10.54
Classifier
15.73
9.53
Critical
14.45
12.15
Not critical
17.11
8.06
Middle
17.43
10.94
Not middle
11.97
9.03
28


Reviews:
[{"review_id": "ARzEd8thfX", "reviewer": "Reviewer_VziZ", "summary": "This paper studies the idea of combining neural networks by averaging only certain layers of the model through the lens of the loss landscape.  If averaging the layers does not result in a model with error significantly above the average error of the two initial models, the the two networks are said to be layer-wise linearly mode connected (LLMC). The paper studies for several different networks/datasets the evolution of the error barrier of the full model as well as the error barrier in each layer.  (Note that layerwise error barrier is not \"symmetric\" as you can use either model 1 or model 2 as the base model; the paper presents plots for both.)  The general conclusion is there is often no layerwise error barrier even when there is a barrier for the full model.\n\nThe paper then makes several further contributions.  LLMC is studied for groups of layers instead of just individual layers, and a toy example of a deep linear network is given to demonstrate how there can be no layerwise error barrier when there is a barrier in the full model.  Some attempts at further understanding this phenomenon are presented by looking at the effect of random perturbations in the loss landscape and perturbations in directions related to the training subspace.  Finally, some results are presented in relation to averaging models in federated learning.", "questions": "Repeating the questions from the review for emphasis:\n\n* The LLMC is the barrier with respect to *the average error of the two full models* not the average error of one full model and that model with a specified layer swapped, correct?\n\n* In Table 1, the results are for different groups of layers the authors think would be interesting, not groups selected by layerwise barriers?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The results on LLMC for single layers and groups of layers across models and datasets is thorough and provides generally interesting results.  I did want to confirm that LLMC is the barrier with respect to *the average error of the two full models* not the average error of one full model and that model with a specified layer swapped.  The latter would not be very informative if this significantly increased the error, but Definition 2 seems to imply the authors used the former.  I would emphasize this before the definition in the paper because readers coming in with intuition from LMC will assume it's the average of the two end points of the interpolation; a figure might also be clarifying, similar to Figure 3 but illustrating how the barrier is computed.\n\nSection 4 and 5.2 are also strong sections in my opinion.  They do a good job of emphasizing the idea that certain layers may provide better directions in the loss surface as compared to averaging the full model (e.g. Fig 4).", "weaknesses": "**Section 5.1:** I found the conclusions of this section hard to parse.  I assume in Fig. 4 the intent is to compare the rows of the left 2 plots to the right 2 plots.  One would then see that the full model has less curvature along the averaging direction as compared to a random perturbation, but a few of the layers have the opposite trend.  The paper states:\n\n```\nMoreover, the networks are much more robust to random perturbations compared to the direction of interpolation between models. This suggests that averaging directions are special in the sense of having much higher curvature than random ones.\n```\nAre the authors referring to just the layerwise results here?  If so this should be made clear because my reading of the plot is this does not hold for the full model.\n\nSecond, is the random perturbation for each layer the size of the norm of the layer-wise interpolation for that particular layer?  This raised the question for me if the results on LLMC were actually better averaging directions in the loss landscape or simply much smaller perturbations to the model. (Also the build up of the error barrier over groups of layers could be increasing the size of the perturbation).  The results in Fig. 4 seem to point to the latter, and *I think this would be an important baseline to include for the other results, i.e. what would the barrier be for a random perturbation of the same size as $\\alpha =0.5$.*  If you get the same result from randomly perturbing the layer there seems to be minimal benefit to doing the averaging.\n\n**Section 6:** In Table 1, these are just the results for different groups of layers, not layers chosen based on the LLMC results.  Is the takeaway just evidence that averaging a subset of layers can be more effective?\n\nIn an application, how would you suggest choosing the layers to average? It seems it would be inefficient in some applications to compute the each layer's barrier every few epochs throughout training.", "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": "", "extra_scores": {}}, {"review_id": "wH1SqhetR5", "reviewer": "Reviewer_zQV8", "summary": "The paper \"Layer-wise Linear Mode Connectivity\" presents an in-depth exploration of barriers on the loss surface observed during model averaging, introducing the concept of layer-wise linear mode connectivity (LLMC). The authors investigate the behavior of individual layers, concluding that the averaging barrier at this level is consistently minor compared to the full model. A significant finding is the propensity of middle layers to create cumulative averaging barriers, suggesting potential connections to existing studies on the neural network training process. The research also examines personalized federated averaging, highlighting the inherent challenges in distinguishing between layers that carry local knowledge versus those that carry common knowledge.", "questions": "- The reviewer is unclear as to why Fig. 2 indicates that \"neither the shallowest nor deepest layers cause the barrier, but the middle ones do\". In Fig. 2 (a), both middle and deep layer averaging seem to result in high losses, and Fig. 2 (a) implies high losses upon averaging both middle and shallow layers.\n- The layer-wise convexity appears straightforward since the single-layer interpolation $\\boldsymbol{X} \\boldsymbol{W}^{(1)} \\ldots\\left(\\alpha \\boldsymbol{W}^{(k)}+(1-\\alpha) \\boldsymbol{W}^{\\prime(k)}\\right) \\ldots \\boldsymbol{W}^{(L)}$ is a linear function of $\\alpha$. The reviewer questions whether interpolating two layers would render $L(\\alpha)$ non-convex.\n- Does \"non-iid\" in Table 1 represent less severe non-i.i.d., and does 'path.' signify pathological non-i.i.d? Why does partial averaging perform better in the pathological non-i.i.d setting than in the less severe non-i.i.d. setting? Shouldn't higher degrees of non-IID reduce performance?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The narrative construction is good, and the literature review is comprehensive and informative.\n- The reviewer appreciates the bold conjectures made throughout the paper. Although these may not always be rigorous, such daring speculations can be beneficial in stimulating further research.\n- Certain experimental outcomes are intriguing, for instance, the most sensitive layers of ViTs are the early attention and fully-connected weights; averaging directions exhibit a peculiar characteristic of having a much higher curvature than random ones of the same norm.", "weaknesses": "### Weaknesses\n- The reviewer acknowledges some insightful observations in this paper. However, from the reviewer's perspective, while the findings are interesting, they may not introduce profound novelty.\n- The presented phenomenon of smaller layer-wise barriers may not be surprising. For instance, if we consider a situation where all layers are created equally, the averaging of only $\\frac{1}{s}$ layers would typically result in approximately $\\frac{1}{s}$ of the loss increase induced by averaging all layers. It is plausible that the observation of Layer-wise Linear Mode Connectivity (LLMC) is due to the averaging of a smaller number of layers rather than a unique layer-wise structure. If the loss increase from averaging $\\frac{1}{s}$ layers is significantly less than $\\frac{1}{s}$ of the loss increase induced by averaging all layers, then this observation would lend more credence to the LLMC conjecture.\n- The theory is for single layer interpolation, whereas the experiments have been conducted on a subset of layers.\n\n### Minor issues\n- There seems to be an error in the author name of the paper \"Which layer is learning faster? A systematic exploration of layer-wise convergence rate for deep neural networks\". The surname of the first author could possibly be Zhou instead of Yixiong Chen.\n- The citations for Singh & Jaggi (2020); Ainsworth et al. (2022); Jordan et al. (2022) on page 3 should use \\citep instead of \\citet.\n- There's a typographical error on page 9: 'to be less communication' doesn't make sense in the context.", "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": "", "extra_scores": {}}, {"review_id": "Tk7XO6pQii", "reviewer": "Reviewer_GMLC", "summary": "This paper studies layer-wise linear mode connectivity (LLMC) and finds that in most cases, there are no barriers in LLMC. The authors also study the learning dynamics behind LLMC  from the perspectives of robustness and loss landscape. Also, implications are given regarding the partial personalization in federated learning.", "questions": "See the above weaknesses for details. I suggest the authors provide additional results according to the weaknesses and I am happy to raise my scores once my concerns are relieved.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* The motivation for studying layer-wise linear mode connectivity (LLMC) is novel and interesting. It contributes new ideas in the community of linear mode connectivity.\n* The authors study LLMC from different perspectives which are thorough.\n* The experiments are solid. The ViTs and LLMs are also studied to show the prevalence of the findings across a large range of model architectures.\n* The findings and takeaway insights are intriguing.", "weaknesses": "Despite of the strengths listed above, I think this paper should be improved in the following aspects.\n* The analysis of cumulative LLMC (LLMC about the group of layers) needs further investigation. The authors should study the LLMC of $l$ consecutive layers on different parts of the models. The authors can conduct an experiment with a moving window of $l$ layers to show the group-layer-wise connectivity. For instance, given a 20-layer network with $l=5$; the experiments should be conducted: LLMC of 1-5 layers, LLMC of 2-6 layers, LLMC of 3-7 layers, ..., LLMC of 16-20 layers. Current experiments cannot jump to the conclusion that the middle part of layers will cause barriers because you didn't control the variable of the number of aggregated layers $l$. \n* For the convexity of LLMC, I think the current Theorem is not enough and is far from the practice. Theorem 4.1 didn't consider the non-linearity of the neural networks, i.e., the activation functions. I think a more solid theorem should be derived to verify the convexity of LLMC.\n* The discussions on personalized federated learning should be given more focus and I think this is an important point for this paper to have a broader audience and be more applicable to practices. However, I think current implications, experiments, and discussions are not enough. \n    * The methods of personalized layers or layer-wise aggregation in personalized federated learning should be implemented and compared.\n    * Based on the findings of this paper, a new method should be devised, but the paper didn't showcase the applicability by proposing a simple method.\n    * Experiments on feature shift non-iid in federated learning should be conducted to verify the claim that personalized-layer-based techniques can work in feature shift federated learning.", "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": "", "extra_scores": {}}, {"review_id": "ewoEMHypIC", "reviewer": "Reviewer_G1sx", "summary": "The presented paper proposes a new concept, layer-wise linear mode connectivity, and shows that it's possible to explain it from a robustness perspective.", "questions": "- In Section 5.1, how is \"robustness to random perturbations\" exactly calculated? \n- There are two competing definition of LLMC, namely the single layer one and the cumulative one. Which one is used in Section 5.1?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The presented paper is well-written and is easy to follow. \n- The proposed phenominon is interesting and worth furthur study. \n- The paper provides specific practice guidance for federated learning.\n\n------\nAfter reading the authors' responce, I decide to keep my current rating unchanged.", "weaknesses": "I don't see how robustness is possible to explain the proposed LLMC, since the LLMC is concerning the relationship of two models while robustness concerns more about the relationship of a model with a random perturbation. The authors claimed in Section 5.1 that \"the networks are much more robust to random perturbations compared to the direction of interpolation between models\". This is not surprising to me, because when togetherly trained with sufficiently many epochs, the two models should be in the same basin (as the original LMC suggests), and the direction of interpolation between models is towards the basin while a random direction might not. I don't see how the proposed robustness explaination goes furthur than the same-basin explanation in original LMC paper.", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)