Paper: Layer-wise linear mode connectivity (ICLR-2024)

Reviewer comment (Theory):
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.

Author response:
**I don't see how robustness is possible to explain the proposed LLMC**: We fully agree with you that robustness alone does not explain LLMC, but our results show that the full picture is complicated. First of all, in our experiments networks are trained so that they are not in the same basin, evident by the large barrier of full averaging. This differs from the original LMC paper, where a single network is trained up to a certain point, from which on two networks are finetuned - in this case, one would expect them to be in the same basin. If networks are not in the same basin, then the barrier size is influenced to some degree by robustness: the more robust a model, and thus the more flat the loss surface around it, the lower the barrier when perturbing in any direction. However, our conclusion from our robustness experiments was that layer-wise perturbations in the direction of averaging can increase the loss significantly more than a layer-wise perturbation in a random direction (see also our answer to Reviewer VziZ). This result is indeed surprising, and we do not claim to fully understand the interplay between robustness (in particular noise magnitude), perturbation direction, and barrier size. We believe that our paper highlights this issue and provides some valuable insights into this complex and interesting phenomenon.

Reviewer comment (Theory):
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.

Author response:
**I don't see how robustness is possible to explain the proposed LLMC**: We fully agree with you that robustness alone does not explain LLMC, but our results show that the full picture is complicated. First of all, in our experiments networks are trained so that they are not in the same basin, evident by the large barrier of full averaging. This differs from the original LMC paper, where a single network is trained up to a certain point, from which on two networks are finetuned - in this case, one would expect them to be in the same basin. If networks are not in the same basin, then the barrier size is influenced to some degree by robustness: the more robust a model, and thus the more flat the loss surface around it, the lower the barrier when perturbing in any direction. However, our conclusion from our robustness experiments was that layer-wise perturbations in the direction of averaging can increase the loss significantly more than a layer-wise perturbation in a random direction (see also our answer to Reviewer VziZ). This result is indeed surprising, and we do not claim to fully understand the interplay between robustness (in particular noise magnitude), perturbation direction, and barrier size. We believe that our paper highlights this issue and provides some valuable insights into this complex and interesting phenomenon.

Reviewer comment (Novelty):
I don't see how the proposed robustness explanation goes further than the same-basin explanation in the original LMC paper.

Author response:
**I don't see how robustness is possible to explain the proposed LLMC**: We fully agree with you that robustness alone does not explain LLMC, but our results show that the full picture is complicated. First of all, in our experiments networks are trained so that they are not in the same basin, evident by the large barrier of full averaging. This differs from the original LMC paper, where a single network is trained up to a certain point, from which on two networks are finetuned - in this case, one would expect them to be in the same basin. If networks are not in the same basin, then the barrier size is influenced to some degree by robustness: the more robust a model, and thus the more flat the loss surface around it, the lower the barrier when perturbing in any direction. However, our conclusion from our robustness experiments was that layer-wise perturbations in the direction of averaging can increase the loss significantly more than a layer-wise perturbation in a random direction (see also our answer to Reviewer VziZ). This result is indeed surprising, and we do not claim to fully understand the interplay between robustness (in particular noise magnitude), perturbation direction, and barrier size. We believe that our paper highlights this issue and provides some valuable insights into this complex and interesting phenomenon.

Reviewer comment (Evaluation):
In Section 5.1, how is "robustness to random perturbations" exactly calculated?

Author response:
1) Robustness to random perturbations is computed in the following way: (i) we compute the distance between the two networks in the layer of interest (ii) we generate random vectors of the same norm (iii) we move the layer according to this random change (iv) we compute the loss of the obtained network. This is performed 5 times and the results are averaged.

Reviewer comment (Writing):
There are two competing definition of LLMC, namely the single layer one and the cumulative one. Which one is used in Section 5.1?

Author response:
2) In Section 5.1 we were using only single layer aggregation. Cumulative aggregation is investigated in Section 3 in order to improve our understanding of which layers actually lead to the barrier between networks.

Reviewer comment (Experiments):
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 \).

Author response:
1) **The authors should study the LLMC of consecutive layers on different parts of the models.**: We thank the reviewer for this idea of an interesting experiment. We performed several runs of it  and included the results in the appendix of the paper (App. 3). The results indicate that a small sliding window allows the same conclusion as our initial experiments (i.e. the barrier is produced by the middle layers), but the larger the sliding window is, the more uniform the picture becomes. We want to emphasize that we also performed random cumulation (i.e. averaging layers in random order) and minimal cumulation (i.e. averaging group of the layers that seem not to cause barrier) as well (see Figure 7 in the appendix) in order to check our conjecture about middle layers. Our results, including this new experiment, support this conjecture.

Reviewer comment (Theory):
For the convexity of LLMC, 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. A more solid theorem should be derived to verify the convexity of LLMC.

Author response:
2) **Theorem 4.1 didn't consider the non-linearity of the neural networks, i.e., the activation functions.**: The goal of Theorem 4.1 was not to explain the behavior in practical neural networks, but to provide a useful and interesting intuition about the difference between layer-wise and full aggregation. We do not draw any conclusion about the convexity of layer-wise averaging in non-linear neural networks from this theorem. We agree with you that investigating this challenging theoretical question further will deepen our understanding of the phenomenon, and makes for interesting future work.

