Task Arithmetic in the Tangent Space: Improved Editing of Pre-Trained Models

TL;DR

The paper investigates why task arithmetic—combining fine-tuned weights to edit pre-trained models—works, focusing on vision-language models. It identifies weight disentanglement as the key mechanism, showing that distinct weight directions map to localized input regions, enabling independent, task-specific edits. It further demonstrates that linearizing models in the tangent space amplifies this disentanglement and improves task addition and negation, linking the phenomenon to localized NTK eigenfunctions and pre-training dynamics. The work combines theoretical framing with extensive experiments across CLIP-based architectures, CNNs, and NLP models, offering a practical, scalable approach to model editing via NTK-inspired linearization. Overall, it reveals that pre-training shapes weight-space locality, and that tangent-space editing provides a more reliable path for multi-task customization of large 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.

Figures (6)

• Figure 1: Illustration of weight disentanglement, where distinct directions in the weight space, $\bm{\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 $\bm{\tau}_t$'s to a pre-trained checkpoint $\bm{\theta}_0$.
• Figure 2: Non-linear advantage. Single-task accuracies of non-linearly fine-tuned models $f(\cdot\,;\,\bm{\theta}^\star)$ and their post-hoc linearization $f_{\rm lin}(\cdot\,;\,\bm{\theta}^\star)$. Markers represent different ViTs.
• 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.
• Figure 4: Conceptual illustration of the different approaches we use to edit a pretrained model $f(\cdot\,;\,\bm{\theta}_0)$. Here $\mathcal{N}$ represents the space of neural network functions $f$, non-linearly parameterized by $\bm{\theta}\in\bm\Theta$; and $\mathcal{K}$ its tangent space, given by the space of linearized functions $f_{\text{lin}}$.
• Figure 5: Single-task accuracies of non-linearly FT, $f(\cdot\,;\,\bm{\theta}^\star)$ and linearly FT, $f_{\rm lin}(\cdot\,;\,\bm{\theta}_{\rm lin}^\star)$, models.

Theorems & Definitions (1)

• Proposition 1: Simplified