Paper: Training-time Neuron Alignment for Improving Linear Mode Connectivity and Model Fusion (ICLR-2024)

Reviewer comment (Evaluation):
It is a bit hard to judge from the experimental results for what networks/settings one would prefer this method over other methods.

Author response:
Thank you for the comment. We would like clarify that our main contribution may lie in the new perspective and insights regarding improving training-time linear mode connectivity (instead of the post-training). The TNA-PFN is just a showcase and the verification of our hypothesis/perspective. 

Moreover, TNA-PFN has practical implications in the following aspects: 
- a. _Enable efficient re-basin and improve re-basin._ We note that in some cases, TNA-PFN + weight matching (WM) can achieve even better interpolated generalization compared with WM alone (Table 1 of the paper: MLP_h2_w200 for CIFAR-10 and MLP_h5_w200 for MNIST); also, it costs less computation of WM. 
- b. _Extensions to federated learning._ We extend TNA-PFN in federated learning to improve the generalization of the global model after model fusion. The proposed FedPFN and FedPNU are light-weight and effective.

Reviewer comment (Experiments):
Lack of experiments at scale (e.g. ImageNet).

Author response:
Thanks for this helpful comment. Due to the time constraint of rebuttal period, we may not be able to implement the experiments on full ImageNet, as an alternative, we have conducted experiments on Tiny ImageNet [2], a subset of ImageNet.

Tiny ImageNet contains 100000 64×64 colored images of 200 classes (500 for each class), which can be regarded as more large-scale. The results are in the following Table A, also added to the revised paper. It is observed that TNA-PFN can also reduce the barriers on Tiny ImageNet by up to 23%.

**Table A. Linear mode connectivity on Tiny ImageNet.** The $\rho$ for CNN is 0.4 and the $\rho$ for ResNet18 is 0.3. The learning rate is 0.08.

| **Models** | **Metrics**    | **TNA-PFN**                        | **Vanilla Train** |
|------------|----------------|------------------------------------|-------------------|
| CNN        | Avg. Acc.      | $11.4\pm0.6$                   | $9.85\pm0.3$   |
|            | Interp. Acc.   | $2.91\pm0.9$                   | $1.4\pm0.2$     |
|            | Acc. Barrier   | $0.75\pm0.07$ (12.8%↓)         | $0.86\pm 0.03$  |
|            | Loss Barrier   | $0.75\pm0.09$ (10.4%↓)         | $0.84\pm0.08$   |
| ResNet20   | Avg. Acc.      | $31.6\pm0.4$                   | $31.8\pm0.3$    |
|            | Interp. Acc.   | $12.5\pm2.1$                   | $6.86\pm1.8$   |
|            | Acc. Barrier   | $0.60\pm0.07$ (23%↓)           | $0.78\pm0.06$   |
|            | Loss Barrier   | $1.2\pm0.09$ (22.2%↓)          | $1.6\pm0.2$     |

Reviewer comment (Experiments):
The method is conjectured to perform better for wider networks and worse for deeper networks. It would be interesting to have a more quantitative result i.e. plot the performance for varying width and depth.

Author response:
Thank you for the comment. Acutally, in Figures 3 and 8 of the initial submission, we have conducted experiments regarding the performances for varying width and depth. It seems that our method works well under different depths and widths.

Reviewer comment (Experiments):
For the setting of a single network and multi-model fusion why is there only a comparison to vanilla training and not other methods for improving LMC?

Author response:
Since it might be the first paper to improve linear mode connectivity (LMC) from the training time and previous methods are post-training, other methods are orthogonal to ours, as a result, we only compare with the vanilla training. But we have included experiments of previous post-hoc matching methods in Table 1 of the paper, showing our method is compatible with other post-hoc methods and can have advantages in efficiency.

Reviewer comment (Theory):
As mentioned in the paper, the MNIST results in the second plot of Figure 3 and the experimental results of Entezari, et al. (2022) show that the barrier actually decreases as the network width increases. According to the theorem presented in the paper, the barrier may decrease as the learned dimension decreases, but it does not apply well to these experiments.

