PaperlayerView exact revision on arXiv

Source-linked AI summary

Editing Models with Task Arithmetic

Gabriel Ilharco, Marco Tulio Ribeiro, Mitchell Wortsman, Suchin Gururangan, Ludwig Schmidt, Hannaneh Hajishirzi, Ali Farhadi

arXiv:2212.04089v3cs.LGcs.CLcs.CV

TL;DR

Editing pretrained models is useful for improving downstream performance, mitigating unwanted behavior, and updating models, but conventional retraining is costly. The paper constructs task vectors from fine-tuning weight changes and edits models through vector arithmetic. Across vision and language tasks, negation, addition, and task analogies support targeted forgetting, multi-task performance, and improvement on data-scarce tasks without extra inference-time cost.

  • Problem

    Pretrained models often need post-training edits to improve downstream tasks, mitigate biases or unwanted behavior, align with preferences, or incorporate new information.

  • Method

    The paper subtracts pretrained weights from task-fine-tuned weights to form task vectors, then combines them through negation, addition, or task-analogy arithmetic.

  • Results

    Across vision and NLP models, task arithmetic supports targeted forgetting, strong multi-task models, and improved performance on tasks with little or no data.

  • Takeaways & Limitations

    Task arithmetic provides a simple, efficient way to reuse or transfer capabilities from fine-tuned models without additional training or inference-time memory and compute.

  • Takeaways & Limitations

    Task vectors require models with the same architecture and, in the experiments, the same pretrained initialization; models with added parameters require future work or shared-weight merging.

Abstract

from arXiv · show

Changing how pre-trained models behave -- e.g., improving their performance on a downstream task or mitigating biases learned during pre-training -- is a common practice when developing machine learning systems. In this work, we propose a new paradigm for steering the behavior of neural networks, centered around \textit{task vectors}. A task vector specifies a direction in the weight space of a pre-trained model, such that movement in that direction improves performance on the task. We build task vectors by subtracting the weights of a pre-trained model from the weights of the same model after fine-tuning on a task. We show that these task vectors can be modified and combined together through arithmetic operations such as negation and addition, and the behavior of the resulting model is steered accordingly. Negating a task vector decreases performance on the target task, with little change in model behavior on control tasks. Moreover, adding task vectors together can improve performance on multiple tasks at once. Finally, when tasks are linked by an analogy relationship of the form ``A is to B as C is to D", combining task vectors from three of the tasks can improve performance on the fourth, even when no data from the fourth task is used for training. Overall, our experiments with several models, modalities and tasks show that task arithmetic is a simple, efficient and effective way of editing models.

1 INTRODUCTION

The paper introduces task vectors as a modular way to edit pretrained models and shows that arithmetic operations can remove behaviors, combine capabilities, and transfer performance to data-scarce tasks.

  • Task vectors encode task-relevant information by subtracting pretrained weights from weights after fine-tuning on that task.
  • Negating task vectors reduces target-task performance while causing little change on control tasks, supporting targeted forgetting and behavior mitigation.
  • Adding task vectors can produce a single model that performs well across multiple tasks or improves performance on an individual task.In image models, adding two task vectors maintains 98.9% of accuracy, while average performance increases as more vectors are added.
  • Task analogies combine vectors from three related task-and-data combinations to improve a fourth task even when little or no target-task data is available.
  • Task arithmetic operates on model weights without extra inference-time memory or compute and can reuse existing fine-tuned or publicly available models without additional training data.

2 TASK VECTORS

Task vectors are weight-space differences that can be applied and combined element-wise to edit models. The framework supports negation, addition, and analogy operations, with scaling selected using held-out validation data.

  • A task vector is the element-wise difference between a model fine-tuned on a task and the corresponding pretrained model.
  • Task vectors can be added element-wise to parameters from the same architecture, with an optional scaling term λ selected using held-out validation sets.
  • Adding one task vector to a pretrained model with λ = 1 reproduces the model fine-tuned on that task.
  • The approach focuses on open-ended models that can be fine-tuned without introducing new parameters, while merging shared weights for models with added parameters remains future work.
  • The paper studies element-wise negation, addition, and analogy combinations of task vectors to edit model behavior.

3 FORGETTING VIA NEGATION

Negating task vectors enables targeted forgetting: it substantially reduces performance or undesirable generations on the edited task while largely preserving behavior on control tasks.

  • Negating task vectors reduces target-task performance without substantially hurting performance elsewhere, supporting forgetting of unwanted capabilities or behaviors.
  • Negative task vectors reduce ViT-L/14 accuracy on eight image-classification targets by 45.8 percentage points with little loss on ImageNet.
  • Negation is more effective than random vectors or gradient ascent for reducing target accuracy while preserving control-task performance.Random vectors have little target impact, whereas gradient ascent severely deteriorates control-task performance.
  • The text-generation experiment fine-tunes on highly toxic Civil Comments data, then compares negation with gradient ascent, random vectors, and non-toxic fine-tuning.
  • Negative task vectors reduce toxic GPT-2 generations from 4.8% to 0.8% while keeping control-task perplexity within 0.5 points of the pretrained model.

