ZipIt! Merging Models from Different Tasks without Training

We thank you for your thoughtful review and suggestions. We have collected our response below.

W1 / Q1. What is the motivation for this? When would I find myself training from scratch on two different tasks?
This is a good question. The assumption that we and prior work make in this setting isn’t that you would train these models yourself, but that they would be given to you (e.g., downloaded from the internet), or trained in a federated learning setting where due to, e.g., privacy reasons, not all the data is available at once. In those cases, we would like to be able to merge the models without any additional training, in the former because it may be too expensive, and in the latter because we are not allowed to access the original data.

We only train these models ourselves to explicitly control the evaluation setting. Note that some of the models in the multi-dataset experiment were not trained by us.

W2. More transparent claims.
Thank you for this comment. We’ve updated the draft to tone down the language in our PDF abstract and Fig 1 caption.

W3. Figure 1a/b is not exactly clear with respect to related work.
We appreciate the reviewer for pointing this out and we completely agree. We have updated our draft to specify methods with different initializations in the caption of that figure.

Q2. Does this improve on same data merging settings?
We match the performance of baselines on the same data settings, but we focus on the multi-dataset case.

Q3. I would be curious to hear if the authors would expect the same trend with a different method to increase FLOPs on the x-axis, such as depth scaling, or joint scaling, or training iterations.
We did not conduct this search, but it is unlikely that we would see the same magnitude of gains as from scaling width. Mode connectivity heavily relies on the width of models, not necessarily the depth or other possible scaling components. Though, ZipIt is less constrained because of Theorem 1, so it is likely that it would still outperform the other methods in those cases.

We thank you for your review and feedback. We would like to clarify a couple points about our setting and the comparisons to Git-Rebasin:

Note 1. Our Setting.
To clarify, we work in the same setting of [2,4,5], where the user would like to merge two models that were trained independently (i.e., different initialization) without full access to the original data and without having to train—with the added stipulation of the models being trained on different tasks. To the best of our knowledge, we are the first to tackle the setting from this angle. Solving this task would allow us to, for instance, merge foundation models for which we do not have the compute to train, or improve federated learning techniques.

In this setting, the ensemble is the true upper bound, as training is not allowed so the perfect merging would exactly match the output of the original models. Additionally, the extension of working on different tasks / datasets broadens the possible uses of ZipIt! beyond these prior work. Allowing training (like for instance, STE in Git-Rebasin) would increase our performance further, but is out of scope for this work.

Note 2. Which version of Git-Rebasin?
We conduct all our Git-Rebasin [2] experiments in their official public code release [6], which for all ResNet models employs weight-matching (WM). Note that [2] find weight-matching actually performs competitively (and better) to their activation-matching version, though our permute baseline performs better than both. We’ve updated our draft to mention that we use the weight-matching version.

Note 3. The Permute baseline.
We noticed there may be a misunderstanding about this baseline. The permute method is an ablation of ZipIt!. We use our algorithm, but without merging-within, and utilizes the optimal Hungarian algorithm to find permutations. Thus, it is a strong baseline that already outperforms prior work.

We have also included our response to each of the reviewer’s concerns below.

W1. My understanding is that Theorem 1 is positioned to support merging on different datasets + within merging. But this doesn’t make sense to me.
We clarify that we only position Theorem 1 to motivate why merging within models is better than permutation-based methods when redundancy exists within models. We agree that Theorem 1 does not explicitly provide justification for why we would expect strong performance when using within-merging to merge models on different tasks. Instead, we demonstrate this empirically through our extensive experiments (Tables 1/2/3, Figure 6b in the main paper, and more in the Appendix).

W2. Merged model is inaccurate on imagenet1k and multi-dataset unless partial zip is used.
It is correct that ZipIt! does not achieve the same performance on ImageNet1k as it does on CIFAR. Yet, we note that all our baselines including git-rebasin (WM) perform significantly worse than ZipIt!, illustrating the task difficulty when retraining is not allowed (note 1). However, we attribute this result to model capacity. It is likely that our ResNet50 is not large enough to contain sufficient redundancy for which ZipIt! (or the baselines / prior work) can make use of. Based on Figure 6b, we expect performance to dramatically improve with scale.