Author response:
We apologize that the previous theorem may cause potential misleadings and we have refined the theorem to make it more solid. 

- The previous theorem may cause a misleading that lowering the network dimension by changing the network architecture can improve linear mode connectivity. Actually, this is not what we try to convey and is also inconsistent with the empirical observations. 
- After our revision, the revised theorem is more aligned with our hypothesis and method. It is proved that given a network architecture (without changing the dimension, i.e., width and depth), if we fix a proportion of parameters and train the remaining (like what TNA-PFN does), the linear mode connectivity of two models after training would be improved. 
- The revised theorem is Theorem 3.2 in blue on page 3, and the detailed proof is in the appendix. We hope the improved theorem can relieve the concerns and the potential misleading.

Reviewer comment (Evaluation):
In many experiments, the absolute barrier is too high when using TNA-PFN alone to be an LMC. Of course, it lowers the barrier compared to the vanilla train with no training, but many experiments show a non-negligible barrier.

Author response:
Thanks for your comment. Actually, we didn't claim that we have completely _eliminated_ the barriers, and what we said is we largely _decreased_ the barriers. To the best of our knowledge, we might be the first to attempt to discover the potential of improving linear mode connectivity (LMC) during training, which is very difficult and challenging. 

SGD randomness results in SGD noise, and as a consequence, causes an LMC barrier after training. We make the preliminary attempt to fix a proportion of weights to reduce the permutation effects of SGD noise. However, it is impossible to eliminate the SGD noise when the model is independently trained, and as long as SGD noise remains, the LMC barrier will not be completely eliminated during training. 

We note even if the barrier may still exist after applying TNA-PFN, we have already achieved up to 70% reduction during training time, and it is promising enough for promoting the theoretical and empirical studies.

Reviewer comment (Experiments):
It's good to have a variety of experiments, but there are no experiments on large datasets like ImageNet. It would have been nice to see some experiments on larger datasets, as they generally have different characteristics than MNIST and CIFAR10.

Author response:
Thanks for this helpful comment. Due to the time constraint of rebuttal period, we may not be able to implement the experiments on full ImageNet, as an alternative, we have conducted experiments on Tiny ImageNet [2], a subset of ImageNet.

Tiny ImageNet contains 100000 64×64 colored images of 200 classes (500 for each class), which can be regarded as more large-scale. The results are in the following Table A, also added to the revised paper. It is observed that TNA-PFN can also reduce the barriers on Tiny ImageNet by up to 23%.

**Table A. Linear mode connectivity on Tiny ImageNet.** The $\rho$ for CNN is 0.4 and the $\rho$ for ResNet18 is 0.3. The learning rate is 0.08.

| **Models** | **Metrics**    | **TNA-PFN**                        | **Vanilla Train** |
|------------|----------------|------------------------------------|-------------------|
| CNN        | Avg. Acc.      | $11.4\pm0.6$                   | $9.85\pm0.3$   |
|            | Interp. Acc.   | $2.91\pm0.9$                   | $1.4\pm0.2$     |
|            | Acc. Barrier   | $0.75\pm0.07$ (12.8%↓)         | $0.86\pm 0.03$  |
|            | Loss Barrier   | $0.75\pm0.09$ (10.4%↓)         | $0.84\pm0.08$   |
| ResNet20   | Avg. Acc.      | $31.6\pm0.4$                   | $31.8\pm0.3$    |
|            | Interp. Acc.   | $12.5\pm2.1$                   | $6.86\pm1.8$   |
|            | Acc. Barrier   | $0.60\pm0.07$ (23%↓)           | $0.78\pm0.06$   |
|            | Loss Barrier   | $1.2\pm0.09$ (22.2%↓)          | $1.6\pm0.2$     |

Reviewer comment (Experiments):
In the experiments in the paper, the mask ratio was set to 0.4; was this an experimentally tuned value, or is there some underlying theoretical basis?

