PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Guillermo Ortiz-Jimenez, Alessandro Favero, Pascal Frossard

arXiv:2305.12827v3cs.LGcs.CV

TL;DR

Task arithmetic can edit pre-trained models scalably, yet its effectiveness and underlying principles are not fully understood. The paper studies this problem in vision-language models, formalizes the role of weight disentanglement, and linearizes models in tangent space. It finds that linearization amplifies disentanglement and improves task arithmetic, while linking the mechanism to localized NTK eigenfunctions and pre-training.

  • Problem

    The effectiveness, internal mechanisms, and reliability conditions of task arithmetic remain insufficiently understood, limiting confidence in its adoption.

  • Method

    The paper systematically studies task arithmetic in CLIP models, analyzes weight disentanglement and NTK eigenfunction localization, and proposes tangent-space linearization.

  • Results

    Linearized models improve task arithmetic across benchmarks and models, achieving up to 5.8 points more accuracy on task addition and 13.1 points less on task negation.

  • Takeaways & Limitations

    Weight disentanglement is the crucial task-arithmetic mechanism, emerges during pre-training, and is amplified by linearized fine-tuning.

  • Takeaways & Limitations

    The practical trade-offs between computational cost and performance, along with faster linearization techniques, remain open areas for future work.

Abstract

from arXiv · show

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

Task arithmetic offers a scalable way to edit pre-trained models, but its mechanisms and reliability remain poorly understood. This study identifies weight disentanglement and tangent-space linearization as central to improving task arithmetic.

  • Motivation: Task arithmetic edits pre-trained models by adding or subtracting fine-tuned weight components to improve or forget task-specific behavior.These operations can avoid costly joint fine-tuning across multiple tasks.
  • Study scope: The study examines task arithmetic systematically in contrastively pretrained vision-language models and tests whether linearization explains its effectiveness.The authors specifically investigate CLIP models and the neural tangent kernel.
  • Mechanism: Weight disentanglement is identified as the condition enabling task arithmetic: distinct weight directions govern disjoint, localized input regions.This localization permits independent manipulation of task-related behavior.
  • Linearization: Linearizing models in their tangent space amplifies weight disentanglement and improves task arithmetic across benchmarks and models.The effect is reported for both task addition and task negation.
  • Results: Up to 5.8 points more accuracy on task addition and 13.1 points less accuracy on task negation are achieved on several vision-language benchmarks.The reported changes compare linearized with nonlinear models.
  • Implications: The findings connect weight disentanglement in linearized models to spatial localization of neural tangent kernel eigenfunctions and identify it as emergent from pre-training.The work presents this as a basis for more reliable and scalable model editing.

2 Notation and problem statement

The paper formalizes tasks, task vectors, task arithmetic, and neural tangent-space approximations for analyzing how weight changes affect localized task behavior.

  • Task notation: Each task is represented by a data support Dt, an input distribution μt with supp(μt)=Dt, and a target function f⋆t.Training examples are sampled from μt and paired with target outputs.
  • Task notation: Fine-tuning a pre-trained model from θ0 on task t produces fine-tuned weights θ⋆t.The task training set is used to obtain these weights.
  • Task vectors: A task vector is the fine-tuning update τt=θ⋆t−θ0 relative to the pre-trained checkpoint.Task arithmetic operates on these relative weight components.
  • Task arithmetic: Property 1 defines task arithmetic for task vectors associated with non-intersecting supports, requiring additions to preserve outputs outside the selected support.The property formalizes independent modification of task-specific behavior.
  • Neural tangent kernel: The neural tangent approximation linearizes network outputs around θ0 using the parameter gradient and induces the neural tangent kernel.The approximation becomes exact as network width approaches infinity.
  • Training regimes: At finite width, the linear approximation may fail, whereas minimal parameter movement during fine-tuning can keep pre-trained models in the tangent space.These cases correspond to nonlinear and linear training regimes, respectively.

