LiNeS: Post-training Layer Scaling Prevents Forgetting and Enhances Model Merging

We thank reviewer tFK1 for their constructive feedback. We are pleased that they found our paper clearly structured and well-written, our evaluation thorough, and that our work demonstrates a strong performance tradeoff between target and control domains. We respond to their comments and questions below.

Theoretical analysis (W1 + Q1)

LiNeS is motivated by prior research and substantiated by strong and diverse empirical evidence. We acknowledge the reviewer's request for theoretical analyses but note that, given the scale of the models we consider (e.g., as large as LLaMa-7B), there are currently no formal mathematical tools or theoretical frameworks suitable for providing rigorous guarantees. In general, the field of model merging has yet no theoretical guarantees since the assumptions would likely be too restrictive to apply in the present practical settings.

Therefore, a theoretical analysis is beyond the scope of this paper. However, we are open to considering any specific theoretical directions from the reviewer for future work. It is also important to recognize that our method is well-motivated, following a long line of work in understanding neural networks and the effect of each layer, see paper references: Neyshabur et al., 2020; Yosinski et al., 2014; Raghu et al., 2019b, as well as multi-task model merging works, such as Ilharco et al., 2023; Yadav et al., 2023, Wang et al., 2024, among others. Given that this line of work does not have any theory and given our strong empirical evidence across an extensive and diverse list of settings, we kindly encourage the reviewer to reconsider their score.

Clarification on Task Combination (W2 + Q2)

For Table 1, we used a round-robin approach, testing all possible combinations of target and control tasks. Whenever one of the 8 tasks is selected as the target task, the remaining 7 tasks are set as control tasks. The final results in Table 1 are averaged over the 8 task combinations, covering all possible combination scenarios. It is exactly our intention to demonstrate the robustness and general applicability of our method by testing it on all possible combinations. We have added a sentence to the paper to emphasize this point.

Learned Scaling (W3+Q3)

The original submission already includes a detailed comparison with methods that learn scaling coefficients, which eventually display consistently increasing behaviors, as discussed in Section 6. Specifically, we evaluated Ada-Merging, an unsupervised approach that leverages entropy minimization to learn these coefficients. On the 8-task benchmark, Ada-Merging reported 81.1% averaged multi-task accuracy. However, this comes with significant trade-offs in resource requirements and unreasonable assumptions. In particular, Ada-Merging assumes that all task vectors are available at test time, breaking the model merging rationale. Moreover, it loads all task vectors simultaneously, resulting in a linear increase in memory consumption with the number of tasks. In contrast, LiNeS leverages insights on the hierarchical nature of learned features, constraining shallow layers to preserve general features. This allows our method to scale efficiently to large models like LLaMA (Section 5.4) or a large number of tasks (Section 5.2). Furthermore, defining an unsupervised loss introduces additional complexity, limiting its broader applicability, especially in language settings. In fact, Ada-Merging has been only applied to computer vision classification tasks.

Moreover, we also considered aTLAS, a contemporaneous work that learns scaling factors in a supervised manner using cross-entropy, which suffers from computational limitations similar to those of Ada-merging. Remarkably, as shown in Figure 6, the learned scalings for both methods demonstrate a clearly increasing trend, which is very close to those obtained without training by LiNeS.

We thank reviewer McJs for their constructive feedback. We are pleased that they found our proposed method novel, effective through extensive experiments, versatile across domains, and compatible with existing model merging baselines. We respond to their comments and questions.

Continual Training

We thank the reviewer for suggesting this potential area for exploration. In this paper, we focused on a static setting and demonstrated the effectiveness of LiNeS in mitigating forgetting of general knowledge after fine-tuning on specific tasks. Besides, as appreciated by the reviewer, we have conducted extensive experiments in demonstrating the effectiveness in many diverse problems in the literature, including OOD generalization, multi-task model merging, improving model soups, and merging LLMs fine-tuned with different rewards. Extending LiNeS to address long-term model maintenance and adaptation to evolving data distributions is beyond the scope of this paper while indeed an important direction for future work.

LLM experiments

We respectfully disagree on the lack of NLP results. In particular, we already tested against all the benchmarks explored by the field [1-3], containing tasks such as QASC, QuaRTz, PAWS, Story Cloze, WikiQA, Winogrande, WSC, CosmosQA, QuAIL, ROPES, and SocialIQA using the T5-large language model which is widely used in this literature [1-4]. Furthermore, beyond these previously explored tasks, we are among the first works in the model merging field to experiment on LLaMA-2 (with 7B parameters), following the setup of Rewarded Soups [5], and for which we fine-tune the checkpoints using PPO. Specifically, we apply LiNeS to improve the merging of two LLMs fine-tuned with 2 different rewards, leading to a merged model with a dominating Pareto-Front on both rewards compared to the model obtained with rewarded soups. The experiments in the paper are extensive, covering different applications and different modalities. We hope this clarifies the reviewer’s comment, and therefore, we encourage them to reassess their score.

LoRA and Fine-tuning baselines