Reviewer comment (Experiments):
The discussions on personalized federated learning should be given more focus, as this is an important point for the paper to have a broader audience and be more applicable to practices. However, the current implications, experiments, and discussions are not enough.

Author response:
3) **The discussions on personalized federated learning**: We agree with the reviewer that the personalized federated learning experiments and discussion is not in the focus of the paper. We want to emphasize that PFL was one of the examples of an application where LLMC can be of use - see general comment - but our results do not provide a practical method, yet. Nevertheless, (i) we do compare to the existing approaches with partial averaging, such as aggregation of the classifier layer only or of the feature extraction layers (i.e. everything except for classifier), which are the most prominent partial averaging approaches in the existing literature. (ii) We try to conjecture an approach that will be useful in PFL using both our robustness experiments insights (averaging only critical or only non-critical layers) and cumulative aggregation insights (averaging only middle layers or only non-middle ones). (iii) In addition to our empirical results for this approach, we performed an additional experiment on DomainNet (a dataset with feature shift). The experiment and its results are described in the general comment.

Reviewer comment (Experiments):
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.

Author response:
3) **The discussions on personalized federated learning**: We agree with the reviewer that the personalized federated learning experiments and discussion is not in the focus of the paper. We want to emphasize that PFL was one of the examples of an application where LLMC can be of use - see general comment - but our results do not provide a practical method, yet. Nevertheless, (i) we do compare to the existing approaches with partial averaging, such as aggregation of the classifier layer only or of the feature extraction layers (i.e. everything except for classifier), which are the most prominent partial averaging approaches in the existing literature. (ii) We try to conjecture an approach that will be useful in PFL using both our robustness experiments insights (averaging only critical or only non-critical layers) and cumulative aggregation insights (averaging only middle layers or only non-middle ones). (iii) In addition to our empirical results for this approach, we performed an additional experiment on DomainNet (a dataset with feature shift). The experiment and its results are described in the general comment.

Reviewer comment (Novelty):
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.

Author response:
1) **not introduce profound novelty**: We kindly disagree with the position that the findings do not introduce profound novelty. Our novel contributions are as following: (i) describing the surprising phenomenon of LLMC, that was not introduced and researched before, (ii) investigating the subspaces of the loss landscape and showing that averaging direction stands out even in the training subspace, (iii) applying the knowledge about layer-wise behavior to personalized federated learning.

Reviewer comment (Experiments):
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.

Author response:
2) **phenomenon of smaller layer-wise barriers may not be surprising**: The initial investigation for this work was exactly based on the intuition that each layer aggregation should contribute its part to the full barrier and thus the less layers are aggregated the less is the barrier and each layer separately has its share of the overall barrier value. Instead, we were very surprised that throughout our empirical investigation we found evidence that this is not the case. First, the barrier for each layer is nearly non-existent which is significantly smaller than the full barrier (Figure 1). Second, cumulative aggregation shows that different groups of layers have different “thresholds” of the amount needed to create a barrier. Third, in the experiments with LLMs we observed a barrier in the final layer that is even larger than the full barrier (together with absence of barriers in all other layers) and analogously in Figure 4 (2nd plot) `layer-0-0-norm` with alpha=0.5 has a higher loss (1.0) than for full interpolation (0.9). All this together falsifies the conjecture that the reason for LLMC is simply a proportionally smaller barrier than the full one.

Reviewer comment (Theory):
The theory is for single layer interpolation, whereas the experiments have been conducted on a subset of layers.

Author response:
3) **theory is for single layer interpolation**: The theory provided serves as one possible approach to understanding LLMC in a minimalistic setup of linear networks. We do not claim that multiple layers or non-linear layers have the same property of convexity as shown in the theorem, but we investigate these setups in the experiments (see also our reply to your question 2) for a broader understanding of the phenomenon.

Reviewer comment (Writing):
There's a typographical error on page 9: 'to be less communication' doesn't make sense in the context.

Author response:
We thank the reviewer for the correction, we fix it now.

Reviewer comment (Presentation):
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.

Author response:
1) We are sorry for the confusion that our “cumulative averaging” concept caused. We will try to clarify it. The idea was to increase the number of layers that are averaged one by one starting either from the deepest or most shallow layer. That is, for shallow cumulation we average only the first layer in the first row of the plot, first two layers in the second row, and like this till the full networks are averaged in the last row. Analogously, with deep cumulation we start with only the classifier layer in the last row and end with the full network in the first row. From these experiments we  concluded that the barrier starts “cumulating” once the middle layers are added - and after this adding more layers has the same level of barrier.

Reviewer comment (Theory):
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.

Author response:
2) We want to emphasize that our theoretical result about layer-wise convexity is provided mainly for explaining the phenomenon of LLMC from all possible perspectives. Regarding the question about two layers, Figure 3 demonstrates that even in the simplest setup averaging more than one layer is not necessarily convex. Proving any further results is a challenging task, as we also mention in the general comment. Our empirical results indicate that for individual layers the loss surface is (close to) convex for realistic neural networks, too. We are excited to continue our theoretical investigation of this challenging question in future work.

