Scalable Model Merging with Progressive Layer-wise Distillation

We would like to express our sincere gratitude to the reviewer for their valuable feedback. The provided suggestions are extremely helpful and constructive, and we will revise the paper accordingly. We address the reviewer's questions as follows.

Q1: The experiments focus on vision and NLP tasks. Have you tested ProDistill on other domains, such as speech or reinforcement learning? If not, do you believe the method would generalize well to these areas?

A1: Thank you for the great question. Most existing works on model merging primarily evaluate their methods on vision and NLP tasks, with limited exploration in other domains. Actually, we can not even find a commonly used model merging benchmark for speech or reinforcement learning in our pretrain-finetune paradigm. Constructing such a benchmark would require significant effort in dataset selection, model training, and evaluation setup, which is unfortunately beyond the scope of the limited time available for the rebuttal.

However, we believe ProDistill can generalize well across different domains. The key reason is that ProDistill is inherently domain-agnostic—it does not leverage specific properties of vision or language data for its effectiveness. The only requirement is a layer-wise network structure, which is a common standard in modern deep learning.

Besides, we would like to highlight that the LLMs used in our experiments, such as WizardMath-13B, are actually trained using RLHF-based algorithms. This can partly give evidence that ProDistill applies to RL trained models.

Q2: Will you release the code and implementation details for ProDistill?

A2: Yes, the code and implementation details are already provided in the zip file of the Supplementary Material part on the openreview console.

Q3: While layer-wise distillation reduces memory overhead, are there any trade-offs in terms of performance or training stability compared to end-to-end distillation?

A3: Thank you for the question. We compare ProDistill with its end-to-end distillation counterpart, DistillMerge, in Appendix C.2. In our experiments, we do not observe degradation in performance or training stability when using layer-wise distillation.

Finally, we thank the reviewer once again for the effort in providing us with valuable and helpful suggestions. We will continue to provide clarifications if the reviewer has any further questions.

We would like to express our gratitude for the reviewer's helpful and positive comments. The suggestions provided have been instrumental in refining our work, and we will incorporate the necessary revisions accordingly. Below, we address each of the reviewer’s questions in detail.

Q1: Honesty, considering the success of data-agnostic methods and their empirical performance, I don’t think that these results are particularly motivating. ... Honestly, I think that these ablations and the results are much more insightful than the theoretical motivation. I suggest expanding them and finding space in the main paper.

A1: Thank you for your thoughtful suggestion. We acknowledge that our theoretical results focus on worst-case analysis and may not fully capture real-world practice. However, we find these results delicate and interesting, as they highlight aspects that are not widely recognized in the community, which is why we included them in the main text.

We agree the ablations in the appendix convey much information that we would like to share with the readers. We will restructure the paper to incorporate these results into the main text in future revisions.

Q2: Missing comparisons in table 1,2 for data-agnostic methods is TIES merging [1], while for the data-driven approaches, it is important to include also MaTS [2].

A2: Thanks for the suggestions. We add additional experiments to compare with TIES and MaTS. For TIES, we use the same hyperparameter grid as in our LLM experiments. For MaTS, we use the RegMean Objective and RegMean Initialization, and choose the same hyperparameter grid as that of RegMean in our paper.

The results are given below. Our method ProDistill outperforms the new baselines. An interesting finding is that TIES merging has a bad performance on NLP tasks, which coincides with previous findings in [1].

Q3: I would have expected that DistillMerge would be an upper bound on the performance of ProDistill. Could you elaborate more on why it is not so? See Fig.5.

A3: Thanks for pointing this out. We also find this result very intriguing. We hypothesize that ProDistill decomposes the overall training objective of DistillMerge, into more fine-grained layer-wise objectives. This decomposition makes it easier to minimize the objectives sequentially rather than optimizing the overall objective in one step.

Q4: Also in the right figure (Roberta exp) caption is DirectDistill. If this is not a typo, the caption should be updated.

A4: Thanks for pointing this out. We will correct the typo in the revision.

Q5: Could you please give more details on DirectDstill? Am I understanding correctly that this baseline is just distillation where you fine-tune the model parameters of the merged model, instead of the lambda coefficients?

