Title: Task Arithmetic in the Tangent Space: Improved Editing of Pre-Trained Models
Venue: NeurIPS 2023
Accepted: yes
Keywords: model editing

transfer learning

neural tangent kernel

vision-language pre-training

deep learning science
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

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.

1 Introduction

Pre-trained models play a 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 human-feedback [64], which constrain scalability and democratization. Furthermore, enhancing downstream task performance typically degrades the model's pre-training performance or zero-shot accuracy [28, 58, 87].

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 arithmetic 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

on two separate tasks results in a new multi-task model with improved performance on both tasks. Similarly, subtracting a 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 limited. Specifically, a comprehensive understanding of how these techniques affect a model's internal representations and the necessary conditions to make it reliable is lacking. This knowledge gap can undermine the adoption of these techniques, as it erodes their trustworthiness 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 pretrained vision-language models ( i.e. , CLIP [69]), offering novel insights into its underlying mechanisms and introducing new approaches which enhance the performance of pre-trained models edited through task arithmetic. Specifically, we probe the hypothesis presented in Wortsman et al. [87] that task arithmetic is possible thanks to the fact that models inherently operate in a 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 performance with respect to their nonlinear counterparts (see Tables 1 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 a model to 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 in 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, we achieve up to 5 . 8 points more and 13 . 1 points less in accuracy on task addition and task negation, respectively, on several vision-language benchmarks.

We link weight disentanglement in 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, facilitating 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 a broader range of tasks.

2 Notation and problem statement

Let f : X × Θ →Y be a neural network taking inputs x ∈ X and parameterized by a set of weights θ ∈ Θ . We will assume X ⊆ R d , Θ ⊆ R m and Y ⊆ R c . Consider T tasks, with every task t consisting of a triplet ( D t , µ t , f ⋆ t ) where D t ⊆ X is a data support ( e.g. , ImageNet [21] images), µ t an input distribution such that supp( µ t ) = D t , and f ⋆ t : D t →Y a target function ( e.g. , labels). In practice, each task is identified with a training set { ( x ν , f ⋆ t ( x ν )) } ν ∈ [ n t ] with x ∼ µ t , that is used to fine-tune the models starting from the pre-trained weights θ 0 and obtain the fine-tuned weights θ ⋆ t .

Task arithmetic. Let the task vector of task t be the difference between the fine-tuned and the pretrained weights, i.e. , τ t = θ ⋆ t -θ 0 . 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 = { τ t } t ∈ [ T ] with associated nonintersecting task supports D = {D t ⊂ X} t ∈ [ T ] , i.e., ∀ t, t ′ , if t = t ′ then D t ∩ D t ′ = ∅ . We say a network f satisfies the task arithmetic property around θ 0 with respect to T and D if with ( α 1 , . . . , α T ) ∈ A ⊆ R T .

In short, a model satisfies Property 1 if adding τ t does not modify the output of the model outside D t .

Neural tangent kernel. Around the initialization weights θ 0 , a neural network can be approximated with a first-order Taylor expansion:

This approximation is equivalent to a kernel predictor with a kernel known as the neural tangent kernel (NTK) [41], k NTK ( x , x ′ ) = ∇ θ f ( x ; θ 0 ) ⊤ ∇ θ f ( x ′ ; θ 0 ) , 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, as the evolution of parameters during 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 a 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 to similar output function combinations. However, we will now demonstrate that CLIP 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 ( · ; θ 0 ) 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 approximated by its first-order Taylor expansion, i.e., f ( x ; θ ⋆ ) -f ( x ; θ 0 ) ≈ ( θ ⋆ -θ 0 ) ⊤ ∇ θ f ( x ; θ 0 ) .

In simple terms, the approximation of the network in the tangent space around initialization must hold after fine-tuning. To test this, we evaluate the performance of the post-hoc linearized version of f , f lin . That is, we apply the fine-tuned task vectors τ = θ ⋆ -θ 0 to the linear approximation of f at θ 0 , i.e. , and we check whether f lin ( · ; θ ⋆ ) performs similarly to f ( · ; θ ⋆ ) 2 .

The results in Figure 2 indicate that CLIP models do not exhibit a kernel behavior. Specifically, we fine-tune (FT) several CLIP pre-trained Vision Transformers (ViTs) [24] of different sizes following the same setup as Ilharco et al. [39] on 8 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 f lin ( · ; θ ⋆ ) is significantly lower than that of f ( · ; θ ⋆ ) for ViTs of all sizes. This non-linear advantage [26] is a clear sign that fine-tuning has not happened in a linear regime as expected by Wortsman et al. [87].

Yet, this observation is not enough to rule out that task arithmetic can be explained by linearizing the network function. Indeed, even if the non-linear components are important for

50

60

70

80

90

100

Post-hoc linearization accuracy (%)

50

60

70

80

90

100

Non-linear fine-tuning accuracy (%)

Non-linear

advantage

ViT-B/32

ViT-B/16

ViT-L/14

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.

We now show this is also not the case, as doing task arithmetic with the non-linearly fine-tuned task vectors over f lin significantly decreases performance. To show this, we employ the benchmark proposed in Ilharco et al. [39] to evaluate the task arithmetic ability of a pre-trained model, which consists of the 8 tasks described before and two sub-benchmarks:

Task addition : The sum of the task vectors τ = ∑ t τ t is added to a 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 are shown in Table 1.

Task negation : A task vector is subtracted from the pre-trained checkpoint to forget a task while retaining performance on a control task (ImageNet). The success of this benchmark is measured in terms of the maximum drop in accuracy on the forgetting task that retains the performance on 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 ViTs from before, and use a single mixing coefficient α = α 1 = · · · = α T optimized separately for the non-linear and post-hoc linearized models to ensure a fair comparison. We provide all details of this experiment in Appendix A.

The results in Table 1 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 in 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 in non-linear models leverages the non-linear components of f , as well.

Although these results reject the linear hypothesis, it is still remarkable that the post-hoc linearized models do better at task negation than the non-linear ones (see Table 2). Furthermore, even in task addition (see Table 1) they achieve higher normalized accuracies (see definition in Appendix A). Indeed, as 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.

4 Weight disentanglement

If the linear regime is not necessary to 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 to the set of fine-tuning tasks.

Property 3 (Weight disentanglement) . A parametric function f : X × Θ →Y is weight disentangled with respect to a set of task vectors T = { τ t } t ∈ [ T ] and the corresponding supports D = {D t } t ∈ [ T ] if where g t ( x ; α t τ t ) = 0 for x / ∈ D t and t = 1 , . . . , T , and g 0 ( x ) = 0 for x ∈ ⋃ t ∈ [ T ] D t .

