Foldable SuperNets: Scalable Merging of Transformers with Different Initializations and Tasks

Authors: Edan kinderman, Itay Hubara, Haggai Maron, Daniel Soudry
Primary area: transfer learning, meta learning, and lifelong learning
Keywords: Model merging, Knowledge Distillation, Deep Learning

Abstract:
Many recent methods aim to merge neural networks (NNs) with identical architectures trained on different tasks to obtain a single multi-task model. Most existing works tackle the simpler setup of merging NNs initialized from a common pre-trained network, where simple heuristics like weight averaging work well. This work targets a more challenging goal: merging large transformers trained on different tasks from distinct initializations. First, we demonstrate that traditional merging methods fail catastrophically in this setup. To overcome this challenge, we propose Foldable SuperNet Merge (FS-Merge), a method that optimizes a SuperNet to fuse the original models using a feature reconstruction loss. FS-Merge is simple, data-efficient, and capable of merging models of varying widths. We test FS-Merge against existing methods, including knowledge distillation, on MLPs and transformers across various settings, sizes, tasks, and modalities. FS-Merge consistently outperforms them, achieving SOTA results, particularly in limited data scenarios.

Decision: reject
Reviewer scores: 6, 5, 5, 6

Reviews and discussion:
<thread>
creation date:
1730455834285

thread_type:
official_review

<thread_object>
creator: reviewer_bvwy

<content>
summary:
This paper explores the method of merging large transformers trained on different tasks from distinct initializations, deviating from the traditional setting where models are required to originate from the same checkpoint. The authors build upon and enhance the concept of folding from Zipit![1], successfully extending its application to transformers, thereby broadening its original scope. Furthermore, the authors employ data augmentation to address the constraints of limited data settings. This approach has achieved remarkable results on models such as MLPs and ViTs.

strengths:
1. Expanding the scope of model merging is a highly significant direction.
2. Very detailed introduction and guidance.
3. Thorough ablation studies and large-scale experimental results and discussions.

weaknesses:
1. Concerns about application scenarios: While the paper takes a step forward in expanding the scope of model merging, it does not significantly enhance the applicability of model merging scenarios. The target scenario (same model architecture with identical pre-training data but different initialization, followed by fine-tuning on different tasks) is not common in practice.
2. I noticed that data augmentation techniques are also used and discussed in Zipit. I suggest the authors address this point when discussing the connections with Zipit, highlighting similarities and differences in the methods used.

questions:
1. There is not much discussion about the next steps in the paper. What are the authors' thoughts on merging truly similar architectures but trained entirely from scratch on different datasets?
2. I am a little curious about the potential effects if the method were applied to the more common setting we have now (originating from the same checkpoint, fine-tuning on different tasks).

Typos:
Table 22, last row.

[1] ZIPIT! MERGING MODELS FROM DIFFERENT TASKS without Training
</content>
</thread_object>

</thread>
<thread>
creation date:
1730619784470

thread_type:
official_review

<thread_object>
creator: reviewer_esd4

<content>
summary:
The paper tackles the problem of merging two Transformers with different initializations and target tasks. To this problem, the paper proposes an extension of the previous idea of folding two weight matrices to the Transformer architecture, and proposes to optimize the merging/unmerging layers for folding by knowledge distillation on unlabeled data.

strengths:
1. The paper is well-structured and easy to follow.
2. The paper polishes the idea of merging by folding from [1], i.e., merging two weights by inserting merging/unmerging layers, and extends it to the specific architecture of Transformer. Also the paper shows it works well when combined with knowledge distillation on unlabeled data.

weaknesses:
1. Although the extension of the previous idea to Transformer is a technical contribution, the novelty of the proposed method is still limited since it just applies the feature-level knowledge distillation to the merging/unmerging layer which is originally proposed in [1].
2. While the title suggests that this paper addresses the general problem of merging Transformers with different initializations, but the experiments are performed only with the Vision Transformer with a few downstream tasks. Also, since there are various pre-trained Vision Transformers available, it should be tested with other initializations rather than just ImageNet-1K pre-trained one.
3. The proposed approach heavily relies on optimization with unlabeled data, while the previous works (SLERP, RegMean, ZipIt, Opt) are designed to be applied without any optimization (but also can be with additional finetuning). If we allow such optimization, the problem can also be reduced to multi-task learning or distillation (with unlabeled data in this case), which now has a plenty of existing approaches ([2,3,4,5] for e.g.) to solve it. Thus, the paper should discuss more on the relationship to such works.
4. There is a concern about the computational/memory inefficiency in the optimization phase particularly when the size of models to be merged increases, because the knowledge distillation part involves large matrices. Since model merging is typically used with a low-end GPU, the proposed approach may be not promising.
5. The number of optimized parameters reported in Table 3,4,5 may be (possibly intentionally) too misleading. It apparently suggests that the proposed method is more than 10x efficient compared to the vanilla knowledge distillation, but the actual time for merging reported in Table 6 of Appendix shows this is not the case. Rather, the proposed method seems more than 2x inefficient, possibly due to the above weakness.

