Layer-wise linear mode connectivity

We thank the reviewer for emphasizing the relevance of our paper, and its practical usefulness. In the following we address their questions:

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.

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.

Weakness:

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.

I thank the authors for their responce and clarification. I think I have some misunderstanding in my original comment. I agree with the authors' claim that the robustness result is surprising and is different from the same-basin case. I think the current overall organization of this paper is kind of confusing. Based on this, I decide to keep my current rating to this paper as weakly accept and encourage the author to reorganize the paper and make it clearer in their future revision, especially in the following aspects:

• Besides proposing the definition of LLMC, clarify when will this phenominon happen.
• Make every term clearly defined. For example the "robustness to random perturbations". Although the authors clarified this term but I don't think it is nowhere defined in the paper.
• As there are multiple definitions about LLMC, it's better to clarify in each experiment is it looking at which version of LLMC.

We thank the reviewer for appreciating our work and are happy to improve it according to the proposals made. In the following we address the questions raised:

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.

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.

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.

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.

4. an error in the author name of the paper: We used the publicly available version of the paper "Which layer is learning faster? A systematic exploration of layer-wise convergence rate for deep neural networks" from https://openreview.net/forum?id=wlMDF1jQF86, where Yixiong Chen is indicated as the first author. We apologize if there exists another version of the paper, but we were not able to find it. Could you please point us towards the version you have in mind?

5. Thank you for your detailed comments on the text. We will fix these issues in the manuscript.

We thank the reviewer for the appreciation of our conjectures and experimental results. In the following we address the questions raised by the reviewer.

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.

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.

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.

Addressing the weaknesses:

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.

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.

The reviewer appreciates the clear explanation of cumulative averaging and layer-wise barriers. The work on LLMC is indeed novel and intriguing. Despite the room for further enhancement in the theory part, the reviewer decides to vote for acceptance.

Minor corrections:

• The surname of the first author of the referenced paper appears to be inconsistent. While the citation presents it as (Yixiong Chen, 2023), the reference section lists it as Zhou.

We thank the reviewer for the positive feedback and are happy to clarify the questions listed.

1. 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.

2. 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.

To the weaknesses:

1. 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.

2. 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.

3. 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.

Dear Seok-Ju,

Thank you for your interest in our work.

Indeed, we discuss personalized federated learning in Section 2 and Section 6 of our paper, including some works that use feature matching for enforcing better connectivity predating your work, namely FedMA from 2020, and the more tightly connected to our research works of Singh et al. 2020 and Ainsworth et al. 2022. In Section 6 we specifically discuss partial averaging and not connectivity inducing algorithms. We are not exhaustively discussing existing work in personalized federated learning here, since our research investigates layer-wise linear mode connectivity in standard training scenarios.

We understand your concern with improving the visibility of your approach, but it is out of scope of this work. We will definitely consider your algorithm as a benchmark in case we continue this research into an applicational algorithm for personalized federated learning.

Best regards,

The authors