Author response:
It is an experimentally appropriate value. It can be seen from Figures 1 and 11 that TNA-PFN is not sensitive in terms of $\rho$. TNA-PFN could be effetive to improve the interpolated accuracy under a wide range of $\rho$, mainly from 0.1 to 0.6. As $\rho$ is not sensitive, so we use the value of 0.4 which is moderate to balance the accuracy-connectivity tradeoff in most cases.

Reviewer comment (Reproducibility):
Which of the three algorithms presented in Ainsworth et al. (2023) was used for weight matching (WM)? What is the approximate computational cost per iteration?

Author response:
In Ainsworth et al. (2023) [3], they proposed three algorithms, activation matching, weight matching (WM), and straight-through estimators matching. The one we used is actually weight matching, which is data-independent. The computational cost of WM relies on the model's parameters, more parameters will cost more computation.

Reviewer comment (Experiments):
Do FedPFN or FedPNU work in general learning situations other than federated learning, and can they be used for model ensembling?

Author response:
We appreciate the intersting comment and have conducted a preliminary experiment on model ensembling. For two indepedently trained networks, we validate the a) _Avg. Acc._: averaged accuracy of two independent models; b) _Ensemble Two_: accuracy of ensembling the two models; c) _Ensemble Four_: accuracy of ensembling the two models and two linearly interpolated models (fusion $\alpha \in \{0.3, 0.7\}$). The results are shown in the following Table B. It is found that ensembling linearly interpolated models for TNA-PFN can improve the generlization but the advantage might by marginal. We reckon it could be an interesting future research direction that studies how training-time neuron alignment improves model ensembling.

**Table B. Model ensembling experiments.** The model is CNN and the dataset is CIFAR-10.

|  **Metrics**    | **TNA-PFN**                        | **Vanilla Train** |
|----------------|------------------------------------|-------------------|
|  Avg. Acc.     | $65.10\pm0.57$         | $62.78\pm0.90$  |
| Ensemble Two   | $70.45\pm0.12$         | $68.21\pm0.82$  |
| Ensemble Four  | $70.53\pm0.17$         | $68.33\pm0.79$  |

Reviewer comment (Presentation):
I'm not sure why the LoRA section was added to Appendix C.1 and what it is trying to say.

Author response:
Thanks for the comment. It may show a potential way to verify our "subspace hypothesis" for improving training-time neuron alignment, but further investigations on the equivalent learning rates across LoRAs with different ranks are needed, so it is out of the paper's scope and we appreciate it as future work.

Reviewer comment (Novelty):
The authors propose a mask-based training scheme, which has already been explored in the literature. While there are variations in the masking details, such as random weight initialization for untrained weights, the overall approach resembles dropout regularization, where different subsets of neurons are masked during each iteration.

Author response:
> 1. Response to "I have a few concerns about the novelty and results of this paper. Firstly, the authors propose a mask-based training scheme, which has already been explored in the literature. While there are variations in the masking details, such as random weight initialization for untrained weights, the overall approach resembles dropout regularization, where different subsets of neurons are masked during each iteration."
> 

Thanks for this comment. We would like to kindly clarify our contributions. 
- **1)** *New insights into linear mode connectivity and permutation invariance from training time.* Our novelty lies in the **training-time** "subspace hypothesis" perspective while previous works only focus on **post-training** connectivity via permutations symmetries (re-basin). The proposed TNA-PFN is just a showcase and verification of the hypothesis/perspective. The contributions are alike [1] in that its main contribution is the hypothesis of linear mode connectivity and permutation invariance and the simulated annealing algorithm (which may not be very effective) is just the verification. We believe our new perspective can inspire more future work in the community. 
- **2)** *The applications of training-time neuron alignment in federated learning.* While linear mode connectivity is somewhat theoretical, we propose two algorithms in federated learning to showcase how improving connectivity during training can help practices. The proposed methods are lightweight and effective.

