Task Arithmetic in the Tangent Space: Improved Editing of Pre-Trained Models
Authors: Guillermo Ortiz-Jimenez
Alessandro Favero
Pascal Frossard
Venue: NeurIPS 2023 oral (2023)
Keywords: model editing
transfer learning
neural tangent kernel
vision-language pre-training
deep learning science
TL;DR: We present a comprehensive study of task arithmetic and find that linearizing models before fine-tuning improves their performance after editing.

Abstract: Task arithmetic has recently emerged as a cost-effective and scalable approach to edit pre-trained models directly in weight space: By adding the fine-tuned weights of different tasks, the model's performance can be improved on these tasks, while negating them leads to task forgetting. Yet, our understanding of the effectiveness of task arithmetic and its underlying principles remains limited. We present a comprehensive study of task arithmetic in vision-language models and show that weight disentanglement is the crucial factor that makes it effective. This property arises during pre-training and manifests when distinct directions in weight space govern separate, localized regions in function space associated with the tasks. Notably, we show that fine-tuning models in their tangent space by linearizing them amplifies weight disentanglement. This leads to substantial performance improvements across multiple task arithmetic benchmarks and diverse models. Building on these findings, we provide theoretical and empirical analyses of the neural tangent kernel (NTK) of these models and establish a compelling link between task arithmetic and the spatial localization of the NTK eigenfunctions. Overall, our work uncovers novel insights into the fundamental mechanisms of task arithmetic and offers a more reliable and effective approach to edit pre-trained models through the NTK linearization.

Task Arithmetic in the Tangent Space: Improved Editing of Pre-Trained Models
Guillermo Ortiz-Jimenez" EPFL, Lausanne, Switzerland
Alessandro Favero" EPFL, Lausanne, Switzerland alessandro wtw.:ndrr...aaeeeep...... ch
guillermo. fealleemo.orriijaii.. orti2ji
Pascal Frossard EPFL, Lausanne, Switzerland pascal frossard@epfl......m ch
Abstract
Task arithmetic has recently emerged as cost-effective and scalable approach to edit pre-trained models directly in weight space: By adding the fine-tuned weights of different tasks, the model's performance can be improved on these tasks, while negating them leads to task forgetting. Yet, our understanding the effectiveness of task arithmetic and its underlying principles remains limited. We present com- prehensive study of task arithmetic in vision-language models and show that weight disentanglement is the crucial factor that makes it effective. This property arises dur- ing pre-ttrining and manifests when distinct directions in weight space govern sepa- rate, localized regions in function space associated with the tasks. Notably, show that fine-tuning models their tangent space by linearizing them amplifies weight disentanglement This leads substantial performance improvements across multi- ple task arithmetic benchmarks and diverse models. Building on these findings, we provide theoretical and empirical analyses of the neural tangent kernel (NTK) of these models and establish compelling link between task arithmetic and the spatial localization of the NTK eigenfunctions. Overall, our work uncovers novel insights into the fundamental mechanisms of task arithmetic and offers a more reliable and effective approach to edit pre-trained models through the NTK linearization.
Introduction
Pre-trained models play pivotal role in contemporary machine learning systems. However, to enhance performance on downstream tasks [38, 39, 92], align them with human preferences [31, 50, 64, 74], and increase robustness [63, 76, 87], they often necessitate further editing. Traditional model editing methods rely on costly joint fine-tuning across multiple tasks [92] and which constrain scalability and democratization. Furthermore, enhancing downstream task perfor- mance typically degrades the model's pre-training performance or zero-shot accuracy [28, 58, 87).
[64],
Recent research has introduced cost-effective and scalable model editing techniques that try to preserve the pre-trained model behavior by acting on the model weights through task arithmetic or weight interpolation techniques [3, 23, 27, 38-40, 46, 57, 71, 72, 79, 86, 87, 89], thus circumventing expensive joint fine-tuning over multiple tasks. These methods hinge on the observation that arith- metic operations between fine-tuned weights often produce analogous functional behaviors [39]. For example, summing the relative weight components of a model between pre- training and fine-tuning
Equul contribution.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Weight space
Weight spane
Weight space
[1:: osr))
[(a; Oy + ogg)))
f(x: Bu + onty + togrs)
8
6
8
D2
D2
D2
Figure 1: Illustaation of weight disentanglement, where distinct directions in the weight space, Ta, are associated with localized areas of the input space, Dr. This allows a model, f. to manipulate these areas independently by adding linear combinations of T; 's to pre-trained checkpoint On-
on two separate tasks results in a new multi-task model with improved performance on both tasks. Similarly, subtracting task's relative component can lead to the model forgetting that task.
Despite these advancements, the understanding of task arithmetic's underlying principles and its general effectiveness remains lirited.. Specifically, a comprehensive understanding of how these tech- niques affect a model's internal representations and the necessary conditions to make it reliable is lack- ing. This knowledge gap can undermine the adoption of these techniques, as it erodes their trustworthi- ness in real-world applications. In addition, reducing this gap could help us improve them even further.
To address these challenges, we present a systematic study of task arithmetic in contrastively pre- trained vision-language models (i.e., CLIP [69]), offering novel insights into its underlying mecha- nisms and introducing new approaches which enhance the performance of pre-trained models edited through task arithmetic. Specifically, we probe the hypothesis presented in Wortsman al. [87] that task arithmetic is possible thanks to the fact that models inherently operate in linear regime, where their behavior is dictated by the finite-width neural tangent kernel (NTK) [16, 41].
Our study reveals that linearized CLIP models exhibit significantly improved task arithmetic perfor- mance with respect to their nonlinear counterparts (see Tables and 2), but also that the NTK cannot fully account for the task arithmetic abilities of pre-trained non-linear models. Indeed, we show that the sole requirement for task arithmetic is actually weight disentanglement, where distinct directions in weight space correspond to changes of the network in disjoint regions of the input space (see Figure 1). This allows model perform task arithmetic by independently manipulating these weight directions.
Notably, we show that fine-tuning models in their tangent space by linearizing them amplifies weight disentanglement, leading to substantial performance improvements across multiple task arithmetic benchmarks and models. However, although weight disentanglement is stronger in the tangent space, it is also present in non-linear models. We demonstrate that weight disentanglement of semantically meaningful tasks is an emergent property of pre-training, as it is absent at random initialization.
In particular, our main contributions are as follows:
We formalize the notion of task arithmetic introduced in Ilharco et al. [39] as Property 1, allowing us to reason quantitatively about it.
We show that task arithmetic non-linear models cannot be explained by their NTK, and introduce the concept of weight disentanglement as the necessary condition to enable it.
We propose to linearize models as a way to enhance weight disentanglement and improve task arithmetic. Doing so, achieve up 5.8 points more and 13.1 points less in accuracy on task addition and task negation, respectively, on several lasnuanlanguage benchmarks.
We link weight disentanglement linearized models to spatial localization of the kernel eigenfunctions and validate this prediction numerically in pre-trained transformer models.
Finally, we show that weight disentanglement is an emergent property of pre-training
Overall, our work delivers new insights into the fundamental mechanisms of task arithmetic, facil- itating more reliable and scalable model editing. Our findings suggest that linearized fine-tuning of pre-trained models warrants further investigation, with the potential for substantial impact on effective model editing. These insights can foster the development of more efficient and precise model editing techniques, empowering practitioners to adapt pre-trained models to broader range of tasks.

---

