$C^2M^3$: Cycle-Consistent Multi-Model Merging

We thank the reviewer for the insightful comments and spot-on questions. We now proceed to address the raised concerns and answer the questions to the best of our capabilities.

• W1: model merging foundations based on prior work. We agree with the reviewer that our work builds on top of existing work and partially leverages the same foundations. This is true, however, for each permutation-based model merging work, e.g. [1, 2, 3, 4, 5, 6], and we believe that sharing a set of common assumptions is inevitable in a subfield. We would also like to remark that none of the existing works provably guarantees cycle consistency, with the problem itself being overlooked. Therefore, while we acknowledge owing the existence of this research to prior work that has laid the foundations of the field, we argue that our work takes it a step further and, being theoretically grounded, has the potential to be itself a foundation for future research.

• W2: pros and cons of global vs local optimization. We thank the reviewer for bringing this up. The answer is double-fold: i) global optimization is required to enforce cycle consistency, as it is an inherently global property; ii) global optimization removes the arbitrariness in the layer iteration order, which, as we report in Figure 4 and Table 6, results in a marked variance in the results.
  More in detail, it is not clear how cycle consistency, a global property of the model-level transformation, can be achieved using layer-local optimization. For this reason, we do not only choose global optimization due to the deterministic nature of its results but also because it is the only way to achieve our main and foremost objective — that of having cycle-consistent permutations between models. Agreeing with the reviewer about the importance of these considerations, we intend to use the extra page in the camera-ready version to clarify them.

• Q1: why conditional gradient descent. We thank the reviewer for the question. Conditional gradient descent (Frank-Wolfe) is a particularly suitable algorithm for the problem we are trying to solve as it only requires computing the gradient, for which we derived a closed form, and enjoys two advisable properties: i) monotonicity of the objective function (see also Figure 11 in appendix) ii) guaranteed convergence rates, for which we added a proof in the appendix. Regarding the difficulty of the process, the monotonicity of the loss makes it arguably stable, even if the step size is decreasing in an alternating manner. As briefly explained in Appendix A. 5, we believe this behavior to reflect a fixed pattern in the optimization rather than an instability in the overall process. Regarding the scalability of the approach with larger networks, we report here its wall-clock time when merging n=2, 3 ResNet20 models having 1x, 2x, 4x, 8x, and 16x width, together with their number of parameters.

  	1x	2x	4x	8x	16x
  # params	292k	1.166m	4.655m	18.600m	74.360m
  C2M3 time n=2	33.4s	33.5s	40.5s	80.8s	367.8s
  C2M3 time n=3	32.9s	83.18s	91.0s	162.0s	715.8s
  MergeMany time n=2	0.24s	0.4s	3.4s	8.9s	59.4s
  MergeMany time n=3	1.2s	4.1s	19.5s	105.8s	892.3s

  As can be inferred from the table, the scaling laws depend on the complexity of the resulting matching problem and cannot be predicted merely from the number of parameters, with a 4-fold increase in parameters resulting in no increase in runtime for the first three columns, a double increase in the second-last column and a 5-fold increase in the last. Compared to MergeMany, our approach enjoys a milder increase in running time when increasing the number of parameters. We also included a rigorous proof determining the convergence rate of the algorithm in the rebuttal document.

• Q2: Efficiency of the method for n>2. We thank the reviewer for the question. To answer it, we computed the matching and merging time when merging $n=2, \dots, 10$ ResNet20 models with 4x width, as can be seen in Figure 1 of the rebuttal document. Compared to the pairwise baseline that maps all the models to a fixed one, our approach incurs a significantly steeper cost. However, Figure 2 (rebuttal document) shows that the latter also suffers from a performance decrease when increasing $n$ which is much more pronounced, making it not advisable for the task.

• Q3: Merging factor alpha. We thank the reviewer for pointing out this missing detail. In all our experiments, we have used $\alpha = \frac{1}{n}$, where $n$ is the number of models to merge. This is indeed usually the hardest case, as it defines a point that is the furthest from the individual basins.

• Q4: cycle consistency. We thank the reviewer for allowing us to clarify this aspect. Luckily, cycle consistency, i.e., the property for which applying a cycle of such transformations leads back to the starting point, holds in this case as long as each transformation is orthogonal. Looking at Figure 1 in the manuscript, it can be appreciated how mapping $A$ to another model $B$ means always passing through $U$ and then to $B$ with $P_B$, and that mapping $B$ to another model $C$ means passing through $U$ with $P_B^T$. This means that, as long as $P_B P_B^T = I$, the effect of mapping $A$ to $B$ is nullified. Since the proposed algorithm returns permutations (a subspace of orthogonal transformations), cycle consistency always holds for the resulting permutations.