The mask-based methods are previously studied in the literature but none of them have been discovered to improve the linear mode connectivity during training, and that's where our new insights come from. We deem that new insights, new perspectives, and new findings can also be viewed as novelty. Also, as claimed in "Discussion on gradient/model masks" on page 4, we have elaborated on the differences between our mask and the previous masks, and it is clear that they are algorithmically different.

The reviewer mentioned that our method resembles dropout, but we think this might not be accurate. 
- Dropout means pruning some weights/neurons during training, however, in Figure 2, we have shown that our TNA-PFN is different from pruning and may have better results. 
- But we need to acknowledge that TNA-PFN shares similar regularization effects with dropout/pruning, that it can improve generalization in some cases but may reduce the performances if the mask ratio is too high.
- Another difference between dropout and our method lies in the applications of federated learning. Dropout makes the neurons dead, whereas our masking methods enable the neurons/weights to be deactivated in one round and activated in another round, which improves training-time connectivity without hurting the structure of networks.

Reviewer comment (Novelty):
The results presented in the paper seem somewhat expected. The authors themselves acknowledge in Theorem 3.2 that improving Linear Mode Connectivity (LMC) is achieved by reducing the number of neurons, denoted by $d$. This essentially means using a simpler model, and the neural network is not necessarily overparameterized. Consequently, the number of minima is reduced, leading to a reduction in barriers between different solutions.

Author response:
> 2. Response to "Furthermore, the results presented in the paper strike me as somewhat expected. The authors themselves acknowledge in Theorem 3.2 that improving Linear Mode Connectivity (LMC) is achieved by reducing the number of neurons, denoted by $d$. This essentially means using a simpler model, and the neural network is not necessarily overparameterized. Consequently, the number of minima is reduced, leading to a reduction in barriers between different solutions."
> 

Thanks for this helpful comment. We realize that previous theorem may cause potential misleadings and we have refined the theorem to make it more solid. 

- The previous theorem may cause a misleading that lowering the network dimension by changing the network architecture can improve linear mode connectivity. Actually, this is not what we try to convey and is also inconsistent with the empirical observations. 
- After our revision, the revised theorem is more aligned with our hypothesis and method. It is proved that given a network architecture (without changing the dimension, i.e., width and depth), if we fix a proportion of parameters and train the remaining (like what TNA-PFN does), the linear mode connectivity of two models after training would be improved. 
- The revised theorem is Theorem 3.2 in blue on page 3, and the detailed proof is in the appendix. We hope the improved theorem can relieve the concerns and the potential misleading.

Reviewer comment (Theory):
Assume weight matrix $W$ are randomly initialized as IID standard Gaussian, which is a common practice in neural network initialization. Then $W$ follows a Matrix Normal distribution, i.e., $W\sim MN(0, I, I)$. By applying permutation matrices $P$ and $Q$ to $W$, the result distribution has the same Matrix Normal form, i.e., $PWQ \sim MN(0, I, I)$. Hence, the permutated weight matrix maintains the same distribution as the original $W$. This suggests that permuting the weight matrix is equivalent to reinitializing it. Given this equivalence, it raises the question of how to consider all permutations, given that the search space is uncountable.

Author response:
> 3. Response to "Assume weight matrix $W$ are randomly initialized as IID standard Gaussian, which is a common practice in neural network initialization. Then $W$ follows a Matrix Normal distribution, i.e., $W \sim MN(0, I, I)$. By applying permutation matrices $P$ and $Q$ to $W$, the result distribution has the same Matrix Normal form, i.e., $PWQ \sim MN(0, I, I)$. Hence, the permutated weight matrix maintains the same distribution as the original $W$. This suggests that permuting the weight matrix is equivalent to reinitializing it. Given this equivalence, it raises the question of how to consider all permutations, given that the search space is uncountable."
> 