2 Notation and problem statement
Let f X' x Q - y' be a neural network taking inputs x: E X' and parameterized by of weights E O. We will assume X' S RR, θ _ I and y S R Consider 7 tasks, with every task t consisting of triplet (Di, j4. ft) where D, C X is data support (e.g., ImageNet [21] images), P an input distribution such that supp(μ+) practice, each task is identified with a training set {(Ey, fax))) ] VVelng] with æ ~ jt,, that is used to fine-tune the models starting from the pre-trained weights On and obtain the fine-tuned weights 0,
Dr, and ft Dr + y a target function (e.g., labels). In
Task arithmetic. Let the task vector of task be the difference between the fine-tuned and the pre- trained weights, i.e., Ti = 07² -Og. The following property formalizes the notion of task arithmetic introduced in Ilharco et al. [39], where the authors observed that the accuracies of pre-trained models on different datasets can be modified independently through the addition or removal of task vectors. Property 1 (Task arithmetic). Consider a set of task vectors T = {Ti (TET with associated non- intersecting task supports D = (D, C hetr,, i.e., Vt, t', if & t' then Dj D., = Ø. We say network I satisfies the task arithmetic property around Ou with respect to T and D if
[[a;; 0o or T1) Ov)
EET Dr x & U Dr
f xx 00
Σ
04 Ti
(I)
with (ai,...., aT) E AC...
In short, model satisfies Property 1 if adding T. does not modify the output of the model outside D,
Neural tangent kernel. Around the initialization weights 0g neural network can be approximated with first-order Taylor expansion:
f(x; 0) se f(x; 0o) + (0 -0) 0o) ( (ef(<; 0o).
(2)
This approximation is equivalent to a kernel predictor with a kernel known as the neural tangent kernel (NTK) [41], kxTk(æ, x²) = VoJ fn) Vof(='; 0o). and defines a neural tangent space in which the relationship between weights and functions is linear. Remarkably, as the network width approaches infinity, Eq. (2) becomes exact and remains valid throughout training [4, 41, 45].
However, this linear approximation is often invalid at finite widths, the evolution of parameters dur- ing training is inadequately captured by Eq. (2). In such cases, training occurs in a non-linear regime. Conversely, often during fine-tuning, parameter evolution in many pre-trained models is frequently minimal, meaning that training does not exit the tangent space and Eq. (2) closely approximates the network behavior [22, 53, 62, 90, 91]. In such cases, training occurs in linear regime.
3 Task arithmetic is not a consequence of linear fine-tuning
The objective of this work is to understand the conditions that enable task arithmetic in deep neural networks. Previous studies hypothesized that task arithmetic results from fine-tuning in the linear regime [39, 86, 87], as linear weight combinations correspond similar output function combinations. However, we will now demonstrate that CLIIP models do not fine-tune in the linear regime and we therefore need other ways to explain task arithmetic.
In general, if a pre-trained network f(: 0o) demonstrates kernel behavior during fine-tuning - i.e., fine-tuning occurs in the linear regime- the following property must be satisfied [53]: Property 2 (Post-hoc linearization). The change in the network output after training can be approxi- mated by its first-order Taylor expansion, i.e., [x:; 0*) f(x 0o) (0- 0p) Vof(a; 0p).
In simple terms, the approximation of the network in the tangent space around initialization must hold after fine-tun tuning. To test this, we evaluate the performance of the post-hoc linearized version of f. fun That is, we apply the fine-tuned task vectors 7 0 Ou to the linear approximation of 0o, ie..
fz;( 0, + -) f(x; 0.) T Vof 6o),
and we check whether fin(:; 0*) performs similarly to ʃ(: 0²)
'The code to reproduce our experiments can be found at https: /// htt///gttthuubb http://gttuuuuumogrrreeeee comm/oorizji/tangntt task arririitii thmetic.
3

---

Table 1: Task addition. Average absolute (%) and normalized accuracies (%) of different CLIP ViTs edited by adding the sum of the task vectors of 8 tasks. We report results for the non-linea and linearized models of Sections and normalizing performance by their single-task accuracies.
ViT-B/32
ViT-B/16
ViT-L/14
Method
Abs. (1) Norm. (1)
Abs. (f)
Norm. (t)
Abs. (t)
Norm, (f)
Pre-trained f(::
48.4
55.2
64.4
Non-lin, FT f(:; 0, + +T)
71.4
76.5
75.5
80.0
85.1
88.8
Post-hoe lin. fin(-; 9o + F)
57.1
81.9
65.0
85.2
75.2
90.0
Linear. FT (:i(-:: 0o + Tim)
76.5
85.4
81.3
86.0
88.5
93.5
Table 2: Task negation. Minimum accuracy (%) of different CLIP ViT' edited by negating task vec- tor from a target task while retaining 95% of their performance on the control task. We report average performances over eight tasks non-linear and linearized models as introduced in Sections 3 and 5.
ViT-B/32
ViT-B/16
VYT-L/14
Method
Targ. (1)
Cont. (1)
Targ. (1)
Cont. (t)
Targ. (L)
Cont. (1)
Pre-trained f(:; 0u)
48.4
63.4
55.2
68.3
64.4
75.5
Non-lin. FT f(: O -T)
24.0
60.7
19.2
64.6
18.0
72.5
Post-hoc lin. (i ; Ou T)
14.8
60.3
10.8
64.8
12.1
71.8
Linear. FT fual:; (-; Ou- Tim)
10.9
60.8
11.3
64.8
7.9
72.5
The results in Figure 2 indicate that CLIP mod- els do not exhibit kernel behavior. Specifically,
100
we fine-tune (FT) several CL.P pre-trained
Non-lmoor finestuning accurmey
Vision Transformers (ViTs) [24] of different
Notrlineer advinatuge
sizes following the same setup Ilharco
80-
et al. [39] on tasks: Cars [43], DTD [20].
SUN397 [88], EuroSAT [33], GTSRB [80].
MNIST [44], SVHN [60] and RESISC45 [15].
We observe that the single-task performance of fim(-; 0*) is significantly lower than that of f(::0" for ViTs of all sizes. This non-linear advantage [26] is a clear sign that fine-tuning has not happened in linear regime expected by Wortsman et al. [87].
60-
VTT-B/32
WT-B/16
50
WTT,114
50
70
80
100
Post-hoc linearization accuracy (%)
Figure 2: Non-linear advantage. Single-task accuracies
Yet, this observation is not enough to rule out that task arithmetic can be explained by linearizing the network function. Indeed, even the non-linear components are important for single-task performance, they might not be used during task arithmetic,, which is the objective of this study. That is, the projection of f onto the tangent space could be the only useful component.
non-linearly fine-tuned mod- els f(:: O*) and their post-hoc lincarization
fun(: 0*). Markers represent different ViTs.
We now show this is also not the case, doing task arithmetic with the non-linearly fine-tuned task vectors over fin significantly decreases performance. To show this, we employ the benchmark proposed in Ilharco et al. [39] to evaluate the task arithmetic ability of pre-trained model, which consists of the tasks described before and two sub-benchmarks:
1. Task addition:
The sum of the task vectors T Σ, T. is added pre-trained checkpoint to produce a multi-task model. The success of this benchmark is measured in terms of the maximum average accuracy over the different tasks. Results shown in Table 1.
2. Task negation: A task vector is subtracted from the pre-trained checkpoint to forget a task while retaining performance on control task (ImageNet). The success of this benchmark is measured in terms of the maximum drop accuracy on the forgetting task that retains the performance the control task. Results are averaged over tasks and shown in Table 2.

---

To obtain the task vectors, we use the fine-tuned weights of the different ViT' from before, and use single mixing coefficient a aj =.... at optimized separately for the non-linear and post-hoe lin- earized models to ensure fair comparison. We provide all details of this experiment in Appendix A.
The results in Table confirm that task arithmetic in CLIP models does not stem from the combination of their linear components only. Specifically, we observe a significant drop in absolute task addition accuracy the post-hoc linearized models compared to the non-linear ones. This decrease in performance is consistent across tasks (see Appendix D.2) and highlights that task arithmetic non-linear models leverages the non-linear components of f, as well.
Although these results reject the linear hypothesis, is still remarkable that the post-hoc linearized models do better at task negation than the linelr ones (see Table 2). Furthermore, even in task ad- dition (see Table 1) they achieve higher normalized accuracies (see definition in Appendix A). Indeed, we formalize in Section 4, this observation suggests that linearized models are more consistent with Property 1. In Section 5, we will use this fact to devise a new way to enhance task arithmetic.
Weight disentanglement
If the linear regime is not necessary explain task arithmetic, what are the necessary conditions that allow it? In this section, we argue that the only necessary condition to perform task arithmetic with a model f is that the model is weight disentangled with respect the set of fine-tuning tasks.
Property 3 (Weight disentanglement). A parametric function funccttooi f X x θ + Y' is weight disentangled with respect to set of task vectors = (Te) and the corresponding supports D = (De)eE() if
(
T
T
x;
0o +
Σ
THT
=
Σ
De (ar (..TT) + go(ar),
(4)
1=1
1=1
where (.())))TT)) = 0 for x E De and 1. 1..... T, and go = 0 for x E UrE!T Dr-
In essence, this definition captures the idea that the function f can be decomposed as a sum of spatially- localized components, i.e., vanishing outside a spatial region, whose functional variation is entirely captured by each Ti (see Figure 1). Morcover, is trivial to see that satisfying weight disentanglement is equivalent to satisfying Property 1 on task arithmetic one always write Eq. (1)
T
T
f
1x;00
+
Σ
OuTT
=
Σ
[(x; 0o aZT)) E D.) + 0o);
x &
RUUZ LE[T]
Dr
(5)
and identify (e(2;aaT) = f(x; + a E D:) and (z) f(x; 1o) & Dz). It is
important to highlight, however, that this additive decomposition does not imply linearity, as the local functions {g. he are not required be linear with respect the parameters.
Furthermore, note that weight disentanglement is a property of the predictors and not related to the performance different tasks. That is, model could be weight disentangled with respect to a set of task vectors and still perform poorly on a task, e.g., if f(: 0, + aT) does not generalize for some a. More generally, we can visualize the level of weight disentanglement of model by measuring its dis- crepancy with Eq. (4). To do so, given two tasks, one can check the disentanglement error of a model,
2
((01,02)
Σ Eunapi (dist (a Oo + OuTT)) (ax; Og + a1 + (222)))
where dist denotes any distance metric between output vectors. As we dealing with classification tasks, in what follows we use the prediction error dist(y) y2) = (1) * y2) as the distance metric. In general, the smaller the value of E(a1, the more weight disentangled a model is at (a1, a2).
Figure displays the disentanglement error of a CLIP ViT-B/32 model concerning several task vector pairs. We observe that the CLIP model exhibits a minimal disentanglement eтror within small region surrounding Og. which enables task arithmetic. However, for an 02 > 1, the error increases, indicating a high degree of interaction between tasks. This explains why task arithmetic performs better in small neighborhood of 0g task arithmetic is more effective when fine-tuning with small learning rates and few training steps [39]- with the optimal value of typically being less than 1.
5

---

EuroSAT 9VHN
DTD SUN397
Cuz- RESISC45
3.0
3.0
3.0
100%
Not- Lin-li
1.0-
10-
1.-
-3.0
100
-3.0
-30
1.0
055
3.0
3.0
3.0
100%
Post-
62
0
-1.0
-..
-1.0
-3.0
(2g
Figure 3: Visualization of weight disentanglement. The heatmaps show the disentanglement error E(a1,a)) of a non-linear CLIP ViT-B/32 (top) and its post-hoc linearization (bottom) on different example task pairs. The light regions denote areas of the weight space where weight disentanglement is stronger. The red box delimits the search space used to compute the best a in all our experiments.
Comparing the disentanglement error of the non-lineat models and their post-hoc linearization reveals interesting finding: linearized models exhibit greater disentanglement than their non-linear counterparts. This evident from the more extensive regions with low disentanglement errors in Figure 3 (bottom). This explains why the post-hoe linearized models achieve higher normalized accuracies via task addition (cf. Table 1) and manage to forget more through task negation (cf. Table 2). Paradoxically, however, although the greater disentanglement of linearized models allows them to retain more of their relative performance when edited with task arithmetic, they still perform worse in absolute terms due the great advantage of the non-linear models in single-task accuracy (cf. Figure 2). This suggests that closing the single-task performance gap between linearized and non-linear models could be a way enhance task arithmetic, We leverage this idea the next section.
5 Enhancing task arithmetic via lincarization
We have seen that linearized models are more weight-disentangled than non-linear ones. How- ever, post-hoc linearization degrades single-task performance. We now demonstrate that enforc- ing models to fine-tune in the tangent space to their pre-trained initialization significantly im- proves task arithmetic by reducing the single- task accuracy gap.
(ol(On)
44
K
(0
f(0o)
fun (0))
Specifically, rather than applying the non- linearly fine-tuned task vectors T = 0* - 0, to fün as Section we propose directly ob- tain the task vectors through explicit fine-tuning in the tangent space as illustrated in Figure That is, given a model f. we directly fine-tune tune its linear approximation fim around O, [26). The fine-tuning process can follow the same protocols used before but with the network parameterization dictated by Eq. (3). Due to the
N
f(0))
Figure 4: Conceptual illustration of the differ- approaches
to edit pretrained model f(: 0b). Here N represents the space of neural network functions non-linearly parameterized by E O and IC its tangent space, given by the space of linearized functions fii-

---

linear connection between the weight-space and function- spacon defined in Eq. (3), fine-tuning fun essentially the same as training a kernel predictor with kemee kNTK. As a result, we obtain the fine-tuned weights fm of the linearized model for each task, which allows us to construct the cor- responding task vector Tin = fi.. 0g. We provide further details of this procedure Appendix
Moreover, as the considered models do not inherently exhibit linear fine-tuning (see Sec- tion 3), this approach yields significantly different results compared to post-hoc lineariza- tion, i.e., finm (2; 0y + Tha) * fiin(:; 0g + T). In particular, although both models share the same kernel bNTK (z, x'). the task vectors Thi have been explicitly optimized to maximize the performance of such linearized models. Conse- quently, by construction, linearized fine-tuning outperforms post-hoc linearization. Indeed, in Figure 5, we observe that linearized fine-tuning significantly reduces the non-linear advantage of non-linear models, as in most cases the performance of fian( Ou + Tin is very similar to the one of f(: 0g + T) (cf. Figure 2).
100
Non-limout fine- fine-tuning necureyy (%)
90
Non-llaeet abventaage
80
VTT-B/32
70
WT-B/166
WTT/114
000
80
90
100
Linewrized fine-tuiing accuracy (%)
Figure 5: Single-task accuracies of non-linearly FT,
f(:: 0*) and linearly FT, fin(-: Oi.)). models.
Remarkably, as we show in Appendix D.4, this increase in single-task performance does not com- promise weight disentanglement, which remains as high as for the post-hoc linearized models in Figure 3. As result, linear fine-tuning allows for improved task arithmetic compared to standard non-linear fine-tuning. In particular, Tables and 2 in their last rows show that linearized fine-tuned models significantly outperform their non-linear counterparts and achieve state-of-the-ard results on the task addition and negation benchmarks [39]. The linearized fine-tuned models achieve higher multi-task accuracies through task addition (up to 5.8 points more) and can forget more through task negation (up to 13.1 points more) while maintaining a similar level of accuracy on the control task. Additionally, observe that the advantage of the linearized models over the non-linear ones is higher for the smaller ViT-B/32 and progressively diminishes as the model size increases up to ViT-L/14
In general, thanks to the efficiency of the Jacobian-vector product implementations in most deep learning frameworks [61], training and inference in linearized neural networks only require an ((1) increase in computational costs with respect to their non-linear counterparts (see Appendix B for longer discussion). this regard, the superiority task arithmetic of linearized models can make this technique appealing for practical applications. Identifying the right trade-offss between computational cost and performance, as well as faster linearization techniques, is an exciting avenue for future work.
Towards understanding task arithmetic
We conclude by providing further fundamental insights that can aid our understanding of task arithmetic. In particular, we ask whether any kemel can satisfy Property 1, and we establish a connection between task arithmetic and the spectral properties of the NTK. Then, argue that weight disentanglement and task arithmetic are emergent properties of pre-training.
6.1 Eigenfunction localization
Generally, a kernel & admits decomposition in terms of family of eigenfunction-eigenval pairs (p)p heNN which implies that & can only represent functions of the form (*x) = Ex (x) with a finite kernel norm, i.e., Ilf 1|2 = Ee <t/>, < +oo. Specifically, the coefficients {cp bet constitute a representation of the function fr the kernel basis.
Consider 7 tasks {Li heery supported in their respective non-intersecting domains {Dtherr) Fur- thermore. let } be an orthogonal basis of cigenfunctions that diagonalizes the kernel on the union of all D, s. The following proposition provides a sufficient condition on the representation of the tasks in this basis to ensure the task arithmetic property:
In Appendix D.3, observe that larger models more weight disentangled.
7

---

Proposition (Simplified). Suppose that {ff hEe(T) can be represented by the kernel k. The kemnel k is capable of performing task arithmetic with respect to {fi herr and (erh if. for each task t, there exists a subset of localized eigenfunctions such that i) supp(ø) _ for each & in the subset, and ii) the representation of ft only involves these basis functions.
The proof and formal statement are deferred to Appendix C. Intuitively, if each task is represented with eigenfunctions that vanish outside the spatial region identified by the task support, the functions corresponding to different tasks do not interfere. Based on Proposition 1, it is natural to examine whether the NTK of CLIP models displays eigenfunctions localized in each task domain and it rep- resents the different tasks using these functions. According to the representer theorem of kernels [77], after linear fine-tuning on task with a training set { f (ar)) (2r))))eend and By ~ le,, the CLIP's predictor evaluated at a new point æ E X can be expressed as a linear combination of its kernel kNTK evaluated on æ and the training data, i.e. (i)(() = f(:; 0o) + Eve(r By ANTK( x).
To explore whether CLIP models use localized eigenfunctions for task arithmetic, we diagonalize the matrix (K &NTK(<;; 2z)) with 2, E D:. i.e., the task on which we trained, and 2y E Dj Dyr, where De is the support of a control task. If the eigenfunctions used to represent (x) are localized, then the power of the cigenvectors of KNTK must be concentrated in the points belonging to the dataset used for training. To measure this concentration, we introduce the local energy (x)) = Ep (zr), which sums the power of all the eigenfunctions Φp at a given point æ.
In Figure 6, we plot this metric for a ViT-B/32 CLIP model trained on RESISC45 with Cars as control. We provide results for other task pairs in Appendix D.8. Notably, the local energy of the eigenfunctions that the predictor uses to rep- resent the RESISC45 task is significantly higher for points belonging to the training dataset. This confirms the presence of eigenfunctions localized across the different data domains and the fact that task arithmetic occurs thanks the use of those. Indeed, thanks to this localization, CLIP models can effectively separate the rep- resentation of different tasks and carry out task- specific operations without interference. We believe that further investigation into this intrigu- ing localization phenomenon holds the potential to deepen our understanding of these models.
(x))
x E RESISC45
Cars x
Figure 6: Eigenfunction localization. Estimated support of the eigenfunctions of the NTK of a VIT- B/32 CLIP model trained on RESISC45. The plot shows the sum of the local energy of the eigenfunc- tions over random subset of the training and con- trol supports (RESISC45 and Cars, respectively).
Remark. While we have shown that localized eigenfunctions can play a crucial role in task arithmetic, it is important to note that they are not always necessary. In fact, the task arithmetic property can hold even if the eigenfunctions used to represent single task cancel outside the corresponding domain. Indeed, although eigenfunctions are linearly independent on the union of the domains, they are not necessarily linearly independent when evaluated on single domain and, in general, can cancel out. However, if the eigenfunctions maintain their linear independence on each of the domains, i.e., they are locally linear independent, then the existence of localized eigenfunctions becomes a necessary con- dition for task arithmetic. This means that if the eigenfunctions are locally linearly independent and not localized, task arithmetic is not possible, We provide some analytical examples of the latter case in Appendix C, including the NTKs of fully-connected and convolutional networks at initialization.
6.2 Weight disentanglement emerges during pre-training
Task arithmetic is not exclusive to CLIP models. In fact, task arithmetic can also be performed on pre-trained text transformers [39, 86), such as GPT-2 [68] or T5 [70] and convolutional neural networks [38] as we also show in Appendices D.5 and D.7. However, it is still unclear if the origin of weight disentanglement comes from pre-training, or if it is general property of deep networks.
To investigate this, we replicate the task addition experiments but employ randomly initialized ViT's instead of pre-trained ones. The results in Table 3 reveal that task arithmetic is not achievable on randomly initialized VIT'. Indeed, adding task vectors obtained from a random initialization ord does not result in significant improvements in multi-task accuracy over random chance. This
8

---