3 Task arithmetic is not a consequence of linear fine-tuning

CLIP fine-tuning does not occur in the linear regime, and task arithmetic cannot be explained by linear components alone. However, linearized models retain stronger relative editing behavior in task addition and negation.

  • Testing the linear hypothesis: Post-hoc linearization approximates fine-tuned outputs with a first-order Taylor expansion around the pretrained parameters.The approximation is f(x; θ0) + τ^T∇θf(x; θ0), where τ = θ⋆ − θ0.
  • Testing the linear hypothesis: Single-task performance of post-hoc linearized models is significantly lower than non-linearly fine-tuned models across ViT sizes and eight tasks.The tasks are Cars, DTD, SUN397, EuroSAT, GTSRB, MNIST, SVHN, and RESISC45.
  • Benchmarks: Task addition combines task vectors for eight tasks, whereas task negation subtracts one vector while retaining performance on ImageNet.Addition is evaluated by maximum average accuracy; negation by maximum forgetting-task drop subject to control-task retention.
  • Task arithmetic beyond linear components: Task addition accuracy significantly decreases when non-linearly fine-tuned task vectors are applied to post-hoc linearized models.This decrease is consistent across tasks, showing that non-linear components contribute to task arithmetic.
  • Task arithmetic beyond linear components: Despite lower absolute task-addition accuracy, post-hoc linearized models perform better at task negation and achieve higher normalized task-addition accuracies.These results indicate greater consistency with the task-arithmetic property despite weaker single-task accuracy.

4 Weight disentanglement

Weight disentanglement is the necessary condition identified for task arithmetic: task-vector directions govern localized, separable regions of input space. Linearization strengthens this property, especially near the pretrained model.

  • Definition: Weight disentanglement means that distinct task-vector directions produce functional changes confined to corresponding, disjoint input regions.The resulting function decomposes into spatially localized components, without requiring linear dependence on parameters.
  • Definition: Weight disentanglement is equivalent to satisfying the task-arithmetic decomposition, but it is independent of task performance or generalization.A model may be disentangled yet perform poorly if a task-vector update fails to generalize.
  • Measuring disentanglement: Disentanglement error compares a single-task output with the output after combining two task vectors; smaller values indicate stronger disentanglement.For classification, the distance metric is prediction disagreement.
  • Empirical behavior: CLIP ViT-B/32 has minimal disentanglement error near the pretrained checkpoint, while error rises for α1, α2 > 1 as task interactions increase.This supports stronger task arithmetic in a small neighborhood of θ0, where the optimal α is typically below 1.
  • Empirical behavior: Post-hoc linearized models show broader low-error regions than nonlinear models, explaining their stronger normalized addition and negation performance.They retain more relative edited performance but remain worse in absolute terms because nonlinear models have higher single-task accuracy.

5 Enhancing task arithmetic via linearization

The paper improves task arithmetic by fine-tuning models directly in the tangent space, preserving the stronger weight disentanglement of linearized models while recovering single-task performance. This approach substantially improves task addition and negation.

  • 5 Enhancing task arithmetic via linearization: Tangent-space fine-tuning directly optimizes the linearized model around the pretrained initialization to obtain task vectors.This differs from applying nonlinear task vectors after post-hoc linearization.
  • 5 Enhancing task arithmetic via linearization: Linearized fine-tuning produces task vectors distinct from post-hoc linearization because they are optimized for the linearized model.The two approaches use the same NTK but generally yield different predictions and task vectors.
  • 5 Enhancing task arithmetic via linearization: Linearized fine-tuning substantially reduces the nonlinear advantage in single-task accuracy while retaining strong task arithmetic.Its single-task performance is often similar to nonlinear fine-tuning.
  • 5 Enhancing task arithmetic via linearization: Up to 5.8 points more multi-task accuracy through task addition and up to 13.1 points more forgetting through task negation are achieved while maintaining similar control-task accuracy.The advantage is larger for smaller ViT-B/32 models and diminishes toward ViT-L/143.
  • 5 Enhancing task arithmetic via linearization: Training and inference in linearized neural networks require only an O(1) increase in computational costs relative to nonlinear counterparts.The paper identifies the trade-off between computational cost and performance as an avenue for future work.

