Fine-Tuning Attention Modules Only: Enhancing Weight Disentanglement in Task Arithmetic
Authors: Ruochen Jin
Bojian Hou
Jiancong Xiao
Weijie J Su
Li Shen
Venue: ICLR 2025 Poster (2025)
Keywords: task arithmetic
weight disentanglement
merging models
multi-task learning
Primary area: transfer learning, meta learning, and lifelong learning

Abstract: In recent years, *task arithmetic* has garnered increasing attention. This approach edits pre-trained models directly in weight space by combining the fine-tuned weights of various tasks into a *unified model*. Its efficiency and cost-effectiveness stem from its training-free combination, contrasting with traditional methods that require model training on large datasets for multiple tasks. However, applying such a unified model to individual tasks can lead to interference from other tasks (lack of *weight disentanglement*). To address this issue, Neural Tangent Kernel (NTK) linearization has been employed to leverage a ''kernel behavior'', facilitating weight disentanglement and mitigating adverse effects from unrelated tasks. Despite its benefits, NTK linearization presents drawbacks, including doubled training costs, as well as reduced performance of individual models. To tackle this problem, we propose a simple yet effective and efficient method that is to finetune the attention modules only in the Transformer. Our study reveals that the attention modules exhibit kernel behavior, and fine-tuning the attention modules only significantly improves weight disentanglement. To further understand how our method improves the weight disentanglement of task arithmetic, we present a comprehensive study of task arithmetic by differentiating the role of the representation module and task-specific module. In particular, we find that the representation module plays an important role in improving weight disentanglement whereas the task-specific modules such as the classification heads can degenerate the weight disentanglement performance.

Published as a conference paper at ICLR 2025
FINE-TUNING ATTENTION MODULES ONLY: ENHANCING WEIGHT DISENTANGLEMENT IN TASK ARITHMETIC
Ruochen Jin .1.2, Bojian Hou". Jiancong Xiao" H
Weijie J. Su², and Li Shen
1 East China Normal University, Shanghai, China
2 University of Pennsylvania, Philadelphia, PA, USA
(ruochenj, bojiahh jcxiao)eupenn. edu, www.eewaarto..ueeeeee 1 isseneupenn.co edu
ABSTRACT
In recent years, task arithmetic has garnered increasing attention. This approach edits pre-trained models directly in weight space by combining the fine-tuned weights of various tasks into a unified model. Its efficiency and cost-effectiveness stem from its training-free combination, contrasting with traditional methods that require model training on large datasets for multiple tasks. However, applying such unified model to individual tasks can lead interference from other tasks (lack of weight disentanglement). To address this issue. Neural Tangent Kernel (NTK) linearization has been employed to leverage a "kernel behavior", facilitat- ing weight disentanglement and mitigating adverse effects from unrelated tasks. Despite its benefits, NTK linearization presents drawbacks, including doubled training costs, as well reduced performance of individual models. To tackle this problem, we propose a simple yet effective and efficient method that is to finetune the attention modules only in the Transformer. Our study reveals that the attention modules exhibit kernel behavior, and fine-tuning the attention modules only significantly improves weight disentanglement. To further understand how our method improves the weight disentanglement of task arithmetic, we present comprehensive study of task arithmetic by differentiating the role of the rep- resentation module and task-specific module. In particular, we find that the rep- resentation module plays an important role in improving weight disentanglement whereas the task-specific modules such as the classification heads can degrade the weight disentanglement performance.
INTRODUCTION
The emergence of large pre-trained models in the open-source community has significantly ex- panded the potential to enhance performance on downstream tasks (Ilharco al. 2023 2022
Zhuang et al. 2020). align with human preferences (Lu et al. 2022 Glaese et al. 2022 Xiao et al. 2024 Li et al. 2025 Wang al., 2025). and improve robustness (Hou et al.) 2017 Ortiz-Jiménez
et al. 2021 Santurkar et al. 202 Tancik et al. 2020). However, traditional methods often in- volve expensive joint fine-tuning across multiple tasks (Zhuang al., 2020 and rely heavily on human feedback (Ouyang et al. 2022), which limits their scalability and broad adoption. Moreover, optimizing performance for specific downstream tasks usually compromises the model's initial pre-
training performance or zero-shot accuracy (French) 1999 McCloskey & Cohen 1989).
In light of these challenges, the necessity of task arithmetic in multitask learning has become increas- ingly evident. Task arithmetic offers a cost-effective and efficient alternative by enabling training- free combinations the weight space of pre-trained models without sacrificing the model's original capabilities IIlhacco al. 2023). Central to this approach is the concept of task vector, which represents a set of weight adjustments specifically calibrated for given task through fine-tuning,
"Equal Contribution
'The code is available at httt:a://gittbb.coo//tttet////
c_tangent.

---

Published as a conference paper at ICLR 2025
Table 1: Task arithmetic performance comparison between different methods. The task arith- metic performance (the average accuracy of the unified model over all the tasks) of our method outperforms the state-of-the-art due to both good performance of individual models and good kernel behavior (weight disentanglement).
Performance of Individual Models
Kernel Behavior
Task Arithmetic Performance
(Weight Disentanglement)
Ilharco et al. 2023
70.00%
Jitiz-Jiennz/ al. 2024)
76.26%
Ours
78.37%
Task Arithmetic
Weight Disentanglement
Individual Model
Nearly Equal Performance
Finetuning
0o
0g +
Task
Pre-trained Model
Unified Model
0g +
r
0o
00 +
Σ
ocT
6g + 72
Task
1-1
On
8
Task r
Figure I: Illustration of the concepts of task arithmetic and weight disentanglement. On the left-hand side, in task arithmetic, we first finetune the pre-trained model Og and get the finetuned individual model 0 + ayT where Tę is the tth task vector. We eventually obtain the unified model by adding all the task vectors to the pre-trained model: 9o + Σ Tr On the right-hand side, weight disentanglement means that the prediction of the unified model specific task will not be affected by other tasks.
obtained by subtracting the task-specific weights from the original pre-trained weights (Tharco et al., 2023). Each task vector encodes unique representational signature tailored to particular task.
As illustrated in Figure | (left), Oo is the pretrained model and Tt, t 1,...., is the ith task vector. The individual model on each task is derived by 0g + deTr, Task arithmetic is to add all the task vectors the pre-trained model to obtain the unified model Og + Σ 0yTe. A main goal in task arithmetic is to achieve weight disentanglement as shown in Figure [ (right) which means the prediction performance of the unified model such as the accuracy on a specific task will not be affected by other tasks. In other words, the unified model has nearly equal performance with the individual model.
The reason that we aim to achieve weight disentanglement is because task arithmetic primarily focuses on competing tasks rather than synergistic ones (Ilharco et al.) 2023), as the challenge lies in balancing and optimizing performance across tasks that may have conflicts or require different model behaviors. This emphasis on competing tasks crucial for developing robust multi-task models that can effectively handle diverse range of applications without performance degradation.
However, weight disentanglement remains the most formidable challenge in task arithmetic. Recent research (Ortiz-Jimenez et al., 2024) has shown that constraining models to fine-tune within the tan- gent space significantly improves weight disentanglement, thanks to the inherent kernel behavior of models during early fine-tuning stages. This kernel behavior, formalized by the Neural Tangent Ker- nel (NTK) theory (Jacot et al. 2018), refers to neural networks updating primarily around are-trained parameter initializations. While NTK linearization is effective, it compromises the performance of
2

---

