Leveraging Submodule Linearity Enhances Task Arithmetic Performance in LLMs

TL;DR

The paper addresses merging multi-task fine-tuned LLMs without retraining by exploiting linearity at the submodule level. It formalizes a linearity criterion and a Non-linearity Score, demonstrates that layers, self-attention, and MLPs are more linear than full models, and derives a closed-form, per-module merging solution that uses minimal data. Empirically, layer-level and Attn/MLP-level merging outperform baselines across model scales and tasks (math, coding, translation), with notable gains when merging three models. This training-free, data-efficient approach offers a practical pathway to robust, multi-task LLM merging and prompts further exploration of submodule-level linearity for model composition.

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.

Figures (6)

• Figure 1: (a) Comparison of non-linearity in full models and submodules for fine-tuned models, measured using Non-linearity Score defined in Definition \ref{['din1']}. A lower value indicates better linearity, based on data from three fine-tuned models detailed in Section \ref{['model']}. (b) Comparison of non-linearity in full models and submodules within the context of model merging, assessed via Eq.\ref{['proj']}ctl. A lower value indicates better linearity, based on data from three merging combinations outlined in the caption of Figure \ref{['fig:coef']}. Error bars represent standard deviation.
• Figure 2: Comparison of Non-linearity Score in full models and submodules for models fine-tuned from tasks of Math, Coding, and Translate in Llama-2-7B and Llama-2-13B. A lower value indicates better linearity. Error bars represent standard deviation across different modules.
• Figure 3: Comparison of Projection Distances in full models and submodules within the context of model merging. A lower value indicates better linearity. For each backbone, we computed the results across all possible combinations of three fine-tuned models, alongside the corresponding merging weights $\alpha_1$ and $\alpha_2$ (and $\alpha_3$). In the case of merging two models, there are three combinations of fine-tuned models, yielding $25$ possible configurations for merging weights $\alpha_t \in [0.2, 0.4, 0.6, 0.8, 1]$ where $t\in [1, 2]$. For merging three models, there's only one combination of fine-tuned models, resulting in $27$ configurations for merging weights $\alpha_t \in [0.3, 0.5, 0.7]$ with $t\in [1, 2, 3]$. Error bars represent standard deviations.
• Figure 4: Comparison of Cosine Similarity in full models and submodules within the context of model merging. A lower value indicates better linearity. For each backbone, we computed the results across all possible combinations of three fine-tuned models, alongside the corresponding merging weights $\alpha_1$ and $\alpha_2$ (and $\alpha_3$). In the case of merging two models, there are three combinations of fine-tuned models, yielding $25$ possible configurations for merging weights $\alpha_t \in [0.2, 0.4, 0.6, 0.8, 1]$ where $t\in [1, 2]$. For merging three models, there's only one combination of fine-tuned models, resulting in $27$ configurations for merging weights $\alpha_t \in [0.3, 0.5, 0.7]$ with $t\in [1, 2, 3]$. Error bars represent standard deviations.
• Figure 5: A visual representation of the merging parameters obtained at different levels during the actual merging process is presented. Here, we showcase a result from the merging of three fine-tuned models. Additionally, only a portion of the head's merging parameters is displayed, while the distribution of the merging parameters for the other heads is similar to those presented.