6 Towards understanding task arithmetic

The paper connects task arithmetic to localized NTK eigenfunctions and shows that weight disentanglement is acquired during pre-training. It also establishes that localized eigenfunctions provide a sufficient condition, but are not always necessary, for task arithmetic.

  • 6 Towards understanding task arithmetic: A kernel can support task arithmetic when each task is represented using localized eigenfunctions confined to its task support.Such representations prevent functions for different tasks from interfering.
  • 6.1 Eigenfunction localization: The NTK localization analysis measures eigenfunction energy on a training task and a control task to test whether representations concentrate on the trained domain.The local-energy metric sums the power of eigenfunctions at each point.
  • 6.1 Eigenfunction localization: For ViT-B/32 CLIP trained on RESISC45, eigenfunction local energy is significantly higher on RESISC45 than on Cars, supporting localization across data domains.The paper links this localization to the separation of task representations.
  • 6.1 Eigenfunction localization: Localized eigenfunctions are not always necessary: task arithmetic can also hold when eigenfunctions cancel outside a task domain.If eigenfunctions remain locally linearly independent, however, localization becomes necessary for task arithmetic.
  • 6.2 Weight disentanglement emerges during pre-training: Task arithmetic is not achieved by adding task vectors from randomly initialized ViTs, whose multi-task accuracy does not significantly exceed random chance.This holds for both nonlinear and linearized task vectors.
  • 6.2 Weight disentanglement emerges during pre-training: The findings support the conclusion that task arithmetic is acquired during pre-training, alongside disentanglement of weights governing semantic task regions.The paper distinguishes this from semantic feature disentanglement alone.

7 Related work

Prior work studies weight interpolation and linearized models, while this paper distinguishes weight disentanglement from feature disentanglement in task arithmetic.

  • Weight interpolation can improve single-task performance, sometimes surpassing fine-tuning accuracy, and has also been explored for multi-task model manipulation.
  • Linear vs non-linear regime: Linearized models have been applied to predict fine-tuning generalization and training speed, enhance calibration, and study neural-network dynamics.
  • Feature disentanglement: Feature disentanglement separates distinct data variation factors along different feature-space directions, whereas this work studies weight disentanglement in task arithmetic.
  • Feature disentanglement: Weight disentanglement connects weight-space changes to localized function-space transitions rather than treating disentanglement only as a static representation property.
  • Feature disentanglement: The paper reports natural emergence of task-specific weight disentanglement for semantically meaningful tasks during CLIP pre-training.

8 Conclusion

The conclusion identifies linearized models and weight disentanglement as central to effective task arithmetic, while leaving the origins of disentanglement during pre-training open.

  • Linearized models governed by the NTK outperform nonlinear counterparts in task arithmetic, offering a more effective approach for model editing.
  • Weight disentanglement supports task arithmetic because distinct weight-space directions correspond to localized function-space regions.
  • Weight disentanglement emerges during pre-training, but how it arises remains an open question for future work.
  • The paper identifies algorithms that enhance disentanglement and tangent-space editing of other pre-trained models as research directions.

A Experimental details