In essence, this definition captures the idea that the function f can be decomposed as a sum of spatiallylocalized components, i.e. , vanishing outside a spatial region, whose functional variation is entirely captured by each τ t (see Figure 1). Moreover, it is trivial to see that satisfying weight disentanglement is equivalent to satisfying Property 1 on task arithmetic as one can always write Eq. (1) as and identify g t ( x ; α t τ t ) = f ( x ; θ 0 + α t τ t ) ✶ ( x ∈ D t ) and g 0 ( x ) = f ( x ; θ 0 ) ✶ ( x / ∈ D t ) . It is important to highlight, however, that this additive decomposition does not imply linearity, as the local functions { g t } t ∈ [ T ] are not required to be linear with respect to the parameters.

Furthermore, note that weight disentanglement is a property of the predictors and not related to the performance on different tasks. That is, a 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 + α τ ) does not generalize for some α . More generally, we can visualize the level of weight disentanglement of a model by measuring its discrepancy with Eq. (4). To do so, given two tasks, one can check the disentanglement error of a model,

̸

Figure 3 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 error within a small region surrounding θ 0 , which enables task arithmetic. However, for α 1 , α 2 > 1 , the error increases, indicating a high degree of interaction between tasks. This explains why task arithmetic performs better in a small neighborhood of θ 0 - 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 .

where dist denotes any distance metric between output vectors. As we are dealing with classification tasks, in what follows we use the prediction error dist( y 1 , y 2 ) = ✶ ( y 1 = y 2 ) as the distance metric. In general, the smaller the value of ξ ( α 1 , α 2 ) the more weight disentangled a model is at ( α 1 , α 2 ) .

-3.0

-1.0

1.0

3.0

α

1

-3.0

-1.0

1.0

3.0

α

2

θ

0

τ

1

τ

2

Non-linear FT

EuroSAT - SVHN

-3.0

-1.0

1.0

3.0

α

1

-3.0

-1.0

1.0

3.0

α

2

θ

0

τ

1

τ

2

DTD - SUN397

-3.0

-1.0

1.0

3.0

α

1

-3.0

-1.0

1.0

3.0

α

2

θ

0

τ

1

τ

2

Cars - RESISC45

-3.0

-1.0

1.0

3.0

α

1

-3.0

-1.0

1.0

3.0

α

2

θ

0

τ

1

τ

2

Post-hoc linear.

-3.0

-1.0

1.0

3.0

α

1

-3.0

-1.0

1.0

3.0

α

2

θ

0

τ

1

τ

2

-3.0

-1.0

1.0

3.0

α

1

-3.0

-1.0

1.0

3.0

α

2

θ

0

τ

1

τ

2

0%

100%

ξ

(

α

1

, α

2

)

0%

100%

Comparing the disentanglement error of the non-linear models and their post-hoc linearization reveals an interesting finding: linearized models exhibit greater disentanglement than their non-linear counterparts. This is evident from the more extensive regions with low disentanglement errors in Figure 3 (bottom). This explains why the post-hoc 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 to 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 to enhance task arithmetic. We leverage this idea in the next section.

5 Enhancing task arithmetic via linearization

We have seen that linearized models are more weight-disentangled than non-linear ones. However, post-hoc linearization degrades single-task performance. We now demonstrate that enforcing models to fine-tune in the tangent space to their pre-trained initialization significantly improves task arithmetic by reducing the singletask accuracy gap.

Specifically, rather than applying the nonlinearly fine-tuned task vectors τ = θ ⋆ -θ 0 to f lin , as in Section 3, we propose to directly obtain the task vectors through explicit fine-tuning in the tangent space as illustrated in Figure 4. That is, given a model f , we directly fine-tune its linear approximation f lin around θ 0 [26]. The fine-tuning process can follow the same protocols used before but with the network parameterization dictated by Eq. (3). Due to the

Vef(00)

K

f(00)

fin(0*)

N

linear connection between the weight-space and function-space defined in Eq. (3), fine-tuning f lin is essentially the same as training a kernel predictor with kernel k NTK . As a result, we obtain the fine-tuned weights θ ⋆ lin of the linearized model for each task, which allows us to construct the corresponding task vector τ lin = θ ⋆ lin -θ 0 . We provide further details of this procedure in Appendix B.

̸

Moreover, as the considered models do not inherently exhibit linear fine-tuning (see Section 3), this approach yields significantly different results compared to post-hoc linearization, i.e. , f lin ( x ; θ 0 + τ lin ) = f lin ( x ; θ 0 + τ ) . In particular, although both models share the same kernel k NTK ( x , x ′ ) , the task vectors τ lin have been explicitly optimized to maximize the performance of such linearized models. Consequently, 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 f lin ( · ; θ 0 + τ lin ) is very similar to the one of f ( · ; θ 0 + τ ) (cf. Figure 2).

60

70

80

90

100

Linearized fine-tuning accuracy (%)

60

70

80

90

100

Non-linear fine-tuning accuracy (%)

Non-linear

advantage

ViT-B/32

ViT-B/16

ViT-L/14

Remarkably, as we show in Appendix D.4, this increase in single-task performance does not compromise weight disentanglement, which remains as high as for the post-hoc linearized models in Figure 3. As a result, linear fine-tuning allows for improved task arithmetic compared to standard non-linear fine-tuning. In particular, Tables 1 and 2 in their last rows show that linearized fine-tuned models significantly outperform their non-linear counterparts and achieve state-of-the-art 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, we 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 3 .

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 O (1) increase in computational costs with respect to their non-linear counterparts (see Appendix B for a longer discussion). In this regard, the superiority of task arithmetic of linearized models can make this technique appealing for practical applications. Identifying the right trade-offs between computational cost and performance, as well as faster linearization techniques, is an exciting avenue for future work.

6 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 kernel can satisfy Property 1, and we establish a connection between task arithmetic and the spectral properties of the NTK. Then, we argue that weight disentanglement and task arithmetic are emergent properties of pre-training.

6.1 Eigenfunction localization

Generally, a kernel k admits a decomposition in terms of a family of eigenfunction-eigenvalue pairs { ( ϕ ρ , λ ρ ) } ρ ∈ N ; which implies that k can only represent functions of the form f ⋆ ( x ) = ∑ ∞ ρ =1 c ρ ϕ ρ ( x ) with a finite kernel norm, i.e. , ∥ f ⋆ ∥ 2 H = ∑ ∞ ρ =1 c 2 ρ / λ ρ < + ∞ . Specifically, the coefficients { c ρ } ρ ∈ N constitute a representation of the function f ⋆ in the kernel basis.

Consider T tasks { f ⋆ t } t ∈ [ T ] supported in their respective non-intersecting domains {D t } t ∈ [ T ] . Furthermore, let { ϕ ρ } ρ ∈ N be an orthogonal basis of eigenfunctions that diagonalizes the kernel on the union of all D t 's. The following proposition provides a sufficient condition on the representation of the tasks in this basis to ensure the task arithmetic property:

Proposition 1 (Simplified) . Suppose that { f ⋆ t } t ∈ [ T ] can be represented by the kernel k . The kernel k is capable of performing task arithmetic with respect to { f ⋆ t } t ∈ [ T ] and {D t } t ∈ [ T ] if, for each task t , there exists a subset of localized eigenfunctions such that i) supp( ϕ ) ⊆ D t for each ϕ in the subset, and ii) the representation of f ⋆ t 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 if it represents the different tasks using these functions. According to the representer theorem of kernels [77], after linear fine-tuning on task t with a training set { ( x ν , f ⋆ t ( x ν )) } ν ∈ [ n t ] and x ν ∼ µ t , the CLIP's predictor evaluated at a new point x ∈ X can be expressed as a linear combination of its kernel k NTK evaluated on x and the training data, i.e. , f lin ( x ) = f ( x ; θ 0 ) + ∑ ν ∈ [ n t ] β ν k NTK ( x ν , x ) .

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 represent 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 to the use of those. Indeed, thanks to this localization, CLIP models can effectively separate the representation of different tasks and carry out taskspecific operations without interference. We believe that further investigation into this intriguing localization phenomenon holds the potential to deepen our understanding of these models.

x

0

1

x

RESISC45

x

Cars

E

loc

(

x

)

To explore whether CLIP models use localized eigenfunctions for task arithmetic, we diagonalize the matrix ( K NTK ) ij = k NTK ( x i , x j ) with x i ∈ D t , i.e. , the task on which we trained, and x j ∈ D t ∪ D t ′ , where D t ′ is the support of a control task. If the eigenfunctions used to represent f ⋆ ( x ) are localized, then the power of the eigenvectors of K NTK must be concentrated in the points belonging to the dataset used for training. To measure this concentration, we introduce the local energy E loc ( x ) = ∑ ρ ϕ 2 ρ ( x ) , which sums the power of all the eigenfunctions ϕ ρ at a given point x .

∈

∈

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 a 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 a 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 condition 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 a general property of deep networks.

To investigate this, we replicate the task addition experiments but employ randomly initialized ViTs instead of pre-trained ones. The results in Table 3 reveal that task arithmetic is not achievable on randomly initialized ViTs. Indeed, adding task vectors obtained from a random initialization θ rd 0 does not result in significant improvements in multi-task accuracy over random chance. This

holds true for both non-linear task vectors, τ rd , and linearized ones, τ rd lin . In 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 observation 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 disentanglement 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 interpolations between model weights and task arithmetic to manipulate and enhance the capabilities of pre-trained models. In particular, several studies have shown that interpolating between a model's finetuned weights and its 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 models [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 interpolations extend to models trained from scratch, as long as they are 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 as 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 understanding, 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 fewshot 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 ideal 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 modeling [14, 34, 73] and self-supervised learning [6, 13, 48, 69] strive to achieve this goal. Our investigation, however, explores a distinct aspect: weight disentanglement within the framework of task arithmetic. Departing from the static perspective of feature disentanglement, weight disentanglement connects weight space and function space transitions, thereby enriching our understanding of disentanglement in neural networks from a functional standpoint. Several studies have previously attempted to exploit a similar notion by inducing the learning of task-specific subnetworks within a larger network [32, 36, 54-56, 83, 85]. To the best of our knowledge, our work is the first to demonstrate the natural emergence of such phenomena in specific semantically meaningful tasks during CLIP pre-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 a more effective approach for model editing. Crucially, we revealed that weight disentanglement plays a vital role in the success of task arithmetic, as distinct directions in weight space correspond to localized areas in the function space; and that it is an emergent property of pre-training.

A fascinating open question consists in understanding how weight disentanglement arises during pretraining 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 a 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.

References

Alessandro Achille and Stefano Soatto. Emergence of invariance and disentanglement in deep representations. Journal of Machine Learning Research (JMLR) , 2018. http://jmlr.org/ papers/v19/17-646.html . (page 9).

Alessandro Achille, Aditya Golatkar, Avinash Ravichandran, Marzia Polito, and Stefano Soatto. LQF: Linear quadratic fine-tuning. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , 2021. https://arxiv.org/abs/2012.11140 . (page 9).

Samuel K Ainsworth, Jonathan Hayase, and Siddhartha Srinivasa. Git re-basin: Merging models modulo permutation symmetries. In International Conference on Learning Representations (ICLR) , 2023. https://arxiv.org/abs/2209.04836 . (pages 1 and 9).

Sanjeev Arora, Simon S. Du, Wei Hu, Zhiyuan Li, Ruslan Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. In Advances in Neural Information Processing Systems (NeurIPS) , 2019. https://proceedings.neurips.cc/paper/2019/ hash/dbc4d84bfcfe2284ba11beffb853a8c4-Abstract.html . (page 3).

Sanjeev Arora, Simon S. Du, Zhiyuan Li, Ruslan Salakhutdinov, Ruosong Wang, and Dingli Yu. Harnessing the power of infinitely wide deep nets on small-data tasks. In International Conference on Learning Representations (ICLR) , 2020. https://arxiv.org/abs/1910. 01663 . (page 9).

Philip Bachman, R. Devon Hjelm, and William Buchwalter. Learning representations by maximizing mutual information across views. In Advances in Neural Information Processing Systems (NeurIPS) , 2019. https://proceedings.neurips.cc/paper/2019/hash/ ddf354219aac374f1d40b7e760ee5bb7-Abstract.html . (page 9).

Gregor Bachmann, Seyed-Mohsen Moosavi-Dezfooli, and Thomas Hofmann. Uniform convergence, adversarial spheres and a simple remedy. In International Conference on Machine Learning (ICML) , 2021. http://proceedings.mlr.press/v139/bachmann21a.html . (page 9).

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 Statistics (AISTATS) , 2021. http: //proceedings.mlr.press/v130/baratin21a.html . (page 9).

Atilim Gunes Baydin, Barak A. Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. Automatic differentiation in machine learning: a survey. Journal of Machine Learning Research (JMLR) , 2017. http://jmlr.org/papers/v18/17-468.html . (page 19).

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://arxiv.org/abs/1206.5538 . (page 9).

Blake Bordelon, Abdulkadir Canatar, and Cengiz Pehlevan. Spectrum dependent learning curves in kernel regression and wide neural networks. In International Conference on Machine Learning (ICML) , 2020. http://proceedings.mlr.press/v119/bordelon20a.html . (page 18).

Francesco Cagnetta, Alessandro Favero, and Matthieu Wyart. What can be learnt with wide convolutional neural networks?, 2022. https://arxiv.org/abs/2208.01003 . (page 20).

Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International Conference on Machine Learning (ICML) , 2020. https://proceedings.mlr.press/v119/chen20j/chen20j.pdf . (page 9).