[1] Stoica et al., "ZipIt! Merging Models from Different Tasks without Training"

[2] Li et al., "An Unsupervised Multiple-Task and Multiple-Teacher Model for Cross-lingual Named Entity Recognition"

[3] Nguyen-Meidine et al., "Unsupervised Multi-Target Domain Adaptation Through Knowledge Distillation"

[4] Park and Kwak, "Feature-level Ensemble Knowledge Distillation for Aggregating Knowledge from Multiple Networks"

[5] Shi et al., "Data-free Model Fusion with Generator Assistants"

questions:
See Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730637953121

thread_type:
official_review

<thread_object>
creator: reviewer_hgax

<content>
summary:
This work addresses the challenging task of merging large transformers trained on different tasks from distinct initializations. The authors first demonstrate that traditional merging methods fail catastrophically in this scenario. To tackle this, they propose Foldable SuperNet Merge (FS-Merge), a method that optimizes a SuperNet to fuse the original models using a feature reconstruction loss. FS-Merge is straightforward, data-efficient, and capable of merging models with varying widths. The method is evaluated against existing approaches, including knowledge distillation, on MLPs and transformers across diverse settings, sizes, tasks, and modalities. FS-Merge consistently achieves state-of-the-art (SOTA) results, particularly in data-limited scenarios.

strengths:
1.	The paper tackles an intriguing problem: merging large transformers trained on different tasks from distinct initializations into a single model.
2.	The proposed method is simple and easy to follow.

weaknesses:
1.	FS-Merge is a training-based merging method, which can be costly compared to other model merging techniques.
2.	The figures, such as Figure 3, are low resolution, and the overall writing quality of the paper is not very professional, requiring significant improvement and refinement.
3.	The datasets used, such as MNIST and SVHN, are relatively small, and the performance improvements appear marginal. Also, the experimental setup seems unique to this paper and not aligned with standard practices in prior literature.
4.	The paper is largely empirical, lacking an in-depth discussion on why the proposed method effectively combines models trained on different domains.

questions:
1.	What is the training cost associated with different datasets, such as FLOPs?
2.	The proposed approach is similar to adapter-based methods. Could the authors discuss this similarity?
3.	It would be beneficial to include results on larger-scale datasets, such as the VTAB-1K benchmark and even ImageNet-1K, which are widely used in model merging or domain adaptation research.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730689601488

thread_type:
official_review

<thread_object>
creator: reviewer_2f4v

<content>
summary:
This paper addresses the challenging problem of merging large transformers trained on different tasks from distinct initializations, where prior works typically rely on models that share a common pretrained initialization. The proposed method, FS-Merge, utilizes a feature reconstruction loss to merge the original models effectively.

strengths:
- The method is data-efficient, requiring only an unlabeled subset of the training data for optimization, which is advantageous when full access to data is limited.
- The paper presents comprehensive experimental results across various model architectures and data scenarios, demonstrating the scalability and effectiveness of FS-Merge in merging models of different scales and tasks.

weaknesses:
- While the method is designed for models trained from scratch, it would be insightful to investigate its performance when applied to pretrained models that are fine-tuned on different sources. Specifically, it would be beneficial to explore potential challenges or advantages this application might present compared to merging models trained from scratch. This analysis could provide a broader understanding of the method's applicability and limitations.
- The paper could benefit from a discussion on robustness to distribution shifts, similar to what is explored in the WiSE-FT paper, "Robust fine-tuning of zero-shot models".
- It would be helpful to compare this method to Mixture of Experts (MoE) approaches, such as "Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM", to provide context on how FS-Merge differs or could complement these strategies. I encourage the authors to provide a comparison between FS-Merge and MoE approaches, with a focus on key differences such as computational efficiency, model size, and how FS-Merge could potentially complement MoE methods. This comparison would contextualize the distinct contributions of FS-Merge and highlight its unique strengths.

