SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Guillermo Ortiz-Jimenez* EPFL, Lausanne, Switzerland guillermo.ortizj Jimenez@epfl.ch
Alessandro Favero* EPFL, Lausanne, Switzerland alessandro.favero@epfl.ch
Pascal Frossard  
EPFL, Lausanne, Switzerland  
pascal.frossard@epfl.ch
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.
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
Figure 1: Illustration of weight disentanglement, where distinct directions in the weight space,  $\pmb{\tau}_{t}$ , are associated with localized areas of the input space,  $\mathcal{D}_t$ . This allows a model,  $f$ , to manipulate these areas independently by adding linear combinations of  $\pmb{\tau}_{t}$ 's to a pre-trained checkpoint  $\theta_0$ .
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 Wortman 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:\mathcal{X}\times \Theta \to \mathcal{V}$  be a neural network taking inputs  $\pmb {x}\in \mathcal{X}$  and parameterized by a set of weights  $\theta \in \Theta$ . We will assume  $\mathcal{X}\subseteq \mathbb{R}^d$ $\Theta \subseteq \mathbb{R}^{m}$  and  $\mathcal{V}\subseteq \mathbb{R}^c$ . Consider  $T$  tasks, with every task  $t$  consisting of a triplet  $(\mathcal{D}_t,\mu_t,f_t^\star)$  where  $\mathcal{D}_t\subseteq \mathcal{X}$  is a data support (e.g., ImageNet [21] images),  $\mu_{t}$  an input distribution such that  $\mathrm{supp}(\mu_t) = \mathcal{D}_t$ , and  $f_{t}^{\star}:\mathcal{D}_{t}\rightarrow \mathcal{V}$  a target function (e.g., labels). In practice, each task is identified with a training set  $\{(x_{\nu},f_t^\star (x_\nu))\}_{\nu \in [n_t]}$  with  $\pmb {x}\sim \mu_t$ , that is used to fine-tune the models starting from the pre-trained weights  $\pmb{\theta}_{0}$  and obtain the fine-tuned weights  $\pmb{\theta}_{t}^{\star}$ .
Task arithmetic. Let the task vector of task  $t$  be the difference between the fine-tuned and the pretrained weights, i.e.,  $\pmb{\tau}_{t} = \pmb{\theta}_{t}^{\star} - \pmb{\theta}_{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  $\mathcal{T} = \{\pmb{\tau}_t\}_{t\in [T]}$  with associated non-intersecting task supports  $\mathcal{D} = \{\mathcal{D}_t\subset \mathcal{X}\}_{t\in [T]}$ , i.e.,  $\forall t,t^{\prime}$ , if  $t\neq t^{\prime}$  then  $\mathcal{D}_t\cap \mathcal{D}_{t'} = \emptyset$ . We say a network  $f$  satisfies the task arithmetic property around  $\theta_0$  with respect to  $\mathcal{T}$  and  $\mathcal{D}$  if
$$
f \left(\boldsymbol {x}; \boldsymbol {\theta} _ {0} + \sum_ {t = 1} ^ {T} \alpha_ {t} \boldsymbol {\tau} _ {t}\right) = \left\{ \begin{array}{l l} f (\boldsymbol {x}; \boldsymbol {\theta} _ {0} + \alpha_ {t} \boldsymbol {\tau} _ {t}) & \boldsymbol {x} \in \mathcal {D} _ {t} \\ f (\boldsymbol {x}; \boldsymbol {\theta} _ {0}) & \boldsymbol {x} \notin \bigcup_ {t = 1} ^ {T} \mathcal {D} _ {t} \end{array} \right. \tag {1}
$$
with  $(\alpha_{1},\ldots ,\alpha_{T})\in \mathcal{A}\subseteq \mathbb{R}^{T}$
In short, a model satisfies Property 1 if adding  $\pmb{\tau}_t$  does not modify the output of the model outside  $\mathcal{D}_t$ .
Neural tangent kernel. Around the initialization weights  $\theta_0$ , a neural network can be approximated with a first-order Taylor expansion:
$$
f (\boldsymbol {x}; \boldsymbol {\theta}) \approx f (\boldsymbol {x}; \boldsymbol {\theta} _ {0}) + \left(\boldsymbol {\theta} - \boldsymbol {\theta} _ {0}\right) ^ {\top} \nabla_ {\boldsymbol {\theta}} f (\boldsymbol {x}; \boldsymbol {\theta} _ {0}). \tag {2}
$$
This approximation is equivalent to a kernel predictor with a kernel known as the neural tangent kernel (NTK) [41],  $k_{\mathrm{NTK}}(\pmb{x}, \pmb{x}^{\prime}) = \nabla_{\pmb{\theta}}f(\pmb{x}; \pmb{\theta}_{0})^{\top}\nabla_{\pmb{\theta}}f(\pmb{x}^{\prime}; \pmb{\theta}_{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(\cdot ;\theta_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(\pmb{x}; \pmb{\theta}^{\star}) - f(\pmb{x}; \pmb{\theta}_0) \approx (\pmb{\theta}^{\star} - \pmb{\theta}_0)^{\top} \nabla_{\pmb{\theta}} f(\pmb{x}; \pmb{\theta}_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_{\mathrm{lin}}$ . That is, we apply the fine-tuned task vectors  $\tau = \theta^{\star} - \theta_0$  to the linear approximation of  $f$  at  $\theta_0$ , i.e.,
$$
f _ {\mathrm {l i n}} (\boldsymbol {x}; \boldsymbol {\theta} _ {0} + \boldsymbol {\tau}) = f (\boldsymbol {x}; \boldsymbol {\theta} _ {0}) + \boldsymbol {\tau} ^ {\top} \nabla_ {\boldsymbol {\theta}} f (\boldsymbol {x}; \boldsymbol {\theta} _ {0}), \tag {3}
$$
and we check whether  $f_{\mathrm{lin}}(\cdot ;\pmb{\theta}^{\star})$  performs similarly to  $f(\cdot ;\pmb{\theta}^{\star})^{2}$
The code to reproduce our experiments can be found at https://github.com/gortizji/tangent_task_arithmetic.
Table 1: Task addition. Average absolute  $(\%)$  and normalized accuracies  $(\%)$  of different CLIP ViTs edited by adding the sum of the task vectors of 8 tasks. We report results for the non-linear and linearized models of Sections 3 and 5 normalizing performance by their single-task accuracies.
Method	ViT-B/32	ViT-B/16	ViT-L/14
Abs. (↑)	Norm. (↑)	Abs. (↑)	Norm. (↑)	Abs. (↑)	Norm. (↑)
Pre-trained	f(·; θ0)	48.4	-	55.2	-	64.4	-
Non-lin. FT	f(·; θ0 + τ)	71.4	76.5	75.5	80.0	85.1	88.8
Post-hoc lin.	flin(·; θ0 + τ)	57.1	81.9	65.0	85.2	75.2	90.0
Linear. FT	flin(·; θ0 + τlin)	76.5	85.4	81.3	86.0	88.5	93.5
Table 2: Task negation. Minimum accuracy (\%) of different CLIP ViTs edited by negating a task vector from a target task while retaining  $95\%$  of their performance on the control task. We report average performances over eight tasks on non-linear and linearized models as introduced in Sections 3 and 5.
Method	ViT-B/32	ViT-B/16	ViT-L/14
Targ. (↓)	Cont. (↑)	Targ. (↓)	Cont. (↑)	Targ. (↓)	Cont. (↑)
Pre-trained	f(·; θ0)	48.4	63.4	55.2	68.3	64.4	75.5
Non-lin. FT	f(·; θ0−τ)	24.0	60.7	19.2	64.6	18.0	72.5
Post-hoc lin.	flin(·; θ0−τ)	14.8	60.3	10.8	64.8	12.1	71.8
Linear. FT	flin(·; θ0−τlin)	10.9	60.8	11.3	64.8	7.9	72.5
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_{\mathrm{lin}}(\cdot ; \theta^{\star})$  is significantly lower than that of  $f(\cdot; \theta^{\star})$  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
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_{\mathrm{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:
1. Task addition: The sum of the task vectors  $\pmb{\tau} = \sum_{t}\pmb{\tau}_{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.
2. 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.
tors  $\tau = \sum_{t}\tau_{t}$  is added to a pre-trained checkpoint. Process of this benchmark is measured in terms of the different tasks. Results are shown in Table 1.
lated from the pre-trained checkpoint to forget a task or job task (ImageNet). The success of this benchmark drop in accuracy on the forgetting task that retains results are averaged over tasks and shown in Table 2.
Figure 2: Non-linear advantage. Single-task accuracies of non-linearly fine-tuned models  $f(\cdot; \theta^{\star})$  and their post-hoc linearization  $f_{\mathrm{lin}}(\cdot; \theta^{\star})$ . Markers represent different ViTs.
and during task arithmetic, which is the objective of a tangent space could be the only useful component.
To obtain the task vectors, we use the fine-tuned weights of the different ViTs from before, and use a single mixing coefficient  $\alpha = \alpha_{1} = \dots = \alpha_{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: \mathcal{X} \times \Theta \to \mathcal{Y}$  is weight disentangled with respect to a set of task vectors  $\mathcal{T} = \{\pmb{\tau}_t\}_{t \in [T]}$  and the corresponding supports  $\mathcal{D} = \{\mathcal{D}_t\}_{t \in [T]}$  if
$$
f \left(\boldsymbol {x}; \boldsymbol {\theta} _ {0} + \sum_ {t = 1} ^ {T} \alpha_ {t} \boldsymbol {\tau} _ {t}\right) = \sum_ {t = 1} ^ {T} g _ {t} (\boldsymbol {x}; \alpha_ {t} \boldsymbol {\tau} _ {t}) + g _ {0} (\boldsymbol {x}), \tag {4}
$$
where  $g_{t}(\pmb {x};\alpha_{t}\pmb{\tau}_{t}) = \mathbf{0}$  for  $\pmb {x}\notin \mathcal{D}_t$  and  $t = 1,\dots ,T$  , and  $g_0(\pmb {x}) = 0$  for  $\pmb {x}\in \bigcup_{t\in [T]}\mathcal{D}_t$
In essence, this definition captures the idea that the function  $f$  can be decomposed as a sum of spatially-localized components, i.e., vanishing outside a spatial region, whose functional variation is entirely captured by each  $\tau_{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
$$
f \left(\boldsymbol {x}; \boldsymbol {\theta} _ {0} + \sum_ {t = 1} ^ {T} \alpha_ {t} \boldsymbol {\tau} _ {t}\right) = \sum_ {t = 1} ^ {T} f \left(\boldsymbol {x}; \boldsymbol {\theta} _ {0} + \alpha_ {t} \boldsymbol {\tau} _ {t}\right) \mathbb {1} \left(\boldsymbol {x} \in \mathcal {D} _ {t}\right) + f (\boldsymbol {x}; \boldsymbol {\theta} _ {0}) \mathbb {1} \left(\boldsymbol {x} \notin \bigcup_ {t \in [ T ]} \mathcal {D} _ {t}\right), \tag {5}
$$
and identify  $g_{t}(\pmb{x};\alpha_{t}\pmb{\tau}_{t}) = f(\pmb{x};\pmb{\theta}_{0} + \alpha_{t}\pmb{\tau}_{t})\mathbb{1}(\pmb{x}\in \mathcal{D}_{t})$  and  $g_{0}(\pmb{x}) = f(\pmb{x};\pmb{\theta}_{0})\mathbb{1}(\pmb{x}\notin \mathcal{D}_{t})$ . It is important to highlight, however, that this additive decomposition does not imply linearity, as the local functions  $\{g_t\}_{t\in [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(\cdot ;\pmb {\theta}_0 + \alpha \pmb {\tau})$  does not generalize for some  $\alpha$ . 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,
$$
\xi \left(\alpha_ {1}, \alpha_ {2}\right) = \sum_ {t = 1} ^ {2} \mathbb {E} _ {\boldsymbol {x} \sim \mu_ {t}} \left[ \operatorname {d i s t} \left(f (\boldsymbol {x}; \boldsymbol {\theta} _ {0} + \alpha_ {t} \boldsymbol {\tau} _ {t}), f (\boldsymbol {x}; \boldsymbol {\theta} _ {0} + \alpha_ {1} \boldsymbol {\tau} _ {1} + \alpha_ {2} \boldsymbol {\tau} _ {2})\right) \right], \tag {6}
$$
where  $\text{dist}$  denotes any distance metric between output vectors. As we are dealing with classification tasks, in what follows we use the prediction error  $\text{dist}(y_1, y_2) = \mathbb{1}(y_1 \neq y_2)$  as the distance metric. In general, the smaller the value of  $\xi(\alpha_1, \alpha_2)$  the more weight disentangled a model is at  $(\alpha_1, \alpha_2)$ .
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  $\theta_0$ , which enables task arithmetic. However, for  $\alpha_{1}, \alpha_{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  $\theta_0$  - task arithmetic is more effective when fine-tuning with small learning rates and few training steps [39] - with the optimal value of  $\alpha$  typically being less than 1.
Figure 3: Visualization of weight disentanglement. The heatmaps show the disentanglement error  $\xi (\alpha_{1},\alpha_{2})$  of a non-linear CLIP ViT-B/32 (top) and its post-hoc linearization (bottom) on different example task pairs. The light regions denote areas of the weight space where weight disentanglement is stronger. The red box delimits the search space used to compute the best  $\alpha$  in all our experiments.
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 single-task accuracy gap.