Xi Chen, Yan Duan, Rein Houthooft, John Schulman, Ilya Sutskever, and Pieter Abbeel. InfoGAN: Interpretable representation learning by information maximizing generative adversarial nets. In Advances in Neural Information Processing Systems (NeurIPS) , 2016. https://proceedings.neurips.cc/paper/2016/hash/ 7c9d0b1f96aebd7b5eca8c3edaa19ebb-Abstract.html . (page 9).

Gong Cheng, Junwei Han, and Xiaoqiang Lu. Remote sensing image scene classification: Benchmark and state of the art. Proceedings of the IEEE , 2017. https://ieeexplore.ieee. org/document/7891544 . (page 4).

Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in Neural Information Processing Systems (NeurIPS) , 2019. https://proceedings.neurips.cc/paper/2019/file/ ae614c557843b1df326cb29c57225459-Paper.pdf . (page 2).

Leshem Choshen, Elad Venezian, Noam Slonim, and Yoav Katz. Fusing finetuned models for better pretraining, 2022. https://arxiv.org/abs/2204.03044 . (page 9).

Ole Christensen and Khadija L Christensen. Linear independence and series expansions in function spaces. The American Mathematical Monthly , 2006. https://www.tandfonline. com/doi/abs/10.1080/00029890.2006.11920343 . (page 20).

Kurtland Chua, Qi Lei, and Jason D Lee. How fine-tuning allows for effective meta-learning. In Advances in Neural Information Processing Systems (NeurIPS) , 2021. https://proceedings.neurips.cc/paper_files/paper/2021/ file/4a533591763dfa743a13affab1a85793-Paper.pdf . (page 9).

Mircea Cimpoi, Subhransu Maji, Iasonas Kokkinos, Sammy Mohamed, and Andrea Vedaldi. Describing textures in the wild. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , 2014. https://openaccess.thecvf.com/content_cvpr_2014/html/Cimpoi_ Describing_Textures_in_2014_CVPR_paper.html . (page 4).

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , 2009. https://ieeexplore.ieee.org/abstract/document/5206848 . (page 3).

Aditya Deshpande, Alessandro Achille, Avinash Ravichandran, Hao Li, Luca Zancato, Charless C. Fowlkes, Rahul Bhotika, Stefano Soatto, and Pietro Perona. A linearized framework and a new benchmark for model selection for fine-tuning, 2021. https://arxiv.org/abs/2102. 00084 . (pages 3 and 9).

Shachar Don-Yehiya, Elad Venezian, Colin Raffel, Noam Slonim, Yoav Katz, and Leshem Choshen. Cold fusion: Collaborative descent for distributed multitask finetuning, 2022. https: //arxiv.org/abs/2212.01378 . (pages 1 and 9).

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 on Learning Representations (ICLR) , 2021. https://openreview.net/forum?id=YicbFdNTTy . (page 4).

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 (NeurIPS) , 2021. https://proceedings.neurips.cc/paper_files/ paper/2021/hash/4e8eaf897c638d519710b1691121f8cb-Abstract.html . (page 20).

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://proceedings.neurips.cc/paper/2020/ file/405075699f065e43581f27d67bb68478-Paper.pdf . (pages 4, 6, and 9).

Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin. Linear mode connectivity and the lottery ticket hypothesis. In International Conference on Machine Learning (ICML) , 2020. https://proceedings.mlr.press/v119/frankle20a.html . (pages 1 and 9).

Robert M French. Catastrophic forgetting in connectionist networks. Trends in Cognitive Sciences , 1999. https://www.sciencedirect.com/science/article/pii/ S1364661399012942 . (pages 1 and 9).

Amnon Geifman, Meirav Galun, David Jacobs, and Basri Ronen. On the spectral bias of convolutional neural tangent and gaussian process kernels. Advances in Neural Information Processing Systems (NeurIPS) , 2022. https://arxiv.org/abs/2203.09255 . (page 20).

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. https://iopscience.iop.org/article/10.1088/1742-5468/abc4de . (page 9).

Amelia Glaese, Nat McAleese, Maja Trebacz, John Aslanides, Vlad Firoiu, Timo Ewalds, Maribeth Rauh, Laura Weidinger, Martin Chadwick, Phoebe Thacker, Lucy Campbell-Gillingham, 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, Iason Gabriel, William Isaac, John Mellor, Demis Hassabis, Koray Kavukcuoglu, Lisa Anne Hendricks, and Geoffrey Irving. Improving alignment of dialogue agents via targeted human judgements, 2022. https://www.deepmind.com/blog/building-safer-dialogue-agents . (page 1).

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. https://openreview.net/forum?id=OGg9XnKxFAH . (page 9).

Patrick Helber, Benjamin Bischke, Andreas Dengel, and Damian Borth. Eurosat: A novel dataset and deep learning benchmark for land use and land cover classification. Journal of Selected Topics in Applied Earth Observations and Remote Sensing , 2019. https://arxiv. org/abs/1709.00029 . (page 4).

Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-VAE: Learning basic visual concepts with a constrained variational framework. In International Conference on Learning Representations (ICLR) , 2017. https://openreview.net/forum?id=Sy2fzU9gl . (page 9).

Irina Higgins, David Amos, David Pfau, Sébastien Racanière, Loïc Matthey, Danilo J. Rezende, and Alexander Lerchner. Towards a definition of disentangled representations, 2018. http: //arxiv.org/abs/1812.02230 . (page 9).

Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, 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://openreview. net/forum?id=nZeVKeeFYf9 . (page 9).

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/10.5281/ zenodo.5143773 . (page 18).

Gabriel Ilharco, Mitchell Wortsman, Samir Yitzhak Gadre, Shuran Song, Hannaneh Hajishirzi, Simon Kornblith, Ali Farhadi, and Ludwig Schmidt. Patching open-vocabulary models by interpolating weights. In Advances in Neural Information Processing Systems (NeurIPS) , 2022. https://arXiv.org/abs/2208.05592 . (pages 1, 8, 9, and 18).

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://arxiv.org/abs/2110. 08207 . (pages 1, 2, 3, 4, 5, 7, 8, 9, 18, and 25).

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. https://arxiv.org/abs/1803.05407 . (pages 1 and 9).

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in Neural Information Processing Systems (NeurIPS) , 2018. https://proceedings.neurips.cc/paper_files/paper/2018/ file/5a4be1fa34e62bb8a6ec6b91d2462f5a-Paper.pdf . (pages 2 and 3).

Tushar Khot, Peter Clark, Michal Guerquin, Peter Jansen, and Ashish Sabharwal. Qasc: A dataset for question answering via sentence composition, 2020. https://arxiv.org/abs/ 1910.11473v2 . (page 25).