4 LEARNING VIA ADDITION

Task-vector addition combines capabilities from multiple fine-tuned models into one model, supporting multi-task performance and improvements on individual target tasks. Across image and language tasks, the resulting models can remain competitive with specialized fine-tuned models and improve target-task accuracy.

  • Task-vector addition reuses knowledge from fine-tuned models without additional training or access to training data.
  • IMAGE CLASSIFICATION: 98.9% normalized accuracy makes a single model competitive with two specialized models when adding pairs of image-classification task vectors.
  • IMAGE CLASSIFICATION: 91.2% average normalized accuracy is achieved when adding task vectors across all eight image-classification tasks.
  • NATURAL LANGUAGE PROCESSING: Adding task vectors downloaded from Hugging Face Hub can improve fine-tuned T5 models on four GLUE text-classification target tasks.

5 TASK ANALOGIES

Task analogies combine vectors from related tasks or subpopulations to improve a fourth target task when labeled target data is scarce or unavailable. The approach supports domain generalization, new-category recognition, and subpopulation transfer.

  • Task analogies combine vectors from three tasks to improve a fourth task even when little or no target-task data is available.
  • DOMAIN GENERALIZATION: Labeled auxiliary data and unlabeled data from both domains improve sentiment-analysis accuracy on target datasets without labeled target examples.
  • SUBPOPULATIONS WITH LITTLE DATA: The subpopulation experiment constructs four populations from 125 overlapping ImageNet and human-sketch classes, pairing shared classes and styles.
  • SUBPOPULATIONS WITH LITTLE DATA: 3.4 percentage points of average accuracy improvement is obtained across four target subpopulations and three CLIP models.
  • KINGS AND QUEENS: Task analogies yield large gains over pre-trained models for a new image category despite having no training data for that category.

6 DISCUSSION

The discussion examines task-vector geometry, fine-tuning choices, and computational efficiency, while identifying an architectural and initialization constraint. Task vectors are often near-orthogonal, can stabilize early in fine-tuning, and require care with learning rates.

  • SIMILARITY BETWEEN TASK VECTORS: Task vectors are typically close to orthogonal, which the authors speculate enables addition with minimal interference.
  • Applying one task vector can improve accuracy on another task without data for that second task, such as MNIST improving SVHN accuracy.
  • THE IMPACT OF THE LEARNING RATE: Larger learning rates degrade accuracy for task-vector use and fine-tuning, with a more gradual decrease for individual models.
  • THE EVOLUTION OF TASK VECTORS: Intermediate task vectors rapidly converge toward the final direction, while added-vector accuracy saturates after a few hundred steps.
  • LIMITATIONS: Task vectors require models with the same architecture and, in these experiments, the same pre-trained initialization.

7 RELATED WORK

The paper situates task arithmetic within work on weight interpolation, model interventions, and task embeddings. Its distinction is using task vectors primarily to steer pre-trained model behavior rather than only represent task relationships.

  • THE LOSS LANDSCAPE AND INTERPOLATING WEIGHTS: Prior work finds that weight interpolation can preserve high accuracy when neural networks share part of their optimization trajectory.
  • THE LOSS LANDSCAPE AND INTERPOLATING WEIGHTS: Averaging weights from models fine-tuned on different tasks from the same initialization can improve accuracy on those tasks.
  • MODEL INTERVENTIONS: Model-intervention research develops efficient post-training modifications because retraining is prohibitively expensive in many circumstances.
  • TASK EMBEDDINGS: Task-embedding methods represent tasks to predict similarities, transferability, or taxonomic relations, whereas task vectors are mainly used here to steer models.

8 CONCLUSION

The paper presents task arithmetic as an efficient paradigm for editing models by adding or subtracting task vectors. Across vision and NLP, these operations support multi-task performance, behavior removal, task forgetting, and adaptation to data-scarce domains.

  • Adding multiple specialized task vectors produces a single model that performs well across target tasks and can sometimes improve a single task.
  • Negating task vectors removes undesirable behaviors or forgets specific tasks while retaining performance elsewhere.
  • Task analogies use existing data to improve performance on domains or subpopulations where data is scarce.
  • Arithmetic operations over task vectors are efficient to compute and produce a single model of the same size without extra inference cost.

A THE LOSS LANDSCAPE, WEIGHT AVERAGING AND ENSEMBLES

Task vectors are motivated by the behavior of linear paths between pretrained and fine-tuned models and by connections between weight combinations and prediction ensembles. The experiments compare task-vector edits with gradient-ascent and random-vector baselines.

  • Shared optimization trajectories can make linear interpolation between pretrained and fine-tuned weights preserve performance, motivating task-vector edits.
  • Task-vector arithmetic is equivalent to linear combinations of the pretrained model and the fine-tuned models used to generate the vectors.
  • Linear combinations of model weights are connected to ensembling predictions, with exact equivalence under the stated linear-loss approximation.
  • Adding two task vectors with λ = 0.5 yields the uniform average of the corresponding fine-tuned model weights.
  • The experiments compare task-vector edits against gradient-ascent fine-tuning and same-magnitude random vectors.

