Leveraging Submodule Linearity Enhances Task Arithmetic Performance in LLMs

Authors: Rui Dai, Sile Hu, Xu Shen, Yonggang Zhang, Xinmei Tian, Jieping Ye
Primary area: other topics in machine learning (i.e., none of the above)
Keywords: Large Language model, Task Arithmetic

Abstract:
Task arithmetic is a straightforward yet highly effective strategy for model merging, enabling the resultant model to exhibit multi-task capabilities. Recent research indicates that models demonstrating linearity enhance the performance of task arithmetic. In contrast to existing methods that rely on the global linearization of the model, we argue that this linearity already exists within the model's submodules. In particular, we present a statistical analysis and show that submodules (e.g., layers, self-attentions, and MLPs) exhibit significantly higher linearity than the overall model. Based on these findings, we propose an innovative model merging strategy that independently merges these submodules. Especially, we derive a closed-form solution for optimal merging weights grounded in the linear properties of these submodules. Experimental results demonstrate that our method consistently outperforms the standard task arithmetic approach and other established baselines across different model scales and various tasks. This result highlights the benefits of leveraging the linearity of submodules and provides a new perspective for exploring solutions for effective and practical multi-task model merging.

Decision: accept (poster)
Reviewer scores: 6, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1729594973483

thread_type:
official_review

<thread_object>
creator: reviewer_kqv8

<content>
summary:
This work is inspired by models demonstrating linearity that can improve the performance of task arithmetic, and hypothesize that linearity already exists within the model’s submodules. Linearity refers to the linear relationship between the differences in model weights and the differences in output features caused by fine-tuning. A model framework that independently merges these submodules is proposed to enhance task arithmetic and hence lead to superior multi-task models. The author justify their motivation by comparing the non-linearity scores between the full model and the submodules. Using submodules instead full models requires computation of the optimal merging weights after splitting the models into submodules that exhibit linear characteristics. Experiments are conducted on fine-tuned models based on Llama-2 and the results indicate superior merging performance with task arithmetic.

strengths:
The paper propses an interesting idea of decomposing the models into submodules that exhibit linear characteristics. The non-linearity score is proposed and merging the weights this way is novel. In particular, Figure 1 and Figure 3 clearly illustrate the motivation behind this paper and the weight merging is well-defined. The results and the ablation study support the author's claim well. The contribution of Layer level and Attention/MLP level weight merging is convincing.

weaknesses:
One relevant topic concerning linearity is the non-linear activation function. This paper focused on the submodules and weight merging, but lacks a thorough discussion on activation functions. Figures show that the non-linearity score for all different submodeules are near zero. Why is that? Have you discarded all non-linear activation functions in submodules? If that is not the case, how do they affect the linearity assumptions and why do they not contribute to a larger non-linerarity score? Please include a discussion of activation functions.

questions:
1. Merging weights in head level appears to give unstable outcomes in Table 3, and the hypothesis here is that the representation functionalities can be averaged out. Is there any potential method to avoid that?

2. The fine-tuned models are based on Llama-2, and the 7B model tends to show stronger non-linearity than the 13B model. Have you considered the influence of the varying depth of the model and data size, and the implications with repect to universal approximation theorem?

3. Is there any family of activation function more suitable for your model framework than others (e.g. ReLU, GELU, Swish) and what is their impact on the linearity of submodules and overall merging performance?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730620315381

thread_type:
official_review

<thread_object>
creator: reviewer_g3he

<content>
summary:
The paper tackles task vector-based model merging, following a line of approaches that obtain a multi-task model by adding to a pretrained model the differences between fine-tuned models and their pretrained base.  In particular, the work builds on top of previous work relating the effectiveness of task vectors to a particular kind of linearity such that if a model is taken on the line connecting pretrained and fine-tuned, its output features should be on the line connecting the features of the pretrained and those of the fine-tuned. The treatment starts by defining a non-linearity score that measures how much this linearity property is violated, showing this score to be much lower for individual layers than for the whole model. The authors then leverage this finding and propose merging individual layers, deriving a closed form for the layer weights based on the linearity requirement.  Experiments show the approach to be slightly beneficial for two LLAMA-based architectures when merging a model with math, coding, and translating capabilities.