Jonathan Krause, Michael Stark, Jia Deng, and Li Fei-Fei. 3D Object representations for finegrained categorization. In International Conference on Computer Vision Workshops (ICCVw) , 2013. https://www.cv-foundation.org/openaccess/content_iccv_workshops_ 2013/W19/html/Krause_3D_Object_Representations_2013_ICCV_paper.html . (page 4).

Yann LeCun. The MNIST database of handwritten digits, 1998. http://yann.lecun.com/ exdb/mnist/ . (page 4).

Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl-Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. In Advances in Neural Information Processing Systems (NeurIPS) , 2019. https://proceedings.neurips.cc/paper_files/paper/2019/ file/0d1a9651497a38d8b1c3871c84528bd4-Paper.pdf . (page 3).

Margaret Li, Suchin Gururangan, Tim Dettmers, Mike Lewis, Tim Althoff, Noah A Smith, and Luke Zettlemoyer. Branch-train-merge: Embarrassingly parallel training of expert language models, 2022. https://arxiv.org/abs/2208.03306 . (pages 1 and 9).

Zhuang Liu, Hanzi Mao, Chao-Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. A ConvNet for the 2020s. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , 2022. https://arxiv.org/abs/2201.03545 . (page 25).

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) , 2019. http://proceedings.mlr.press/v97/locatello19a.html . (page 9).

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations (ICLR) , 2019. URL https://openreview.net/ forum?id=Bkg6RiCqY7 . (page 18).

Ximing Lu, Sean Welleck, Liwei Jiang, Jack Hessel, Lianhui Qin, Peter West, Prithviraj Ammanabrolu, and Yejin Choi. QUARK: Controllable text generation with reinforced unlearning. In Advances in Neural Information Processing Systems (NeurIPS) , 2022. https: //arxiv.org/abs/2205.13636 . (page 1).

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 Computational Linguistics (ACL) , 2011. http://www.aclweb.org/anthology/P11-1015 . (page 25).

Wesley Maddox, Shuai Tang, Pablo G. Moreno, Andrew Gordon Wilson, and Andreas Damianou. Fast adaptation with linearized neural networks. In International Conference on Artificial Intelligence and Statistics (AISTATS) , 2021. https://proceedings.mlr.press/v130/ maddox21a/maddox21a.pdf . (page 9).

Sadhika Malladi, Alexander Wettig, Dingli Yu, Danqi Chen, and Sanjeev Arora. A kernel-based view of language model fine-tuning, 2022. https://arxiv.org/abs/2210.05643 . (pages 3 and 9).

Arun Mallya and Svetlana Lazebnik. Packnet: Adding multiple tasks to a single network by iterative pruning. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , 2018. https://arxiv.org/abs/1711.05769 . (page 9).

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. https://arxiv.org/abs/1801.06519 .

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://doi.org/10.1073/pnas.1803839115 . (page 9).

Michael Matena and Colin Raffel. Merging models with fisher-weighted averaging. In Advances in Neural Information Processing Systems (NeurIPS) , 2021. https://arxiv.org/abs/2111. 09832 . (pages 1 and 9).

Michael McCloskey and Neal J Cohen. Catastrophic interference in connectionist networks: The sequential learning problem. In Psychology of Learning and Motivation . Elsevier, 1989. https: //www.sciencedirect.com/science/article/abs/pii/S0079742108605368 . (pages 1 and 9).

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/abs/2004.05529 . (page 9).

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 in Neural Information Processing Systems (NeurIPS) Workshops , 2011. https://storage.googleapis. com/pub-tools-public-publication-data/pdf/37648.pdf . (page 4).

Roman Novak, Jascha Sohl-Dickstein, and Samuel S Schoenholz. Fast finite width neural tangent kernel. In International Conference on Machine Learning (ICML) , 2022. https: //arxiv.org/abs/2206.08720 . (page 7).

Guillermo Ortiz-Jiménez, Apostolos Modas, Seyed-Mohsen Moosavi-Dezfooli, and Pascal Frossard. What can linearized neural networks actually say about generalization? In Advances in Neural Information Processing Systems (NeurIPS) , 2021. https://proceedings.neurips.cc/paper_files/paper/2021/file/ 4b5deb9a14d66ab0acc3b8a2360cde7c-Paper.pdf . (page 3).

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 of the IEEE , 2021. https://ieeexplore.ieee. org/document/9348948 . (page 1).

Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll L Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback, 2022. https://arxiv.org/abs/2203.02155 . (page 1).

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://arxiv.org/abs/2007.11471 . (page 9).

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. In Advances in Neural Information Processing Systems (NeurIPS) , 2019. https://proceedings.neurips.cc/paper_files/paper/ 2019/file/bdbca288fee7f92f2bfa9f7012727740-Paper.pdf . (page 18).

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 (NeurIPS) , 2022. https://proceedings.neurips.cc/paper_files/ paper/2022/file/3d3a9e085540c65dd3e5731361f9320e-Paper-Conference.pdf . (page 9).

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language Models are Unsupervised Multitask Learners, 2019. https://openai.com/blog/ better-language-models/ . (page 8).

Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning (ICML) , 2021. https://arxiv.org/abs/2103. 00020 . (pages 2, 9, and 19).

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/papers/v21/20-074.html . (pages 8 and 25).

Alexandre Ramé, Kartik Ahuja, Jianyu Zhang, Matthieu Cord, Léon Bottou, and David LopezPaz. Model ratatouille: Recycling diverse models for out-of-distribution generalization. In International Conference on Machine Learning (ICML) , 2022. https://arxiv.org/abs/ 2212.10445 . (pages 1 and 9).

Alexandre Ramé, Matthieu Kirchmeyer, Thibaud Rahier, Alain Rakotomamonjy, Patrick Gallinari, and Matthieu Cord. Diverse weight averaging for out-of-distribution generalization. In Advances in Neural Information Processing Systems (NeurIPS) , 2022. https: //arxiv.org/abs/2205.09739 . (pages 1 and 9).

Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International Conference on Machine Learning (ICML) , 2014. http://proceedings.mlr.press/v32/rezende14.html . (page 9).

Marco Tulio Ribeiro and Scott Lundberg. Adaptive testing and debugging of nlp models. In Annual Meeting of the Association for Computational Linguistics (ACL) , 2022. https: //aclanthology.org/2022.acl-long.230/ . (page 1).

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 (NeurIPS) , 2019. https://proceedings.neurips.cc/paper/2019/ file/5ac8bb8a7d745102a978c5f8ccdb61b8-Paper.pdf . (page 20).

Shibani Santurkar, Dimitris Tsipras, Mahalaxmi Elango, David Bau, Antonio Torralba, and Aleksander Madry. Editing a classifier by rewriting its prediction rules. In Advances in Neural Information Processing Systems (NeurIPS) , 2021. https://proceedings.neurips.cc/paper_ files/paper/2021/file/c46489a2d5a9a9ecfc53b17610926ddd-Paper.pdf . (page 1).