Published as a conference paper at ICLR 2025
NTK linearization (Ortiz-Jimenezz al. 2024)
Better task arithmetic
Better performance individual model
Proves
Altention module exhibits kernel behavior
Fine tuning attention modules only
-inspires>>
Our work
Shows>
Enables
-Result
Kernel behavior leads to weight disentanglement
Greater weight disenlanglement
More efficient
Figure 2: Logic flow of our work.
individual models and demands two to three times more computational resources, conflicting with task arithmetic's original efficiency goals.
To tackle this problem, we hypothesize that a combination of non-linear fine-tuning and weight dis- entanglement is necessary. This hypothesis motivates investigate the following question: Does there exist a sub-module within neural networks where non-linear fine-tuning can exhibit weight dis- entanglement? We demonstrate that attention modules transformer models exhibit kernel behavior (see Figure 3 which is crucial for weight disentanglement. This finding is supported by empirical evidence using post-hoc linearization (Ortiz-Jimene" et al. 2024), technique that approximates the change in network output after training using first-order Taylor expansion. Building upon the insights above, we propose to fine-tune the attention modules only, which will improve the per- formance of individual models and efficiency while maintaining the weight disentanglement. This logic flow is illustrated in Figure
By focusing fine-tuning only the attention modules, approach significantly improves weight disentanglement and accuracy of the unified model compared to non-linear fine-tuning and NTK linearization (see Table 2). while substantially reducing computational burden and memory usage. This method offers a balanced alternative that maintains strong performance of individual models while enhancing weight disentanglement capabilities, providing practical solution for improving task arithmetic performance without sacrificing efficiency or accuracy.
To further understand how method improves the weight disentanglement of task arithmetic, present a study by differentiating the role of the representation module and task-specific module, while existing literature (Ortiz-Jimenez al. 2024) formulated task arithmetic using single model without clearly differentiating them. We conduct a comprehensive study of task arithmetic on pre- trained Vision Transformer (ViT) models like the Contrastive Language-Iman Image Pre- Training (CLIP) Radford et al. (021), providing new insights into its fundamental mechanisms and proposing novel methods to improve the performance of pre-trained models through task arithmetic. Specifically, we illustrate that the representation module plays an important role in improving weight disentangle- ment whereas this has been constrained by task task-specific modules, such as classification heads.
In particular, our main contributions are as follows:
We propose a simple yet effective and efficient method that only fine-tunes attention mod- ules, which improves weight disentanglement and the average accuracy of the unified model on all the tasks up to 2.38% improvement compared to the ofate-of-the-art meth- ods and 8.37% over the non-linear bascline on several
benchmarks.
We demonstrate that the attention module exhibits kernel behavior, suggesting that focusing fine-tuning these modules could enhance the weight disentanglement capabilities in task arithmetic while maintaining efficiency.
We reformulate the architecture of task arithmetic by separating the representation module from task-specific modules, revealing that while weight disentanglement mostly comes from the representation module, the effectiveness of task arithmetic constrained by task- specific components like classification heads.
2 PRELIMINARIES: TASK ARITHMETIC AND WEIGHT DISENTANGLEMENT
We begin by introducing the necessary mathematical notations. Let F: X x + y be a neural network taking inputs x X' and parameterized by set of weights i E O, which consists of
3

---

Published as a conference paper at ICLR 2025
representation module f(: 0) and task-specific module g(::w) where & = [0,w]. We assume X C R., 6 C In and } _ R. Consider T tasks, with every task consisting of a triplet (Di, Mt. F?), where D, S X is a data support (c.g., ImageNet (Deng al. 2009) images), 24 an imput distribution such that supp( = D., and F? Dr y a target function (e.g.. labels). In practice each task is identified with a training set {(re, F (xx)) ve(nn where F2 (.+) o(f(xv) 0f :)) with ft. that is used to fine-tune the representation modules starting from the pre-trained weights Pg and to obtain the fine-tuned weights 0r. while the task-specific modules are fixed at w/.
Having established this context, we can now introduce the key concepts of this work: task vectors and task arithmetic.
Definition (Task Vector and Task Arithmetic) Let Og denote the parameters of the pre-trained representation modules and 0g denote the parameters after fine-tuning on task t. The task vector To
for task is defined as: Ti = 0: 00- Task arithmetic is an operation of adapting a pre-trained model to T different tasks by modifying the pre-trained parameters 8g to the unified parameters aonitied as follows:
Ounified 0o +
Tr
where are scalar coefficients.
To distinguish from the concept of a unified model, we refer to the model fo, fine-tuned on specific task as an individual model. It is worth noting that adding a single task vector тą to a pre-trained yields model equivalent to this individual model fo.
model with coefficient
Non-linear Fine-Tuning, The initial idea of task arithmetic was introduced by Ilharco et al. (023). They demonstrated that performance of the unified model on multiple tasks could be im- proved simultaneously through task arithmetic. However, the performance of the model still lagged behind that of individual models specifically fine-tuned for particular tasks. In the remainder of this paper, we refer to this baseline method as non-linear fine-tuning, as the task vectors are obtained through this approach.
Accuracy Gap. We define the accuracy gap as the difference in accuracy on task t between the unified model and the corresponding individual model. A plausible hypothesis for this accuracy gap is that the task vectors for different tasks exhibit implicit conflicts with one another. To address the limitations non-linear fine-tuning, Ortiz-Jimenez et al. (2024) proposed that weight disentangle- ment is an important and potentially necessary condition for effective task arithmetic.
Definition 2 (Weight disentanglement) Given T different tasks and their corresponding supports D = (Dihtetry We say a set of lask vectors T {Te heary is weight disentangled with respect to parametric function f X x θ + y and the initial weights Oo if
T
T
cc On
arTt
=
E(x; (x E D.) + fxx; x x$UD. [E[T]
D.D.)
Dz
(I)
where f is the representation module.
Remark. Although the concept of weight disentanglement was originally proposed by Ortiz- Jimenez et al. al. (2024), our Definition differs from theirs in several key aspects.
Firstly, our definition characterizes weight disentanglement as a property of set of task vectors w.r.t. the function f, whereas it was originally defined as a property of the function f w.r.t. the task vectors. There are two reasons for reversed formulation. 1) The term "weight disentanglement" corresponds to T different weights (i.c., task vectors), rather than f. 2) Both Ortiz-Jimenee et al. (2024'ss 's and our approach aim to find better task vectors T, rather than finding better f.
Secondly, our definition applies specifically to the representation module, whereas the original def- inition encompasses the entire neural network. This is primarily motivated by the fact that task arithmetic is performed exclusively on the representation module. We will show later in Section 4 that this definition is well-defined: weight disentanglement emerges from the representation module.
4

---

Published as a conference paper at ICLR 2025
NTK Linearization Fine-tuning. Inspired by studies of NTK showing that very wide neural net- works behave similarly to linear functions around their initialization Bb. Ortiz-Jimménzz et al. (2021) proposed tne-tuning the task vectors post-hoe linear function, denoted fim- of 0o:
(an 8 7) [æ; Bu) (e
The obtained task vectors are denoted as Tim- The authors further demonstrated empirically that this approach yields an important property: Tiin is weight disentangled w.r.t. fin at Po-
To conclude this section, we define good task arithmetic performance as characterized by three key factors: high accuracy, computational efficiency. and good weight disentanglement.
3 TASK ARITHMETIC IN ATTENTION MODULES
3.1 MAIN CHALLENGE OF TASK ARITHMETIC
Based on our previous definition, the performance of the unified model on each task can be simply decomposed as follows:
Accuracy of unified models = Accuracy of individual models Accuracy Gap.
Unfortunately, existing studies reveal a trade-off between these factors. Non-linear fine-tuning can achieve high accuracy for individual models, but the resulting task vectors are less weight disentan- gled, leading to a larger accuracy gap in the unified model. Conversely, linear fine-tuning guarantees weight disentanglement and thus lower accuracy gap, but the linearly fine-tuned individual models tend to be less accurate, as linear approximation brings error to the models. This trade-off presents significant challenge in task arithmetic.
To this aim, we hypothesize that a combination of non-linear fine-tuning and weight disentangle- ment is necessary. This hypothesis motivates us investigate the following question: Does there exist a sub-module within neural networks where non-linear fine-tuning can exhibit weight disen- tanglement? Our investigation reveals that the attention module is promising candidate.
3.2 ACCURACY GAP: KERNEL BEHAVIOR AND WEIGHT DISENTANGLEMENT OF ATTENTION MODULE
As per the study of NTK studies and the dis- cussion in Ortiz-Jimenez al. (2024). if pre-
1.0
Post-hoc Finetuning Accuracy(%)
trained network f(: 0o) exhibits kernel behav-
0.9
ior during fine-tuning that is, if the neu- ral network updates primarily around its pre-
0.8
trained parameter initialization and can be ap- proximated by its first-order Taylor expansion then the resulting task vectors are weight dis- entangled. Directly examining whether kernel behavior can be challenging; however, it can approximated by the following test.
0.7
0.6
0.5
Post-hoc Non-linear (VIT)
Post-hoc vs. Non-linear (Attention) Equal Accuracy Line
0.4
0.4
0.5
0.6
0.7
0.8 0.9 1.0
Kernel Behavior Test. Given a function f with initial parameters do. and task with task vector Tr, we define the Kernel Behavior Test as follows. If the equation
Non-linear Finetuning Accuracy(%)
Figure Accuracy of non-linear and post-hoc models by tasks. The diagonal dashed line indi- cates post-hoe performance meets non-linear.
[xz 0o TH) (in( On + Tt),
holds for all in the dataset of task we say that f(: exhibits kernel behavior during fine-tuning using the given approach on task f, or more simply, that the given approach exhibits kernel behavior.
In experiments, compare the average accuracy across T tasks of the non-linear function (f()) and the host-hoc línear function (fin), fune-tuning on (1) all the parameters and (2) only on the attention modules. Specifically, we fine-tune several CLIIP pre-trained VITs (Dosovitskiy al.
5

---

Published as a conference paper at ICLR 2025
Trainable Parameters
Fixed Parameters
Linearized Part
Layer Norm
Layer Norm
Layer Norm
Feed Forward
Feed Forward
Feed Forward
Layer
Layer Norm
Layer
Attention
Attention
Attention
W
W
(9)
[D]
Figure 4: Three types of fine-tuning paradigms. (a) Non-linear fine-tuning where all the parame-
ters will be updated. (b) Full-model linearization. (c) Attention modules only fine-tuning where only
Wg. We, We and Wo will be updated. In this paper, we explore attention modules only fine-tuning.
2021 of different sizes following the same setup as ITharco et al. (0023) on 8 tasss: Cars Krause et al. 2013), DTD (Cimpoi et al. 2014), SUN397 (Xiao et al. 2016), EuroSAT (Helber et al. 2019).
GTSRB (Stallkamp et al. 2011). MNIST (LeCun 1998). SVHN Netzer et al. 2011) and RESISC45
(Cheng et al. 2017).
The results Figure indicate that the attention module demonstrates kernel behavior. The triangle dots show the comparison of the kernel behavior test between the attention modules (yellow) and the whole models (blue), respectively. The proximity of dots to the diagonal dashed line indicates kernel behavior. The post-hoc of attention module, represented by yellow dots, consistently appears closer to the diagonal dashed line than the whole ViT (blue dots), suggesting superior performance.
This indicates that attention modules demonstrate stronger kernel behavior compared the full model, suggesting that focusing on fine-tuning these modules could enhance the weight disentan- glement, finally resulting in low accuracy gap.
3.3 ACCURACY OF INDIVIDUAL MODELS WITH FINE- TUNING ATTENTION MODULES
Based on the kernel behavior of attention modules, we propose focusing on fine-tuning only the attention modules. The comparison of our fine-tuning paradigms with the non-linear fine-tuning paradigm and the NTK linearization fine-tuning paradigm is demonstrated Figure 4 Next, we will show that fine tineituning attention modules also achieves high accuracy for individual models.
Non-Linear Advantage. We will first introduce a crucial concept referred to as non-linear advan- tage. For given approach, the non-linear advantage is defined as the difference in accuracy of individual models between non-linear fine-tuning and the approach in question. Since non-linear fine-tuning typically achieves the highest accuracy for individual models, the non-linear advantage is always non-negative, i.e., non-linear advantage 2 0.
Accuracy of Individual Models. In Figure 5 we demonstrate that fine-tuning attention mod- ules can reduce the non-linear advantage-essentiaal improving the accuracy of individual mod- els cls-coepared to NTK lincarization fine-tuning The figure presents two comparisons: 1) Round markers represent the comparison between our method and non-linear fine-tuning. 2) Triangular markers show the comparison between NTK linearization and non-linear fine-tuning.
The proximity of markers to the diagonal dashed line indicates a non-linear advantage equal to zero. Our method, represented by round markers, consistently appears closer to the diagonal dashed line than the NTK linearization (triangular markers), indicating a smaller non-linear advantage. This visual representation demonstrates that our approach of fine-tuning attention modules achieves per-
Please Appendix for performance on each task.

---

Published as a conference paper at ICLR 2025
Table 2: Comparison of performance for task arithmetic across various visual models. This table presents the average accuracy (%) and normalized accuracy (%) of various ViTs after incor- porating the sum of task vectors from eight different tasks. We report results for the non-linear fine-tuning and NTK linearized models normalizing performance by their single-task accuracy.
ViT-B-32
ViT-B-16
ViT-L-14
Method
Abs.(f)
Norm.(t)
Abs.(t)
Norm.(t)
Abs.(1)
Norm.(t)
Pre-trained
48.40
55.25
66.40
Non-linear Fine-tuning
70.00
77.04
74.75
80.59
84.40
89.47
NTK Linearization 76.26
85.82
79.01
86.32
85.53
91.44
Ours
78.37
87.42
80.44
87.25
87.91
93.66
Table 3: Comparison of performance for task arithmetic across different language models.
This table presents the average accuracy (%) of Flan-TT models on GLUE benchmark.
glue-cola glue-mnli glue-mrpe glue-qqp glue-rte glue-sst2 glue-stsb avg
Non-linear Fine-tuning
79.87
80.94
59.31
82.19
50.54
89.33
70.55
73.25
NTK Lincarization
75.93
83.19
76.72
87.88
62.09
92.09
66.51
77.77
Ours
80.63
86.25
86.76
89.69
72.20
93.92
87.67
85.30
formance closer to that of non-linear fine-tuning compared to NTK linearization, thus reducing the non-linear advantage more effectively.
3.4 ACCURACY OF UNIFIED MODELS WITH FINE-TUNING ATTENTION MODULES
We have demonstrated that our method has achieved both high accuracy for individual models and solved the accuracy gap. Then we will validate that our method has achieved high accuracy for unified models in terms of average accuracy and normalized accuracy.
To obtain the task vectors, use the fine-tuned weights of different ViTs from before and use the to ensure fair comparison with
same mixing coefficient for all tasks, i.e., a 01 Ortiz-Jimenez al. (2024). We provide all the details of this experiment in Appendix
Normalized Accuracy. The normalized accuracy calculated by the individual accuracy achieved by the model fine-tuned on each task. Mathematically,
OurINTK lineeriation Finetuning Accuracy(%)
1.00
0.95
Jace (f(2;00 + TTT:))]
0.90
Normalized Acc- =
aac (f(x; 0o + a(T:))
Ours vs. Non-linear (VIT-B-32)
0.85
Non-linear (VIT-B-16)
Ours Nari linear (VIT-L-14)
Accuracy of Unified Models. We employ the benchmark proposed by Ilharco et al. (2023)
0.80
NTK Linearization vs. Non linear (VTTB332)
NTK Linearization vs. Non-linear (VIT-8-16)
evaluate the task arithmetic ability of pre-trained model, which consists of the tasks described in Section 3.2 In particular, Table 2 shows that non-linear
NTK Linearization vs. Non-linear (VIT-L-14
Equal Accuracy
0.70
0.70
0.75
0.80
0.85
0.90
0.95
1.00
Nan-linear finetuning Accuracy(%)
our method significantly outperforms counterparts (Ilharco et al., (023) and achieves state- of-the-art results on the task addition benchmarks. Our model achieves higher accuracy of the unified model through task addition (up to 2.38%). Addi- tionally, method not only outperforms averaged accuracy but also normalization accuracy.
Figure Averaged accuracy of non-linear and linear models. The diagonal dashed line indicates linear fine-tuning performance meets non-linear.
For Natual Language Processing (NLP) tasks, we utilize the Flan- T5 (Chung et al. 2022) as our pre-trained language model. For fine-tuning, employ the Flan-T5-bar models on seven tasks derived from the General Language Understanding Evaluation (GLUE) benchmark (Wang et al, 2018) with the same random seed 42 initialize the models. These tasks CoLA, MNLI, MRPC.

---

Published as a conference paper at ICLR 2025
Table 4: Efficiency comparison in terms of parameters, size, and training time. The NTK linearization requires two to three times more computational resources and doubles the memory footprint compared to its non-linear counterpart. Our method outperforms NTK linearization in accuracy with only quarter of the costs.
Computational Cost (MB)
Total Params Trainable Params
Training Time (Min)
Non-linear Fine-tuning
222
222 M
891.614
17:38
NTK Linearization
470 M
222 M
1,881.928
40:59
Ours
222
63.7
891.614
10:15
QQP, RTE, SST2, and STSB. We report accuracy for all tasks and an average accuracy in Table 3 method outperforms both lonerrn fine-tuning and NTK linearization in vision and NLP tasks.
Ablation Study. To investigate whether we enhance weight disentanglement perfor- mance by fine-tuning the Multiple-Layer Per- ceptron (MLP) modules in addition to atten- tion modules, we conduct an ablation exper- iment with four paradigms: (1) fine-tuning only attention weights (Q, K, V. and O projections) (ours), (2) tine-tuning attention
Table 5: Ablation study. Single-task and perfor- mance unified model (%) on different settings.
Single-task Accuracy(1)
Multi-task
Paradigm
Abs.(1)
Norm.(1)
(1)
89.55
78.37
87.42
weights and biases, (3) fine-tuming both at- tention and MI.P weights, and (4) fine-tuning attention and MI.P weights along with bi- ases. Remarkably, all four paradigms outper- formed NTK lincarization in terms of both performance and weight disentanglement, indicating that VIT models exhibit strong kernel behavior within the attention modules and MLP. However, per- formance varied based whether bias parameters were fine-tuned, with the best results aligning closely with settings used LoRA. This suggests that further exploration of these configurations could yield valuable insights into optimizing task arithmetic.
(2)
89.48
77.71
86.79
88.95
76.52
86.11
89.43
77.80
86.93
Efficiency Comparison. Besides superior accuracy, our method is much more efficient than non- linear fine-tuning and NTK linearization shown in Table to fine-tuning fewer parameters.
Our fine-tuning method significantly enhances the appeal of task arithmetic for practical applica- tions. By improving the performance of individual models, our approach demonstrates the superior- ity of task arithmetic in achieving good accuracy of the unified model efficiently. Additionally, we have observed that kernel behavior within attention modules fosters greater task disentanglement. In the subsequent section, we will delve deeper into this concept, exploring its implications and potential for future advancements.
4 ROBUSTNESS OF TASK ARITHMETIC WITH RESPECT TO COEFFICIENT
In previous sections, we discuss how to find a good T in task arithmetic. Yet the robustness of task arithmetic (i.e., the effect of a)) on weight disentanglement has not been explored. To this end, in the following section, we first propose a metric to evaluate the weight disentanglement performance called "disentanglement error" on the representation module and prove that weight disentanglement emerges from the representation module. Then, we illustrate that our method has great weight disentanglement for wider choice of for both representation and classification modules, which demonstrates the robustness of the task arithmetic of our method.
4.1 WEIGHT DISENTANGLEMENT EMERGES FROM REPRESENTATION MODULE
To explore the robustness task arithmetic and its effect on weight disentanglement, we propose a metric called "disentanglement error" to evaluate weight disentanglement performance. Unlike pre- vious work that focused on task-specific modules, we investigate whether the representation module satisfy Definition (Weight Disentanglement) without relying on task-specific components.
8

---

Published as a conference paper at ICLR 2025
Non-linear Fine-tuning
Lincarization
Ours
Repres entocton
0. 1.125
0.00
0.075
050
Classification
1.25
0.75
-1-
0.25
Figure 6: Visualization of weight disentanglement. The heatmaps show the disentanglement error E(01, (see Eq. (2) of a single representation module CLIP ViT-B/32 (top) and combination of representation module and classification module DTD SUN397 task pair. We use prediction error for classification task as Ortiz-Jimenez did (Ortiz-Jimenez et al. 2024). Three fine-tuning paradigms are used from left right: non-linear fine-tuning, NTK linearization, and ours. The blue
regions denote areas of the weight space where weight disentanglement is stronger.
Our hypothesis is that pre-trained models can demonstrate task arithmetic properties independently of downstream tasks, maintaining consistent representations through task arithmetic. By focusing on the representation module alone, we aim show that the inherent properties of pre-trained mod- els are sufficient to support task arithmetie, potentially simplifying the process and broadening its applicability.
To visualize the level of weight disentanglement, we measure the discrepancy with Eq. using the disentanglement error (Ortiz-Jimene) et al. 2024):
2
E(on, (2)
Σ Eawp, (dist (f(2; + aatt), + am (222)))
(2)
where "dist" denotes any distance metric between output vectors. As we are dealing with represen- tation distributions, in what follows we the Kullback-Leibhe divergence as the distance metric. In general, the smaller the value of E(a1, (()))02) the more weight disentangled a model is at (o1, Q2).
4.2 WEIGHT DISENTANGLEMENT RESULTS
Figure o displays the disentanglement error of CLIP VIT-B/32 model concerning several task vector pairs from different fine-tuning paradigms. We observe a minimal disentanglement error within a small region surrounding Og. which enables task arithmetic. Different from disentanglement error downstream tasks, remains relatively small even for a1 02 > 1, which indicates the power of task arithmetic has been limited by the performance -pecifcc modules (classification head).
Disentanglement Error Comparison. Our method demonstrates greater weight disentanglement than its counterparts, as evidenced by the more extensive regions with low disentanglement errors in Figure 6 (right). This explains the higher normalized accuracy achieved (cf. Table 2) when fine-tuning attention modules only. The combination greater weight disentanglement and better performance of individual models results in higher performance of the unified model.

---

Published as a conference paper at ICLR 2025
These results demonstrate that our method achieves great weight disentanglement for wider choice of a for both representation and classification modules, illustrating the robustness of task arithmetic in our approach.
5 RELATED WORK
Existing model merging techniques can be broadly categorized into two main types (Yang et al,, 2024]: (i) Pre-Merging Methods: These methods focus on enhancing the conditions necessary for effective model merging by optimizing the fine-tuning process of individual models. (ii) During Merging Methods: These approaches address task conflicts and interference through various strate- gies before executing the parameter merging operations.
Pre-Merging Methods. To establish better conditions for model merging, significant body of work has focused on the fine-tuning processes of independent models. For instance, Ortiz-Jimenez jtt al. (2024) propose fine-tuning within the tangent space of the pre-trained model, leveraging the NTK framework to enhance performance during the fine-tuning stage. However, fine-tuning all parameters in the linearized model can be computationally intensive compared to nonlinear fine- tuning. To mitigate this issue, some studies advocate for selectively linearizing only certain layers; for example, Tang et al. (2023) suggest partially linearizing Adapter modules prior to merging them. Additionally, TAFT (Liu et 2024) introduces an efficient linearization method specifically for Transformer architectures, deriving closed-form linearized solutions that facilitate smoother inte- gration of models. Overall, fine-tuning in the tangent space aids in disentangling both input and weight spaces, thereby reducing potential interference during subsequent model merging.
During Merging Methods. In the context of multi-task learning (MTL), model merging can be effectively achieved by employing various strategies to resolve task conflicts and perform parameter merging operations. Traditional methods often involve averaging or combining weights from multi- ple models to create a unified system, as demonstrated in prior works (Garipov et al.) [018]] [Iharco] jet al. 2023 Wortsman et al., 2022). However, these basic merging techniques frequently underper- form, particularly when tasks interfere with one another. Advanced methods have been developed to address this challenge by incorporating weighted -based strategies that assign different importance levels to task vectors during merging (Matena & Raffel 2021 Ainsworth et al.. [0223] Stoica al.) 2023 Yang et al. 2023). J. Furthermore, some approaches transform models into sparse subspaces be- fore merging, effectively mitigating task interference and allowing for the removal of insignificant neurons from individual models while enabling the combination of multiple sparse models within a parameter subspace (Yadav et al. 22023] Tam al. [023] Li et al.) 2023 Zhang al. 2023 Huh et al. 2024 Huang et al. 2024). This innovative perspective opens new avenues for model merging,
enhancing overall performance and flexibility in multi-task applications.
Our method falls into the Pre-Merging category, focusing on the fine-tuning process and achieving superior performance in task arithmetic with high efficiency.
6 DISCUSSION
In this work, we conducted a comprehensive analysis of task arithmetic in deep neural networks, un- covering its fundamental mechanisms and enhancing its performance. Our findings reveal that atten- tion modules exhibit kernel behavior, leading to improved weight disentanglement when fine-tuned exclusively, without compromising individual accuracy or efficiency. Crucially, we demonstrated that weight disentanglement emerges primarily from the representation module, while task-specific modules can limit the effectiveness of task arithmetic. This insight opens up new possibilities for applying task arithmetic in unsupervised learning scenarios and broadens its potential applications.
While our approach significantly advances the field of task arithmetic, several limitations and op- portunities for future research remain. Current task vectors are constrained to models with identical architectures and initializations due their reliance on element-wiss weight operations. Future stud- ies could explore integrating task arithmetic with partial fine-tuning techniques, focusing on varying numbers of attention blocks. Additionally, investigating the relationship between the sparsity of attention modules and their kernel behavior may provide insights into learnable tasks. Understand- ing the nuanced impact of fine-tuning bias on model performance and weight disentanglement also presents important avenue for future research. These investigations could lead to more robust and efficient methods for adapting pre-trained models to various tasks, significantly enhancing their applicability and effectiveness in real-world scenarios.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Samuel K Ainsworth, Jonathan Hayase, and Siddhartha Srinivasa. Git re-basin: Merging mod- els modulo permutation symmetries. International Conference on Learning Representations (ICLR), 2023.
Gong Cheng, Junwei Han, and Xiaoqiang Lu. Remote sensing image scene classification: Bench- mark and state of the art. Proceedings of the IEEE, 2017. URL https://ieeexplore leee. org/ http///cmmmnttttt5/....
Mehdi Cherti, Romain Beaumont, Ross Wightman, Mitchell Wortsman, Gabriel Ilharco, Cade Gor- don, Christoph Schuhmann, Ludwig Schmidt, and Jenia Jitsev. Reproducible scaling laws for contrastive language-imann image learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 2818-22882 2829, June 2023.
Hyung Won Chung, Hou, Shayne Longpre, Barret Zoph, Tay. William Fedus, Yunxuan Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, Albert Webson, Shixiang Shane Gu, Zhuyun Dai, Mirac Suzgun, Xinyun Chen, Aakanksha Chowdhery, Alex Castro-Ros, Maric Pel- lat, Kevin Robinson, Dasha Valter, Sharan Narang, Gaurav Mishra, Adams Yu, Vincent Zhao, Yanping Huang, Andrew Dai, Hongkun Yu, Slav Petrov, Ed H. Chi, Jeff Dean, Jacob Devlin, Adam Roberts, Denny Zhou, Quoc V. Le, and Jason Wei. Scaling cnrrrcctionfinntuned language models. arXiv preprint arXiv:2210. 1416, December 2022. http://: /ttp:///aaavvvvvi org/ abs/ 2210.11416
Mircea Cimpoi, Subhransu Maji, lasonas Kokkinos, Sammy Mohamed, and Andrea Vedaldi. De- scribing textures in the wild. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2014. URL http:accs///ecccssssssss thecvf. com/ content cww_evep____e evpr_2014/
html /Cimpoi hta.errrbbng__eeetrree________ 2014 2atpee___ii_21A___________P
.htm
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In IEEE Conference on Computer Vision and Pattern Recogni- tion (CVPR), 2009. URL
5206848
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recogni- tion at scale. In International Conference on Learning Representations (ICLR), 2021. URL https: http:r//////eeeeveeeee net. wtt.:e/ooooorrn-d//iiddddaa
Robert M French. Catastrophic forgetting in connectionist networks. Trends in Cognitive Sci- ences, 1999. URL http://wwwwwciieeeeeeeee.c....oc wtw.://eeeeeeecteeooccooooooooee
9134466139011299
Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry Vetrov, and Andrew Gordon Wilson. Loss surfaces, mode connectivity, and fast ensembling of dnns. In Conference Uncertainty in Artificial Intelligence (UAI), 2018. URLtttta:://arae :/tppii:///xxi....// org/ abs///b//// 1802. 10026
Amelia Glaese, Nat McAleese, Maja Trebacz, John Aslanides, Vlad Firoiu, Timo Ewalds, Mari- beth Rauh, Laura Weidinger, Martin Chadwick, Phoebe Thacker, Lucy Camphell-Gillinghan,, Jonathan Uesato, Po-Sen Huang. Ramona Comanescu, Fan Yang. Abigail See, Sumanth Dathathri, Rory Greig, Charlic Chen, Doug Fritz, Jaume Sanchez Elias, Richard Green, Sona Mokra, Nicholas Fernando, Boxi Wu, Rachel Foley, Susannah Young, Iason Gabriel, William Isaac, John Mellor, Demis Hassabis, Koray Kavukcuoglu, Lisa Anne Hendricks, and Geoffrey Irving. Improving alignment of dialogue agents via targeted human judgements. https: // w/www. deepmind. com/'
homm/bloo lolooo////o ooooogg/bi bbug//uiii lenildiingg swwer-aaeeeeeeee www.lgialllleeooo agents 2022.
Patrick Helber, Benjamin Bischke, Andreas Dengel, and Damian Borth. Eurosat: A novel dataset and deep learning benchmark for land use and land cover classification. Journal of Selected Topics in Applied Earth Observations and Remote Sensing. 2019. htthts:s: / //arxivv org//bgg/a.
T709. 00029
11

---

Published as a conference paper at ICLR 2025
Bo-Jian Hou, Lijun Zhang, and Zhi-Hua Zhou. Learning with feature evolvable streams. Advances in Neural Information Processing Systems, 30, 2017.
Chenyu Huang, Peng Ye, Tao Chen, Tong He, Xiangyu Yue, and Wanli Ouyang. Emr-merging: Tuning-free high-performance model merging. arXiv preprint arXiv:2405. May 2024.
Minyoung Huh, Brian Cheung, Jeremy Bernstein, Phillip Isola, and Pulkit Agrawal. Training neural networks from scratch with parallel low-rank adapters. arXiv preprint arXiv:2402.16828, Febru- ary 2024. Submitted on 26 Feb 2024, last revised 26 Jul 2024.
Gabriel Ilharco, Mitchell Wortsman, Samir Yitzhak Gadre, Shuran Song, Hannanch Hajishirzi, Si- Kornblith, Ali Farhadi, and Ludwig Schmidt. Patching open-vocabulary models by interpo- lating weights. In Advances in Neural Information Processing Systems (NeurIPS), 2022. URL
https:/ /apps://axxii: org/ abs 2208 05592
Gabriel Ilharco, Marco Túlio Ribeiro, Mitchell Wortsman, Suchin Gururangan, Ludwig Schmidt, Hannaneh Hajishirzi, and Ali Farhadi. Editing models with task arithmetic. In International Conference on Learning Representations (ICLR), 2023. URL
abs/
2110 0. 08207
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in Newral Information Processing Systems (NeurlPS), 2018. URL URL.ttps::/prrceeedieee
newrips cc/ wwtppper.p.paappeeeeeee.o happppapppp__a httppper_fflleepppeeeeeee Paper.
Jonathan Krause, Michael Stark, Jia Deng, and Li Fei-Fei. 3d object representations for fine-grained
categorization. In International Conference on Computer Vision Workshops (ICCVw), 2013. URL
https:/ http:////wwwww//m cv- www.ccccoooooaiooooooooccc org/ wwwpaccces/ccontttccooee content _ecpeettttiiccceeeeeeeeeeeeeeeee wwr.shoss.sss.....
2013/W1 http://mmm//aaueee3311 3D |r000//hhhhaaaaarnegooooooooeoee html
2tti:ooonnno011_ee 20C3_ICCV__________ orrp paper.
Yann LeCun. The mnist database of handwritten digits, 1998. URL http::/yan..
com/ hon//edb/mnist.
Weishi Li, Yong Peng, Miao Zhang, Liang Ding, Han Hu, and Li Shen. Deep model fusion: A
survey, 2023. arXiv preprint arXiv:2309.15698
Ziniu Li, Congliang Chen, Tian Xu, Zeyu Qin, Jiancong Xiao, Ruoyu Sun, and Zhi-Quan Luo. Entropic distribution matching for supervised fine-tuning of LLMSs Less overfitting and better diversity. In The Thirteenth International Conference on Learning Representations, 2025. URL https: wwwppeeeeeeew.eeeeeeeeee net/ www.inee//tttttn.e DarrumiddNNEee72
Tian Yu Liu, Aditya Golatkar, and Stefano Soatto. Tangent transformers for composition, pri- vacy and removal. In Proceedings of the International Conference on Learning Representations (ICLR), 2024.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Con- ference on Learning Representations (ICLR), 2019. URL
Sorum??d=BRgg6iio.
Ximing Lu, Sean Welleck, Liwei Jiang, Jack Hessel, Lianhui Qin, Peter West, Prithviraj Am- manabrolu, and Yejin Choi. Quark: Controllable text generation with reinforced unleam- ing. In Advances in Neural Information Processing Systems (NeurIPS), 2022. URL https: 7/arxiv. org/ abs 2205. 13636
Michael Matena and Colin Raffel. Merging models with fisher weightee averaging. Advances in Neural Information Processing Systems (NeurIPS), 2021. URL https::/aaxxi.... org/ abs/
09832
Michael McCloskey and Neal Cohen. Catastrophic interference in connectionist networks: The sequential learning problem. In Psychology of Learning and Motivation. Elsevier, 1989. URL https:/ /www. www.sscciiieeciceeeooee.eee wwt.cett.cccmeeeeeeeeeeeeeee http://ddeeee.cccccccccce//aatta ptttt/ccllccaaspp///
S0079 02800338
12

---

Published as a conference paper at ICLR 2025
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Ng. Reading digits in natural images with unsupervised feature learning. In Advances in Neural Information Processing Systems (NeurIPS) Workshops, 2011. URL https: storagge googleapis. ccomppb--ee tool s-publisse suublic-pulle c-publ cat ion data pttpp//ff3333/pppmmp/
G. Ortiz-Jimenez, A. Favero, and Frossard. Task arithmetic in the tangent space: Improved editing of pre-trained models. In Advances in Neural Information Processing Systems, volume 36, 2024.
Guillermo Ortiz-Jiménez, Apostolos Modas, Seyed-Mohsen Moosavi-Dezzooli, and Pascal Frossard. Optimism in the face of adversity: Understanding and improving deep learning through adversarial robustness. Proceedings of the IEEE, 2021. URL httts::///eeeeepe/eeeeeee
ieee.
Orc//oocuent/9934555
Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll L Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback, 2022. URL
org/ abse:/bs/2200000 2ttp//2223.02555555
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agar- wal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and llya Sutskever. Learning transferable visual models from natural language supervision. In Interna- tional Conference on Machine Learning (ICML), 2021. URL
org/ abs/
2103. 00020
Shibani Santurkar, Dimitris Tsipras, Mahalaxmi Elango, David Bau, Antonio Tor- ralba, and Aleksander Madry. Editing classifier by rewriting its prediction rules. In Advances in Neural Information Processing Systems (NeurIPS), 2021. URL
httpps://proceedingg www.innn...nuuuuuueeeeeeeeeeeeee neur wttpipppp.c.ppppppeeeeeeeeeeppse htcc/paaer__llesappeeeeeeee c4 2668a200aaa3oooo0331776668803332
Paper. ppd
Johannes Stallkamp, Mare Schlipsing, Jan Salmen, and Christian Igel. The german traffic sign recog- nition benchmark: A multi-class classification competition. In International Joint Conference Neural Networks (IJCNN), 2011. URL htttps://eeexxploree... feee. org document.
6033395
George Stoica, Daniel Bolya, Jakob Bjorner, Taylor Hearn, and Judy Hoffman. Zipit! merging models from different tasks without training. In International Conference on Learning Represen- tations (ICLR), 2023. URL https: httpvvaiivvv..oa http://abb233330000
Derek Tam, Mohit Bansal, and Colin Raffel. Merging by matching models task subspaces. arXiv preprint 23X22223/2.04339, December 2023.
Matthew Tancik, Pratul P. Srinivasan, Ben Mildenhall, Sara Fridovich-Keil Nithin Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan T. Barron, and Ren Ng. Fourier features let net- works learn high frequency functions in low dimensional domains. In Advances in Neural In- formation Processing Systems (NeurIPS), 2020. URL httpp::/prooceeiggg...... neurips.. cc/ cc/paper////ccc ht/paper/2222) http:///22220/111115555333333333 21t.1110000502253333333333331110
A. Tang, I.. Shen, Y. Luo, Y. Zhan, H. Hu, B. Du, and D. Tao. Parameter-efficient multi-task model fusion with partial linearizeation. In The Twelfth International Conference on Learning Representations, October 2023.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, pp. 353-355, Brussels, Belgium, 2018. Association for Computational Linguistics. doi: 11116553v1/W18--5
htLp: web. erg/ ant ho ogy / /1 5446
Mingzhi Wang, Chengdong Ma. Qizhi Chen, Linjian Meng. Yang Han, Jiancong Xiao, Zhaowei Zhang, Jing Huo, Weijie J Su, and Yaodong Yang. Magnetic mirror descent self-play preference optimization. In The Thirteenth International Conference on Learning Representations, 2025.
URL httpp:::ps............ ptttt:::://operreieeee net wwt.eeewweeeeoooorrroooooenesae
13

---

Published as a conference paper at ICLR 2025
Mitchell Wortsman, Gabriel Ilharco, Samir Yitzhak Gadre, Rebecca Roelofs, Raphael Gontijo- Lopes, Ari Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Korbblith, et al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without in- creasing inference time. In International Conference on Machine Learning (ICML), 2022. URL
htt.ps://arxiis org/ abs / 2203 05482
Jiancong Xiao, Ziniu Li, Xingyu Xie, Emily Getzen, Cong Fang, Qi Long, and Weijie Su. On the algorithmic bias of aligning large language models with rlhf: Preference collapse and matching regularization. arXiv preprint arXiv:2405. 16455, 2024.
Jianxiong Xiao, Krista A Ehinger, James Hays, Antonio Torralba, and Aude Oliva. Sun database: Exploring large collection of scene categories. International Journal of Com- puter Vision (IJCV), 2016. URL (ottaa:tt///ingg springer. ccm.aatticle/10.100077
11126-014-0748-3
Prateek Yadav, Derek Tam, Leshem Choshen, Colin Raffel, and Mohit Bansal. Resolving interfer- ence when merging models. In Advances in Neural Information Processing Systems (NeurIPS), 2023. URL/tttss:/aarxt //rxivv org/ abs /2000.2306.222 01708
E. Yang, Z. Wang. L. Shen, S. Liu, G. Guo, Wang, and D. Tao. Adamerging: Adaptive model merging for multi-task learning. In The Twelfth International Conference Leurning Represen- tations, October 2023.
Enneng Yang, Li Shen, Guibing Guo, Xingwei Wang, Xiaochun Cao, Jie Zhang, and Dacheng Tao. Model merging in Ilms, milms, and beyond: Methods, theories, applications and opportunities. arXiv preprint arXiv:2405. 174611 May 2024.
Jinghan Zhang, Shiqi Chen, Junteng Liu, and Junxian He. Composing parameter-efficient modules with arithmetic operations. In Advances in Neural Information Processing Systems (NeurIPS), 2023. URL w/wp/aaaavvv////////////aa///.aa org/ /as////////aaaa 22../.33......
Fuzhen Zhuang, Zhiyuan Qi, Keyu Duan, Dongbo Xi, Yongchun Zhu, Hengshu Zhu, Hui Xiong, and Qing He. A comprehensive survey on transfer learning. Proceedings of the IEEE, 2020. URL https://iit https::/ieeexplore leee org /t////mcmuuent/133300
14

---

Published as a conference paper at ICLR 2025
A EXPERIMENTAL DETAILS
All our experiments are performed using the same hardware consisting of four 3090 NVIDIA GPUs with 24GB of memory each, which can be reproduced in less than 150 GPU hours. The details of each experiment are the following.
Datasets. We evaluate task arithmetic on a set of popular benchmark datasets from various domains. The dataset collection includes:
SVHN (Netzer et al. 2011]: The Street View House Numbers dataset a real-world image dataset for developing machine learning and object recognition algorithms with minimal requirements on data preprocessing and formatting.
MNIST (LeCun 1998): A database of handwritten digits, with 60,000 training images and 10,000 testing images.
EuroSAT (Helber al., 2019); A dataset based on Sentinel-2 satellite images covering 13
spectral bands, with 10 classes and total of 27, 27,000 labeled and geo-referenced images.
RESISC45 (Cheng et al.) 2017): The remote sensing image scene classification dataset, consisting of 31,500 images in 45 scene classes.
Cars (Krause et al. 2013): This dataset contains images of cars categorized into various fine-grained classes. is widely used for fine tineigraiedd image classification tasks, providing a rich set of vehicle images for training and evaluation.
DTD (Describable Textures Dataset) (Cimpoi et al.) 2014): This dataset is designed for texture recogmition and categorization. It consists of texture images organized into 47 cate- gories, each labeled with attributes describing the texture patterns. It is commonly used to evaluate texture recognition algorithms.
SUN397 (Xiao al., 2016): The Scene UNderstanding (SUN) dataset is large-scale dataset for scene recognition, containing 397 categories with a total of over 100,000 images. It is used evaluate scene understanding models and benchmark scene classification algorithms.
GTSRB (German Traffic Sign Recognition Benchmark) (Stallkamp et al. 2011 : This dataset comprises images of German traffic signs, classified into over 40 categories. It is used to develop and evaluate traffic sign recognition systems, particularly in the context of autonomous driving and intelligent transportation systems.
Fine-tuning. All the fine-tuning experiments follow the same training protocol specified in Ilharco et al. (Ihhacco et al., 2022) with minor modifications to the training code to use linearized models when needed. In particular, we fine-tune all datasets starting from the same CLIP pre-trained check- point downloaded from the open_clip repository (Cherti et al. 2023). We fine-tune for 2,000 iter- ations with a batch size of 128, learning rate of 10 and a cosine annealing learning rate schedule with 200 warm-up steps and the AdamW optimizer (Loshchilov & Hutter 2019). As introduced in Ilharco et al. (Ilharco et al.] 2022), during fine-tuning, we freeze the weights of the classification layer obtained by encoding standard set of zero-shot template prompts for each dataset. Freezing this layer does not harm accuracy and ensures that no additional learnable parameters are introduced during fine-tuning (Ilharco et al. 2022). We use this exact same protocol to fine-tune the non-linear and linearized models.
Tuning of a in Task Arithmetic Benchmarks. As in Ilharco et al. (Ilaacco et al. 2022 we use a single coefficient to tune the size of the task vectors used to modify the trantedined models. or in Eq. [ In the task addition benchmarks, after
This is equivalent to setting a = a1 fine-tuning, we evaluate different scaling coefficients E (0.0, 0.1......... 1.0) and choose the value that achieves the highest target metric on a small held-out proportion of the training set as specified in Ilharco et al. (Ilhareo et al.) 2022). Namely, maximum normalized average accuracy, and minimum target accuracy on each dataset that still retains at least 95% of the accuracy of the pre- trained model on the control task. The tuning of a is done independently for non-linear fine-tuning, linearized fine-tuning, and post-hoe linearization.
Disentanglement Error. To produce the weight disentanglement visualizations of Figure 6 compute the value of E(a), a2) on a 15 x 15 grid of equispaced values in [-2, x [-2, 2T. To estimate the disentanglement error, use random subset of 2,048 test points for each dataset.
15

---

Published as a conference paper at ICLR 2025
B FURTHER EXPERIMENTAL RESULTS
We now present additional experiments that expand the findings discussed in the main text.
we report the single-task accuracy achieved by different CLIP models after fine-tuning with different approaches (referred as non-linear, NTK linearization, and our method).
Fine-tuning Accuracy. In Figure 7
Weight Disentanglement on Different Task Pairs. ment on different task pairs.
In Figure g
we illustrate weight disentangle-
Non-linear Fine-tuning
NTK Linearization
Ours
0.200
0.355
Representation
0.75
0.550
0.525
Classification
0.5
8
0.5
0.1
-1
Figure 8: Visualization of weight disentanglement. The heatmaps show the disentanglement error E(a1100 02) of a single representation module of CL.IP ViT-B/32 (top) and combination represen- tation module and classification module (bottom) Cars RESISC45 task pair. Three fine-tuning paradigms are used from left to right: non-linear fine-tuning, NTK lincarization, and ours. The light regions denote areas of the weight space where weight disentanglement is stronger.
Parameter Sensitivity Analysis. In our experiments, we replicate the experimental setup used by Ortiz-Jimenez et al. (2024) to evaluate the impact of varying coefficients on model performance. The results are summarized Figure 9 which demonstrates that our method exhibits greater robust- ness across different choices of a compared to both non-linear fine-tuning and NTK lincarization. As illustrated, our method consistently outperforms both non-linear fine-tuning and NTK lineariza- tion across wide range of values, indicating its robustness in maintaining performance even with varying coefficients.
Similarity between Task Vectors. Figure 10 shows the cosine similarity between task vectors from VIT for three types of fine-tuning (ef. Figure image classification tasks. Vectors from at- tention modules only fine-tuning are closer to orthogonal than those from both non-linear fine-tuning and NTK linearization, indicating that models fine-tuned with full parameters
more independent. This finding aligns with discussions in (Tharcoo et al.) 2023 Tang et al. 2023) and supported by the experimental results in Table 2 The experimental details described in Appendix A
16

---

Published as a conference paper at ICLR 2025
0.8
Training Accurecy
0.0
0.2
0.0
Cars
0T0
EuroSAT
GTSRB
MNIST
RESEC4S
SKW
SUN397
Dataset
(a)
VIT-B/32
Training Accuracy Dataset and Model
1.0
Non linear
NTK
Ours
0.8
Teining Accuracy
0.2
0.0
DTD
EuroSAT
GTSRe
MNST
RESISCAS
SVIN
SUN397
Dataset
(b) ViT-B/16
Non-lnear
0.8
Trainingg Accuracy
0.6
0.4
0.2
Cars
0T0
EurSsA
orsna
MUsT
RESISC45
Swint
SUN997
Dataset
(c) VIT-L/14
Figure 7: Single-task accuracy of different models obtained using different strategies on each of the tasks.
17

---

Published as a conference paper at ICLR 2025
Non-linear fine tuning NTK Unnarizattio
Ours
60
Performance
10
0.00 0.15
0.30 0.35
0.50
0.70 0.85 0.90 1.00
Figure 9; Performance comparison across different methods with varying a values.
Non-linear Fine-tuning
Linearization
Ours
0:03
9.35
0.33
300
3:55
0.5
0:0
0,30
3,66
0.00
2.0
0.01
8.00
606
0.21
0.22
HNNT
8.63
3.44
607
PAOT:
NESSONS 0,00
8:04
0,7
3.84
6:4
2.5
00
0.34
R250005
844
0,0
3
0.04
1.00
033
E.44
PASMUS
040
813
1.0
RADI, 003 281 6/6
924
6.16
021
0.2
043
3.5
6.99
037
6.0
5.0
9.90
1,00
0.0
14
640
0.34
130
303
6.00
6.4
0C
105
C41
R21
0.00
Figure 10: Similarity heatmaps. These figures show heatmaps of the cosine similarity between task vectors from task-specific CLIP models (Radford et al. (022)) fine-tuned different tasks. Three fine-tuning paradigms from left to right: non-linear fine-tuning, NTK linearization, and Attention modules only fine-tuning (Ours).
18

---

Published as a conference paper at ICLR 2025
C
IMPACT STATEMENT
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which feel must be specifically highlighted here.
19


Reviews:
[{"review_id": "l63QIsoUYy", "reviewer": "Reviewer_S2KF", "summary": "This paper presents a method that fine-tunes only the attention modules of models to enhance weight disentanglement in task arithmetic. Task arithmetic is a training-free approach that enables new tasks by linearly combining the weights of a task-specific model with those of the original model. Building on NTK studies on kernel behaviour, the proposed method introduces an approximate kernel behaviour test to verify that attention modules exhibit this characteristic. The authors validate their hypothesis by fine-tuning only the attention modules in CLIP models, showing that this approach outperforms other fine-tuning strategies as well as the original NTK-initialization method.", "questions": "See the weaknesses.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well-organized and clearly written. The authors thoughtfully introduce relevant prior work, providing clear definitions and context for task arithmetic, weight disentanglement, and their associated challenges. Each argument is well-supported, with compelling visualizations that reinforce the method's performance. The progression from motivation to methodology is logical and engaging.", "weaknesses": "The paper has no prominent weaknesses. Here are some other open questions:\n \n1. Contribution of Individual Attention Modules. \nWhile the method targets attention modules to achieve weight disentanglement, it remains unclear whether each module contributes equally to task performance. In large models, which may have over 100 billion parameters, understanding the relative contribution of each attention module could make the approach more efficient, possibly allowing for sparse fine-tuning of only the most impactful modules. An analysis or ablation study on the importance of individual modules could provide further insights into how weight disentanglement varies within the attention layers.\n\n2. Generalizability to Non-Attention-Based Architectures. \nWhile the method demonstrates promising results on attention-based models, it would be useful to discuss the potential for generalizing this approach to non-attention architectures, such as convolutional or Mamba-based models. Would the same strategy apply effectively, or would alternative fine-tuning techniques be required for architectures without attention mechanisms? Providing an outlook on this point could help extend the paper's relevance across a wider variety of model architectures.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WRtsGVmwm1", "reviewer": "Reviewer_CQUv", "summary": "The paper addresses task interference issues in task arithmetic, which arise when combining weights from multiple tasks into a single unified model. To tackle this issue, it uses Neural Tangent Kernel (NTK) linearization and trains only the attention layers in the transformer to reduce training costs. Through several experiments, the paper demonstrates the efficiency of the proposed approach and explains task arithmetic by distinguishing the roles of the representation module and the task-specific module.", "questions": "Refer to the weaknesses section.", "limitations": "", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 1, "strengths": "1. The paper is easy to follow and provides clear explanations of essential concepts, making it easier to understand both the proposed approach and the fundamentals of task arithmetic.", "weaknesses": "1. The paper appears to be a naive combination of prior work, specifically [1] and [2], with modifications focused on training parameter selection (e.g., training only the attention layers in transformers), which limits its novelty.\n\n[1] Multitask prompted training enables zero-shot task generalization (ICLR 2023)\n[2] Task arithmetic in the tangent space: Improved editing of pre-trained models (NeurlPS 2024)\n\n2. The main contribution claimed by the authors—asserting that weight disentanglement primarily arises from the representation module while the effectiveness of task arithmetic is constrained by task-specific components—is not well-supported. Additional ablation studies, analysis, or theoretical support are needed to substantiate this claim. Depending on interpretation, the proposed contribution may be a trivial insight already known in the fields of task arithmetic or multi-task learning.\n\n3. The experiments are quite limited, as many related works cited by the authors are missing from the experimental results, particularly in Tables 2 and 3. This omission reduces the validity of the proposed method.\n\n4. The paper's theoretical contributions rely heavily on [2], which limits its originality and overall contribution.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xNSMjQaAGf", "reviewer": "Reviewer_6SAc", "summary": "This paper proposes a new method to enhance weight decoupling in task arithmetic by fine-tuning only the attention module in the Transformer model. This method improves the weight decoupling ability of the unified model while maintaining the performance of individual models, thereby improving the performance of task arithmetic.", "questions": "While the paper mentions some related work, the relationship to existing methods could be discussed in more detail, especially the latest advances in model fusion and weight interpolation.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "Innovation: The proposed method of fine-tuning only the attention module is innovative in the field of task arithmetic and can effectively solve the weight decoupling problem.\nExperimental design: The experimental part covers multiple data sets and tasks, enabling a comprehensive evaluation of the effectiveness of the proposed method.\nThe results are remarkable: the method proposed in the paper achieves better performance than existing techniques on multiple benchmarks.", "weaknesses": "Parameter sensitivity analysis: The paper mentions the effect of the choice of α coefficient on performance, but does not provide a detailed sensitivity analysis. It is recommended to add experiments on the effect of the choice of α coefficient on model performance.\nComputational resource consumption: The paper mentions the efficiency advantage of the proposed method, but does not provide a direct comparison with existing methods in terms of computational resource consumption. It is recommended to add analysis in this regard, especially memory and time consumption in actual deployment.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0Ey3qB2Bif", "reviewer": "Reviewer_Rsur", "summary": "The authors propose to improve the performance of weight-interpolation-based model merging methods such as Task Arithmetic by finding a sub-module within the deep neural networks. I strongly agree with this idea because, for existing work based on NTK linearization, the model's performance on individual downstream tasks has been compromised in order to achieve better performance in multi-task model merging. This trade-off is usually unacceptable in practical applications, where ensuring performance on downstream tasks is paramount for fine-tuning. However, there are some typos and minor mistakes in this manuscript, which need to be addressed to improve the overall clarity and professionalism of the document.\n\nTherefore, I would like to assign this manuscript an overall score of 6, indicating that it requires minor revisions. However, **I am open to raising the score if the authors address the specific concerns and suggestions outlined in my detailed review.**", "questions": "1. The differences between Definition 2 in the manuscript and the original definition in Ortiz-Jimenez et al. (2024) remain unclear. Could the authors elaborate further on these differences?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The motivation of this study is clear. The authors seek to identify a subset of parameters that can undergo non-linear fine-tuning while retaining linear characteristics, such as those observed in NTK linearization. Consequently, the attention modules, particularly the linear layers within them, are selected for this purpose.\n2. The experimental results presented in Figure 3 demonstrate the validity of the proposed method.\n3. Models of varying sizes (ViT-B-32, B-16 and L-14) are employed for experimental evaluation.\n4. The authors conduct experiments on both the vision domain and NLP tasks.", "weaknesses": "1. The task arithmetic presented in Definition 1 (lines 167-175) deviates from standard practice. In the original paper [1], task arithmetic is defined as $\\theta_{unified} = \\theta_0 + \\lambda \\sum_{i=1}^T \\tau_i$, employing only a single scaling factor. Otherwise, it should be explicitly stated as task-wise AdaMerging [2].\n2. Missing scaling factor of task vectors in lines 52-53 and lines 353-357.\n3. Is the summation symbol missing on the right-hand side of Eq.(1) (lines 196-199)? \n$$ RHS = \\sum_{t=1}^T \\dots $$\n4. In lines 400-401, the \"disentanglement error\" was first proposed by Ortiz-Jimenez et al. (2024).\n5. Recent work on deep model fusion is missing from the related literature section. The authors may refer to these surveys [3,4,5] for additional information.\n6. A minor suggestion: The results for NLP tasks (Table 4) should be included in the main text.\n7. The ablation study in Appendix B.4 is not mentioned in the main text. Consider including it in the main text.\n\n[1] Ilharco et al. Editing models with task arithmetic. In International Conference on Learning Representations. ICLR 2023. \n[2] Yang et al. AdaMerging: Adaptive Model Merging for Multi-Task Learning. ICLR 2024.\n[3] Tang et al. FusionBench: A Comprehensive Benchmark of Deep Model Fusion. \n[4] Yadav et al. A Survey on Model MoErging: Recycling and Routing Among Specialized Experts for Collaborative Learning.\n[5] Yang et al. Model Merging in LLMs, MLLMs, and Beyond: Methods, Theories, Applications and Opportunities.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)