strengths:
- The paper is generally clear and easy to read. The experiments make intuitive sense and are well explained, with enough details to be reproduced.
- The layer-wise linearity perspective is novel. Overall, I find layer-wise approaches to better leverage the structure of the network when compared with standard vanilla task vectors. The proposed non-linearity score is novel, well-motivated, and intuitive.

weaknesses:
- First and foremost, I am not convinced that the results warrant the added complexity: overall, the performance of the approach is in the same ballpark as vanilla task arithmetic while not inheriting its intuitiveness and simplicity. While it exhibits some slight improvements in a portion of the (limited) evaluation settings, these are not substantial nor consistent enough.
- Still on the evaluation, the paper only considers two LLAMA-based architectures and only a subset of the evaluation settings considered by the previous works. For instance, Task Arithmetics also considers the GLUE benchmark and a suite of computer vision benchmarks, while DARE considers GLUE and OpenLLM. An exhaustive evaluation is necessary to rule out possibly lucky choices of models and datasets that do not guarantee the effectiveness of the approach in general.
- The baselines are also somewhat limited, some possible works to consider are *e.g.* [1], [2] and [3].
- The main finding of the paper, that is the fact that modules show more linearity than the whole model, obtained as a composition of all those intermediate modules, is not surprising. In fact, if the intermediate modules have non-zero non-linearity scores, the overall non-linearity score will grow as these ones accumulate.
- Experiments have shown that merging three models exhibits a lower projection distance than merging just two, highlighting the potential of merging even more models. However, the experiments only merge up to three models. Can the trend be confirmed? More experiments would help clarify this aspect.

[1] Yadav, Prateek, et al. "Ties-merging: Resolving interference when merging models." *Advances in Neural Information Processing Systems* 36 (2024).

[2] Davari, MohammadReza, and Eugene Belilovsky. "Model breadcrumbs: Scaling multi-task model merging with sparse masks." ECCV 2024.

[3] Wang, Ke, et al. "Localizing Task Information for Improved Model Merging and Compression." *Forty-first International Conference on Machine Learning*.

questions:
- The requirement of a small set of data per task is reasonable in practice, but what if no data is available? Can you still do something?
- Maybe add a few more compatible tasks so that merging more than three models reveals some interesting trends as the number of tasks increases. Maybe your framework works even better compared to the baselines when the number of tasks is raised? Or perhaps the advantage becomes less pronounced?
- Minor suggestion: In Property 2, I would really avoid using footnotes over equations as these look like powers instead.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730657193175

thread_type:
official_review

<thread_object>
creator: reviewer_f2k5

<content>
summary:
This paper conducted a statistical analysis to examine the linearity of model-level and decomposed submodule-level components. The authors discovered that submodules (e.g., layers, self-attentions, and MLPs) demonstrate significantly higher linearity compared to the overall model. Based on this finding, they introduced a new model merging strategy that independently merges these submodules. Experimental results indicate that the proposed method surpasses the standard task arithmetic approach and other baselines across various model scales and tasks.

strengths:
1. The paper is elegantly written and easy to comprehend. The motivation, proposed method, and experiments are all well elucidated. 
2. The authors focus on an important issue in model merging, and proposed a principled approach to address this issue.
3. The authors offer detailed explanations of the proposed method and experimental outcomes in the appendix, which is highly commendable.

weaknesses:
1. The observation that submodules exhibit higher linearity than the overall model may be deemed rather obvious. Since the overall model is constructed from these submodules, the nonlinearity of each submodule contributes to the overall model's nonlinearity. Consequently, the comparison results presented in Figure 1 and Figure 2 are unsurprising. The linearity of the model is fundamentally inferior to that of its submodules.

2. Building upon the first comment, the proposed method could be interpreted as a straightforward extension of previous global linearization approaches. The authors introduce a layerwise linearization version, which has been explored in prior studies, such as Git Re-Basin [1], Zipit! [2], and MuDSC [3]. It would be beneficial if the authors could provide comparisons with these works.