questions:
- Why does distillation perform poorly in the last task (C, M, C100, E) as shown in Table 4, while it performs well in the earlier tasks? An analysis or hypothesis explaining this discrepancy would strengthen the paper's results section. Discussing potential contributing factors such as task similarity, dataset characteristics, or interactions between specific models being merged would add depth to the findings and clarify this discrepancy.
- Given the use of the “first” initialization, how sensitive is the method to changes in the order of models? Would changing the order significantly affect the merged model’s performance, and if so, why? The authors should consider conducting an ablation study on the impact of model ordering during initialization. Reporting performance metrics for different orderings and discussing any observed patterns or practical implications would provide valuable insights into the method's robustness and guide its real-world application.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733042316414

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We deeply appreciate the time and effort the reviewers have dedicated to reviewing our work, and thank them for the constructive feedback.

We are encouraged that the reviewers found our paper well-written and easy to follow (HgAx , eSD4) and recognized that it addresses a challenging and intriguing problem (2F4V, HgAx, bVWy). The reviewers also highlighted the data efficiency of the proposed method (2F4V, HgAx) and its ability to achieve SOTA results across various scenarios, architectures, model sizes, datasets, and modalities (2F4V, HgAx, bVWy).

Following the responses to our rebuttal, we see several questions remaining. We would like to respond to these questions:

> 1. Isn’t the comparison between traditional merging methods and FS-Merge or the KD baseline “unfair”, due to the latter methods requiring significantly more resources?

This comparison was meant only to demonstrate that traditional merging methods perform poorly in our setting — across all the tasks, modalities, and model sizes we evaluated. This motivated using KD and developing FS-Merge, to address this challenging scenario.

However, we can understand the sentiment that if we could somehow add more resources to the traditional local methods, this would make the comparison more “fair”. Interestingly, we already did something very similar: in Appendix H.1 we examined the (more resource-intensive) local version of FS-Merge and found it is ineffective in our setting. This suggests that traditional methods are ineffective because they rely on local objectives (which are inherently unsuitable for this challenging task) — not only because they use fewer resources. Therefore, we are not sure what we can do to make the comparison more fair.

> 2. Is the KD method a strong baseline?

Yes, since it performs better than existing merging methods both in: 

(a) The standard ‘same init’ setting (Table 24)

(b) Our ‘different init’ setting (most other tables in our paper)

Moreover, we tested various KD approaches for model merging (Appendix G.2), including different initializations, inner feature usage (Appendix H.3), and augmentations. The KD method that we identified as the strongest was used as a baseline (and it was still outperformed by FS-Merge).

Thus, we do not see a better baseline than the KD method we used. However, merging transformers from different initializations is a relatively new direction, and we hope future works will introduce additional baselines.

> 3. Are suggested methods too resource-intensive to be practical? 

No, since there is no better option. Specifically, in our challenging setting:

(a) No other method works at all, other than KD and FS-merge.

(b) FS-merge (or FS-merge seq.) has a comparable computational footprint as KD, yet achieves better accuracy. 
This should not come as a surprise: many times in the past computationally cheap methods worked well in an ‘easy’ setting but became inadequate in more challenging cases, and this led to more computationally intensive methods (e.g., neural networks replacing linear models).
</content>
</thread_object>

</thread>


Metareview:
metareview:
Summary: The paper introduces FS-Merge, a method for merging large transformers trained on different tasks with distinct initializations. It extends the concept of folding from previous work and uses feature reconstruction loss with unlabeled data for optimization. The method shows promise in data-limited scenarios and is effective across various model architectures and tasks.

Strengths:

FS-Merge addresses the challenging problem of merging models trained from scratch, offering a data-efficient solution that does not require labeled data for optimization.

The paper provides a comprehensive evaluation, demonstrating FS-Merge's effectiveness and scalability across different model sizes, tasks, and data scenarios.


Drawbacks:

The method's reliance on knowledge distillation makes it computationally expensive compared to other model merging techniques, which could limit its practical applicability.

The paper lacks a comparison to Mixture of Experts (MoE) approaches, which could offer insights into how FS-Merge differs or complements these strategies.

There is a concern about the robustness of FS-Merge to distribution shifts and its sensitivity to the order of models during initialization, which could affect the merged model's performance.

Based on the above points, I must reject this work due to concerns about its computational efficiency, lack of comparison to existing methods like MoE, and potential robustness issues.