We thank the reviewer for their suggestion. We note that LiNeS operates as a post-training editing method, making it orthogonal and compatible with standard fine-tuning methods such as LoRA. We have shown LiNeS’s compatibility and orthogonality with LoRA in the experiment for merging LLaMA-7B models (Section 5.4), where they are fine-tuned exactly with LoRA. We made this point more clear in the revised version of the paper.

Moreover, following the reviewer’s suggestion, we have added additional results for models fine-tuned with LoRA in the updated manuscript (Appendix C.8.4), with a focus on the mitigation of catastrophic forgetting. We have considered different hyper-parameter settings for LoRA. The results in Table 11 show that the models fine-tuned with LoRA also demonstrate degraded performance on control tasks to as much as 6.4%. Nevertheless, we further show that LiNeS can be applied to models after fine-tuning with LoRA to mitigate catastrophic forgetting. For example, for the models fine-tuned with LoRA with rank=128, LiNeS effectively improves the control task performance from 41.6% to 46.3% while preserving almost full performance in target tasks (87.2% out of 87.5%). The additional results demonstrate LiNeS’s orthogonality and compatibility with parameter-efficient fine-tuning methods like LoRA.

For convenience, we add the table regarding experiments with LoRA below, where $r$ indicates the different ranks:

Finally, as suggested by other reviewers, we have also added experiments with different fine-tuning schemes beyond standard and LoRA. Specifically, we considered layer-wise learning rate decays, linear and exponential, as well as freezing early layers of the network. The results are also presented in Appendix C.8 of the updated manuscript. It is important to note that modifying the fine-tuning process introduces multiple hyperparameters, and each experiment is computationally expensive and time-consuming. In contrast, LiNeS is flexible and efficient while achieving the least forgetting, as shown by the added results.

Thank you for your response. I have decided to maintain my rating score.

We thank reviewer KCXu for their constructive feedback. We are pleased that they find our paper well-motivated and easy to follow, and our results impressive. We respond to each weakness and question below. The additional experiments requested by the reviewer can be found in the updated manuscript in Appendix C.8.

OOD Experiments

The evaluation of OOD performance (see Section 5.1) is performed on test sets without searching for hyper-parameters. Thus, there is no influence from a validation set. It is important to note that the benchmark and evaluation pipeline are not our own and are widely used in the literature and, more importantly, our baseline [1] CVPR 2022 oral. Additionally, our paper does not solely focus on OOD generalization but showcases an issue found in many settings that our method resolves. Beyond OOD generalization, we also include experiments on multi-task model merging (vision and NLP), single-task model merging, and merging foundation models fine-tuned on different rewards.

Beyond Transformer Architectures

In all our experiments, we follow the experimental protocols established by prior works that only use transformer architectures, which currently dominate both the CV and NLP landscape. Specifically, §5.1 comes from [1], §5.2 from [3-5], §5.3 from [2], and §5.4 from [6] and [7]. The inclusion of so many diverse settings where our method improves performance highlights the generality of LiNeS.

Given the reviewer's suggestion and to further validate the robustness of our findings, we have explored different architectures, such as CNNs. ResNet-based architectures have been shown to not work well on task arithmetic scenarios (see, e.g., [9]), probably due to a lack of sufficient weight disentanglement [8]. Yet, we have added new results on ConvNext, following [8], in Appendix C.8.2. Remarkably, LiNeS improves performance in both mitigating catastrophic forgetting (similar to §3) as well as multi-task model merging scenarios (similar to §5.2). Given these new strong results addressing the reviewer's concern, we kindly ask them to reconsider their score.

Sensitivity analysis

We thank the reviewer for the suggestion. We have added new experiments focusing on the sensitivity w.r.t. the hyperparameters in Appendix C.8.1.

We emphasize that for multi-task merging, $\alpha$ is not tuned as a hyper-parameter but fixed via a heuristic to keep the computational cost the same as the baseline methods, as explained in §5.2. Therefore, we focus on the hyperparameter $\beta$ – the slope of the line – and compare its sensitivity to the uniform scaling hyperparameter $\lambda$ used by prior work such as Task Arithmetic, Ties, and Consensus. Our results, in Figure 14, show that LiNeS is less sensitive compared to the corresponding constant scaling coefficient of vanilla Task Arithmetic while improving performance.

Furthermore, we also conducted an experiment to test the sensitivity of the multi-task performance to both $\alpha$ and $\beta$, treating them both as hyper-parameters and presenting the results as a 2D heatmap in Figure 15, Appendix C.8.1. The heatmap clearly demonstrates the necessity for applying a layer-increasing scaling schedule for multi-task merging, as the optimal performances are obtained with $\beta>0$ for all three merging methods. In these new results, we also indicate the optimal configuration, which validates the heuristics previously introduced in the paper (cf. Eq. 2) to restrict hyperparameter search to a single parameter, i.e., the slope. Notably, with the latter choice, LiNeS maintains the computational complexity of the baseline methods.

Typos