Bernhard Schölkopf and Alexander Johannes Smola. Learning with Kernels: support vector machines, regularization, optimization, and beyond . Adaptive computation and machine learning series. MIT Press, 2002. https://mitpress.mit.edu/9780262536578/ learning-with-kernels/ . (page 8).

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. https://arxiv.org/abs/2111. 02114 . (page 25).

Sidak Pal Singh and Martin Jaggi. Model fusion via optimal transport. In Advances in Neural Information Processing Systems (NeurIPS) , 2020. https://proceedings.neurips. cc/paper/2020/hash/fb2697869f56484404c8ceee2985b01d-Abstract.html . (pages 1 and 9).

Johannes Stallkamp, Marc Schlipsing, Jan Salmen, and Christian Igel. The german traffic sign recognition benchmark: a multi-class classification competition. In International Joint Conference on Neural Networks (IJCNN) , 2011. https://ieeexplore.ieee.org/document/ 6033395 . (page 4).

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://proceedings.neurips.cc/ paper/2020/file/55053683268957697aa39fba6f231c68-Paper.pdf . (page 9).

Nikhil Vyas, Yamini Bansal, and Preetum Nakkiran. Limitations of the NTK for understanding generalization in deep learning, 2022. https://arxiv.org/abs/2206.10012 . (page 9).

Yeming Wen, Dustin Tran, and Jimmy Ba. Batchensemble: an alternative approach to efficient ensemble and lifelong learning. In International Conference on Learning Representations (ICLR) , 2020. https://arxiv.org/abs/2002.06715 . (page 9).

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 transformers: State-of-the-art natural language processing, 2019. https://arxiv.org/abs/1910. 03771 . (page 25).

Mitchell Wortsman, Vivek Ramanujan, Rosanne Liu, Aniruddha Kembhavi, Mohammad Rastegari, Jason Yosinski, and Ali Farhadi. Supermasks in superposition. In Advances in Neural Information Processing Systems (NeurIPS) , 2020. https://proceedings.neurips.cc/ paper/2020/file/ad1f8bb9b51f023cdc80cf94bb615aa9-Paper.pdf . (page 9).

Mitchell Wortsman, Gabriel Ilharco, Samir Yitzhak Gadre, Rebecca Roelofs, Raphael GontijoLopes, Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, et al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In International Conference on Machine Learning (ICML) , 2022. https://arxiv.org/abs/2203.05482 . (pages 1, 3, 8, and 9).

Mitchell Wortsman, Gabriel Ilharco, Mike Li, Jong Wook Kim, Hannaneh Hajishirzi, Ali Farhadi, Hongseok Namkoong, and Ludwig Schmidt. Robust fine-tuning of zero-shot models. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , 2022. https: //arxiv.org/abs/2109.01903 . (pages 1, 2, 3, 4, and 9).

Jianxiong Xiao, Krista A Ehinger, James Hays, Antonio Torralba, and Aude Oliva. Sun database: Exploring a large collection of scene categories. International Journal of Computer Vision (IJCV) , 2016. https://link.springer.com/article/10.1007/s11263-014-0748-y . (pages 4 and 22).

Prateek Yadav, Derek Tam, Leshem Choshen, Colin Raffel, and Mohit Bansal. Resolving interference when merging models. In Advances in Neural Information Processing Systems (NeurIPS) , 2023. https://arxiv.org/abs/2306.01708 . (pages 1 and 9).

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 Pattern Recognition (CVPR) , 2022. https://doi.org/10.1109/CVPR52688. 2022.01863 . (pages 3 and 9).

Luca Zancato, Alessandro Achille, Avinash Ravichandran, Rahul Bhotika, and Stefano Soatto. Predicting training time without training. In Advances in Neural Information Processing Systems (NeurIPS) , 2020. https://proceedings.neurips.cc/paper/2020/hash/ 440e7c3eb9bbcd4c33c3535354a51605-Abstract.html . (pages 3 and 9).

Fuzhen Zhuang, Zhiyuan Qi, Keyu Duan, Dongbo Xi, Yongchun Zhu, Hengshu Zhu, Hui Xiong, and Qing He. A comprehensive survey on transfer learning. Proceedings of the IEEE , 2020. https://ieeexplore.ieee.org/document/9134370 . (page 1).

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 et 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 a batch size of 128 , learning rate of 10 -5 and a cosine annealing learning rate schedule with 200 warm-up steps and the AdamW 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-shot template prompts for each dataset. Freezing this layer does not harm accuracy and ensures that no additional learnable parameters are introduced during fine-tuning [38]. We use this exact same protocol to fine-tune the 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 et al. [39] we use a single coefficient α to tune the size of the task vectors used to modify the pre-trained models. This is equivalent to setting α = α 1 = . . . α T in Eq. (1). Both in the task addition and task negation benchmarks, after fine-tuning, we evaluate different scaling coefficients α ∈ { 0 . 0 , 0 . 05 , 0 . 1 , . . . , 1 . 0 } and choose the value that achieves the highest target metric on a small held-out proportion of the training set as specified in Ilharco et al. [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 control task; for task addition and negation, respectively. The tuning of α is done independently for non-linear FT, linearized FT, and post-hoc linearization.

Normalized accuracies in task addition. Table 1 shows the normalized accuracies after editing different models by adding the sum of the task vectors on 8 tasks τ = ∑ t τ t . Here, the normalization is performed with respect to the single-task accuracies achieved by the model fine-tuned on each task. Mathematically,

Disentanglement error. To produce the weight disentanglement visualizations of Figure 3 we compute the value of ξ ( α 1 , α 2 ) on a 20 × 20 grid of equispaced values in [ -3 , 3] × [ -3 , 3] . To estimate the disentanglement error, we use a random subset of 2 , 048 test points for each dataset.

NTKeigenfunction estimation. Weuse the finite-width NTK implementation from the functorch sublibrary of PyTorch [66] to compute the K NTK matrices described in Section 6.1. In particular, we use a random subset of 200 training points for each dataset and compute the singular value decomposition (SVD) of K NTK to estimate the entries of ϕ ρ on each dataset. As described in Bordelon et al. [11], and to avoid a high memory footprint, we estimate a different set of singular vectors for each output class, equivalent to estimating one kernel matrix per output logit. Figure 6 shows the values of E loc ( x ) 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 of 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 a trainable linearized version of it around its initialization. We give a simple example of this in

Listing 1, where we see that the resulting LinearizedModel 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 is frozen in our 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 text embeddings. This normalization does not 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 a clearly positive effect in single-task accuracy both for the non-linear and linearized models.

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 a linearized model usually takes around two or 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 the 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 ∇ θ f lin ( x ; θ ) is also roughly twice the cost of obtaining the gradient of its non-

linear counterparts ∇ θ f ( x ; θ ) . 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 arithmetic and provide analytical examples with exactly diagonalizable NTKs to support our discussion.

Proposition 2 (Formal version of Proposition 1) . Suppose that the task functions { f ⋆ t } t ∈ [ T ] belong to the RKHS of the kernel k and their coefficients in the kernel eigenbasis are { ( c ⋆ t,ρ ) ρ ∈ N } t ∈ [ T ] . If ∀ t, ρ , either c ⋆ t,ρ = 0 or supp( ϕ ρ ) ⊆ D t , then the kernel k has the task arithmetic property with respect to { f ⋆ t } t ∈ [ T ] and {D t } t ∈ [ T ] .

Proof. The task arithmetic property requires that ∀ t ′ ∈ [ T ] , ∀ x ∈ D t ′ , ∑ t ∈ [ T ] f ⋆ t ( x ) = f ⋆ t ′ ( x ) . Representing the task functions in the kernel basis, we have

This condition can be rewritten as

̸

If, for each t , the eigenfunctions corresponding to non-zero coefficients are supported within a subset of D t and all domains D t 's are disjoint, then all the summands inside the integral in Eq. (9) become zero inside D t ′ , 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 { f ⋆ t } t ∈ [ T ] belong to the RKHS of the kernel k and their coefficients in the kernel eigenbasis are { ( c ⋆ t,ρ ) ρ ∈ N } t ∈ [ T ] . Furthermore, let the kernel eigenfunctions be either zero or linearly independent over each domain D t . The kernel k has the task arithmetic property with respect to { f ⋆ t } t ∈ [ T ] and {D t } t ∈ [ T ] if and only if ∀ t, ρ , either c ⋆ t,ρ = 0 or supp( ϕ ρ ) ⊆ D t .

̸

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 D t ′ , for this integral to be zero, we have only two possibilities: either i) all coefficients { ( c ⋆ t,ρ ) ρ ∈ N } t ∈ [ T ] , t = t ′ must be zero or ii) the eigenfunctions corresponding to non-zero coefficient c ⋆ t,ρ ( t = t ′ ) must be zero in D t ′ . Since the proposition is valid for any set of functions, condition i) is not feasible. Therefore, condition ii) must hold. Furthermore, since Eq. (9) is valid ∀ t ′ ∈ [ T ] , it follows that the eigenfunctions used to represent each task t ′ are zero in D t ′ = ⋃ t ∈ [ T ] , t = t ′ D t . Consequently, these eigenfunctions are only supported in D t ′ or a 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 a 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 a property of pre-training and is not inherently present at initialization, as we numerically demonstrated for CLIP models in Section 6.2.