For low-capacity models, partial zipping is necessary to avoid merging features in later layers that have very low correlation with one another, thus avoiding catastrophic performance when not allowed to train. See our response to Q6 for more details.

Q1. Is git-rebasin activation matching considered equivalent to zipit without within merge? I would like to know the difference between git rebasin, permute and zipit.
ZipIt! within-merging utilizes our own algorithm which achieves near optimal performance (Table 4) whilst being significantly faster and more scalable. The permute method is an ablation of ZipIt!. We use our algorithm, but without merging-within, and utilizes the Hungarian algorithm to find permutations. Git-rebasin with activation matching (AM) is different from both these methods. AM finds permutations based on the pairwise dot-products of neuron activations between models and it appears worse than weight-matching (WM) (Figure 2 of [2]). However, both our permute and ZipIt! significantly outperform WM in every experiment.

Notably, ZipIt! is the only method capable of jointly aligning >2 models into a shared space, as opposed to aligning all models to the space of a chosen anchor model. All prior work (including Git-Rebasin) does the latter, and it is very sensitive to the choice of anchor model.

We thank you for your review and feedback. We have responded to your concerns below.

Clarification: The Permute baseline.
We noticed there may be a misunderstanding about this baseline. The permute method is an ablation of ZipIt!. We use our algorithm, but without merging-within, and utilize the optimal Hungarian algorithm to find permutations. Thus, it is a strong baseline that already outperforms prior work.

W1. “The method seems not to be practical for real-world usage, while lacks insights and theoretical analysis of the model properties”.
Our main goal is to show that it is possible to obtain good performance by merging models from different tasks without retraining, which is something no prior work has been able to achieve. Thus, to show this we present merged accuracy in many different settings with many more evaluations in the appendix. While a theoretical analysis of the resulting models might be interesting, it is not the core focus of our work.

However, we do include extensive analysis of different aspects of the resulting models in Fig. 6/7/8 and Sec. 6, as well as provide rigorous theoretical motivation for our method in Eq. 4 and Appendix G.

W2. “The method works fine in the CIFAR dataset, but doesn’t have remarkable advantages in large datasets such as ImageNet, and disjoint tasks.”
We would like to emphasize that ZipIt! significantly outperforms prior work on both CIFAR-10 and CIFAR-100 and gets close to the maximum possible performance in these settings (i.e., ensemble accuracy). For much harder settings like ImageNet, ZipIt! still outperforms prior work and is able to perform significantly better when using partial unzipping. Note that the “permute baseline” is a strong baseline we introduce that already outperforms prior work.

W3. Moreover, the paper lacks comparisons between its method and the model trained on both datasets directly, which limits its usage scenario.
To clarify, we work in the same setting of [1,2,3], where the user would like to merge two models that were trained independently (i.e., different initialization) without full access to the original data and without having to train—with the added stipulation of the models being trained on different tasks. To the best of our knowledge, we are the first to tackle the setting from this angle. Solving this task would allow us to, for instance, merge foundation models for which we do not have the compute to train, or improve federated learning techniques.

In this setting, the ensemble is the true upper bound, as training is not allowed, so the perfect merging would exactly match the output of the original models. Additionally, the extension of working on different tasks / datasets broadens the possible uses of ZipIt! beyond these prior work. Allowing training would increase our performance further, but is out of scope for this work.

Q1. Permute stronger than Git Rebasin?
We conducted all our Git-Rebasin experiments in their official public code release [4], which for all ResNet models employs weight-matching. Note that [1] find weight-matching actually performs competitively (and better) to the activation-matching version. Our permute baseline is simply our ZipIt! algorithm with the matching algorithm replaced by an optimal permutation finding algorithm. Thus, it is a strong baseline and already outperforms prior work.