The experiments use a common CLIP fine-tuning setup, tune task-vector scaling on held-out data, and evaluate task arithmetic with normalized accuracies and disentanglement measurements.

  • All experiments use four V100 NVIDIA GPUs with 32GB of memory each and require less than 350 GPU hours to reproduce.
  • Fine-tuning: Fine-tuning starts from the same CLIP checkpoint and uses 2,000 iterations, batch size 128, learning rate 10^-5, cosine annealing, and 200 warm-up steps.
  • Tuning of α in task arithmetic benchmarks: A single coefficient α scales task vectors, with α selected from 0.0 through 1.0 using a held-out training-set proportion and the highest target metric.
  • Normalized accuracies in task addition: Task-addition results normalize accuracy against single-task accuracies from models fine-tuned separately on each task.
  • Disentanglement error: Disentanglement error is estimated on a 20 × 20 grid over [-3, 3] × [-3, 3] using 2,048 random test points per dataset.
  • NTK eigenfunction estimation: NTK eigenfunctions are estimated from KNTK matrices computed on 200 random training points per dataset, followed by singular value decomposition.

B Implementation aspects of linearized models

The implementation wraps any neural network as a trainable linearized model using functorch, computes outputs through Jacobian-vector products, and incurs higher computational costs.

  • Practical implementation: The implementation creates a trainable linearized version of any nn.Module around its initialization using functorch.
  • CLIP-specific implementation: Only the CLIP image encoder is linearized because the text encoder remains frozen, and image outputs are normalized before the text-embedding inner product.
  • Implementation: The implementation uses PyTorch nn.Module, functional conversion, and forward-mode automatic differentiation through functorch.
  • Model construction: The model stores frozen initialization parameters and separate trainable current parameters for linearized training.
  • Forward computation: Linearized outputs add the Jacobian-vector contribution from parameter differences to the initialization output.
  • Computational complexity: Inference with linearized models usually costs two or three times more, while training has roughly twice the gradient cost and double memory footprint.

C Spectral analysis of linearized models

The spectral analysis connects task arithmetic to localized NTK eigenfunctions under specific independence conditions. Experiments show that linearization strengthens disentanglement and task-arithmetic performance, while pre-training is necessary for the observed localization.

  • Theoretical spectral conditions: Task arithmetic follows when task-relevant kernel eigenfunctions are supported within disjoint task domains.Under local linear independence across domains, localization becomes necessary; without that independence, it is generally sufficient but not always necessary.
  • Theoretical spectral conditions: Fourier eigenfunctions on rings and tori are not localized, illustrating settings where the spectral condition cannot explain task arithmetic.The analysis uses analytically diagonalizable NTKs whose Fourier atoms are linearly independent on intervals but not localized.
  • Empirical spectral findings: Linearized task arithmetic achieves the highest accuracies for all tested models and datasets except SUN397.The performance gap between linear and nonlinear task vectors decreases as model size increases, consistent with stronger kernel behavior in larger models.
  • Empirical spectral findings: Linearized CLIP models exhibit significantly more weight disentanglement than their nonlinear counterparts across multiple dataset pairs.This pattern appears for both direct tangent-space fine-tuning and post-hoc linearization.
  • Empirical spectral findings: Weight disentanglement is absent after fine-tuning from random initialization, supporting its emergence through pre-training.Randomly initialized models show high disentanglement error, whereas pre-trained models develop localized task-relevant directions.
  • Empirical spectral findings: Closed-vocabulary supervised models can perform task arithmetic at lower absolute and normalized accuracy than open-vocabulary models.Linearized fine-tuning increases normalized accuracy in the closed-vocabulary model but does not surpass nonlinear addition because its single-task performance is lower.

D.9 Further experiments with randomly-initialized networks

Randomly initialized models can learn individual benchmark tasks, but their task-arithmetic failure is attributed primarily to insufficient weight disentanglement rather than inability to learn.

  • Single-task performance: Randomly initialized CLIP models achieve non-trivial or moderate single-task accuracy on most benchmark tasks.This holds under both nonlinear and linearized fine-tuning dynamics.
  • Interpretation: Task-arithmetic failure in non-pretrained models is attributed to a lack of weight disentanglement, not inability to learn the tasks.The models’ single-task results provide the supporting comparison.
  • Interpretation: The randomly initialized large model generally performs more poorly than the base models.The passage attributes this pattern to greater susceptibility to overfitting as model capacity increases.
Loading 2305.12827v3…