̸

D Further experimental results

We now present additional experiments that expand the findings discussed in the main text.

Cars

DTD

EuroSAT GTSRB

MNIST RESISC45SUN397

SVHN

0

20

40

60

80

100

Single-task accuracy (%)

Non-linear FT

Linear FT

Post-hoc linear.

Zero-shot

(a) ViT-B/32

Cars

DTD

EuroSAT GTSRB

MNIST RESISC45SUN397

SVHN

0

20

40

60

80

100

Single-task accuracy (%)

Non-linear FT

Linear FT

Post-hoc linear.

Zero-shot

(b) ViT-B/16

Cars

DTD

EuroSAT GTSRB

MNIST RESISC45SUN397

SVHN

0

20

40

60

80

100

Single-task accuracy (%)

Non-linear FT

Linear FT

Post-hoc linear.

Zero-shot

(c) ViT-L/14

D.1 Fine-tuning 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 as 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 a non-linear advantage . However, an interesting observation is that the gap between non-linear, linear, and post-hoc 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 a stronger kernel behavior during fine-tuning. As a result. they tend to stay closer to the NTK approximation, closing the gap with linearized models.

D.2 Detailed results on task addition

In addition to the results presented in Table 1 in the main text, we report in Figure 8 the absolute accuracies of different CLIP models on the single tasks before ( zero-shot ) and after performing task addition with different strategies ( non-linear fine-tuning , post-hoc linearization , and linear fine-tuning ).

Cars

DTD

EuroSAT

GTSRB

MNIST

RESISC45

SUN397

SVHN

20

40

60

80

ViT-B/32

Cars

DTD

EuroSAT

GTSRB

MNIST

RESISC45

SUN397

SVHN

20

40

60

80

100

ViT-B/16

Cars

DTD

EuroSAT

GTSRB

MNIST

RESISC45

SUN397

SVHN

20

40

60

80

100

ViT-L/14

Non-linear FT

Linear FT

Post-hoc linearization

Zero-shot

For all models and all datasets, except SUN397 [88], we 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 is better approximated by the NTK description.

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 a low disentanglement error grows with the model scale. One plausible explanation for this is that larger models 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.

-3.0

-1.0

1.0

3.0

α

1

-3.0

-1.0

1.0

3.0

α

2

θ

0

τ

1

τ

2

Non-linear FT

ViT-B-32, DTD - SUN397

-3.0

-1.0

1.0

3.0

α

1

-3.0

-1.0

1.0

3.0

α

2

θ

0

τ

1

τ

2

ViT-B-16, DTD - SUN397

-3.0

-1.0

1.0

3.0

α

1

-3.0

-1.0

1.0

3.0

α

2

θ

0

τ

1

τ

2

ViT-L-14, DTD - SUN397

0%

100%

ξ

(

α

1

, α

2

)

-3.0

-1.0

1.0

3.0

-3.0

-1.0

1.0

3.0

α

2

θ

0

τ

1

τ

2

Linearized FT

ViT-B-32, DTD - SUN397

-3.0

-1.0

1.0

3.0

-3.0

-1.0

1.0

3.0

α

2

θ

0

τ

1

τ

2

ViT-B-16, DTD - SUN397

-3.0

-1.0

1.0

3.0

-3.0

-1.0

1.0

3.0

α

2

θ

0

τ

1

τ

2

ViT-L-14, DTD - SUN397

0%

100%

ξ

(

α

1

, α

2

)

?

α

1

α

1

α

1

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, we showcase the disentanglement error of a 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.

-3.0

-1.0

1.0

3.0

α

1

-3.0

-1.0

1.0

3.0

α

2

θ

0

τ

1

τ

2

Linearized FT

EuroSAT - SVHN

-3.0

-1.0

1.0

3.0

α

1

-3.0

-1.0

1.0

3.0

α

2

θ

0

τ

1

τ

2

DTD - SUN397

-3.0

-1.0

1.0

3.0

α

1

-3.0

-1.0

1.0

3.0

α

2

θ

0

τ

1

τ

2

Cars - RESISC45

0%

100%

ξ

(

α

1

, α

2

)

-3.0

-1.0

1.0

3.0

-3.0

-1.0

1.0

3.0

θ

0

τ

1

τ

2

Random init

EuroSAT - SVHN