Thanking again the reviewer for their effort, we remain available for any further clarification.

References

[1] Ainsworth, Samuel, Jonathan Hayase, and Siddhartha Srinivasa. 2022. “Git Re-Basin: Merging Models modulo Permutation Symmetries.” In The Eleventh International Conference on Learning Representations (ICLR), 2022

[2] Jordan, Keller, Hanie Sedghi, Olga Saukh, Rahim Entezari, and Behnam Neyshabur. 2023. “REPAIR: REnormalizing Permuted Activations for Interpolation Repair.” In The Eleventh International Conference on Learning Representations (ICLR), 2023

[3] Navon, Aviv, Aviv Shamsian, Ethan Fetaya, Gal Chechik, Nadav Dym, and Haggai Maron. 2023. “Equivariant Deep Weight Space Alignment.” in The Forty-first International Conference on Machine Learning (ICML), 2024

[4] Peña, Fidel A. Guerrero, Heitor Rapela Medeiros, Thomas Dubail, Masih Aminbeidokhti, Eric Granger, and Marco Pedersoli. “Re-Basin via Implicit Sinkhorn Differentiation.”, IEEE / CVF Computer Vision and Pattern Recognition Conference (CVPR), 2023

[5] Singh, Sidak Pal, and Martin Jaggi. 2020. “Model Fusion via Optimal Transport.” In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020 (NeurIPS), 2020

[6] Horoi, Stefan, Albert Manuel Orozco Camacho, Eugene Belilovsky, and Guy Wolf. 2024. “Harmony in Diversity: Merging Neural Networks with Canonical Correlation Analysis.”, in The Forty-first International Conference on Machine Learning (ICML), 2024

Meaning of color in Figure 3

The color in Figure 3 represents the value of the loss in the loss landscape given by interpolations of the models. Red values indicate low-loss regions (basins), while blue values indicate high-loss regions. We added this information in the caption.

Thanking again the reviewer for their effort, we remain available for any further clarification.

References

[1] Ainsworth, Samuel, Jonathan Hayase, and Siddhartha Srinivasa. 2022. “Git Re-Basin: Merging Models modulo Permutation Symmetries.” In The Eleventh International Conference on Learning Representations (ICLR), 2022

[2] Jordan, Keller, Hanie Sedghi, Olga Saukh, Rahim Entezari, and Behnam Neyshabur. 2023. “REPAIR: REnormalizing Permuted Activations for Interpolation Repair.” In The Eleventh International Conference on Learning Representations (ICLR), 2023

[3] Qu, Xingyu, and Samuel Horvath. "Rethink Model Re-Basin and the Linear Mode Connectivity." arXiv preprint arXiv:2402.05966 (2024).

[4] Imfeld, Moritz, et al. "Transformer fusion with optimal transport." In The Twelfth International Conference on Learning Representations (ICLR), 2024

[5] Verma, Neha, and Maha Elbayad. "Merging text transformer models from different initializations." arXiv preprint arXiv:2403.00986 (2024).

We thank the reviewer for the insightful comments and spot-on questions. We will now do our best to address each and every critique and question.

Sensitivity to hyperparameters

We thank the reviewer for raising this point. We would like, in fact, to clarify this aspect: when, in the limitations, we refer to merging methods being sensible to hyperparameters, we do not refer to this particular approach, but to the whole linear mode connectivity phenomenon. Different works have, in fact, observed this in practice: among these, Ainsworth et al. [1] when listing the known failure modes of their approaches, mention models trained with SGD and too low learning rate, or ADAM coupled with too high learning rate. Keller et al. [2] show that the chosen normalization layer incredibly affects the accuracy of the resulting merged model, while Qu et al. [3] observe learning rate, weight decay, and initialization method to play a strong role as well. We therefore argue that the sensitivity to hyperparameters is not a peculiarity of our approach, but something we observed throughout all the approaches we used and that has been confirmed in the existing literature. We added this clarification to the main manuscript.

Lack of theoretical guarantees

We thank the reviewer for the comment. As explicitly stated in the limitations of the paper, we share with the reviewer an overarching wish for an improved theoretical ground for model merging and linear mode connectivity. This is mostly due to the whole permutation-based model merging field building upon a conjecture that cannot be proven or disproven due to the huge number of possible neuron permutations. We would like, however, to remark that this is not something that has to do with our research in particular but with the field in general and that, differently from several works in the field, we provide a principled approach that indeed holds some guarantees. While this may be just a brick on a bumpy road, this brick is stable and allows further work to build upon it.

