SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Specifically, rather than applying the nonlinearly fine-tuned task vectors  $\pmb{\tau} = \pmb{\theta}^{\star} - \pmb{\theta}_{0}$  to  $f_{\mathrm{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_{\mathrm{lin}}$  around  $\pmb{\theta}_{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
Figure 4: Conceptual illustration of the different approaches we use to edit a pretrained model  $f(\cdot ;\theta_0)$ . Here  $\mathcal{N}$  represents the space of neural network functions  $f$ , non-linearly parameterized by  $\pmb {\theta}\in \Theta$ ; and  $\kappa$  its tangent space, given by the space of linearized functions  $f_{\mathrm{lin}}$ .
linear connection between the weight-space and function-space defined in Eq. (3), fine-tuning  $f_{\mathrm{lin}}$  is essentially the same as training a kernel predictor with kernel  $k_{\mathrm{NTK}}$ . As a result, we obtain the fine-tuned weights  $\pmb{\theta}_{\mathrm{lin}}^{\star}$  of the linearized model for each task, which allows us to construct the corresponding task vector  $\pmb{\tau}_{\mathrm{lin}} = \pmb{\theta}_{\mathrm{lin}}^{\star} - \pmb{\theta}_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_{\mathrm{lin}}(\pmb{x}; \pmb{\theta}_0 + \pmb{\tau}_{\mathrm{lin}}) \neq f_{\mathrm{lin}}(\pmb{x}; \pmb{\theta}_0 + \pmb{\tau})$ . In particular, although both models share the same kernel  $k_{\mathrm{NTK}}(\pmb{x}, \pmb{x}')$ , the task vectors  $\pmb{\tau}_{\mathrm{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_{\mathrm{lin}}(\cdot; \pmb{\theta}_0 + \pmb{\tau}_{\mathrm{lin}})$  is very similar to the one of  $f(\cdot; \pmb{\theta}_0 + \pmb{\tau})$  (cf. Figure 2).
Figure 5: Single-task accuracies of non-linearly FT,  $f(\cdot ;\pmb{\theta}^{\star})$  and linearly FT,  $f_{\mathrm{lin}}(\cdot ;\pmb{\theta}_{\mathrm{lin}}^{\star})$ , models.
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<sup>3</sup>.
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  $\mathcal{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  $\{(\phi_{\rho}, \lambda_{\rho})\}_{\rho \in \mathbb{N}}$ ; which implies that  $k$  can only represent functions of the form  $f^{\star}(\pmb{x}) = \sum_{\rho=1}^{\infty} c_{\rho} \phi_{\rho}(\pmb{x})$  with a finite kernel norm, i.e.,  $\|f^{\star}\|_{\mathcal{H}}^{2} = \sum_{\rho=1}^{\infty} c_{\rho}^{2} / \lambda_{\rho} < +\infty$ . Specifically, the coefficients  $\{c_{\rho}\}_{\rho \in \mathbb{N}}$  constitute a representation of the function  $f^{\star}$  in the kernel basis.
Consider  $T$  tasks  $\{f_t^\star\}_{t \in [T]}$  supported in their respective non-intersecting domains  $\{\mathcal{D}_t\}_{t \in [T]}$ . Furthermore, let  $\{\phi_\rho\}_{\rho \in \mathbb{N}}$  be an orthogonal basis of eigenfunctions that diagonalizes the kernel on the union of all  $\mathcal{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:
<sup>3</sup>In Appendix D.3, we observe that larger models are more weight disentangled.
Proposition 1 (Simplified). Suppose that  $\{f_t^\star\}_{t \in [T]}$  can be represented by the kernel  $k$ . The kernel  $k$  is capable of performing task arithmetic with respect to  $\{f_t^\star\}_{t \in [T]}$  and  $\{\mathcal{D}_t\}_{t \in [T]}$  if, for each task  $t$ , there exists a subset of localized eigenfunctions such that  $i)$ $\mathrm{supp}(\phi) \subseteq \mathcal{D}_t$  for each  $\phi$  in the subset, and ii) the representation of  $f_t^\star$  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_{\nu},f_t^\star (x_\nu))\}_{\nu \in [n_t]}$  and  $x_{\nu}\sim \mu_t$ , the CLIP's predictor evaluated at a new point  $\pmb {x}\in \mathcal{X}$  can be expressed as a linear combination of its kernel  $k_{\mathrm{NTK}}$  evaluated on  $\pmb{x}$  and the training data, i.e.,  $f_{\mathrm{lin}}(\pmb {x}) = f(\pmb {x};\pmb {\theta}_0) + \sum_{\nu \in [n_t]}\beta_\nu k_{\mathrm{NTK}}(\pmb {x}_\nu ,\pmb {x})$ .
To explore whether CLIP models use localized eigenfunctions for task arithmetic, we diagonalize the matrix  $(K_{\mathrm{NTK}})_{ij} = k_{\mathrm{NTK}}(\pmb{x}_i,\pmb{x}_j)$  with  $\pmb{x}_i\in \mathcal{D}_t$ , i.e., the task on which we trained, and  $\pmb{x}_j\in \mathcal{D}_t\cup \mathcal{D}_{t'}$ , where  $\mathcal{D}_{t'}$  is the support of a control task. If the eigenfunctions used to represent  $f^{\star}(\pmb{x})$  are localized, then the power of the eigenvectors of  $K_{\mathrm{NTK}}$  must be concentrated in the points belonging to the dataset used for training. To measure this concentration, we introduce the local energy  $\mathcal{E}_{\mathrm{loc}}(\pmb{x}) = \sum_{\rho}\phi_{\rho}^{2}(\pmb{x})$ , which sums the power of all the eigenfunctions  $\phi_{\rho}$  at a given point  $\pmb{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 task-specific operations without interference. We believe that further investigation into this intriguing localization phenomenon holds the potential to deepen our understanding of these models.
Figure 6: Eigenfunction localization. Estimated support of the eigenfunctions of the NTK of a ViT-B/32 CLIP model trained on RESISC45. The plot shows the sum of the local energy of the eigenfunctions over a random subset of the training and control supports (RESISC45 and Cars, respectively).
Remark. While we have shown that localized eigenfunctions can play a crucial role in task arithmetic, it is important to note that they are not always necessary. In fact, the task arithmetic property can hold even if the eigenfunctions used to represent 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  $\theta_0^{\mathrm{rd}}$  does not result in significant improvements in multi-task accuracy over random chance. This
Table 3: Task addition from random initialization. We use the same setup as for the experiments in Table 1 but with task vectors obtained from fine-tuning randomly initialized ViTs. Results compare the average single-task accuracy (\%) after fine-tuning and the multi-task accuracy (\%) via task addition.
Method	ViT-B/32	ViT-B/16	ViT-L/14
Sing. (↑)	Multi (↑)	Sing. (↑)	Multi (↑)	Sing. (↑)	Multi (↑)
Random init	f(·; θ0rd)	5.3	-	4.8	-	5.2	-
Non-lin. FT	f(·; θ0rd + τrd)	48.5	5.5	40.6	4.5	18.0	4.8
Linear. FT	flin(·; θ0rd + τrd)	27.8	3.8	24.7	4.0	24.8	6.1
holds true for both non-linear task vectors,  $\tau^{\mathrm{rd}}$ , and linearized ones,  $\tau_{\mathrm{lin}}^{\mathrm{rd}}$ . 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.
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 fine-tuned 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 few-shot performance [5]. Our work serves as another example of the utility of linearized models in certain scenarios where they do not only offer practical benefits but also provide valuable theoretical insights.
Feature disentanglement. The notion of feature disentanglement lies at the heart of representation learning, where 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.
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.
We thank Nikolaos Dimitriadis, Adam Hazimeh, Pau de Jorge, Nikolaos Karalias, Seyed Mohsen Moosavi-Dezfooli, Antonio Sclocchi, Thibault Séjourné, and Matthieu Wyart for helpful feedback and comments. We also thank Gabriel Ilharco for helping set up the code to reproduce their experiments.
[1] 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).
[2] 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).
[3] 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).
[4] 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).
[5] 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).
[6] 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/bitstream/ddf354219aac374f1d40b7e760ee5bbb7-Abstract.html. (page 9).
[7] 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).
[8] Aristide Baratin, Thomas George, César Laurent, R. Devon Hjelm, Guillaume Lajoie, Pascal Vincent, and Simon Lacoste-Julien. Implicit regularization via neural feature alignment. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2021. http://proceedings.mlr.press/v130/baratin21a.html. (page 9).
[87] 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).
[88] 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).
[89] 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).
[90] Gizem Yüce, Guillermo Ortiz-Jiménez, Beril Besbinar, and Pascal Frossard. A structured dictionary perspective on implicit neural representations. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2022. https://doi.org/10.1109/CVPR52688.2022.01863. (pages 3 and 9).
[91] Luca Zancato, Alessandro Achille, Avinash Ravichandran, Rahul Bhotika, and Stefano Soatto. Predicting training time without training. In Advances in Neural Information Processing Systems (NeurIPS), 2020. https://proceedings.neurips.cc/paper/2020/bit/440e7c3eb9bbcd4c33c3535354a51605-Abstract.html. (pages 3 and 9).
[92] 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).
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 openclip 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  $\alpha$  in task arithmetic benchmarks. As in Ilharco et al. [39] we use a single coefficient  $\alpha$  to tune the size of the task vectors used to modify the pre-trained models. This is equivalent to setting  $\alpha = \alpha_{1} = \dots \alpha_{T}$  in Eq. (1). Both in the task addition and task negation benchmarks, after fine-tuning, we evaluate different scaling coefficients  $\alpha \in \{0.0, 0.05, 0.1, \dots, 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  $\alpha$  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  $\pmb{\tau} = \sum_{t}\pmb{\tau}_{t}$ . Here, the normalization is performed with respect to the single-task accuracies achieved by the model fine-tuned on each task. Mathematically,
$$
\text {N o r m a l i z e d} = \frac {1}{T} \sum_ {t = 1} ^ {T} \frac {\operatorname {a c c} \left[ f \left(\boldsymbol {x}; \boldsymbol {\theta} _ {0} + \sum_ {t ^ {\prime}} \boldsymbol {\tau} _ {t ^ {\prime}}\right) \right]}{\underset {\boldsymbol {x} \sim \mu_ {t}} {\operatorname {a c c}} \left[ f \left(\boldsymbol {x}; \boldsymbol {\theta} _ {0} + \boldsymbol {\tau} _ {t}\right) \right]}. \tag {7}
$$
Disentanglement error. To produce the weight disentanglement visualizations of Figure 3 we compute the value of  $\xi (\alpha_{1},\alpha_{2})$  on a  $20\times 20$  grid of equispaced values in  $[-3,3]\times [-3,3]$ . To estimate the disentanglement error, we use a random subset of 2,048 test points for each dataset.