A5: Yes exactly! The DirectDistill baseline uses $\ell_2$ distillation loss to fine-tune the original model parameters, instead of lambda coefficients.

Q6: Could you give more details on the "dual inputs", I don't think I understood if this is something that is particular to your method or it just identifies the input of the merged and the fine-tuned models. The way is introduced is confusing as it seems a technical contribution, but to me this is just an implementation detail of your method.

A6: Thanks for the question. The proposed "dual inputs" is an important implementation detail in our method, which has a big impact in the final accuracy (see Appendix C.4). Additionally, this design choice is often overlooked in activation-matching algorithms, and we highlight it to ensure awareness within the community.

Once again, we appreciate the reviewer's feedback and hope that our responses clarify the questions. We remain committed to improving the quality of our paper and welcome any further feedback.

References

[1] He Y, Hu Y, Lin Y, et al. Localize-and-stitch: Efficient model merging via sparse task arithmetic

We sincerely appreciate the reviewer’s valuable feedback. We follow the reviewer's advice and conduct additional experiments, with the results provided at https://anonymous.4open.science/r/Experiments-for-Reviewer-CuV6-8476. We address the reviewer's speficic questions as follows.

Q1: What happens if the fine-tuned models were trained under different configurations, such as varying architectures? ... How can this approach be extended to accommodate models with different embedding sizes?

A1: Thank you for this insightful question. Our method can accommodate training differences such as varying hyperparameters; however, it is not directly applicable to merging models with different architectures, as activation matching strategies may fail. The broader challenge of merging models with differing architectures, which we highlight in the Related Works section, is a valuable research direction but is beyond the scope of this paper. Notably, most existing methods based on task vectors struggle with such settings, as weight averaging fails when architectures differ.

Q2: Can the authors justify how your method handles alignment before merging models in such scenarios?

A2: Our method actually aligns the representation of each layer, which effectively addresses misalignment in the considered fine-tuning setup. In cases of severe misalignment, e.g. merging independently trained models, our method can be applied in parallel to existing alignment algorithms. Specifically, one can first apply an off-the-shelf alignment algorithm like Git-rebasin [1], and then apply our ProDistill to complete the merging process.

Q3: How does the proposed ProDistill method compare to multi-task learning (MTL) when both are trained on a validation dataset? What are the performance gap between these approaches?

A3: We have already provided such an ablation study in Appendix C.3, where we study an algorithm termed as DirectTrain, which conducts supervised multi-task training on the few-shot validation datasets. Our method significantly outperforms DirectTrain algorithm, as indicated in Figure 6.

Q4: It is worth noting that existing literature, such as the work on model alignment prior to merging, has explored layer-wise alignment techniques.

A4: Good point! The Git-rebasin paper [1] also explores a layer-wise activation matching strategy. The major difference is that

1. they use activation matching to learn a permutation matrix that aligns the neurons
2. we use activation matching to learn the merging coefficients, which directly determine the final merged model.

we have cited [1] and will make further discussions in the revision.

Q5: How do the coefficients for each layer and downstream task evolve across training steps?

A5: Thank you for this insightful question. We conduct additional experiments to track the change of merging coefficients . The results are provided in Figure 1~3 in the link.

We make several observations:

1. The mean of merging coefficients remains stable around its initialization 0.3.
2. The standard deviation of merging coefficicents keep increasing during training.
3. The cosine similarities between coefficients keep decreasing during training.

These results indicate that ProDistill captures the fine-grained (Obs 2) and task-specific (Obs3) variation within each module, rather than merely adjust the overall scaling (Obs 1). Therefore, unlike previous approaches that use scalar coefficients, our element-wise coefficients carry richer information and are unlikely to be easily predictable from data or model alone.

Q6: How do these coefficients relate to the correlation or similarity between downstream tasks?

A6: We provide additional experiments measuring the cosine similarity between merging coefficients across different tasks. The results, provided in Figure 4 in the link, reveal that the coefficients of different datasets have uniformly low cosine similarity of about 0.15. However, some correlation exists; for example, the coefficients for MNIST and SVHN (both digit classification tasks) have a slightly higher cosine similarity of 0.2.

Q7: Is it necessary to compute coefficients for all the layers, can we reuse several close layers' coefficients without computing them?