Small datasets and networks

We thank the reviewer for the suggestion.

As can be seen in the table, we are using the most established set of architectures and datasets considered in all the previous and concurrent literature in the field. This choice stems from two motivations: 1) for the sake of comparison and continuity with respect to previous works; 2) the complexity of the architecture adds additional challenges that must be taken into account and requires additional research that is not immediately relevant to the merging approach. In particular, the architecture suggested by the reviewer is a transformer-based architecture and requires ad hoc mechanisms to handle multi-head attention, positional embeddings, and the vast number of residual connections, in fact motivating stand-alone works to do this [4, 5]. We agree, however, with the reviewer that we are lacking a very large-scale case and, therefore, are also training ResNet50 endpoints on ImageNet. We aim to include these results as soon as they are available. Given our limited computational budget, these may be ready before the end of the discussion phase or for the camera-ready.

Convergence speed

We thank the reviewer for the question. We report here the wall-clock time when merging n=2,3 ResNet20 models having 1x, 2x, 4x, 8x, and 16x width, together with their number of parameters.

As can be inferred from the table, the scaling laws depend on the complexity of the resulting matching problem and cannot be predicted merely from the number of parameters, with a 4-fold increase in parameters resulting in no increase in runtime for the first three columns, a double increase in the second-last column and a 5-fold increase in the last. Compared to MergeMany, our approach enjoys a milder increase in running time when increasing the number of parameters. We also included a rigorous proof determining the convergence rate of the algorithm in the rebuttal document.

Cosine similarity of the models

We thank the reviewer for the suggestion. We report here the cosine similarity between a model and the model obtained by cyclically applying the permutations obtained with git re-basin and C2M3 respectively.

The result is analogous to that observed in the manuscript with the L2 distance, showing the lack of error accumulation in the proposed approach.

Ambiguity in equation 3

We thank the reviewer for pointing out this ambiguity. We confirm their understanding of the equation summing over all pairs of models (1, 1), (1, 2), …, (n-1, n) . For the sake of clarity, we replaced the equation in the manuscript with a triple sum, for each model p=1, ..., n-1, for each model q=p+1, ..., n and each layer l=1, ..., L.

(TBC)

Inconsistencies:

We thank the reviewer for accurately spotting the naming inconsistency, we have promptly fixed the issue in the manuscript by making $n$ lowercase everywhere. Analogously, we added a figure-level caption to Figure 2 which incorrectly had only two subfigure-level captions. The caption states “Existing methods accumulate error when cyclically mapping a model through a series of permutations, while $C^2M^3$ correctly maps the model back to the starting point.”.

Regarding definition 2.1, we used the definition as presented in Git Re-Basin [1]. We believe, however, the reviewer’s comment to be spot on, as the current formulation may hinder clarity. We therefore replaced the $\mathcal{L}(\Theta_A) \approx \mathcal{L}(\Theta_B)$ assumption by instead asking that the two points correspond to weights of neural networks trained to convergence with SGD.

References

[1] Ainsworth, Samuel, Jonathan Hayase, and Siddhartha Srinivasa. 2022. “Git Re-Basin: Merging Models modulo Permutation Symmetries.” In The Eleventh International Conference on Learning Representations (ICLR), 2022

[2] Jordan, Keller, Hanie Sedghi, Olga Saukh, Rahim Entezari, and Behnam Neyshabur. 2023. “REPAIR: REnormalizing Permuted Activations for Interpolation Repair.” In The Eleventh International Conference on Learning Representations (ICLR), 2023

[3] Navon, Aviv, Aviv Shamsian, Ethan Fetaya, Gal Chechik, Nadav Dym, and Haggai Maron. 2023. “Equivariant Deep Weight Space Alignment.” in The Forty-first International Conference on Machine Learning (ICML), 2024

[4] Peña, Fidel A. Guerrero, Heitor Rapela Medeiros, Thomas Dubail, Masih Aminbeidokhti, Eric Granger, and Marco Pedersoli. “Re-Basin via Implicit Sinkhorn Differentiation.”, IEEE / CVF Computer Vision and Pattern Recognition Conference (CVPR), 2023

[5] Singh, Sidak Pal, and Martin Jaggi. 2020. “Model Fusion via Optimal Transport.” In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020 (NeurIPS), 2020

[6] Horoi, Stefan, Albert Manuel Orozco Camacho, Eugene Belilovsky, and Guy Wolf. 2024. “Harmony in Diversity: Merging Neural Networks with Canonical Correlation Analysis.”, in The Forty-first International Conference on Machine Learning (ICML), 2024