-3.0

-1.0

1.0

3.0

-3.0

-1.0

1.0

3.0

θ

0

τ

1

τ

2

DTD - SUN397

-3.0

-1.0

1.0

3.0

-3.0

-1.0

1.0

3.0

θ

0

τ

1

τ

2

Cars - RESISC45

0%

100%

ξ

(

α

1

, α

2

)

D.5 Task arithmetic with a convolutional architecture

We replicate our task addition experiments using a convolutional architecture rather than a ViT. Specifically, we finetune a ConvNeXt [47] pre-trained on LAION-400M using CLIP [78] on 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).

D.6 Task arithmetic with closed vocabulary models

We also replicate our task addition experiments using a ViT-B/16 pre-trained on ImageNet-1k using standard supervised learning. Note that this is a closed vocabulary model, and therefore, when fine-tuning on the different tasks we need to also fine-tune a randomly initialized head. When interpolating the different task vectors we freeze the resulting head to arrive at the final models.

Interestingly, we find that this closed-vocabulary model can do task arithmetic, albeit at a lower performance (both in terms of absolute and normalized accuracy) than the open vocabulary ones (see Table 5). Remarkably, linearized fine-tuning also enhances disentanglement in this model - it yields higher normalized accuracies - but does not perform better than non-linear task addition due to its lower single-task performance.

D.7 Weight disentanglement in other architectures and modalities

To substantiate the generality of weight disentanglement, we conduct a new experiment on a pretrained T5-Base model [70] from Hugging Face Hub [84], fine-tuned on two benchmark Natural Language Processing tasks, i.e., 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 as demonstrated in Ilharco et al. [39], thereby reinforcing the robustness of 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.

-3.0

-1.0

1.0

3.0

α

-3.0

-1.0

1.0

3.0

α

2

θ

0

τ

1

τ

2

ConvNeXt, Cars - RESISC45

-3.0

-1.0

1.0

3.0

α

-3.0

-1.0

1.0

3.0

α

2

θ

0

τ

1

τ

2

Lin. ConvNeXt, Cars - RESISC45

0%

100%

ξ

(

α

1

, α

2

)

-3.0

-1.0

1.0

3.0

α

-3.0

-1.0

1.0

3.0

α

2

θ

0

τ

1

τ

2

0%

100%

ξ

(

α

1

, α

2

)

X

1

1

1

-3.0

-1.0

1.0

3.0

α

1

-3.0

-1.0

1.0

3.0

α

2

θ

0

τ

1

τ

2

Non-linear FT

EuroSAT - SVHN

-3.0

-1.0

1.0

3.0

α

1

-3.0

-1.0

1.0

3.0

α

2

θ

0

τ

1

τ

2

DTD - SUN397

-3.0

-1.0

1.0

3.0

α

1

-3.0

-1.0

1.0

3.0

α

2

θ

0

τ

1

τ

2

Cars - RESISC45

0%

100%

ξ

(

α

1

, α

2

)

-3.0

-1.0

1.0

3.0

α

1

-3.0

-1.0

1.0

3.0

α

2

θ

0

τ

1

τ

2

Linearized FT

EuroSAT - SVHN

-3.0

-1.0

1.0

3.0

α

1

-3.0

-1.0

1.0

3.0

α

2

θ

0

τ

1

τ

2

DTD - SUN397

-3.0

-1.0

1.0

3.0

α

1

-3.0

-1.0

1.0

3.0

α

2

θ

0

τ

1

τ

2

Cars - RESISC45

0%

100%

ξ

(

α

1

, α

2

)

D.8 Localization of eigenfunctions of CLIP's NTK

In Figure 14, we plot the local energy of the NTK eigenfunctions for a pre-trained CLIP ViT-B/32 model evaluated on three different data supports and control data supports. These panels complement the information presented in Figure 6 in the main text, where we observed that the CLIP has eigenfunctions whose energy is concentrated on points belonging to the respective dataset.

In Figure 15, we extend this analysis to a randomly-initialized CLIP ViT-B/32 model. In all panels, we observe a non-trivial but considerably poor degree of eigenfunction localization. This observation aligns with the finding that randomly-initialized 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 to perform task arithmetic operations. In summary, eigenfunction localization offers a complementary perspective on the limitations of randomly-initialized models.

0

1

x

RESISC45

x

Cars

E

loc

(

x

) =

∑

ρ

φ

2

ρ

(

x

)

0

1

x

DTD

x

SUN397

E

loc

(

x

) =

∑

ρ

φ

2

ρ

(

x

)

0

1

x

MNIST

x

EuroSAT

E

loc

(

x

) =

∑

ρ

φ

2

ρ

(

x

)

∈

∈

∈

∈

∈

∈

0

1

x

RESISC45

x

Cars

E

loc

(

x

) =

∑

ρ

φ

2

ρ

(

x

)

0

1

x

DTD

x

SUN397

E

loc

(

x

) =

∑

ρ

φ

2

ρ

(

x

)

0

1

x

MNIST

x

EuroSAT

E

loc

(

x

) =

∑

ρ

φ

2

ρ

(

x

)

∈

∈

∈

∈

∈

∈

D.9 Further experiments with randomly-initialized networks

We conclude by showing, in Figure 16, the absolute single-task accuracy achieved by different CLIP ViT models that were fine-tuned from a 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-trained models in task arithmetic is their lack of weight disentanglement.

Interestingly, the performance of the randomly-initialized large model is generally poorer compared to the base models. This observation can be attributed to the models' tendency to overfit the training data, which is more likely to occur when a model has a larger capacity.

Cars

DTD

EuroSAT GTSRB

MNIST RESISC45 SUN397

SVHN

0

20

40

60

80

Single-task accuracy (%)

Non-linear FT

Linear FT

Random init

(a) ViT-B/32

Cars

DTD

EuroSAT GTSRB

MNIST RESISC45 SUN397

SVHN

0

20

40

60

80

Single-task accuracy (%)

Non-linear FT

Linear FT

Random init

(b) ViT-B/16

Cars

DTD

EuroSAT GTSRB

MNIST RESISC45 SUN397

SVHN

0

20

40

60

80

Single-task accuracy (%)

Non-linear FT

Linear FT

Random init

(c) ViT-L/14

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\".\n\nSTRENGTHS:\nOverall 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.\n\nWEAKNESSES:\n- 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]).\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nWould have been nice to see a limitations section.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\nAs 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nThe authors have adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis 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\nWEAKNESSES:\n* 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\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\n* 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.", "TITLE": "Official Review", "RECOMMENDATION": 9, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\nThese 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.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\nOne 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.\n\nQUESTIONS:\nI’ve seen that in the supplementary material, there are some code snippets, but do the authors plan to release the full codebase upon acceptance?\n\nLIMITATIONS:\nThe authors adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 10, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}