Thanks for this valuable comment. In the response, we will provide the basic background on the role of permutation invariance in linear mode connectivity to relieve the reviewer's concerns. We kindly suggest the reviewer to refer to [1]. Further, to solve the reviewer's confusion, we will provide detailed explanations from the following four aspects:

- I think you may have assumed that $P$ and $Q$ are constant matrices, or $P$ and $Q$ are independent random matrices with respect to the random matrix $W$, which is $W\sim MN(0,I,I)$. Only in such cases would you arrive at the conclusion that the random matrices $W$ and $PWQ$ have the same distribution. In reality, $P$ and $Q$ are not independent random matrices with respect to $W$. In the process of using permutations to find Linear Mode Connectivity, assuming that $W$ and $U$ are the coefficient matrices of two separate models, the distribution of $P$ and $Q$ depends on $W$ and $U$. Specifically, $P, Q = \arg \min_{P,Q} |PWQ-U|$. Using such $P$ and $Q$ to permute $W$ and obtain $PWQ$ generally does not make $PWQ$ follow the same distribution as $W$. So, the understanding that "permuting the weight matrix is equivalent to reinitializing it" is incorrect. Additionally, whether $W$ and $PWQ$ follow the same distribution is not important; what matters is the distribution of $|PWQ-U|$. Even if $PWQ$ and $U$ follow the same distribution, the distribution of $|PWQ-U|$ can still be complex because $PWQ$ and $U$ are not independent. This complexity arises from their joint distribution (if you are still confused about this, please refer to https://en.wikipedia.org/wiki/Wasserstein_metric, which explains the fact that even if $A$ and $B$ follow the same distribution, the distribution of $|A-B|$ can still be complex).

- A simple experiment can illustrate that $PWQ$ does not follow the same distribution as $W$. Suppose $W = \begin{pmatrix} w_{11}&w_{12}\\ w_{21}&w_{22} \end{pmatrix}$ is a 2x2 random matrix, where $w_{ij} \sim N(0,1)$. We define $Q = \begin{pmatrix}1&0\\ 0 &1\end{pmatrix}$, and $P$ follows the conditional distribution with respect to $W$: when $w_{11} \ge w_{12}$, $P = \begin{pmatrix}1&0  \\ 0&1\end{pmatrix}$, and when $w_{11} < w_{12}$, $P = \begin{pmatrix}0&1 \\ 1 &0\end{pmatrix}$. In this case, the probability density function of the element $w_{11}'$ is given by $p(x) = \frac{1}{2\pi}e^{-\frac{x^2}{2}}\int_{-\infty}^xe^{(-\frac{t^2}{2})}dt$, which does not follow the standard normal distribution (in fact, $w_{11}'$ is the order statistic of two independent normal distribution variables). So, even if $W \sim MN(0,I,I)$, it does not necessarily hold that $PWQ \sim MN(0,I,I)$. Moreover, in the context of $P,Q = \arg \min_{P,Q}|PWQ-U|$, the conditional distributions of $P$ and $Q$ with respect to $W$ and $U$ are more complex than the conditional distributions of $P_0$ and $Q_0$ in our example here. Therefore, it is not possible for $PWQ$ to follow the same distribution as $W$.

- Finding when $W$ and $U$ are random matrices, the distribution of $\min|PWQ-U|$ is actually a famous question called Random Euclidean Matching Problem.  Please refer to [3].  
 
- Now, let's consider the issue of how to find $P$ and $Q$ by calculation. Indeed, for a given $W$ and $U$, finding permutations $P$ and $Q$ that satisfy $P, Q = \arg \min_{P,Q}|PWQ-U|$ is a difficult problem. It has been proven in [2] that this is an NP-hard problem. However, for a given $W$ and $U$, the search space for $P$ and $Q$ is finite, not uncountable, because $W$ and $U$ are of finite dimension and the search space for $P$ and $Q$, which is the space of permutation matrices  for $W$ and $U$ is also finite. Since finding $P$ and $Q$ is an NP-hard problem, there is no polynomial-time exact solution formula in the computation process. Typically, we use the random approximation algorithms proposed in [2], such as weight matching/activation matching.