[7] Lacoste-Julien, S. (2016). Convergence Rate of Frank-Wolfe for Non-Convex Objectives. ArXiv, abs/1607.00345.

[8] Imfeld, Moritz, et al. "Transformer fusion with optimal transport." In The Twelfth International Conference on Learning Representations (ICLR), 2024

[9] Verma, Neha, and Maha Elbayad. "Merging text transformer models from different initializations." arXiv preprint arXiv:2403.00986 (2024).

[10] Marczak, Daniel, et al. "MagMax: Leveraging Model Merging for Seamless Continual Learning." ECCV 2024

Theoretical analysis: convergence properties and guarantees.

Following previous literature on the Frank-Wolfe algorithm [7], we know that FW obtains a stationary point at a rate of $\mathcal{O}(1 / \sqrt{t})$ on non-convex objectives with a Lipschitz-continuous gradient. We now prove that the considered objective function

$

\sum_{p=1}^{n-1} \sum_{q=p+1}^{n} \sum_{\ell=1}^L \langle (P_{\ell}^p )^\top W_\ell^p P_{\ell -1}^p, (P_{\ell}^q)^\top W_{\ell}^q P^q_{\ell -1} \rangle

$

Has a Lipschitz-continuous gradient. We recall that, for each layer permutation $P^A = \{P_1^A, P_2^A, \ldots, P_{L}^A\}$ of model $A$, we can define the gradient of our objective function relatively to the model $B$ we are matching towards:

$

f(P_l^A) = \nabla^{\text{rows},P_l^A} + \nabla^{\text{cols},P_l^A} + \nabla^{\text{rows},\leftrightarrows,P_l^A} + \nabla^{\text{cols},\leftrightarrows,P_l^A} = \left[W^A_l P_{l-1}^A (P^B_{l-1})^\top (W^B_{l})^\top + (W^A_{l+1})^\top P_{l+1}^A (P^B_{l+1})^\top W^B_{l+1}\right] P^B_{l} + \left[W^B_{l} P_{l-1}^B(P^A_{l-1})^\top (W^A_{l})^\top + (W^B_{l+1})^\top P_{l+1}^B (P^A_{l+1})^\top W^A_{l+1}\right] P^A_{l}

$

To prove Lipschitz continuity, we need to show there exists a constant $C$ such that $\forall p\in[1,n] l\in[1,L] \lVert f(P_\ell^p) - f(Q_\ell^p) \rVert \leq C \lVert P_\ell^p - Q_\ell^p \rVert$. To simplify passages, we only consider a fixed $l$ and perform a generic analysis. We begin by observing that

The last form of the above equation can be rewritten as a sum of the two sums:

Since the first term does not depend on either $P_l^p$ or $Q_l^p$, we assume as a worst case that its norm is 0. Then, we remove transposes (since $\lVert M \rVert = \lVert M^\top \rVert$) and apply the triangle inequality and the sub-multiplicative property of matrix norms:

Let

\lVert f(P_l^p) - f(Q_l^p) \rVert \leq C \sum_{q\in [1,n]\setminus p} \lVert P_l^p - Q_l^p \rVert = C (n-1) \lVert P_l^p - Q_l^p \rVert

We thank the reviewer for their efforts in providing insightful comments and questions.

Method’s scalability to large models or datasets.

As can be seen in the table, we are using the most established set of architectures and datasets considered in all the previous and concurrent literature in the field. This choice stems from two motivations: 1) for the sake of comparison and continuity with respect to previous works; 2) the complexity of the architecture adds additional challenges that must be taken into account and requires additional research that is not immediately relevant to the merging approach. In particular, transformer-based architectures require ad hoc mechanisms to handle multi-head attention, positional embeddings and the vast number of residual connections, in fact motivating stand-alone works to do this [8, 9]. We agree, however, with the reviewer that we are lacking a very large-scale case and, therefore, are also training ResNet50 endpoints on ImageNet. We aim to include these results as soon as they are available. Given our limited computational budget, these may be ready for the camera-ready.

Real-world applications. We thank the reviewer for the suggestion. Tackling the problem of model merging, we inherit all the applicative domains suggested in relevant prior works, such as federated learning [1, 3], incremental learning [4], and continual learning [10]. While its foundational nature, in our opinion, makes examples of real-world applications out of scope for our paper, we share the curiosity of the reviewer in assessing its effectiveness in such a scenario. Therefore, we are currently setting up a federated learning experiment whose results we hope to share during the discussion phase.