A7: Thanks for the suggestions. We follow the advice and use vanilla task arithmetic for layers where task vectors have smallest $\ell_2$ norm. The results are given in Table 1 in the link. It shows that skipping closer layers degrades performance. However, even when skipping 6 out of the total 12 layers, the average accuracy is still above 80%.

Once again, we sincerely thank the reviewer for their constructive feedback, and we are eager to engage in further discussions to clarify any concerns.

References

[1] Ainsworth S K, Hayase J, Srinivasa S. Git re-basin: Merging models modulo permutation symmetries[J].

We greatly appreciate the reviewer's comments and valuable suggestions. We conduct additional experiments to clarify the reviewer's question, with the results given in https://anonymous.4open.science/r/Experiments-for-Reviewer-ZoNc-9701. We address the reviewer's questions in more detail as follows:

Q1: While the results of ProDistill are consistently strong on vision and small-scale NLP tasks, the performance gains are less clear and more nuanced in large language model (LLM) tasks...Task Arithmetic, which does not require a few-shot validation dataset or additional training—achieves better results than ProDistill on certain tasks. Why does ProDistill not show a clear performance advantage over Task Arithmetic in these cases?

A1: This is a very good question. Task Arithmetic indeed has a good performance in our LLM experiments, which we summarize is due to the following reasons:

1. Task Arithmetic is more effective for larger models. In our vision experiments, its gap to fine-tuned models is over 20% for small ViT-B-32 models but narrows to about 10% for larger ViT-L-14 models. This trend aligns with findings from [1], which suggest that larger models exhibit stronger kernel behavior during fine-tuning and are closer to linear models (see Section D.1 in [1]).
2. Task Arithmetic works better for merging a small number of models. As shown in [2], increasing the number of models leads to greater task interference, reducing the effectiveness of methods that do not explicitly address interference. Our LLM experiments involve merging only two models, whereas our vision and NLP experiments merge up to eight models.
3. Task Arithmetic can have a less balanced accuracy across tasks, which gives a false sense of high performance. For example, when merging WizardMath-13B and Llama-2-13B-Code-Alpaca models (results given in Table 3 in our manuscript), task arithmetic shows slightly better performance than ProDistill on math related benchmarks (0.6467 vs 0.6279 on GSM8k), but significantly worse performance on code related benchmarks (0.0840 vs 0.2239 on MBPP). In other words, although Task Arithmetic may perform better on some tasks, ProDistill provides more consistent improvements, as reflected in the normalized average metric.

Q2: Additionally, regarding Figure 4 (left), which illustrates a clear positive relationship between validation set size and performance in vision tasks, can we expect similar trends to hold consistently for LLM tasks? Specifically, how do we explain or interpret the anomaly in Table 9, where 32-shot performance is higher than 64-shot performance

A2: Thank you for this question. We conducted additional experiments and found that the inconsistency arises from randomness in the sampling of validation data. To confirm this and further investigate data scaling, we repeated the experiments from Table 9 using three different random seeds, and extend the validation shot to 128. The averaged performance of ProDistill is given in the table below. The updated results demonstrate a clear positive correlation between validation shot count and merging performance.

To mitigate the effect of randomness, we rerun all the LLM experiments for ProDistill using three random seeds. The updated results can be found at https://anonymous.4open.science/r/Experiments-for-Reviewer-ZoNc-9701.

Q3: It would be helpful if the authors explained why a different set of baselines was chosen specifically for the LLM tasks compared to the other experiments.

A3: We change the baselines for LLM tasks due to scalability constraints. Many of the baselines in vision and NLP tasks are computation or storage heavy, making them unsuitable for merging LLMs. For example, Adamerging requries storing each of the fine-tuned checkpoints in memory, and Localize-and-Stich requires solving an end-to-end optimization problems. Our method ProDistill is designed to be scalable and efficient, making it well-suited for merging LLMs without such constraints.

We once again thank the reviewer for their valuable feedback and insightful questions. Please let us know if any further clarifications are needed.

Reference

[1] Ortiz-Jimenez G, Favero A, Frossard P. Task arithmetic in the tangent space: Improved editing of pre-trained models

[2] Yadav P, Tam D, Choshen L, et al. Ties-merging: Resolving interference when merging models