We have updated the manuscript to correct the typos, we thank the reviewer for highlighting them.

Q1

This phenomenon is observed across the multi-task model merging literature [3,4,5,7]. The reason is that larger models are more weight-disentangled [8], i.e., the task representations in weight space are more disentangled and models can retain a larger portion of the performance before merging. For instance, vanilla task arithmetic [3] drops from 90.5% fine-tuning accuracy to 69.7% for ViT-B-32 and 8 tasks (see Table 2) while for 8 tasks with ViT-L-14 the drop is from 94% to 84%, and, therefore, any method including ours has less room for improvement over [3] with the larger models.

Importantly, we note that LiNeS outperforms the previous methods with larger models. For example, for ViT-L-14 architecture, Ties-merging improves Task Arithmetic respectively by 1.6%, 0.1%, 1.6% for the 3 vision benchmarks, while LiNeS improves Task Arithmetic by 2.5%, 3.0%, 3.1% respectively, and it further improves Ties-merging by 2.4%, 3.2%, 4.0% respectively.

Dear Reviewer,

Could you kindly respond and indicate whether authors have addressed your concerns?

Thanks, AC

We thank reviewer oncK for their valuable feedback. We are pleased they find our method well-motivated, well-benchmarked, and the paper easy to follow. Below, we address their specific comments and questions.

W1

LiNeS is a post-processing method that can be applied to any single- or multi-task vector, making it orthogonal to existing and future model merging methods, such as Ties and Consensus Merging. This is a major strength of our approach. As demonstrated in Tables 2, 3, and 5, our method delivers performance increases in all 36 benchmarks, covering the combinations of underlying methods, baselines, and models. These improvements are achieved on top of the baseline methods, including Ties and Consensus Merging. Furthermore, LiNeS is simpler than Ties and Consensus and often yields better results than them individually when applied to vanilla task arithmetic.

Our choice of increasing schedule is motivated by prior literature, which indicates that general features reside in the first layers, whereas task-specific features are predominantly found in the deep layers. This insight is corroborated by our empirical findings in Section 3, where we show that linearly scaling updates effectively mitigates catastrophic forgetting. We opted for linear scaling since it is the simplest design choice. While we also explored alternative scaling functions, such as square root and quadratic (Table 4 in the Appendix), we observed comparable performance. Investigating more sophisticated scaling choices – which would likely require more hyperparameters than LiNeS – is an interesting direction that we leave for future research.

W2

LiNeS is motivated by prior research and substantiated by strong and diverse empirical evidence. We acknowledge the reviewer's request for theoretical analyses but note that, given the scale of the models we consider (e.g., as large as LLaMa-7B), there are currently no formal mathematical tools or theoretical frameworks suitable for providing rigorous guarantees. In general, the field of model merging has yet no theoretical guarantees since the assumptions would likely be too restrictive to apply in the present practical settings.

Therefore, a theoretical analysis is beyond the scope of this paper. However, we are open to considering any specific theoretical directions from the reviewer for future work. It is also important to recognize that our method is well-motivated, following a long line of work in understanding neural networks and the effect of each layer, see paper references: Neyshabur et al., 2020; Yosinski et al., 2014; Raghu et al., 2019b, as well as multi-task model merging works, such as Ilharco et al., 2023; Yadav et al., 2023, Wang et al., 2024, among others. Given that this line of work does not have any theory and given our strong empirical evidence across an extensive and diverse list of settings, we kindly encourage the reviewer to reconsider their score.

Typos

Thank you for highlighting these issues – we have updated the manuscript to correct them.

Q1

The origin of the degradation of performance on control tasks is not just an assumption since (1) it is shown comprehensively in our motivating section 3 (Figures 1+2) and (2) is inspired by a long line of work showing similar observations, such as the paper references (Neyshabur et al., 2020; Yosinski et al., 2014; Raghu et al., 2019b).

Q2

On control tasks, we measure normalized accuracy with respect to the zero-shot performance of the model. All fine-tuned models (the blue marker points) in Figures 2 and 7 achieve less than 100% normalized control accuracy. Thus, there are no such cases as the ones mentioned by the reviewer where fine-tuning improves performance on control tasks. In fact, fine-tuning systematically decreases performance on the control tasks, referred to as catastrophic forgetting. It is only when applying our method on the fine-tuned checkpoints (in orange) that the performance is close to (or even surpasses) 100% on both target and control tasks. We also point out that, as suggested by reviewer KCXu, we have added new experiments showing the same results with convolutional architectures.

Q3

We note that all 70 models come directly from the published checkpoints of a previous work, Model Soups, fine-tuned with different hyperparameter configurations, available in their repo. Model 2 was trained with the largest learning rate and longest epochs, resulting in significant differences between the pre-trained and fine-tuned weights. As the weight interpolation literature suggests, these choices can break mode connectivity. Therefore, this issue originates from the specific hyperparameter choice and not our method.

Dear Reviewer,

Could you kindly respond and indicate whether authors have addressed your concerns?

Thanks, AC