B.3 BREAKDOWN PER TASK

Additional experiments examine when task-vector negation forgets image, OCR, and person-identification tasks while preserving control-task performance. Results identify task similarity and fine-tuning gains as important boundaries, and compare negation with gradient ascent and interpolation.

  • Task-vector curves initially reduce target-task performance faster than control-task performance, unlike same-magnitude random vectors.
  • Forgetting is least effective when target-image distributions overlap more closely with the ImageNet control task.
  • Filtering semantically overlapping classes does not produce large differences in the forgetting results.
  • Gradient-ascent fine-tuning reduces treatment-task accuracy but substantially deteriorates control-task accuracy, whereas interpolation improves the trade-off.
  • Fine-tuning gains positively correlate with accuracy drops after subtracting the corresponding task vector.
  • Negating task vectors causes large OCR performance drops with little ImageNet change, while forgetting is less effective for person identification.

C.1 EXPERIMENTAL DETAILS

The appendix details toxicity, multi-task, scaling, seed, and ImageNet experiments. These results show that task-vector edits can reduce toxicity and combine task performance efficiently, while requiring coefficient tuning and leaving a gap to joint multi-task training.

  • Task-vector editing is more effective for larger GPT-2 models, with qualitatively similar trends across GPT-2 sizes.
  • Negative task vectors reduce toxic generations while keeping WikiText-103 perplexity close to the pretrained model.On RealToxicityPrompts, evaluation tracks expected maximum toxicity across 25 generations and WikiText-103 perplexity.
  • A single scaling coefficient is used for sums of task vectors, although separate coefficients could improve performance and become costly to search as task count grows.
  • Different random seeds have little impact on edited-model accuracy in the tested MNIST and EuroSAT setup.
  • 0.994 average normalized performance is achieved by joint multi-task fine-tuning, compared with 0.912 for the best task-vector result.
  • Optimal scaling coefficients vary substantially across experiments, while coefficients from 0.3 to 0.5 are close to optimal in many cases.
  • For smaller task-vector subsets, the combined model more closely matches multiple specialized models, with the gap increasing as subset size grows.
  • Adding an ImageNet task vector to another task vector produces one model with high accuracy on both tasks, often competitively with separate fine-tuned models.

D.6 ADDING PAIRS OF TASK VECTORS FROM NLP TASKS

This section tests whether task-vector addition can compress pairs of fine-tuned NLP models into one multi-task model. Across diverse tasks and metrics, the resulting model retains most specialized-model performance and averages 96.7% normalized performance.

  • Experimental setup: Six public T5 models span sentiment analysis, question answering, summarization, question generation, and constrained text generation.Evaluation uses accuracy, exact match, or ROUGE-2, with each task normalized by its fine-tuned-model performance.
  • Results: Task-vector addition compresses pairs of NLP models into a single multi-task model with little performance loss.The experiments use checkpoints downloaded from the Hugging Face Hub rather than models fine-tuned by the authors.
  • Results: 96.7% is the average normalized performance of the model obtained by adding task vectors across the NLP task set.Performance varies more across task subsets than in the image-classification experiments, but some combinations outperform specialized models.
  • External task vectors: The section evaluates whether adding external task vectors can improve performance on specific target tasks.This investigation follows the pair-compression experiments and uses T5-base models fine-tuned on four GLUE tasks.

E TASK ANALOGIES

The section applies task analogies to improve target-task performance without labeled target data, covering sentiment, image-category, and subpopulation settings. It also shows that independent vector scaling yields a small average gain at substantially higher evaluation cost.

  • Task analogies: Task analogies combine vectors from related tasks to improve four new target tasks without training data and with little change on control settings.The table reports learning by analogy for CLIP models.
  • Sentiment analogies: For sentiment analysis, the target vector combines an auxiliary sentiment vector with language-modeling vectors from the target and auxiliary domains.The approach is evaluated on Yelp and Amazon binary-sentiment tasks and compared with auxiliary-task and target-task fine-tuning.
  • Image-category analogies: In the image-category experiment, each target category uses vectors from the other three categories, while ImageNet measures control-task behavior.The analogy example for king is τqueen + (τman − τwoman).
  • Subpopulation analogies: Task analogies produce large accuracy gains over pretrained models with very little control-task change, and results improve with model scale.These results hold despite having no training data for the target task.
  • Scaling coefficients: 0.7 percentage points is the average improvement from independent scaling coefficients over a single coefficient, requiring 103 instead of 10 evaluations.The explored model is θnew = θ + λCτC + λBτB − λAτA, with coefficients searched over 0.0 to 1.0.
Loading 2212.04089v3…