3. Furthermore, it is conceivable to decompose the model into more finely-grained components, such as different groups, channels, or even neurons. I posit that the linearity of these components is more pronounced than that of layer-level submodules. However, this would necessitate a greater number of merging parameters, similar to the method proposed in this paper. Could the authors engage in discussions regarding this aspect?

4. The three properties outlined in Section 2 are somewhat redundant and could benefit from simplification.

5. The performance improvement is not so impressive. As the proposed method adopts a layerwise merging strategy, it is expected to be much more performant than others. However, experiments demonstrate that the proposed method is still outperformed by prior methods. For example, coding & translate, math & coding in Table 1 and Table 2.


### References
[1] Git Re-Basin: Merging Models modulo Permutation Symmetries, ICLR 2023.

[2] Zipit!: Merging Models from Different Tasks without Training. NeurIPS 2023.

[3] Training-Free Pretrained Model Merging, CVPR 2024.

questions:
Please see the Weaknesses
</content>
</thread_object>

</thread>
<thread>
creation date:
1730661314478

thread_type:
official_review

<thread_object>
creator: reviewer_sec4

<content>
summary:
This paper shows that by identifying submodules that exhibit linearity, more granular and effective merging can be achieved. After defining these submodules, they approximate the true objective with a linear interpolation, which holds when the submodules behave linearly. The objective is defined as minimizing the loss on the joint dataset by adjusting the weighting coefficients $\alpha$. This approach outperforms existing parameter merging methods, such as Task Arithmetic. They also perform insightful ablations to determine the optimal level of merging granularity—whether at the model, MLP/attention, or head level.

strengths:
Clear presentation, and interesting proposition of understanding which set of weights are more amenable to being merged. Novel idea of relaxing the original optimization with a linear approximation, which provides a fast and closed form solution. Good experimentation, and valuable ablation study on granularity of merging. The paper also shows impressive results with only 30 datapoints per task, making it highly efficient in terms of data and computational demands. It performs well across different scales like Llama-2-7B and Llama-2-13B, highlighting scalability for both small and large models. The approach is flexible, enabling merging at different granularities such as layer-level and attention/MLP-level, and delivers consistent performance across these structures.

weaknesses:
The authors overlook referencing Mixture of Experts (MoE) merging approaches, such as in https://arxiv.org/abs/2402.00433, where merging also involves linearly weighting submodules. However, in traditional MoE, this weighting occurs at the activation level, rather than directly on model parameters. Since the authors allow for optimization of the weighting coefficients on the joint dataset, then the idea in this paper can be a useful baseline to run: https://arxiv.org/pdf/2410.13025. 

That said, it’s unclear why the proposed method would surpass a router-based approach in performance, as a trained router can achieve task-specific specialization without approximations to the objective. While the linear interpolation method is efficient, a router's task-based dynamic routing could potentially allow finer-grained adjustments in response to specific inputs or task demands, leveraging the submodule structures without relying on assumptions of linearity. Adding comparisons or addressing these differences in the paper would clarify the potential advantages and trade-offs between the approaches.

questions:
An interesting direction for improvement could be incorporating second-order information into the optimization process rather than updating each weighting parameter independently. Currently, the approach adjusts the coefficients linearly, assuming each weighting parameter can be optimized in isolation. However, second-order information, such as Hessian-based adjustments or other curvature-aware methods, could potentially capture interactions between weighting parameters, especially when submodules exhibit interdependencies. This is why routing methods could be more powerful, as by iteratively optimizing each layer they can better account for submodule interdependency.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732624398400

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response About Revision

comment:
First of all, we would like to express our gratitude to all the reviewers for their time and efforts in reviewing our paper. The valuable comments from the reviewers have greatly contributed to strengthening our manuscript.

We have revised our manuscript based on the reviewers' comments and have uploaded it. Here, we summarize our modifications, following the original order presented on this webpage:

1. We have added content related to WEMoE [1] and CAT [2] and included specific comparative tables with further discussions in Appendix A.5.2. (Thanks to Reviewer SeC4)

2. We have provided a more detailed discussion on second-order information in Appendix A.5.3. (Thanks to Reviewer SeC4)

3. We have appropriately modified the introduction to more clearly clarify the linearly properties studied in this paper. Additionally, we have included a comparative table and related discussions in Appendix A.5.1. (Thanks to Reviewers F2K5, g3he, D)

4. We have incorporated descriptions of Git Re-Basin [3], Zipit! [4], and MuDSC [5] and included specific comparative tables with more detailed discussions in Appendix A.5.4. (Thanks to Reviewer F2K5)

5. We have added more discussions regarding the head level and included a more detailed discussion on finely-grained components in Appendix A.5.5. (Thanks to Reviewer F2K5)

6. We have merged Properties 1 (for the model) and Properties 2 (for the module) into more general Properties 1. (Thanks to Reviewer F2K5)

7. We have added the results and discussions of the ablation study on the number of models to be merged to Appendix A.1.3. (Thanks to Reviewers F2K5, g3he)

8. We have included experimental results involving more datasets (including Winogrande, Cosmos QA, and Ropes) and additional models (Qwen-2.5-7B) in Appendix A.1.2. (Thanks to Reviewers F2K5, g3he)

9. We have added experimental results involving more baselines (including TIES [6], Breadcrumbs [7], and Consensus TA [8]) to Appendix A.1.2. (Thanks to Reviewer g3he)

10. We have included the results and discussions of the ablation study on the number of data in Appendix A.1.4. (Thanks to Reviewers g3he, KQV8)

11. We have moved the footnote from the formula to the description below. (Thanks to Reviewer g3he)

12. We have provided a more detailed discussion on non-linear activation functions in Appendix A.5.6. (Thanks to Reviewer KQV8)

13. We have included a detailed discussion on possible methods to improve head-level performance in Appendix A.5.7. (Thanks to Reviewer KQV8)

14. We have added a detailed discussion on the impact of model depth on linearity in Appendix A.5.8. (Thanks to Reviewer KQV8)

15. We have included a detailed discussion on the phenomenon that the non-linearity score for Qwen-2.5-7B is much higher than others in Appendix A.5.9. (Thanks to Reviewer KQV8)


Finally, we would like to thank all the reviewers once again for taking the time to review our paper and providing us with your invaluable feedback. We sincerely appreciate your constructive comments on our work. 

Please feel free to ask any remaining questions as we would be glad to provide further clarification. We look forward to any further discussions with all the reviewers.


[1] Tang, Anke , et al. Merging Multi-Task Models via Weight-Ensembling Mixture of Experts. (2024)

[2] Prabhakar, Akshara , et al. LoRA Soups: Merging LoRAs for Practical Skill Composition Tasks. (2024)

[3] Git Re-Basin: Merging Models modulo Permutation Symmetries, ICLR 2023.

[4] Zipit!: Merging Models from Different Tasks without Training. NeurIPS 2023.

[5] Training-Free Pretrained Model Merging, CVPR 2024.

[6] Yadav, Prateek, et al. Ties-merging: Resolving interference when merging models. NIPS 2024.

[7] Davari, MohammadReza, and Eugene Belilovsky. Model breadcrumbs: Scaling multi-task model merging with sparse masks. ECCV 2024.

[8] Wang, Ke, et al. Localizing Task Information for Improved Model Merging and Compression. ICML 2024
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper shows that identifying linear submodules can lead to more granular and efficient model merging, and the experimental results demonstrate superiority over existing methods like Task Arithmetic, with valuable ablations on merging granularity. Overall, the reviewers find the presentation is clear, the idea is novel and effective,  and the experimentation and ablation study are of high quality. One concern is that the proposed method could be seen as an extension of prior global linearization work and the improvement is marginal. AC acknowledges the novelty, especially in decomposing models and the novel weight-merging approach. It is thus recommended for acceptance as a poster.