Reviewer comment (Writing):
The "mask ratio" is not precisely defined, and it seems to be related to the ratio of untrained neurons to the total number of neurons.

Author response:
> 4. Response to "The "mask ratio" is not precisely defined, and it seems to be related to the ratio of untrained neurons to the total number of neurons."
> 

Thanks for this helpful comment. The reviewer understands it correctlly that the mask ratio $\rho$ refers to "keeping $\rho$ fraction of the parameters **fixed** after initialization". We have added the specific defition of $\rho$ in the revised paper (page 4, blue sentence).

Reviewer comment (Theory):
It is unclear whether it is possible to express the results in Theorem 3.2 in terms of "loss barrier" or "accuracy barrier".

Author response:
> 5. Response to "Is it possible to express the results in Theorem 3.2 in terms of "loss barrier" or "accuracy barrier"?"
> 

Thank you for the suggestion. Taking your advice in the newly revised Theorem 3.2, we have reformuate $z_{\alpha}$ into the "barrier" form same with [1]. We note that the "barrier" form is depicted by the network's output function, not the loss or accuracy functions, and previous theoretical analysis also adopted this form [1].

Actually, for a convex loss function $\mathcal{L}$ satisfying the $C$-Lipschitz condition, the loss barrier can be bounded by $\left|z(\alpha)\right|$: $B_{loss}(\mathbf{w_1},\mathbf{w_2})\le C\sup_\alpha|z(\alpha)|$, which differs only by the constant factor $C$, and that's why $|z(\alpha)|$ can be used for the upper bound on the barrier instead of $B_{loss}(\mathbf{w}_1,\mathbf{w}_2)$, both in our Theorem 3.2 and in the previous theoretical analysis [1].

Reviewer comment (Novelty):
Some of the claims are a little overhyped, such as ``training in a subspace'' meaning keeping some set of parameters frozen and optimizing the rest.

Author response:
Thanks for pointing it out. We acknowledge that there are some other definitions/meanings of "subspaces" which may cause misunderstanding. To make it more rigorous and clear, we have added a footnote (page 2 in blue) to clarify the definition of "subspaces" within this paper's scope.

Reviewer comment (Novelty):
The algorithmic novelty is also limited, and most of the strengths lie in the empirical studies on the effect of width, etc.

Author response:
Thanks for this comment. We apologize that the previous theorem may mislead the reviewer to understand our novelty and contributions. We would like to kindly clarify our contributions. 
- **1)** *New insights into linear mode connectivity and permutation invariance from training time.* Our novelty lies in the **training-time** "subspace hypothesis" perspective while previous works only focus on **post-training** connectivity via permutations symmetries (re-basin). The proposed TNA-PFN is just a showcase and verification of the hypothesis/perspective. The contributions are alike [1] in that its main contribution is the hypothesis of linear mode connectivity and permutation invariance and the simulated annealing algorithm (which may not be very effective) is just the verification. We believe our new perspective can inspire more future work in the community. 
- **2)** *The applications of training-time neuron alignment in federated learning.* While linear mode connectivity is somewhat theoretical, we propose two algorithms in federated learning to showcase how improving connectivity during training can help practices. The proposed methods are lightweight and effective.

For the width or depth issue, it is worth noting that according to the previous empirical results [1, 2], deeper networks intrinsically have higher barriers in connectivity and are therefore, harder to improve by post-training rebasin, while wider networks have lower barriers. Therefore, for our TNA-PFN, the strength of width and the weakness of depth are rational and aligned with previous understandings. 

However, we emphasize that our TNA-PFN has better results even on the deeper models compared with previous works. In Figure 7 of [1], the proposed simulated annealing algorithm fails to improve the post-training linear mode connectivity when the depths are high, whereas TNA-PFN improves training-time connectivity even under deeper models (Figures 3 and 4 of the paper).