Table 3: Task addition from random initialization. We use the same setup as for the experiments in Table I but with task vectors obtained from finettuning randomly initialized ViTs. Results compare the average single-task accuracy (%) after fine-tuning and the multi-task accuracy (%) via task addition.
ViT-B/32
ViT-B/16
ViT-L/14
Method
Sing. (t)
Multi (t)
Sing. (t)
Multi (1)
Sing. (1)
Multi (1)
Random init
[(; 0
5.3
4.8
5.2
Non-lin. FT f(: 0h + +')
48.5
5.5
40.6
4.5
18.0
4.8
Linear. FT fial-; 0 + mim)
27.8
3.8
24.7
4.0
24.8
6.1
holds true for both non-linear task vectors, 7rd and linearized ones, rd Appendix D.9, we further corroborate these findings by computing the disentanglement error and the NTK spectrum of randomly initialized models.
Therefore, we conclude that task arithmetic is a property acquired during pre-training. This obser- vation goes beyond the traditional representation learning View of pre-training, emphasizing that pre-training not only leads to semantically disentangled feature representations but also to the disen- tanglement of the weights that govern the output on those semantic sets. Investigating the pre-training dynamics that give rise to such disentanglement is another interesting avenue for future research.
7
Related work
Weight interpolation and task arithmetic. A growing body of work is exploring the use of in- terpolations between model weights and task arithmetic manipulate and enhance the capabilities of pre-trained models. In particular, several studies have shown that interpolating between a model's fine- tuned weights and pre-trained initialization can lead to improved performance on single tasks, even surpassing their fine-tuning accuracies [27, 40, 57, 71, 72, 87]. In the multi-task setting, averaging the parameters of multiple fine-tuned models has been proposed to produce superior multi-task mod- els [38, 39, 46, 86, 89] that avoid catastrophic forgetting [28, 58] and even provide a better starting point for subsequent fine-tuning [17, 23). Interestingly, the benefits of weight ensembles and interpola- tions extend to models trained from scratch, as long they properly aligned before merging [3, 79]. This phenomenon has been observed to enhance downstream performance, further emphasizing the potential of weight interpolation and task arithmetic techniques such the ones studied in this work.
Linear vs non-linear regime. Extensive research has been conducted on comparing generalization and dynamical properties of neural networks in linear and non-linear regimes [8, 26, 30, 65, 67, 82] and investigating specific inductive biases [2, 7, 19, 53, 59, 81, 90]. In addition to theoretical under- standing, several studies have applied linearized models for practical purposes, such as predicting fine-tuning generalization [22] and training speed [91], as well as enhancing calibration [52] and few- shot performance [5]. Our work serves as another example of the utility of linearized models in certain scenarios where they do not only offer practical benefits but also provide valuable theoretical insights.
Feature disentanglement. The notion of feature disentanglement lies at the heart of representation learning, where idcal representations are assumed to separate distinct data variation factors along different directions in the feature space [1, 10, 35]. A multitude of approaches in generative model- ing [14, 34, 73] and self-supervised learning [6, 13, 48, 69] strive to achieve this goal. Our investiga- tion, however, explores distinct aspect: weight disentanglement within the framework of task arith- metic. Departing from the static perspective of feature disentanglement, weight disentanglement con- nects weight space and function space transitions, thereby enriching our understanding of disentangle- ment in neural networks from a functional standpoint. Several studies have previously attempted to ex- ploit similar notion by inducing the learning of task-specific subnetworks within larger network [32, 36, 54-56, 83, 85]. To the best our knowledge, our work is the first to demonstrate the natural emergence of such phenomena in specific semantically meaningful tasks during CLIP pre-true training.

---

8 Conclusion
In this work, we conducted a thorough analysis of task arithmetic in deep neural networks, delving into its fundamental mechanisms and enhancing its performance. Our findings demonstrate that linearized models, governed by the NTK, outperform their non-linear counterparts in task arithmetic, thus providing more effective approach for model editing. Crucially, we revealed that weight dis- entanglement plays vital role in the success of task arithmetic, as distinct directions in weight space correspond localized areas in the function space; and that it is emergent property of pre-training.
fascinating open question consists in understanding how weight disentanglement arises during pre- training and finding algorithms that enhance it. Another exciting research direction is investigating the potential of tangent spaces for editing other pre-trained models. In this sense, developing more efficient linearized models would be significant leap forward in this field. These advancements could pave the way for novel approaches to model editing and deepen our understanding of the intricate relationship between weight space and function space in deep learning.
Acknowledgements
We thank Nikolaos Dimitriadis, Adam Hazimeh, Pau de Jorge, Nikolaos Karalias, Seyed Mohsen Moosavi-Dezfooli, Antonio Sclocchi, Thibault Séjourné, and Matthieu Wyart for helpful feedback and comments. We also thank Gabriel Ilharco for helping set up the code to reproduce their experiments.
References
[1] Alessandro Achille and Stefano Soatto. Emergence of invariance and disentanglement in deep representations. Journal of Machine Learning Research (JMLR), 2018. http:// jmmr org/ papers/v19/17-644 html. (page 9).
[2] Alessandro Achille, Aditya Golatkar, Avinash Ravichandran, Marzia Polito, and Stefano Soatto. LQF: Linear quadratic fine-tuning, In IEEE Conference Computer Vision and Pattern /0/b///2/2..2. 11140. (page 9).
Recognition (CVPR), 2021. bttpss://rxxiv.ooo
[3] Samuel K Ainsworth, Jonathan Hayase, and Siddhartha Srinivasa. re-binn::: Merging models modulo permutation symmetries. In International Conference on Learning Representations (ICLR), 2023. https: //t/x//xxvvvv////////////////xaa org/ abs/ 22099 04836. (pages and 9).
[4] Sanjeev Arora, Simon S. Du, Wei Hu, Zhiyuan Li, Ruslan Salakhutdinov, and Ruosong Wang. On exact computation with infinitely wide neural net. In Advances in Neural Information Processing Systems (NeurIPS), 2019. https: /ttt:////ccccceee..p/eo .wwrrise..ppp... cc/paper/2019/ hash
Abstract. html. (page 3).
[5] Sanjeev Arora, Simon Du, Zhiyuan Li, Ruslan Salakhutdinov, Ruosong Wang, and Dingli Yu. Harnessing the power of infinitely wide deep nets small-data tasks. In International Conference on Learning Representations (ICLR), 2020. https://arxiv.orgg 1910. 01663. (page 9).
[6] Philip Bachman, R. Devon Hjelm, and William Buchwalter. Learning representations by maximizing mutual information across views. In Advances in Neural Information Process- ing Systems (NeurIPS), 2019, https://proceeding wewpeeeddinnnn.uniiinne nwwrips....... cctppppe//011/////// 20t9///99///hh///// /hash/ ddf
html. (page 9).
[7] Gregor Bachmann, Seyed-Mohsee Moosavi-Dezfoolli, and Thomas Hofmann. Uniform conver- gence, adversarial spheres and a simple remedy. In International Conference on Machine Learn- ing (ICML), 2021. http://proceedingg mlr. www.p.l.r..///...m......m.. watess/vvv//// h1139/9/aahmnn211.11 htmm. (page 9).
[8] Aristide Baratin, Thomas George, César Laurent, R. Devon Hjelm, Guillaume Lajoie, Pascal Vincent, and Simon Lacoste-Julien, Implicit regularization via neural feature alignment. In International Conference on Artificial Intelligence and Statisties (AISTATS), 2021. http: html. (page 9).
/ //proceedings.r mlr, www.e/.eeseeeeeeeeee.ee httss/vv0//aaaa v130/baratin21a. htmm.
10

---

[9] Atilim Gunes Baydin, Barak A. Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. Automatic differentiation in machine learning: survey, Journal of Machine Learning Research (JMLR). 2017. http://jmii otg/papers/v1e. 1118811746688 htmml. (page 19).
[10] Yoshua Bengio, Aaron C. Courville, and Pascal Vincent. Representation learning: A review and new perspectives. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2013. https://axiiv. org/ aa/aaaba//200 5538. (page 9).
[11] Blake Bordelon, Abdulkadir Canatar, and Cengiz Pehlevan. Spectrum dependent learning curves in kernel regression and wide neural networks. In International Conference on Machine Learn- ing (ICML), 2020. http://proceed:e 18).
mlr wtt../eesssw////rrlllleeoo
html. (page
[12] Francesco Cagnetta, Alessandro Favero, and Matthieu Wyart. What can be learnt with wide 01003. (page 20).
convolutional neural networks?, 2022. http:////axvvvv.oaaareaosee/
[13] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International Conference Machine Learn- ing (ICML), 2020. https://proceeding mml.. mlr. www.ppresssviveee /119// chen20j /an2200//ch/22/ chen20j pdf. (page 9).
[14) Xi Chen, Yan Duan, Rein Houthooft, John Schulman, Ilya Sutskever, and Pieter Abbeel. InfoGAN: Interpretable representation learning by information maximiz- ing generative adversarial nets.
Advances in Neural Information Processing Systems (NeurIPS), 2016. https: https://proceedings pewpocceedingg.nnurppp.. neurips. wwtpppppe/////////aeeeeeppppeeee /2/h///hh////////// 779001199aabt775sacceeeaaaa Attrobb--bbttatt.... html. (page 9).
[15] Gong Cheng, Junwei Han, and Xiaoqiang Lu. Remote sensing image scene classification: Benchmark and state of the art. Proceedings of the IEEE, 2017. https: ieee. org/document/7891555
(page 4).
[16] Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in dif- ferentiable programming. Advances in Neural Information Processing Systems (NeurlPS), 2019.
https://procedans..nnett www.nneuuuc.ppppeeeeee pearripp..c/ppppe///////// 4Bc1666778B010336668777724-8--83 .dff. (page 2).
[17] Leshem Choshen, Elad Venezian, Noam Slonim, and Yoav Katz. Fusing finetuned models for better pretraining, 2022. https://arxis: wwwg/bg///////eeeeeeee 03044. (page 9).
[18] Ole Christensen and Khadija L Christensen. Linear independence and series expansions in function spaces, The American Mathematical Monthly, 2006. https://www. tandf online. com/doi/abs/ 10. 108000000099000 2006. 11 11920343. (page 20).
[19] Kurtland Chua, Qi Lei, and Jason D Lee. How fine-tuning allows for ef- fective meta-learning In Advances in Neural Information Processing Systems (NeurIPS). 2021. https: http::s//proceeddnggg neurips httpper_filleeppper/eeeee fttl://aa333111faaaaaaaaaaaaaaaa
(page 9).
[20] Mircea Cimpoi, Subhransu Maji, lasonas Kokkinos, Sammy Mohamed, and Andrea Vedaldi. De- scribing textures in the wild. In IEEE Conference Computer Vision and Pattern Recognition (CVPR), 2014.
thecvf. com/ content_ cct___cvr____ hrrp__011/htmm/CCaa Describing Beaiinngeextttts_eeeeee 2aah__2011___P___0 paper html. (page 4).
[21] 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 Computer Vision and Pattern Recognition (CVPR), 2009. 3).
ieee. org
(page
[22] Aditya Deshpande, Alessandro Achille, Avinash Ravichandran, Hao Li, Luca Zancato, Char- less C. Fowlkes, Rahul Bhotika, Stefano Soatto, and Pietro Perona. A linearized framework and a new benchmark for model selection for fine-tuning, 2021. https:/ /aapxii///xvv...// /tt///awww...aaaaaaa////aaa/aaae org/abs/2102. 00084. (pages 3 and 9).

---

[23] Shachar Don- Ychiya, Elad Venezian, Colin Raffel, Noam Slonim, Yoav Katz, and Leshem Choshen. Cold fusion: Collaborative descent for distributed multitask finetuning, 2022. https:
/ //arxiv. org/ooggagao org/abs/////eoaooooooggooo.eebrr abs//222...2.2////////////////// 2212. 01378. (pages and 9).
[24] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference Learning Representattins (ICLR), 2021. (page 4).
[25] Alessandro Favero, Francesco Cagnetta, and Matthieu Wyart. Locality defeats the curse of dimensionality in convolutional teacher-student scenarios. In Advances in Neural Information Processing Systems (NeurlPSS) 2021. https: httpps://proceedings neurips wwwrips.ccppapppeepepe cc/paper files/ http://///a///eeee0000/0/0000000 Abstrfftttbbbsrtttt.......cc .ataractt.tmm.......... (page 20).
[26] Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M Roy, and Surya Ganguli. Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel. In Advances in Neural Information Processing Systems (NeurIPS), 2020. https: //poocedigss neurips.. cc/ cc/paper/ sappe/22200 /020/ (T.10440007669100063817777777880 Pappe pdf. (pages 4, and 9).
[27] Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin. Linear mode connectivity and the lottery ticket hypothesis. International Conference on Machine Learning (ICML), 2020. https://proceeddngg .mmiiinnn...........m...m.. www.mmmww...m...mmmmmmmmmmmmme.. ppesss/1199/fraakleeee html. (pages and 9).
[28] Robert M French. Catastrophic forgetting in connectionist networks. Trends in Cognitive Sciences, 1999. https:/ www../ciennee.......ee conn/cciene//rtiicll $1364661399012942. (pages 1 and 9).
[29] Amnon Geifman, Meirav Galun, David Jacobs, and Basri Ronen. On the spectral bias of convolutional neural langent and gaussian process kernels. Advances in Neural Information Processing Systems (NeurIPS), 2022. https://aat //axiv//////v///// org/abs/ /a/aabbb22200 //2/2/222..222..2. (page 20).
[30] Mario Geiger, Stefano Spigler, Arthur Jacot, and Matthieu Wyart. Disentangling feature and lazy training in deep neural networks. Journal of Statistical Mechanics: Theory and Experiment, 2020.
org. www.oorggairiieee hat/arrtticleeeeeee /000 1008 1742-5462//7// 5466//bcddd.. abc4de. (page 9).
[31] Amelia Glaese, Nat McAleese, Maja Trebacz, John Aslanides, Vlad Firoiu, Timo Ewalds, Mari- beth Rauh, Laura Weidinger, Martin Chadwick, Phoebe Thacker, Lucy Camppell-Gillingann, Jonathan Uesato, Po-Sen Huang, Ramona Comanescu, Fan Yang. Abigail See, Sumanth Dathathri, Rory Greig, Charlie Chen, Doug Fritz, Jaume Sanchez Elias, Richard Green, Sona Mokra, Nicholas Fernando, Boxi Wu, Rachel Foley, Susannah Young, lason Gabriel, William Isaac, John Mellor, Demis Hassabis, Koray Kavukcuoglu, Lisa Anne Hendricks, and Geof- frey Irving. Improving alignment of dialogue agents via targeted human judgements, 2022. https:/ bttts::/www..ce www.wwepwwiddd...mmmmmmmmmmmmmm. com/blog..///ooooocooooooo//oooo httgg/bulddig-gdiiieeeeeergoaial safer- -iiee--daalgueeeee.. agents. (page 1).
[32] Marton Havasi, Rodolphe Jenatton, Stanislav Fort, Jeremiah Zhe Liu, Jasper Snoek, Balaji Lakshminarayanan, Andrew Mingbo Dai, and Dustin Tran. Training independent subnetworks for robust prediction. In International Conference on Learning Representations (ICLR), 2021. (page 9).
[33] 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 Applied Earth Observations and Remote Sensing, 2019. https://arxiv org/abs/ /aabs//109... 1709. 00029. (page 4).
[34] Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matteww Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-VAFE Learning basic visual concepts with constrained variational framework. In International Conference on Learning Representations (ICLR), 2017. https:
(page 9).
12

---

[35] Irina Higgins, David Amos, David Pfau, Sébastien Racanière, Loic Matthey, Danilo J. Rezende, and Alexander Lerchner. Towards definition of disentangled representations, 2018. http: //arxiv org/ abaaabas///// 02230. (page 9).
[36] Edward J. Hu, Yelong Shen, Phillip Wallis. Zeyuan Allen-Zih Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations (ICLR), 2022. https: //op:////peeeevee. net/forum? wwweammddddZZZeeeeeeeee (page 9).
[37] Gabriel Ilharco, Mitchell Wortsman, Ross Wightman, Cade Gordon, Nicholas Carlini, Rohan Taori, Achal Dave, Vaishaal Shankar, Hongseok Namkoong, John Miller, Hannaneh Hajishirzi, Ali Farhadi, and Ludwig Schmidt. OpenCLIP, July 2021. https://doi. org /0///0..000 5281/ zenodo. 5143773. (page 18).
[38] Gabriel Ilharco, Mitchell Wortsman, Samir Yitzhak Gadre, Shuran Song, Hannaneh Hajishirzi, Simon Kornblith, Ali Farhadi, and Ludwig Schmidt. Patching interpolating weights. In Advances in Neural Information Processing Systems (NeurlPS), 2022.
models by
nttps://arXin org/ abs/ /208 05592. (pages 1, 8, 9, and 18).
[39] 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. https://arxivvorr//aa atb/////////........... 08207. (pages 1, 2, 3, 4, 5, 7, 8, 9, 18, and 25).
[40] Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. In Conference on Uncertainty in Artificial Intelligence (UAI), 2018. attps://arxinl http:/aaxivvv....aaaa.. arg/abs/ /10/1800...0 05407.
(pages and 9).
[41] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in Neural Information Processing Sys- tems (NeurIPS), 2018. https://prooeedinggs neurips wwtpacpp.ppe..ecee wttps//feeeeeeeee/////// Pff (pages 2 and 3).
[42] Tushar Khot, Peter Clark, Michal Guerquin, Peter Jansen, and Ashish Sabharwal. Qasc: A dataset for question answering via sentence composition, 2020. https: wwwaaaiiiorrrreeeeeeeeeeeaaeaaea abs// 1910. 11473v2.
(page 25).
[43] Jonathan Krause, Michael Stark, Jia Deng, and Li Fei-Fei. 3D Object representations for fine- grained categorization. In International Conference Computer Vision Workshops (ICCWw), 2013. https:/ //wwww//wwwwww cv foundation httg///ppaaccess/cco//s cttp///oottetttcceeeeeeeeeeeeeee wwcccccorsshosss.....o Sttraraauueeeeeeeeeeeeeeeeeeeeee oboject_Rapresentatiinns 2013 _013__001____________e______ ICCV paper html.
(page 4).
[44] Yann LeCun, The MNIST database of handwritten digits, 1998. http://yame http://yann.lecum. lecun com/ exdb/mnist/. (page 4).
[45] Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl-Dickstein, and Jeffrey Pennington. Wide neural networks any depth evolve as lin- models under gradient descent. In Advances in Neural Information Processing Sys- tems (NeurIPS), 2019. https: / //prc//eeono///ssseeo neurips wccppaper_fileet files///ppprr//119///// http://0111111111111381111111111
Paper pdf. (page 3).
[46] Margaret Li, Suchin Gururangan, Tim Dettmers, Mike Lewis, Tim Althoff, Noah A Smith, and Luke Zettlemoyer. Branch-traii--merge: Embarrassingly parallel training of expert language models, 2022. https://arxis:
aas//bbs/22000 03306. (pages 1 and 9).
[47] Zhuang Liu, Hanzi Mao, Chao-Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. ConvNet for the 2020s. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2022. https: // /arxivv.. org/ abs/2201 03545. (page 25).
13

---

[48] Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Rätsch. Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentangled representations. In International Conference on Machine Learning (ICML), waapiiggsmm........
2019. httt:://proceeeieee
hhmml.. (page 9).
[49] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations (ICLR), 2019. URL https://openrertt net/ forun?id-Bkg6RiCGn (page 18).
[50] Ximing Lu, Sean Welleck, Liwei Jiang, Jack Hessel, Lianhui Qin, Peter West, Prithviraj Ammanabrolu, and Yejin Choi. QUARK: Controllable text generation with reinforced un- learning. In Advances in Neural Information Processing Systems (NeurIPS), 2022. https: //arxiv. org/ aa///////// abs/2205. 13636. (page 1).
[51] Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In Annual Meeting of the Association for Compulational Linguisties (ACL), 2011. http://wwrt aclweb. org/ anthology/PI1-1001 1015. (page 25).
[52] Wesley Maddox, Shuai Tang, Pablo G. Moreno, Andrew Gordon Wilson, and Andreas Dami- anou. Fast adaptation with linearized neural networks. International Conference on Artificial Intelligence and Statistics (AISTATS), 2021. https://proceed:ngs mlr wwwps/vessvvee vv100/ maddox21a/maddoxx/ pdf. (page 9).
[53] Sadhika Malladi, Alexander Wettig, Dingli Yu, Danqi Chen, and Sanjeev Arora. kernel-based view of language model fine-tuning, 2022. https: http///axvvvvo///// org/ abs//21...222... 05643, (pages
3 and 9).
[54] Arun Mallya and Svetlana Lazebnik. Packnet: Adding multiple tasks single network by iterative pruning. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2018. /tatps://axiv....o wwtgg///////..////oo/o/
(page 9).
[55] Arun Mallya, Dillon Davis, and Svetlana Lazebnik. Piggyback: Adapting a single network to multiple tasks by learning to mask weights. In European Conference on Computer Vision (ECCV), 2018. attps://arxxs org/ abs/ 1801.06519
[56] Nicolas Y. Masse, Gregory D. Grant, and David J. Freedman. Alleviating catastrophic forgetting using context-dependent gating and synaptic stabilization. Proceedings of the National Academy of Science (PNAS), 2018. https: //tp://////oioo/////////
10007///as... 1803839115. (page 9).
[57] Michael Matena and Colin Raffel. Merging models with fisher-weighted averaging. In Advances in Neural Information Processing Systems (NeurIPS), 2021. https://arxiv.org/aat 09832. (pages 1 and 9).
[58] Michael McCloskey and Neal J Cohen. Catastrophic interference in connectionist networks: The sequential learning problem. In Psychology of Learning and Motivation, Elsevier, 1989. https: / / /wwwww.. sciencedirect. com/ science/article article/ httptc//e/aa//p/////////00000000
(pages
1 and 9).
[59] Fangzhou Mu, Yingyu Liang, and Yin Li. Gradients as features for deep representation learning. In International Conference on Learning Representations (ICLR), 2020. https: //arxiv. org/
wwwg/ba//////eooeeooooooeeeeeeee abs//b// 05529. (page 9).
[60] Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng. Reading digits in natural images with unsupervised feature learning. In Advances Neural In- formation Processing Systems (NeurIPS) Workshops, 2011. https:// /ttp:://sttaaeeee wwworrggggogggeeeeoooooooooooooo googleapis. com/pub-toola. -ullisc-u-ic--- puwiicccpubiiiiiinnmmmmom aattiiiaaiinddat/pppffff /37648 pdf. (page 4).
[61] Roman Novak, Jascha Sohl-Dickstein, and Samuel S Schoenholz. Fast finite width neural tangent kernel. International Conference on Machine Learning (ICML), 2022. https: // //arxivv org/ abs/ /206. 08720. (page 7).
14

---

[62] Guillermo Ortiz-Jiménez, Apostolos Modas, Seyed-Mohsen Moosavi-Dezfooli; and Pascal Frossard. What can linearized neural networks actually say about gen- eralization? In Advances in Neural Information Processing Systems (NeurIPS), 2021.
hew.:eeeeeddnnnnenneerse....ppp www...ps..cppppprrrrrr.... Abbealaaee6Gllaccc3333330oeee Paper. pdd. (page 3).
[63] Guillermo Ortiz-Jiménez, Apostolos Modas, Seyed-Mohsen Moosavi-Dezfooli, and Pascal Frossard. Optimism in the face of adversity: Understanding and improving deep learning through adversarial robustness. Proceedings the IEEE, 2021. https://ieeexplore ieee. org/ htg//document/9348880 (page 1).
[64] Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll L Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, al. Training language models to follow instructions with human feedback, 2022. (page 1).
org/ org/aba/////eoeaooooeogggoeeeebt 2202222233.2 02155.
[65] Jonas Paccolat, Leonardo Petrini, Mario Geiger, Kevin Tyloo, and Matthieu Wyart. Geometric compression of invariant manifolds in neural nets. Journal of Statistical Mechanics: Theory and Experiment, 2021. https: /awxxvvv/vvv/////////////////aa/ org/ hwtg//bb/2000/// 11471. (page 9).
[66] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. PyTorch: An imperative style, high-performance deep learning library. Advances Neural Information Process- ing Systems (NeurlPSS) 2019. https:: ietpreeeedddinneeeneiepes...p cc/paper files/paper/ pdf. (page 18).
[67] Leonardo Petrini, Francesco Cagnetta, Eric Vanden-Eijnden and Matthieu Wyart. Learning sparse features can lead to overfitting in neural networks. In Advances in Neural Information Processing Systems (NeurlPSS) 2022. https: htpppp/prr//ooeeeeaagss mwwpeeddiiiiinnnnnrspeeeeeeeeeep cc/ wecppapr__aa files/ (page 9).
Paper-Conferercce pdf.
[68] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language Models are Unsupervised Multitask Learners, 2019, https://openai.com htt.e/ai..cmm/ooo///// (page 8).
better--language-modee models/.
[69] 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 Ilya Sutskever. Learning transferable visual models from natural language supervision. In Inter- national Conference Machine Learning (ICML), 2021. 00020. (pages 2, 9, and 19).
[70] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research (JMLR), 2020. http://jmlr. org htgg/paperrs/21120-000 074. html. (pages and 25).
[71] Alexandre Ramé, Kartik Ahuja, Jianyu Zhang, Matthieu Cord, Léon Bottou, and David Lopez- Paz. Model ratatouille: Recycling diverse models for out-of-distribution generalization. In International Conference on Machine Learning (ICML), 2022. https: //arxiv org/abs/ 2212. (pages and 9).
[72] Alexandre Ramé, Matthieu Kirchmeyer, Thibaud Rahier, Alain Rakotomamonjy, Patrick Gallinari, and Matthieu Cord. Diverse weight averaging for out-of-distribution generaliza- tion. In Advances Neural Information Processing Systems (NeurIPS), 2022. https: //arxiv. wwg//ra/////a/ee ab2/ 09739. (pages and 9).
[73] Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International Conference on Machine I. Learning (ICMI), 2014. http:// pttp:////ooceeing..mmom mlr. wwwpppreevvvveeee 2e2//eeeddd14 14, .ttll, (page
9).
15

---

[74] Marco Tulio Ribeiro and Scott Lundberg. Adaptive testing and debugging of nlp models. In Annual Meeting of the Association for Computational Linguisties (ACL), 2022. https: //aclaanna /acllanthology org/ wtw.:222222....... acl-long 230/ (page 1).
[75] Basri Ronen, David Jacobs, Yoni Kasten, and Shira Kritchman. The convergence rate of neural networks for learned functions of different frequencies. In Advances in Neural Information Processing Systems (Neur(PS), 2019. attps://proceediin Paper pdf. (page 20).
file/ (55c8008677744510097775555558885
[76] Shibani Santurkar, Dimitris Tsipras, Mahalaxmi Elango, David Bau, Antonio Torralba, and Alek- sander Madry. Editing a classifier by rewriting its prediction rules. In Advances in Neural Infor- mation Processing Systems (NeurlPSS) 2021. https::/proceed::/
neurips. cc/paper. Paper.ddf. (page 1).
[77] Bernhard Schölkopf and Alexander Johannes Smola. Learning with Kernels: support vec- tor machines, regularization, optimization, and beyond. Adaptive computation and ma- chine learning series. MIT Press, 2002. https: http:://mippesss..t mit. edu/9978022538577/ learning-wild with-kernels/ (page 8).
[78] Christoph Schuhmann, Richard Vencu, Romain Beaumont, Robert Kaczmarczyk, Clayton Mullis, Aarush Katta, Theo Coombes, Jenia Jitsev, and Aran Komatsuzaki. Laion-400m: Open dataset of clip-filtered 400 million image-text pairs, 2021. (page 25).
02114.
[79] Sidak Pal Singh and Martin Jaggi. Model fusion via optimal transport. In Advances in Neural Information Processing Systems (Neur[PS), 2020. p neurips. cc/ cc/paper /ttp:////...eeeeee000000000000e0 Abstrattttdddttttt...o..oaoaaace ceaettrccttttl........... (pages
1 andd 9).
[80] Johannes Stallkamp, Marc Schlipsing, Jan Salmen, and Christian Igel. The german traffic sign recognition benchmark: multi-class classification competition. In International Joint Con- ference on Neural Networks (JJCNN), 2011. https://leeexpl.oo ieee. org/ document/ 6033395. (page 4).
[81] 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 networks learn high frequency functions in low dimensional domains. In Advances in Neural Information Processing Systems (NeurIPS), 2020. https: https://procedd: http:///e/eennnee//eeeeee neurips. cc/ Paper pdf. (page 9).
[82] Nikhil Vyas, Yamini Bansal, and Preetum Nakkiran. Limitations of the NTK for understanding generalization deep learning, 2022. https://arxiae httggg////////////e 10012. (page 9).
[83] Yeming Wen, Dustin Tran, and Jimmy Ba. Batchensemble: an alternative approach to efficient ensemble and lifelong learning. In International Conference Learning Representations (ICLR), 2020. / 22g02 06715. (page 9).
[84] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Huggingface's trans- formers: State-of-the-art natural language processing, 2019. https://arxis orr.iv..orr/aaaa wttp/b/////b////////a///g 03771. (page 25).
[85] Mitchell Wortsman, Vivek Ramanujan, Rosanne Liu, Aniruddha Kembhavi, Mohammad Raste- gari, Jason Yosinski, and Ali Farhadi. Supermasks superposition. In Advances in Neural Information Processing Systems (NeurIPS), 2020. httpp://proceeddigg neurips.ccccco cc/ Paper pdf. (page 9).
[86] Mitchell Wortsman, Gabriel Ilharco, Samir Yitzhak Gadre, Rebecca Roelofs, Raphael Gontijo- Lopes, Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Korblith, et al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In International Conference on Machine Learning (ICMI.), 2022. https:/ :/ /arxiviv org/ abs/abb3.. 05482. (pages 1, 3, 8, and 9).
16

---

[87] Mitchell Wortsman, Gabriel Ilharco, Mike Li. Jong Wook Kim, Hannanch Hajishirzi, Ali Farhadi, Hongseok Namkoong. and Ludwig Schmidt. Robust fine-tuning of zero-shot models. In IEEE Conference on Computer Vision and Puttern Recognition (CVPR), 2022. https:
//arxiv.orrg abaaab///u 2b099 01903. (pages 1, 3, 4, and 9).
[88] Jianxiong Xiao, Krista Ehinger, James Hays, Antonio Torralba, and Aude Oliva. Sun database: Exploring large collection of scene categories. International Journal of Computer Vision (IJCV), 2016. https://linm http:://iink.pprineerr com/article, wtriiiceeee.10000. 110....00/111111 -1--11888--1--------- (pages 4 and 22).
[89] Prateek Yadav, Derek Tam, Leshem Choshen, Colin Raffel, and Mohit Bansal. Resolving interference when merging models. In Advances in Neural Information Processing Systems (NeurlPS), 2023. https://arxiss http:///xvvvvor../...... org/ /206/ 01708. (pages and 9).
[90] Gizem Yüce, Guillermo Ortiz-Jiménez, Beril Besbinar, and Pascal Frossard. A structured dictionary perspective on implicit neural representations. In IEEE Conference on Computer Vision and Puttern Recognition (CVPR), 2022. https://dod .rriiiorg//// 1109///R5266..8 2022. 2022.01863, (pages 3 and 9).
[91] Luca Zancato, Alessandro Achille, Avinash Ravichandran, Rahul Bhotika, and Stefano Soatto. Predicting training time without training. In Advances in Neural Information Process- ing Systems (NeurlPS), 2020. c330:3366bbb433333333335555 -at htmml. (pages 3 and 9).
neurips. cctppp.cccc/cc wwwc//ape///a/ 2020/hash/
[92] Fuzhen Zhuang, Zhiyuan Qi, Keyu Duan, Dongbo Xi, Yongchun Zhu, Hengshu Zhu, Hui Xiong, and Qing He. comprehensive survey on transfer learning. Proceedings of the IEEE, 2020. (page 1).
betps://ieeexper ieee.
17

---

A Experimental details
All our experiments were performed using the same hardware consisting of four V100 NVIDIA GPUs with 32GB of memory each and can be reproduced in less than 350 GPU hours. The details of each experiment are the following.
Fine-tuning. All the fine-tuning experiments follow the same training protocol specified in Ilharco al. [39] 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 checkpoint downloaded from the open_ clip repository [37]. We fine-tune for 2, 000 iterations with batch size of 128, learning rate of 10 and cosine annealing learning rate schedule with 200 warm-up steps and the Adam W optimizer [49]. As introduced in Ilharco et al. [38], during fine-tuning, we freeze the weights of the classification layer obtained by encoding a standard set of zero-shof template prompts for each dataset. Freezing this layer does not harm accuracy and ensures that no additional learnable parameters introduced during fine-tuning [38]. We use this exact same protocol to fine-tune the non- non-linear and linearized models and do not perform any form of hyperparameter search in our experiments.
Tuning of in task arithmetic benchmarks. As in Ilharco al. [39] a to tune the size of the task vectors used to modify the pre-trained models. This is equivalent to setting a 01 ... Qr in Eq. (1). Both in the task addition and task negation benchmarks, after fine-tuning, we evaluate different scaling coefficients a E {0..0, 0.05, o.....,, 1.0) and choose the small held-out proportion of the training set as specified in Ilharco et al. [39]. 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
a single coefficient
value that achieves the highest target metric
control task; for task addition and negation, respectively. The tuning of a is done independently for non-linear FT, linearized FT, and post-hoc linearization.
Normalized accuracies in task addition. Table shows the normalized accuracies after editing different models by adding the sum of the task vectors on 8 tasks = E. Tr Here, the normalization is performed with respect to the single-task accuracies achieved by the model fine-tuned on each task. Mathematically.
T
[f(x; 0u + Eg Te)
1 T
Normalized accuracy
w-p acc [f(:; 00 Tc)
ampp
Disentanglement error. To produce the weight disentanglement visualizations of Figure we compute the value of {(a1, a2) on 20 x 20 grid of equispaced values in [-3, 3] [-3,3], 3). To estimate the disentanglement error, we use a random subset of 2, 048 test points for each dataset.
NTK eigenfunction estimation. We use the finite-width NTK implementation from the functorch sublibrary of PyTorch [66] to compute the KNIK matrices described in Section 6.1. In particular, random subset of 200 training points for each dataset and compute the singular value decomposition (SVD) of KNIK to estimate the entries of Φ, on each dataset. As described in Bordelon et al. [11], and to avoid high memory footprint, we estimate different set of singular vectors for each output class, equivalent to estimating one kernel matrix per output logit. Figure
shows the values of (ar) for each class with a different line. However, there is little variability of the NTK among classes, and hence all curves appear superimposed in the figure.
B Implementation aspects of linearized models
We now provide more details of the different implementation aspects of linearized models, including basic code and a discussion on their computational complexity.
Practical implementation. Creating linearized models a neural network is very simple using the functorch sublibrary of PyTorch. Specifically, using the fast Jacobian- vector implementation of this library, we can easily create a custom class that takes any nn. Module as input and generates trainable linearized version of it around its initialization. We give simple example of this in
18

---

Listing I, where we see that the resulting LinearizedHodee can be directly used in any training script as any other neural network.
In our experiments, we linearize the ViT image encoder of CLIP as the text encoder frozen in experiments. In this regard, during training and inference, as it is common in standard CLIP models [69], we normalize the output of the linearized image encoder prior to performing the inner product with the embeddings. This normalization does change the classification decision during inference, but it has a rescaling effect on the loss that can influence training dynamics. In our fine-tuning experiments, we found this standard normalization technique has clearly positive effect in single-task accuracy both for the non- linear and linearized models.
import copy
2 import torch. nn as nn
from functorch import jvp, cake_functionall buffer
clasa LinearizedModel (nn. Module):
unn Creates linearized version of any Module.
The Linearized version of a model is a proper PyTorch model and can be trained as any other Module.
10
Args:
inst_mode (nn.Nodule) The model to linearize. Its parameters used to instialized the linearized model.
def isennt___ init_model):
16
Convert models to functional form.
func, parama0, buffers0 make finke_funtiionl_.. with buffers b (init_model)
Store parameters and forward function.
self funco lambda params, func(params, buffersO, x)
self. params0 = params0
Initialization parameters.
self. params copy deepcopy (params0)
Trainable parameters.
23
24
Freeze initial parameters and unfreeze current parameters. for p0 in self self. paramso: p0. requirreee-gaad False for p self params: p requires_ True
def __lll____selll x):
Compute linearized model output.
30
dparams [p p0 for P. p0 in zip(self, params, self, self. parama0)] out, dp jup(self. funco, (self, paramparrmm,,), (dparams,)) return out + dp
Listing 1: Basic Py PyTorch code to linearize a model.
Computational complexity. Jacobian-vector products can be computed efficiently, at the same marginal cost as a forward pass, using forward-mode automatic differentiation rules [9]. This means that doing inference with linearized model usually takes around two three times more than with its non-linear counterpart, as for every intermediate operation in the forward pass, its derivative also needs to be computed and evaluated.
Training the linearized models, on the other hand, uses backpropagation algorithm which, for every forward pass, requires another backward pass to compute the gradients. In this regard, the computational cost of obtaining the gradient with respect to the trainable parameters of the linearized models Vefunle 0) is also roughly twice the cost of obtaining the gradient of its non-
19

---

linear counterparts Vef(2; 0). Similarly, as the forward-mode differentiation required to compute the forward pass also depends on the values of the derivatives at this step, the final memory footprint of training with the linearized models is also double than the one of training the non-linear ones.
C Spectral analysis of linearized models
In this section, we present the formal statement and proof of Proposition 1. Additionally, we delve deeper into the question of whether eigenfunction localization is a necessary condition for task arith- metic and provide analytical examples with exactly diagonalizable NTKs to support our discussion. Proposition 2 (Formal version of Proposition I). Suppose that the task functions {fr heam belong to the RKHS of the kernel k and their coefficients in the kernel eigenbasis are ((C hem. If of tt either Ct = O or supp(0p) S DD then the kernel k has the task arithmetic property with respect to {st hee(r) and (Dr)
Proof. The task arithmetic property requires that VL' E [T], Va Dr, Eveny fi()) = f7(x).
Representing the task functions in the kernel basis, we have
Vt' E [T], Vo E Dv, ΣΣ ( = Σ (p))))))
(8)
EEN
PEN
This condition can be rewritten as
Σ I I DVy (E(T),, titt' PEN
dx dx=
If, for each I, the eigenfunctions corresponding to non-zero coefficients are supported within a subset of Dj and all domains Dj's are disjoint, then all the summands inside the integral in Eq. (9) become zero inside Dr. and thus the proof is complete.
As we discussed in Section 6.1, eigenfunction localization is generally not a necessary condition to achieve task arithmetic. However, we now show that if the eigenfunctions are locally linear independent across the different task domains, then the localization property becomes a necessary condition for task arithmetic. The proposition presented below formalizes this concept.
Proposition 3. Suppose that the task functions {fi hE(T) belong to the RKHS of the kernel Furthermore, let the kernel
and their coefficients in the kernel eigenbasis are (Ce.p)peN))er eigenfunctions be either zero or linearly independent over each domain Dr. The kernel l has the task arithmetic property with respect {ftt Jee(T) and (Di ( if and only if t. p. either ct.p = 0 or supp(op) S Dr.
Proof. The initial steps of the proofs follow those of the previous proposition. In particular, let's consider the integral in Eq. (9). Due to the linear independence of the non-zero kernel eigenfunctions on Dr. for this integral to be zero, we have only two possibilities: either i) all coefficients
((ct.)) (ent,))en must be zero or ii) the eigenfunctions corresponding to non-zero coefficient C.p (t # t') must be zero in Dyr Since the proposition is valid for any set of functions, condition is not feasible. Therefore, condition it) must hold. Furthermore, since Eq. (9) is valid Vt' e [T]. it follows that the eigenfunctions used to represent each task t' are zero D. UrerT] D.. Consequently, these eigenfunctions are only supported in De or subset thereof.
In order to understand the implications of this proposition, it is useful to examine simple data geometries and architectures for which the NTK can be analytically diagonalized. For instance, when data is uniformly distributed on a ring or torus, the NTK of fully-connected and convolutional neural networks at initialization can be diagonalized with the Fourier series [12, 25, 29, 75]. Fourier atoms are linearly independent on any interval [18] and not localized. Consequently, according to Proposition 3, these architectures cannot perform task arithmetic within such settings. This straightforward calculation aligns with the observation that task arithmetic generally emerges as property of pre-training and is not inherently present initialization, as we numerically demonstrated for CLIP models in Section 6.2.
20

---

D Further experimental results
We now present additional experiments that expand the findings discussed in the main text.
Single-task aceuacy (%)
100
80
60-
40-
20
Cars
DTD
EuTSAT GTSRB
MNIST RESISC45SUN397
SVHN
Non-linear FT
Linear FT
Post-hoe linear.
Zero-shot
VIT-B/32
Single-task accuracy (5)
100
80-
0-
Cazs
DTD
EuroSAT GTSRB
MNIST RESISCA5
SVHN
Non-linear FT
Linear FT
Post-hoe linear.
Zero-shot
(b) VIT-B/16
Single-task accuracy (%)
100
80
60-
40
20-
0-
Cars
DTD
EurOAAT GTSRB
MNIST RESISC45SUNW7
SVHN
Non-lineur PT
Lineur FT
Post-hoc linear.
Zero-shot
(c) VIT-L/14
Figure 7: Single-task accuracies (CLIP). Accuracy of different models obtained using different strategies on each of the tasks.
21

---

D.1 Fine-funing accuracies
In Figure 7, we report the single-task accuracies achieved by different CLIP models before fine-tuning (referred to as zero-shot), after fine-tuning with different dynamics (referred to non-linear FT and linear FT), and after linearizing the non-linearly fine-tuned models (post-hoc linearization).
These results demonstrate that non-linear fine-tuning consistently achieves the highest accuracy, indicating non-linear advantage. However, interesting observation that the gap between non-linear, linear, and post-hoe linearized models diminishes as the model size increases. This trend can be explained by the fact that larger models, which are more over-parameterized, inherently induce stronger kernel behavior during fine-tuning, As result. they tend to stay closer to the NTK approximation, closing the gap with linearized models.
D.2 Detailed results task addition
In addition to the results presented in Table in the main text, we report in Figure the absolute accura- cies of different CLIP models the single tasks before (zero-shot) and after performing task addition with different strategies (non-linear funetunnng, post-hoc linearization, and linear fine-tuning).
VIT-B/32
VIT-B/16
WI-L/114
Fane447
RuNSAT
CTSRB
500
11
500
XNNER
Ces MNIS
Dan
MNTE
3m
RISHSSS4
SUN39T
SUN22
SUNDI
liear FT
Linear
Post hoc Insoriaation
Zero-lutt
Figure 8: Task addition performance. Absolute accuracy (%) of each task after performing task addition with different linear/non-linear strategies and over different CLIP ViT models.
For all models and all datasets, except SUN397 [88], observe that linearized task arithmetic achieves the highest accuracies. Interestingly, as commented in the main text, the gap in performance between linear and non-linear task vectors decreases while increasing the size of the model. This observation aligns with the previous observation that fine-tuning with larger models better approximated by the NTK description.
22

---

D.3 Weight disentanglement and model scale
We note that weight disentanglement can also explain the increased performance of task arithmetic with model scale. As we see in Figure 9, the size of the areas with low disentanglement eπrоr grows with the model scale. One plausible explanation for this is that larger models inherently induce stronger kernel behavior during fine-tuning. Namely, since the models have more parameters, each parameter has to change less to fit the training examples. As a result, they tend stay closer the NTK approximation, closing the gap with linearized models and taking benefit of the better weight disentanglement of the models lying in the tangent space.
VIT B 32, DTD SUN397
VIT-B-10 1G, DTD SUN397
VIT-L-14, DTD SUN397
f(or, og)
3,0
3.0
200%
Non-linear FT
8'
-3.0
-3.0
-3.0
08
-3.0
-1.0
1.0
8.0
-30
-1.0
1.0
3.5
-3.0
-1.0
1.0
30
VIT-B-32, DTD SUN397
VIT-B-10, DTD SUN397
WIT-L-14 DTD SUN397
3.0-
3.0
100%
Linearized FT
1.0
1.0
1.0
6'
0,
-1.0
-1.0
-1.0
-3.0
-3.0
1.0
3.0
-30
-300
Figure 9: Weight disentanglement and model scale. The heatmaps show the disentanglement error $(01, (22) of different non-linear CLIP ViT's (top) and their post-hoc linearizations (bottom) on DTD and SUN397. The light regions denote areas of the weight space where weight disentanglement stronger. The red box delimits the search space used to compute the best in all our experiments.
23

---

D.4 Weight disentanglement of linearized and random models
In Figure 10, we present the disentanglement error of a linearized CLIP ViT-B/32 model across three different dataset pairs. By comparing these results with Figure 3, we can clearly observe that linearized models exhibit significantly more weight disentanglement compared to their non-linear counterparts, similar to the findings obtained for post-hoc linearization.
Conversely, in Figure 11, showcase the disentanglement error of CLIP ViT-B/32 model that was non-linearly fine-tuned starting from a random initialization. In all panels, we observe a high disentanglement error, which supports the claim that weight disentanglement and, consequently, task arithmetic are emergent properties of pre-training.
EunSAT- SVHN
DTD SUN397
Cars RESISC45
((m,,
3.0
3.0
.003
Linearized PT
1.0
1.0
1.0
0
0
-1.0
-1.
-10
06
3.0 -30
-3.0
-1.0
Figure 10: Visualization of weight disentanglement from linearized models. The heatmaps show the disentanglement error E(a1, 0:2) of a VIT-B/32 linearly fine-tuned different example task pairs. The light regions denote areas the weight space where weight disentanglement is stronger. The red box delimits the search space used to compute a in our experiments.
EunSAT
SVHN
DTD SUN397
Cars
RESISC45
((my)
3.
30
3.0
200%
Random
0
-1.0-
-1.0
300
-3.0
-3.0
08
300
-1.0
1.0
3.0
3.0
-1.0
1.0
3.0
3300
-1.0
1.3
30
Figure 11: Visualization of weight disentanglement from random initialization. The heatmaps show the disentanglement error {(01, of ViT-B/32 fine-tuned from random initialization on different example task pairs. The light regions denote areas of the weight space where weight disen- tanglement stronger. The red box delimits the search space used to compute a in our experiments.
24

---

D.5 Task arithmetic with convolutional architecture
We replicate our task addition experiments using a convolutional architecture rather than ViT. Specifically, finetune ConvNeXt [47] pre-trained on LAION-400M using CLIP [78] the 8 tasks from our task addition benchmark. We observe that also for this architecture linearized fine-tuning improves task arithmetic performance (see Table 4).
Table 4: Task addition with CNN. Average absolute (%) and normalized accuracies (%) of CLIP ConvNeXt edited by adding the sum of the task vectors of 8 tasks. We report results for the non-linear and linearized models normalizing performance by their single-task accuracies.
ConvNeXt
Method
Abs. (1) Norm.
Pre-trained
f((b))
57.5
Non-lin. FT 5(05 + +rdd
79.1
83.6
Linear. FT fisn + mim)
81.1
85.7
D.6 Task arithmetic with closed vocabulary models
We also replicate task addition experiments using ViT-B/16 pre-trained ImageNet-IK 1k using standard supervised learning. Note that this a closed vocabulary model, and therefore, when fine-tuning the different tasks need also fine-tune randomly initialized head. When interpolating the different task vectors we freeze the resulting head to arrive at the final models.
Interestingly, find that this closed-vocabular model can do task arithmetic, albeit lower performance (both in terms of absolute and normalized accuracy) than the open vocabulary ones (see Table 5). Remarkably, linearized fine-tuning also enhances discntanglement in this model - it yields higher normalized accuracies but does not perform better than non-linear task addition due to its lower tanle-task performance.
Table 5: Task addition with closed vocabulary model. Average absolute (%) and normalized accuracies (%) of a ViT-B/16 pre-trained on ImageNet-Ik kk edited by adding the sum of the task vectors of 8 tasks. We report results for the non-linear and linearized models normalizing performance by their single-task accuracies.
ViT-B/16 (supervised)
Method
Abs. (t) Norm. (f)
Pre-trained f((5")
2.3
Non-lin. FT f(0g² + Td)
54.9
52.2
Linear. FT fue. (0p' + +mm)
41.8
66.1
D.7 Weight disentanglement other architectures and modalities
To substantiate the generality of weight disentanglement, we conduct new experiment on a pre- trained T5-Base model [70] from Hugging Face Hub [84], fine-tuned two benchmark Natural Language Processing tasks, i.c., sentiment analysis on movie reviews with the IMDB dataset [51] and question answering with the QASC dataset [42]. The results, illustrated in the right panel in Figure 12, show a notable region around the pre-trained checkpoint characterized by low disentanglement error. This finding echoes the ability of T5 to perform task arithmetic demonstrated in Ilharco et al. [39], thereby reinforcing the robustness our conclusions.
Similarly, in Figure 12 we also see that the tangent space of CLIP ConvNeXt models is also more disentangled than the non-linear function space around the pre-trained initialization. The same findings also apply to closed vocabulary models, as shown in Figure 13.
25

---

Figure 13: Weight disentanglement of closed vocabulary models. The heatmaps show the disentan- glement error E(a), a2) of a ViT-B/16 model pre-trained on ImageNet-Ie Ik using supervised training in the non-linear function space (top) and its linearizations (bottom) different task pairs. The light regions denote areas of the weight space where weight disentanglement stronger. The red box CrovNeNs, Caus 455
Lin CreNNet, Cars RESISC44
T5 Bese, IMB QASC
30
56
30-
-30
-50
-.0
-1.0
1.0
2.0
-3.0
3.0
Figure 12: Visualization of weight disentanglement for other architectures and modalities. The heatmaps show the disentanglement g(o1, (r2) of a non-linearly and linearly fine-tuned ConvNeXt on pair of vision tasks (two left panels) and a T5-Base model fine-tuned on a pair of NLP tasks (right).
EumSAT SVHN
DTD SUN397
Cars RESISC45
f os)
3.0
3.0
200%
Non-linen linear FT
0
0
0,7
-1.0
-1.0
-3.0
-30
-3.0
3.0
-.0
1.0
3.0
-30
1.0
1.0
3.0
3.0
-1.0
1.0
EurSWT SVIIN
DTD SUN397
Cars - RESISCA5
(((g-02)
3.0
30-
3.0
200%
Linearized FT
0
0
0
-30
-3.0
-3.0
-3.0
-1.0
1.0
3.0 -3.0
-1.0
1.0
3.3
-3.0
-1.0
1.0
3.0
delimits the search space used to compute the best in all our experiments.
D.8 Localization eigenfunctions of CLIP's NTK
In Figure 14, we plot the local energy of the NTK cigenfunctions for pre-trained CLIP ViT-B/32 model evaluated three different data supports and control data supports. These panels complement the information presented Figure in the main text, where we observed that the CLIP has eigenfunctions whose energy is concentrated points belonging the respective dataset.
In Figure 15, we extend this analysis to a we observe a non-trivial but considerably poor degree of eigenfunction localization. This observation aligns with the finding that randomly-initialize linearized models cannot perform task arithmetic. Indeed, as we showed in the previous subsection, in this case the model's weights are not effectively disentangled, hindering its ability perform task arithmetic operations. In summary, eigenfunction localization offers complementary perspective the limitations of mandomly-initiaizze models.
initialized CLIP ViT-B/32 model. all panels,
26

---

Euiz)= = EB
RESISC4S
Can
DTD
SUNDO7
MINIST
EMeSAT
Figure 14: Eigenfunction localization. Estimated support of the eigenfunctions of the NTK of VIT-B/32 CLIP model trained on different datasets. The plot shows the sum of the local energy of the eigenfunctions over random subset of the training and control supports
Eu] Σ 02(x)
Blx)
RESISCAS
Cans
MNSST
C.C. EuroSAT
Figure 15: Eigenfunction localization. Estimated support of the eigenfunctions the NTK of randomly initialized ViT-B/32 model trained on different datasets. The plot shows the sum of the local energy eigenfunctions over a random subset the training and control supports
D.9 Further experiments with randomly-initialin networks
We conclude by showing, in Figure 16, the absolute single-task accuracy achieved by different CLIP ViT models that were fine-tuned from random initialization. Both the base models achieve non-trivial or moderate accuracy on the majority of benchmark tasks, using both non-linear and linearized fine-tuning dynamics.
These findings reinforce the intuition that non-pretrained models are not failing in task arithmetic due to their inability to learn the task initially. Instead, as argued earlier, the primary reason for the failure of non-pre-traid trained models in task arithmetic is their lack of weight disentanglement.
Interestingly, the performance of the randomly -nnttal-ii large model is generally poorer compared the base models. This observation can be attributed to the models' tendency to overfit the training data, which is more likely to occur when model has larger capacity.
27

---

Single-task accuracy (%)
80
60-
40
20
Cars
DTD
EuroSAT GTSRB MNIST RESISC45 SUN397 SVHN
Non-linear FT
Linear FT
Random init
VIT-B/32
Single-task accuracy (%)
60
40
20
Cans
DTD
EuroSAT GTSRB MNIST RESISC45 SUN:97
SVHN
Non-linear FT
Linear FT
Random init
(b) VIT-B/16
Single-task accuracy (%)
80
40
20
0
Cans
DTD
EuroSAT GTSRB
MNIST RESISC45 SUN:97
SVHN
Non-linear FT
Linear FT
Random init
(c) VIT-L/14
Figure 16: Single-task accuracies (random init). Accuracy of different models obtained using different strategies on each of the tasks.
28


Reviews:
[{"review_id": "MpKbXwmZKw", "reviewer": "Reviewer_qb7d", "summary": "This work challenges the prevailing belief regarding the origin of task arithmetic in CLIP models. While it is commonly attributed to the linear nature of fine-tuning, the authors argue that the critical factor lies in the \"weight disentanglement\" that occurs between tasks during the fine-tuning process. The paper presents a theoretical analysis and extensive empirical validation to support this claim.", "questions": "I’ve seen that in the supplementary material, there are some code snippets, but do the authors plan to release the full codebase upon acceptance?", "limitations": "The authors adequately addressed the limitations.", "rating": 10, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "This paper explores an intriguing topic with promising real-world applications. The authors exhibit strength in establishing a solid theoretical foundation, formulating straightforward research questions, and providing well-justified explanations that strike a remarkable balance between accessibility and rigor.\n\nThe content is commendably clear and understandable, facilitating understanding for readers from diverse backgrounds. Personally, I was able to (at least, I think) grasp portions of the paper (Section 6) which required prerequisite knowledge I was not already familiar with.\n\nNotably, this paper achieves a synergy between its experimental design and the overarching claim that \"Task arithmetic is not merely a result of linear fine-tuning,\" but rather depends on \"weight disentanglement of the model with respect to the fine-tuning task set.\" The authors successfully validate the performance of their proposed method through rigorous experiments (and ablations) while simultaneously providing substantial support for key theoretical hypotheses.\n\nOverall, this paper represents a remarkable contribution, embodying the meticulousness, clarity, and scholarly standards I would expect in top-tier NeurIPS submissions.\n\nI’m starting my recommendation with an 8, given my not-high confidence, I’ll adjust it according to the rebuttal discussion.", "weaknesses": "One aspect that requires attention is the clarification of the paper's applicability beyond CLIP models. Although there are specific references throughout the text (e.g., lines 38-49 in the introduction) emphasizing the focus on CLIP models, the overall messaging may appear too general regarding the broader applicability of the proposed method and study. To address this, I recommend at least modifying the abstract to explicitly state the paper's objective as \"… a comprehensive study of task arithmetic in **CLIP** models…\" instead of using the term \"vision-language models\" which could be misleading. Furthermore, exploring preliminary tests or investigations concerning the weight disentanglement of different pre-trained models would be valuable. For instance, considering the citation of various architectures in Section 6.2 (lines 301-303), extending the analysis to include some of them (e.g., convolutional neural networks) would be advantageous, strengthening the paper's applicability and relevance to a broader range of models and areas.\n\nAdditionally, while Tables 1 and 2 provide average performance across tasks, it would be beneficial to include the standard deviation to provide a more comprehensive understanding of the variability in performance and, therefore, the method's robustness.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the reviewer’s recognition of our work and their engagement to improve it!  In what follows, we address their comments.\n \n**Generality of our results beyond CLIP/ViT models**\n \nWe would like to emphasize that all our theoretical results are directly applicable to any model which satisfies Property 1 (Task arithmetic) regardless of pre-training scheme or architecture. Notably, any model satisfies Property 1 if and only if it is weight disentangled. However, in order to show that our results are robust to the architecture choice and input modality, we have replicated our experimental analysis using convolutional neural networks and large language models:\n\n- We have repeated all our experiments using a CLIP ConvNeXt model pre-trained on LAION-400M. Remarkably, also for this architecture weight disentanglement is stronger in the tangent space to the pre-trained checkpoint (see left panels in Figure R.1 of the *Author Response document* attached [here](https://openreview.net/forum?id=0A9f2jZDGW&noteId=7E6o5YEkJw)) and linearized fine-tuning enhances task arithmetic (see Table R.1 of the *Author Response document*). We will present the complete set of new results in the revised version.\n\n- In addition, in order to investigate the effects of pre-training, we are repeating the same analysis for non-contrastively pre-trained (closed-vocabulary) ViT and ConvNeXt models, and we will add these results to the final version as well. As suggested, we will also make explicit what architectures we consider, both in the abstract and in the main text.\n\n- To substantiate the generality of weight disentanglement, we conducted a new experiment on a pre-trained T5-Base model from Hugging Face, fine-tuned on two benchmark NLP tasks (sentiment analysis on movie reviews and question answering). The results, illustrated in the right panel in Figure R.1 of the *Author Response document*, show a notable region around the pre-trained checkpoint characterized by low disentanglement error. This finding echoes the ability of T5 to perform task arithmetic as demonstrated in Ilharco et al. [39] (Appendix D.6), thereby reinforcing the robustness of our conclusions. We will report this strong result in the paper. Finally, while exploring whether linearization improves weight disentanglement and task arithmetic for different modalities – such as language – is undoubtedly captivating, it goes beyond the current scope of the paper, and we reserve it for subsequent investigations.\n \n**Robustness of our method**\n \nIn regard to the variability in performance in Tables 1 and 2, it is important to consider that the diverse nature of the tasks results in distinct levels of difficulty. Hence, standard deviations are likely more affected by this variability in task difficulty than the robustness of a given method. Nevertheless, we concur with the reviewer that solely looking at averages might not convey the whole picture. To address this, we showed that the improvements in performance are consistent across tasks in Appendix D.2. This clarification will be further stressed in the revised version.\n \n**Code release**\n \nWe confirm that we will release the complete codebase in a public GitHub repository once the work undergoes deanonymization.\n \nWe thank the reviewer for their valuable feedback and remain available to answer further questions or provide more clarifications regarding the previous points.\n"}, {"review_id": "v6Lz3dFhri", "reviewer": "Reviewer_DjkZ", "summary": "This paper presents a comprehensive analysis of task arithmetic using pre-trained CLIP models. It challenges the early hypothesis that task arithmetic arises from linear fine-tuning in the NTK regime and introduces weight disentanglement as a necessary condition for enabling task arithmetic.  Further experiments demonstrate that linearized fine-tuning of pre-trained CLIP models exhibits stronger weight disentanglement and improved task arithmetic compared to standard, non-linear fine-tuning. This result is accompanied by an analysis of the NTK spectrum to facilitate the understanding of weight disentanglement in linearized models. Lastly, empirical evidence suggests that weight disentanglement emerges from large-scale pre-training. The findings of the paper may shed new light on the effective adaptation of foundation models for downstream applications.", "questions": "These questions are likely outside the scope of the paper, yet I think addressing any of them via an empirical analysis could strengthen the work.\n\n- Task arithmetic makes the strong assumption that datasets from different tasks have disjoint support. This is often not true in practice. Does task addition still yield cooperative behavior when datasets overlap? How does linearized fine-tuning compare to standard, non-linear fine-tuning, especially when the assumption of non-overlapping data is violated?\n\n- In linearized fine-tuning, the optimization is restricted to the tangent space of pre-trained weights. Notably, this is conceptually similar to LoRA and adaptor-based fine-tuning, where a task vector exhibits low-dimensional structure. A natural question to ask is whether these fine-tuning approaches also produce favorable weight disentanglement / task arithmetic.", "limitations": "The paper discussed two main limitations. First, linearized fine-tuning is currently implemented using the JVP algorithm, which doubles the cost of a forward pass as compared to standard fine-tuning. Second, the spatial localization of NTK eigenfunctions is a sufficient (yet not necessary) condition for enabling task arithmetic. While the linearized models indeed respond to disjoint spatial regions in the experiments, this is not a must for task arithmetic to hold.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- This paper advances the theoretical understanding of task arithmetic. In particular, it introduces weight disentanglement as a strong indicator of task arithmetic, and demonstrates that linearized fine-tuning in the tangent space of pre-trained weights promotes weight disentanglement. It further analyzes the NTK spectrum of linearized models and presents a sufficient condition for weight disentanglement.\n\n- In the meantime, the findings of the paper are significant on the practical side. The experiment results justify an emerging fine-tuning scheme for adapting foundation models.\n\n- Overall, this paper strikes a good balance between theory (NTK) and practice (fine-tuning pre-trained models). Embracing a broad audience is a key strength of the paper.\n\n- Finally, the paper is very well-written. The flow of presentation is very easy to follow.", "weaknesses": "- Overall, I found the theoretical analysis motivating and the experiment results convincing. That said, all conclusions of the paper are drawn from CLIP fine-tuning, which makes me wonder whether the same findings are valid for other pre-trained models that similarly exhibit task arithmetic. To this end, I encourage the authors to report results on a second pre-trained model that differs in network architecture (e.g., ResNet), learning objective (e.g., MAE) or input modality (e.g., natural language), in order to establish the generality of their findings.\n\n- The authors observed that increasing model size (ViT-B/32->ViT-L/14) closes the gap between linearized and non-linear fine-tuning. Is it a consequence of stronger weight disentanglement because the fine-tuning of larger models approaches the NTK regime? Visualization of weight disentanglement throughout the paper only considers the smallest ViT-B/32 model. Similar visualizations for larger models could be informative.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate that the reviewer recognized that our paper is well-written, our analysis is motivating, and our experiments are convincing and their engagement to improve it. Below, we address their comments.\n \n**Generality of our results beyond CLIP/ViT models**\n \nWe would like to emphasize that all our theoretical results are directly applicable to any model which satisfies Property 1 (Task arithmetic) regardless of pre-training scheme or architecture. Notably, any model satisfies Property 1 if and only if it is weight disentangled. However, in order to show that our results are robust to the architecture choice and input modality, we have replicated our experimental analysis using convolutional neural networks and large language models:\n\n- We have repeated all our experiments using a CLIP ConvNeXt model pre-trained on LAION-400M. Remarkably, also for this architecture weight disentanglement is stronger in the tangent space to the pre-trained checkpoint (see left panels in Figure R.1 of the *Author Response document* attached [here](https://openreview.net/forum?id=0A9f2jZDGW&noteId=7E6o5YEkJw)) and linearized fine-tuning enhances task arithmetic (see Table R.1 of the *Author Response document*). We will present the complete set of new results in the revised version.\n\n- In addition, in order to investigate the effects of pre-training, we are repeating the same analysis for non-contrastively pre-trained (closed-vocabulary) ViT and ConvNeXt models, and we will add these results to the final version as well. As suggested, we will also make explicit what architectures we consider, both in the abstract and in the main text.\n\n- To substantiate the generality of weight disentanglement, we conducted a new experiment on a pre-trained T5-Base model from Hugging Face, fine-tuned on two benchmark NLP tasks (sentiment analysis on movie reviews and question answering). The results, illustrated in the right panel in Figure R.1 of the *Author Response document*, show a notable region around the pre-trained checkpoint characterized by low disentanglement error. This finding echoes the ability of T5 to perform task arithmetic as demonstrated in Ilharco et al. [39] (Appendix D.6), thereby reinforcing the robustness of our conclusions. We will report this strong result in the paper. Finally, while exploring whether linearization improves weight disentanglement and task arithmetic for different modalities – such as language – is undoubtedly captivating, it goes beyond the current scope of the paper, and we reserve it for subsequent investigations.\n\n**Effect of scale on non-linear advantage**\n \nThe reviewer correctly highlights that by scaling the number of model parameters, the performance of linearized fine-tuning becomes closer to the one of standard non-linear fine-tuning. To the best of our knowledge, this is a novel observation.  As commented briefly in Appendix D.1, one plausible interpretation is that larger models, which are more over-parameterized, inherently induce a stronger kernel behavior during fine-tuning. Namely, since the models have more parameters, each parameter has to change less to fit the training examples. As a result, they tend to stay closer to the NTK approximation, closing the gap with linearized models and taking benefit of the better weight disentanglement of the models lying in the tangent space. \n\nIn response to the reviewer’s suggestion, we conducted supplementary experiments that visualize how weight disentanglement varies changing the model scale (see Figure R.2 of the *Author Response document*). Consistently to our results, larger models exhibit stronger weight disentanglement, as highlighted by the larger light region in the right panel of the first row of Figure R.2. Yet, interestingly, the models linearly fine-tuned are always more weight disentangled than the non-linearly fine-tuned ones, highlighting the strength of linearized models for model editing. We will add this discussion and the new results to the paper.\n \n**Open questions**\n \nWe thank the reviewer for raising these stimulating questions, although we agree that these are beyond our current scope. In particular, the disjoint support assumption adopted in our work stems from the way in which task arithmetic was first introduced in Ilharco et al. [39], wherein the datasets are indeed disjoint. Yet, we acknowledge that studying task arithmetic in other settings is an interesting direction (see also the answer Disjoint Task Support Hypothesis to Reviewer [eVrq](https://openreview.net/forum?id=0A9f2jZDGW&noteId=QWXTFZW72i)). Similarly, studying the sparsity of task vectors and the effects of LoRA training are exciting open questions and avenues that warrant separate exploration.\n \nWe thank the reviewer for their valuable feedback and remain available to answer further questions or provide more clarifications regarding the previous points.\n"}, {"review_id": "CDcYUaXFFc", "reviewer": "Reviewer_sE5a", "summary": "This paper theoretically and emperically investigates the reasons why task arithmetic (an emerging technique for editing pre-trained neural networks) works.\nThe paper shows that, contrary to previous hypotheses [39,79,80], lienarity of the fine-tuning on individual tasks is not sufficient to fully explain the success of task arithmetic (Sec 3).\nInstead, the authors propose the related and straight-forward idea of weight disentaglement (Eq 4) as an explanation (Sec 4), also providing a measure of the disentaglment error for two tasks (Eq 6).\nThis idea is leveraged to improve task arithmetic performance by constraining the fine-tuning to linearized models - which importantly is shown to be an improvement over post-hoc linearlization, with marginal additional computational complexity (Sec 5).\nFinally, an additional connection between task arithmetic and the eigen space of the Neural Tangent Kernel is used to argue that weight disentanglement (and hence the capacity for task artihmetic) are a learned properties during pre-training, not inherent properties of parameterization or architecture (Sec 6).\nThe empirical results focus on the CLIP image-text Vision Transformer model family, and many experimental details are reported, inclduing in the additional material in the appendix.", "questions": " 1. Ln 336 on 'Feature disentanglement'; to what degree is feature disentanglement enabled by over-parameterization / very large numbers of parameters? If this is the case, what is the parameter count relative to - i.e. some measure of the complexity of the input space or the function space represented by the data? Is there a way to quantify with a metric the 'capacity' for feature disentaglement of a given NN archiecture? These could be useful areas for investigation.\n 2. Related to the above point, Ln 233 notes that the advantage of the linearized model task arithmetic approach diminishes as the number of model parameters increases. I find this curious and would have expecte the opposite result - do you have any intuition why this is the case?\n 3. Pseudo-code or actual PyTorch code for Eq 6 would be a very helpful (and I assume straight-forward) addition to your supplementary material - this would be of immediate practical help to other researchers working on practical task arithmetic applications.\n 4. Related to the above, Ln 297 - can the degree of local linear independence of the NTK eigenfunctions be measured/computed with a simple metric to test this condition in other models / datasets? This would likewise be a useful contribution of the paper.", "limitations": " * It's not clear to me if there are situations where the linerarization technique won't be readily applicable. E.g. are some activation functions or more esoteric neural network architectures (e.g. recurrent NN's; long-term memory components) going to problematic to linearize? Some negative examples of where NTK linearization can't be used and/or approximations or work-arounds could be a useful addition in this regard.", "rating": 9, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "This paper is the best kind of NeurIPS paper; Beautifully written and a delight to read, the authors have considered a timely and pertinent problem in an emerging machine learning domain, applied a methodological theoretical investigation; clarified prior hypotheses in the literature; used the theoretical findings to propose a simple novel methodology, and compellingly evidenced the subsequent effectiveness with an appropriate evaluation protocol. As a cherry on the cake, the methodology is a 'drop in' method that can be applied easily to existing approaches, and actual code is provided without breaking review anonymity in the supplementary material. Bravo.\n\n### Originality\n\n * There are multiple practical take-aways that are immediately useful; the task disentanglement error metric (Eq 6), intuition for why task vector coefficients should be << 1 (Ln 181), simple drop-in code in the appendix that enables re-production and immediate application of the ideas (Listing 1 in the appendix).\n\n### Quality\n\n * The work appears to be of high quality, the metrics and emprical appear to support the theoretical findings and claims.\n\n### Clarity\n\n * This paper is superbly well written, and the logical argument and flow is well constructed. This paper was a delight to read and think about.\n\n### Signifigance\n\n * Multiple compelling and genuine avenues for future research are identfidied (Ln 239, Ln 286, Ln 315, Ln354)\n * Provides timely and much-needed (i) overview of the emerging and rapidly evolving technique of task arithmetic with pre-trained models, and (ii) the beginnings of a theoretical grounding and understanding for why this technique is possible.\n", "weaknesses": "* The main weakness of the paper in the present form is that the empirical results are limited to ViT models (specifically, CLIP). Within this modality, the empirical results are compelling (e.g. using multiple model families and 8x distinct datasets/tasks), and provide evidence for the theoretical findings, however the paper's authority and strength would be enhanced by demonstrating some of the key findings hold with another architecture or data type. E.g. I would love to see the results in Table 3 (Sec 6.2) replicated, even in a small way, with GPT or BERT on a text dataset. This would provider further strength for the empirical claims around Eigenfunction localization.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We really appreciate the reviewer’s enthusiasm and acknowledgment of the significance of our work and their engagement to improve it!  Below, we address their comments.\n \n**Generality of our results beyond CLIP/ViT models**\n \nWe would like to emphasize that all our theoretical results are directly applicable to any model which satisfies Property 1 (Task arithmetic) regardless of pre-training scheme or architecture. Notably, any model satisfies Property 1 if and only if it is weight disentangled. As suggested, in order to show that our results are robust to the architecture choice and input modality, we have replicated our experimental analysis using convolutional neural networks and large language models:\n\n- We have repeated all our experiments using a CLIP ConvNeXt model pre-trained on LAION-400M. Remarkably, also for this architecture weight disentanglement is stronger in the tangent space to the pre-trained checkpoint (see left panels in Figure R.1 of the *Author Response document* attached [here](https://openreview.net/forum?id=0A9f2jZDGW&noteId=7E6o5YEkJw)) and linearized fine-tuning enhances task arithmetic (see Table R.1 of the *Author Response document*). We will present the complete set of new results in the revised version.\n\n- In addition, in order to investigate the effects of pre-training, we are repeating the same analysis for non-contrastively pre-trained (closed-vocabulary) ViT and ConvNeXt models, and we will add these results to the final version as well. As suggested, we will also make explicit what architectures we consider, both in the abstract and in the main text.\n\n- To substantiate the generality of weight disentanglement, we conducted a new experiment on a pre-trained T5-Base model from Hugging Face, fine-tuned on two benchmark NLP tasks (sentiment analysis on movie reviews and question answering). The results, illustrated in the right panel in Figure R.1 of the *Author Response document*, show a notable region around the pre-trained checkpoint characterized by low disentanglement error. This finding echoes the ability of T5 to perform task arithmetic as demonstrated in Ilharco et al. [39] (Appendix D.6), thereby reinforcing the robustness of our conclusions. We will report this strong result in the paper. Finally, while exploring whether linearization improves weight disentanglement and task arithmetic for different modalities – such as language – is undoubtedly captivating, it goes beyond the current scope of the paper, and we reserve it for subsequent investigations.\n \n**Disentanglement and overparameterization**\n \nWe are not aware of any results showing to what extent feature disentanglement is enabled by overparameterization. Intuitively, achieving a distinct disentanglement in feature space, as well as in weight space in the case of weight disentanglement, demands a sufficient number of model parameters. A systematic study of these quantities as the model size scales lies beyond our current scope. However, we concur that this is a fascinating question for future research.\n\n**Linear vs non-linear**\n \nOur intuition for the observation of the diminishing advantage of linearized-model task arithmetic with increasing model size is that larger models tend to stay closer to the NTK approximation. In other words, they tend to behave linearly without being explicitly constrained to do so (see, e.g., Figure 2). One explanation for this is that larger models, being more over-parameterized, inherently induce a stronger kernel behavior during fine-tuning, i.e., having more parameters, each parameter requires smaller adjustments to fit the training examples. As a result, being closer to their linearized counterparts, larger models have better weight disentanglement and can perform task arithmetic similarly to linearized models (see also the answer Effect of Scale on Non-Linear Advantage to reviewer [DjkZ](https://openreview.net/forum?id=0A9f2jZDGW&noteId=PEwcdQxgup) and the new results in Section 3 of the *Author Response document* displaying weight disentanglement as a function of model size).\n \n**Pseudo-code for the weight disentanglement error**\n \nWe thank the reviewer for their suggestion of providing the PyTorch code for computing the disentanglement error in the Appendix. This addition will enhance reproducibility, and we intend to implement it in the revised version.\n \n**Eigenfunction localization**\n \nApproximating the eigenfunctions of the NTK is a costly operation since it requires computing the kernel matrix and diagonalizing it. Hence, in general, measuring localization or other properties of the eigenfunctions is challenging. Moreover, in practice, all these properties are not displayed exactly, so a sound investigation should take into account this fact as well. All in all, this precludes an exhaustive exploration within the confines of this paper. Yet, we agree and believe that studying the spectral properties of the NTK for different datasets, architectures, and modalities holds promise for future research. We will acknowledge the importance of this avenue in our manuscript, reflecting your input.\n \n**Linearization of other architectures**\n \nLinearization readily works with the majority of architectures currently used, encompassing convolutional and transformer architectures with both smooth and non-smooth activation functions (as shown by our new experiments using ConvNeXts). Although our current framework doesn't explicitly address recurrent architectures, we believe that implementing linearization for recurrent architectures is possible We will explicitly outline the range of applicability of our procedure in Appendix B (Implementation aspects of linearized models).\n \nWe thank the reviewer for their valuable feedback and remain available to answer further questions or provide more clarifications regarding the previous points.\n"}, {"review_id": "CetQ5EK9GW", "reviewer": "Reviewer_9cnX", "summary": "This paper presents a comprehensive theoretical and empirical analysis of task arithmetic for model editing, where adding different task vectors (obtained by taking the difference between fine-tuned and pretrained model checkpoints) could improve the model’s performance on these tasks and vice versa. The authors propose weight disentanglement to investigate the underlying principles of task arithmetic, which involves decomposing the learned model function into a sum of localized components with disjoint supports. Specifically, the authors compare **regular** and **post-hoc linearized** fine-tuned models and find that weight disentanglement significantly contributes to the ability of task arithmetic. Based on these insights, the authors then propose to directly employ the linearized model for fine-tuning, obtain optimized task vectors, and lead to improved task arithmetic performance. Further analyses are conducted to reveal the connection between task arithmetic and weight disentanglement.", "questions": "1. How does model scale impact weight disentanglement, and thus the benefits of linearized fine-tuning?\n2. (Minor) Can the findings about task arithmetic be generalized to natural language tasks as well? By intuition, it seems that the setting of natural language texts is more straightforward for weight disentanglement due to the inherent discreteness of text tokens.", "limitations": "The authors have adequately addressed the limitations.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- The paper is clearly written and well organized.\n- This work presents a neat analysis of task arithmetic based on the use of linearization and neural tangent kernels, which is novel to my knowledge and interesting, offering a fresh perspective on understanding the geometry of pre-trained checkpoints’ weights.\n- The proposed method to further improve task arithmetic is simple yet effective, supported by extensive empirical findings.", "weaknesses": "As the authors already mentioned, one potential weakness is the introduced computational overhead during training. However, since the main focus of this work is a comprehensive analysis of task arithmetic, this is not a significant issue in the context of this work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the reviewer’s recognition of our work and their engagement to improve it!  Below, we address their comments.\n \n**Effect of scale on weight disentanglement**\n \nOur results reveal that by scaling the number of model parameters, the performance of linearized fine-tuning becomes closer to the one of standard non-linear fine-tuning. As commented briefly in Appendix D.1, one plausible interpretation is that larger models, which are more over-parameterized, inherently induce a stronger kernel behavior during fine-tuning. Namely, since the models have more parameters, each parameter has to change less to fit the training examples. As a result, they tend to stay closer to the NTK approximation, closing the gap with linearized models and taking benefit of the better weight disentanglement of the models lying in the tangent space. \n\nAs also suggested by Reviewer [DjkZ](https://openreview.net/forum?id=0A9f2jZDGW&noteId=v6Lz3dFhri), we conducted supplementary experiments that visualize how weight disentanglement varies changing the model scale (see Figure R.2 of the *Author Response document* attached [here](https://openreview.net/forum?id=0A9f2jZDGW&noteId=7E6o5YEkJw)). Consistently to our results, larger models exhibit stronger weight disentanglement, as highlighted by the larger light region in the right panel of the first row of Figure R.2. Yet, interestingly, the models linearly fine-tuned are always more weight disentangled than the non-linearly fine-tuned ones, highlighting the strength of linearized models for model editing. We will add this discussion and the new results to the paper.\n \n**Generalization to NLP tasks**\n \nTask arithmetic and our findings on weight disentanglement can be readily generalized to NLP tasks. In order to show the generality of weight disentanglement, we conducted a new experiment on a pre-trained T5 base model from Hugging Face, fine-tuned on two benchmark NLP tasks (sentiment analysis on movie reviews and question answering). The results, illustrated in the right panel in Figure R.1 of the *Author Response document*, show a notable region around the pre-trained checkpoint characterized by low disentanglement error. This finding echoes the ability of T5 to perform task arithmetic as demonstrated in Ilharco et al. [39] (Appendix D.6), thereby reinforcing the robustness of our conclusions. We will report this result in the revised version.\n"}, {"review_id": "kbSkLPUU32", "reviewer": "Reviewer_eVrq", "summary": "This paper studies the \"task vectors\" framework where the weights of models can be perturbed in specified directions corresponding to tasks which result in improvements on those tasks. They attribute the success of this framework to \"weight disentanglement\" which means that adding a task vector for task i does not change how the network behaves when seeing input for task j != i. In addition, their investigation reveals that task vectors are not a consequence of fine-tuning occurring in the NTK regime (post-hoc linearization underperforms), but explicitly linearizing fine-tuning improves the task vectors framework as it \"amplifies\" \"weight disentanglement\".", "questions": "- Do the authors findings point to any \"better\" way of adding task vectors rather than just adding them for all weights in the network?\n- How does weight disentanglement change with scale?\n- Can the authors think of an experiment which might falsify their hypothesis? E.g., do they think that task vectors might still \"work\" if two tasks share the same input space. Here's a quick thought: get one task vector which corresponds to the problem \"first 5 classes in CIFAR10 or second 5 classes\" and another corresponding to the task \"traditional CIFAR10\". If you apply both these task vectors then look at CIFAR10 performance on just the first 5 classes I think performance would still be as good as just applying the CIFAR10 task vector?\n- Should the community switch to linearized fine-tuning?", "limitations": "Would have been nice to see a limitations section.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Overall I thought this paper was very interesting, it's claims are mostly precise and thoroughly investigated, and it's results are impressive (+ 5.8 for task addition). I though Figure 5 was particularly nice evidence of the authors hypothesis.", "weaknesses": "- Limited experiments for many tasks. I would be very curious to see what happens in e.g. table 1 if you scale up the number of tasks.\n- Limited attempt to falsify the hypothesis -- for instance can you consider two tasks which share the same input images? Would task vectors still work? (See question in questions section).\n- Throughout the paper the authors mention variations of \"Specifically, we probe the hypothesis presented in Wortsman et al. [80] that task arithmetic is possible thanks to the fact that models inherently operate in a linear regime\" but I cannot actually find anywhere in [80] where this hypothesis is stated? I thought that Wortsman et al. [79, 80] were observing that ensembles behaved roughly similar to averages in the fine-tuning regime, and simply used the NTK regime as an example of where this would occur, but noted that there was still differences between the two (which is not predicted by the NTK regime). How do the authors feel about the hypothesis that averaging models (~= task vectors) ~= output-ensembling + terms-which-may-be-small (e.g., Section 4 of [79]).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate that the reviewer reported that our paper is very interesting and has impressive results, and their engagement to improve it. Below, we address their comments.\n \n**Adding more tasks**\n \nFor the paper, we are using the experimental setting of Ilharco et al. [39], where task arithmetic was introduced. Namely, our task-addition experiments already involve a set of 8 different vision tasks. Yet, while the scalability of task arithmetic across a larger number of tasks is an appealing question, it goes beyond the paper’s objective and is left for future investigations. In the revised version of our manuscript, we will also include results on NLP tasks and using other convolutional architectures.\n \n**Disjoint task support hypothesis**\n \nOur theoretical analysis specifically focuses on non-overlapping scenarios, consistent with the case which is most studied in the empirical literature. In fact, the hypothesis of having disjoint task supports is satisfied in both Ilharco et al. [39] and our work. Encompassing overlapping task supports – which might require a slight change in the definition of weight disentanglement – remains an interesting problem. However, we do not see studying only the non-overlapping case as a weakness. Indeed, in the context of vision-language tasks, the input space is the Cartesian product of all images and captions, which is a high-dimensional space in which the overlap between tasks might be minimal. We will ensure to add this justification in the revised version.\n \n**Previous linearity hypothesis**\n \nWortsman et al. [80] observe in Appendix F that even though averaging weights and averaging output functions in practice are not exactly the same, these two methods are not as different as they appear. In particular, they report that these methods are equivalent in the NTK regime and this regime might be an accurate approximation of fine-tuning. Similarly, in Appendix A of Ilharco et al. [39], the authors justify their results on task arithmetic based on the NTK hypothesis for which they also cite Wortsman et al. [80], sharing many of the same authors.\n\nConcerning the question on the hypothesis introduced in [79], we do not have arguments to refute it. In fact, we believe that weight disentanglement might play a role in this setting as well if different models are specializing to different regions of the input space. In that case, the terms-which-may-be-small would be the disentanglement error.\n \n**Effect of scale on weight disentanglement**\n \nOur results reveal that by scaling the number of model parameters, the performance of linearized fine-tuning becomes closer to the one of standard non-linear fine-tuning. As commented briefly in Appendix D.1, one plausible interpretation is that larger models, which are more over-parameterized, inherently induce a stronger kernel behavior during fine-tuning. Namely, since the models have more parameters, each parameter has to change less to fit the training examples. As a result, they tend to stay closer to the NTK approximation, closing the gap with linearized models and taking benefit of the better weight disentanglement of the models lying in the tangent space. \n\nAs also suggested by Reviewer [DjkZ](https://openreview.net/forum?id=0A9f2jZDGW&noteId=v6Lz3dFhri), we conducted supplementary experiments that visualize how weight disentanglement varies changing the model scale (see Figure R.2 of the *Author Response document* attached [here](https://openreview.net/forum?id=0A9f2jZDGW&noteId=7E6o5YEkJw)). Consistently to our results, larger models exhibit stronger weight disentanglement, as highlighted by the larger light region in the right panel of the first row of Figure R.2. Yet, interestingly, the models linearly fine-tuned are always more weight disentangled than the non-linearly fine-tuned ones, highlighting the strength of linearized models for model editing. We will add this discussion and the new results to the paper.\n\n**Adoption of linearized fine-tuning**\n \nLinearized fine-tuning, as shown in our paper, gives multiple benefits for, e.g., ensembling, composition, and forgetting (negation). Notably, for convex losses, this method also enjoys a convex optimization landscape and could be used to provide further theoretical guarantees in the future. However, we want to remark that studying the performance of linear vs non-linear models is an ongoing line of research (see, e.g., references in Related Works – Linear vs non-linear regime). In fact, linearized fine-tuning may not be uniformly superior in all settings and architectures, although it seems to be the case for the models we studied in the present work. All in all, we hope that our findings will motivate further empirical exploration to discern in which cases linearized fine-tuning competes effectively with standard non-linear fine-tuning.\n \n**Limitations**\n \nOur limitations are clearly reported in the paper. In particular, we had a section in the Appendix to the O(1) computational complexity increase of linearized models and discussed the limitations of Proposition 1 as a remark in the main body of the text. Following the reviewer’s suggestion, we will emphasize those points further in the revised version of our work.\n \nWe thank the reviewer for their valuable feedback and remain available to answer further questions or provide more clarifications regarding the previous points.\n"}]

Author rebuttal:
We kindly thank all the reviewers for their time and for providing valuable feedback on our work. We appreciate that reviewers have pointed out that our work is interesting (Reviewer [eVrq](https://openreview.net/forum?id=0A9f2jZDGW&noteId=kbSkLPUU32)), intriguing (Reviewer [qb7d](https://openreview.net/forum?id=0A9f2jZDGW&noteId=MpKbXwmZKw)), and very well written (Reviewers [qb7d](https://openreview.net/forum?id=0A9f2jZDGW&noteId=MpKbXwmZKw), [DjkZ](https://openreview.net/forum?id=0A9f2jZDGW&noteId=v6Lz3dFhri), [sE5a](https://openreview.net/forum?id=0A9f2jZDGW&noteId=CDcYUaXFFc)), and that our results are solid (Reviewer [qb7d](https://openreview.net/forum?id=0A9f2jZDGW&noteId=MpKbXwmZKw)), impressive (Reviewer [9cnX](https://openreview.net/forum?id=0A9f2jZDGW&noteId=CetQ5EK9GW)), and impactful (Reviewers [DjkZ](https://openreview.net/forum?id=0A9f2jZDGW&noteId=v6Lz3dFhri), [qb7d](https://openreview.net/forum?id=0A9f2jZDGW&noteId=MpKbXwmZKw)).

In response to the reviews, we ran a series of **new experiments** to show the generality of our findings. Specifically,
- We have replicated our experimental analysis using a **convolutional architecture**. Our new results, reported in Table R.1 and Figure R.1 of the *Author Response document* (see attached pdf) reveal that also for this architecture weight disentanglement is stronger in the tangent space and linearized fine-tuning enhances task arithmetic.
- We have extended the experimental results on weight disentanglement to language by
demonstrating that a T5-Base model, fine-tuned on two distinct **NLP tasks**, exhibits a region around its pre-trained initialization with low weight disentanglement error (Figure R.1, right panel).
- We have analyzed the **effect of model scale** on weight disentanglement, showing that larger models are more weight disentangled, but not as much as their linearized counterparts (Figure R.2)

We hope that these new results and the clarifications detailed in the individual comments given to each reviewer will effectively address the concerns raised during the review process. We remain available for engaging in any further discussions that may arise, and we thank you once again for your comments.


Meta-review: This paper studies the use of task vectors, i.e. the difference between a fine-tuned model and a pre-trained model, to perform task arithmetic, i.e. adding or removing capabilities to a model. A new theoretical viewpoint is developed that links the success of task vectors to "weight disentanglement", and this perspective is used to develop a new method for performing task arithmetic that outperforms prior methods. Reviewers all appreciated the clearly articulated theory and convincing experiments and unanimously voted to accept the paper.
Decision: Accept (oral)