Reviewer comment (Evaluation):
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?

Author response:
3) Yes, non-iid means using Dirichlet with parameter 3 and path. means pathological distribution with Dirichlet parameter 0.01. Our tentative conjecture about the obtained results is that with pathological non-iid setup there is no reason for models to exchange their parameters, so that not communicating is the optimal solution. Since partial averaging does not change models significantly, their performance is nearly as good as not communicating at all. On the other hand, when non-iid is not severe full averaging leads to the best performing model for each of the local datasets and the quality of partial averaging decreases by reducing the number of layers that are being aggregated, since it removes the advantage of knowledge exchange. Nevertheless, this is an initial exploration of a possible application of LLMC notion, detailed investigation of which is left for future work.

Reviewer comment (Evaluation):
In Section 5.1, the conclusions are hard to parse. It appears that in Fig. 4 the intent is to compare the rows of the left 2 plots to the right 2 plots, showing that the full model has less curvature along the averaging direction compared to a random perturbation, but a few of the layers have the opposite trend. The paper states:  
```
Moreover, 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.
```  
It is unclear whether the authors are referring only to the layerwise results here. If so, this should be made clear, because the reading of the plot suggests it does not hold for the full model.

Author response:
**I assume in Fig. 4 the intent is to compare the rows of the left 2 plots to the right 2 plots.**: In Figure 4 we compare averaging directions (left two plots) to random directions (right two plots). You are absolutely correct that the text in the paper is related to the layer-wise directions, not the full model ones. We are sorry for the confusion and will edit the text accordingly.

Reviewer comment (Experiments):
In Section 5.1, it is not clear whether the random perturbation for each layer is the size of the norm of the layer-wise interpolation for that particular layer. This raises the question of whether the results on LLMC were actually evidence of better averaging directions in the loss landscape or simply the result of much smaller perturbations to the model. 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, suggesting that an important baseline to include for the other results would be: what would the barrier be for a random perturbation of the same size as $\alpha = 0.5$? If the same result occurs from randomly perturbing the layer, there may be minimal benefit to doing the averaging.

Author response:
**is the random perturbation for each layer the size of the norm of the layer-wise interpolation for that particular layer?**: For each particular layer, the random perturbations are generated so that their norm is equal to the norm of the perturbations in the averaging direction for this layer. We had exactly the same question in mind when performing the robustness investigation of LLMC, i.e., if the layer-wise perturbations are just so much smaller than full averaging so that the loss does not grow. Our random directions experiment indicates that it is indeed the case and we agree with the reviewer that checking if the random perturbation of each layer with a larger norm will show different results. The current experiment setup will become impossible in this case though, since (i) the norms with respect to the entire network and a single layer are not comparable due to the vastly different dimensionality, and (ii) having random perturbations with much larger norm requires to also increase the norm of the averaging perturbations, which then would shoot far beyond the other network and thereby arguably loose their meaning. Nevertheless, this would be an exciting independent investigation of the robustness directions. What fascinated us in the current setup is that some averaging directions are actually worse than random ones - even for a small norm.

Reviewer comment (Evaluation):
In Section 6, Table 1 presents results for different groups of layers, not layers chosen based on the LLMC results. It is unclear whether the takeaway is simply evidence that averaging a subset of layers can be more effective.

Author response:
Table 1 includes the following setups: (i) standard (full model) averaging and fully local training, (ii) existing partial averaging approaches, reiterated in multiple works, averaging only body without classifier layer or only classifier layer, (iii) approaches that can be motivated by the robustness perspective of LLMC (Figure 22b in the Appendix), where we either average only the critical (fast changing loss) layers or non-critical layers, and (iv) approaches can be motivated  by cumulative averaging (Figure 2) averaging only middle layers or everything but middle layers. Given a set of critical layers, it remains an open question whether one should average these critical layers, or do exactly the opposite and average all non-critical layers. We are excited to further investigate this direction, though, since on the one hand averaging barrier prone parts means that in the end there will be less barriers and performance should improve, but on the other hand it might lead to too similar models, which is against personalization goals.

Reviewer comment (Reproducibility):
In an application, it is unclear how to choose which layers to average. It seems that computing each layer's barrier every few epochs throughout training would be inefficient in some applications.

Author response:
**how would you suggest choosing the layers to average**: Our general intuition for the application is that the criticality of layers (i.e. which layers  show rapid growth of loss when disturbed in particular directions) is universal with respect to architecture and task (data to model). If this is indeed the case, then the criticality can be identified once and then reused for all possible applications and training setups. There are indications to this (in particular that our robustness experiments rediscovered the same structure as in the work of Zhang et al.), but we cannot prove it yet.

Reviewer comment (Evaluation):
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?

Author response:
We indeed compute the barrier with respect to the original models’ performance. We specify it also in the Definition 2 and Definition 3 in the paper, but only in mathematical notation. We will add a clarifying note in the text as well.