Q2. I’m curious about the comparison between “ZipIt!” method and a single model trained on both datasets.
See our response for W3 above. While it might be an interesting comparison, this is out of scope for our setting as the upper bound is ensemble performance.

Q3. The presentation for “Ensemble” results which uses light grey is a little misleading.
Thank you for your comment. We gray out the ensemble because it requires evaluating both models, which uses significantly more (up to double) the parameters / compute compared to the rest of the table entries. Do you have any suggestions on how we can make this more clear?

Q4. Are there the results of Git Re-Basin method for Multi-Dataset in Table 3?
Git Rebasin in its original code-base (all Git-Rebasin experiments are conducted in the official repository [4]) does not support this type of experiment. However, we expect its performance to be much worse than our permute method just like in Tab. 1 and Tab. 2.

Q5. Did you consider the time cost of computing transformation matrices in your method when comparing the FLOPs?
The transformation matrices only need to be computed once and are very quick to compute (~1s total on a CPU), so it is not included in the model FLOPs. Benchmark results for the actual transformation calculation can be found in Tab. 4.

[1] Git-Rebasin: https://arxiv.org/pdf/2209.04836.pdf
[2] OT-Fusion: https://arxiv.org/pdf/1910.05653.pdf
[3] REPAIR: https://arxiv.org/pdf/2211.08403.pdf
[4] Git-Rebasin Code: https://github.com/samuela/git-re-basin

Thanks for the response. However, the authors have not addressed this problem that "the method seems not to be practical" as "compared to the ensemble, the FLOPs don’t decrease a lot while getting lower accuracy". This indicates that “ZipIt!” may not achieve a good trade-off between performance and training time. Therefore, I would like to maintain this 5 score.

We thank you for your comprehensive review and feedback. We have included responses to your questions and concerns below.

W1 / Q2. Concern about Stability.
Most mode connectivity and model merging work assumes models are well-trained. In the issue referenced, the original models obtained 15% lower accuracy than well-trained models, as they were underfit with too low of a learning rate, and thus themselves underperformed before the merge. This caused all merging algorithms to drop to under 20% in accuracy, so this is not specifically a ZipIt! issue. The merged model’s performance is naturally going to depend on the original models’ performances.

W2. Performance degradation when merging multiple models. What about transformers?
Unfortunately, that is a limitation of merging multiple models trained on different tasks into the size of one without being able to train–-since each model has some task-specific neurons [1]. However, as shown in Fig. 6b, this could likely be alleviated in practice by increasing the network width. As for transformers, at the time of submission no other model merging work has shown results on transformers with different initializations, but this is a great direction for future work.

Q1. What does “ensemble” mean here?
Thanks for pointing this out, we’ve clarified this in the baseline section of the updated draft. For joint accuracy, we softmax the logits of each model separately to place their predictions on the same scale, and then concatenate resulting probabilities from each model and take an argmax. For avg task accuracy, we use the correct model for each task (i.e., just average the per-task accuracies). Note that this is upper bound accuracy because with perfect merging (not allowing training), the merged model's outputs would be identical to the originals.

Q3. Results for smaller models, e.g., MobileNet?
For smaller models, note that ResNet-20x4 (which we use for our CIFAR experiments) has around the same number of parameters as MobileNetV2, and we show scaling performance (both downward and upward) in Fig. 6b. While it’s true that constraining the model size hurts performance, this depends on task difficulty and is true for all other methods as well.

Q4. Can BN be fused beforehand, so it doesn’t need to be recomputed?
Sadly, while it is possible to fuse the BNs into the models beforehand, [2] shows that simply merging the model’s BNs together would result in “variance decay”. So some mitigation strategy would be necessary regardless, and we found resetting the batch norm to be the simplest option.

[1] Simon Kornblith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited. In ICML, 2019.
[2] REPAIR: https://arxiv.org/pdf/2211.08403.pdf