Task Arithmetic

Task arithmetic is a way of editing a pretrained model by adding and subtracting task vectors, the weight changes that fine-tuning on each task produced. Adding the task vectors of several tasks to the pretrained model weights gives one model that handles all of them, subtracting a task vector makes the model worse at that task while leaving others mostly intact, and combining vectors in an analogy transfers a skill to a setting with no labeled data. Every edit is element-wise addition on the weights, so it needs no training data or gradient steps beyond the original fine-tunes. Gabriel Ilharco and colleagues at the University of Washington, Microsoft Research, and the Allen Institute for AI introduced the approach in Editing Models with Task Arithmetic, published at ICLR 2023.

The Three Operations#

All three operations produce new weights θ_new = θ_pretrained + λτ_new, where the scaling coefficient λ is chosen on held-out validation data. Ilharco et al. tested them on CLIP models fine-tuned on eight image classification tasks, such as traffic signs and satellite images, and on GPT-2 and T5 language models.

Negation uses τ_new = −τ to make a model forget a task. Negating the eight CLIP task vectors lowered the ViT-L/14 model's average accuracy on those tasks by 45.8 percentage points while keeping at least 95% of its ImageNet accuracy. Subtracting a vector fine-tuned on toxic comments cut the share of toxic generations from GPT-2 Large from 4.8% to 0.8%, with perplexity on ordinary text within 0.5 points of the original model.

Addition uses τ_new = τ_1 + τ_2 + … to build multi-task models. Adding the vectors of two CLIP tasks gave one model with 98.9% of the accuracy of the two specialized models on average, and adding all eight reached 91.2%.

Analogies apply when tasks relate as "A is to B as C is to D", using τ_D ≈ τ_C + (τ_B − τ_A). Ilharco et al. built a sentiment analysis model for Yelp reviews without labeled Yelp data by adding a vector for Amazon review sentiment to the difference between language-modeling vectors trained on Yelp and Amazon text.

These operations work because task vectors from different tasks point in nearly orthogonal directions, with cosine similarities close to zero, so adding one barely changes what another encodes. Interference still grows with the number of tasks, and model merging methods such as TIES-Merging reduce it by trimming each vector's smallest changes and resolving parameters whose signs disagree before adding.

Adding and Negating Skills in Code#

This PyTorch example pretrains a small classifier briefly on 8x8 handwritten digits, then fine-tunes its first layer for two tasks: task A classifies the digits 0-4 and task B the digits 5-9, each read from its own five frozen output logits, as each CLIP task uses its own frozen classification head in the paper. It then evaluates the pretrained weights plus different combinations of the two task vectors:

import copy

import torch
from sklearn.datasets import load_digits
from torch import nn

torch.manual_seed(0)
X, y = load_digits(return_X_y=True)
X, y = torch.tensor(X / 16, dtype=torch.float32), torch.tensor(y)
tasks = {"A": (0, 5), "B": (5, 10)}  # task A classifies digits 0-4, task B digits 5-9, each with its own 5 logits
train, test = torch.arange(len(X)) < 1000, torch.arange(len(X)) >= 1000


def task_loss(model, task, rows):
    """Cross-entropy and accuracy of one task's 5-way classifier on the given rows."""
    lo, hi = tasks[task]
    keep = rows & (y >= lo) & (y < hi)
    logits, labels = model(X[keep])[:, lo:hi], y[keep] - lo
    return nn.functional.cross_entropy(logits, labels), (logits.argmax(1) == labels).float().mean().item()


def fit(model, params, steps, loss_fn):
    """Run full-batch SGD on the given parameters."""
    optimizer = torch.optim.SGD(params, lr=0.05, momentum=0.9)
    for _ in range(steps):
        optimizer.zero_grad()
        loss_fn().backward()
        optimizer.step()
    return model


# Pretrained model: a few steps on 200 images of all ten digits
base = nn.Sequential(nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 10))
fit(base, base.parameters(), 50, lambda: nn.functional.cross_entropy(base(X[:200]), y[:200]))

# Fine-tune the first layer for each task (heads stay frozen) and keep the task vectors
vectors = {}
for task in tasks:
    tuned = copy.deepcopy(base)
    fit(tuned, tuned[0].parameters(), 100, lambda: task_loss(tuned, task, train)[0])
    vectors[task] = {k: v - base.state_dict()[k] for k, v in tuned.state_dict().items()}

edits = {"Pretrained": {}, "+A": {"A": 1}, "+B": {"B": 1}, "+A +B": {"A": 1, "B": 1}, "-A": {"A": -1}}
for name, coefs in edits.items():
    # Pretrained weights plus a weighted sum of task vectors
    model = copy.deepcopy(base)
    weights = {k: v + sum(c * vectors[t][k] for t, c in coefs.items()) for k, v in base.state_dict().items()}
    model.load_state_dict(weights)
    with torch.no_grad():
        print(f"{name}: task A {task_loss(model, 'A', test)[1]:.3f}, task B {task_loss(model, 'B', test)[1]:.3f}")

On the 797 test digits, the pretrained model scores 0.802 on task A and 0.897 on task B. Adding only task A's vector raises task A to 0.925 but lowers task B to 0.850, the catastrophic forgetting of ordinary fine-tuning, and adding only task B's vector does the reverse. Adding both gives one model at 0.922 and 0.940, close to each specialist on its own task. Subtracting task A's vector drops task A to 0.616 while task B moves only from 0.897 to 0.887.

Requirements and Limits#

Task arithmetic needs every fine-tune to start from the same pretrained weights in the same architecture. When Ultralytics YOLO26 trains from yolo26n.pt on a new dataset, it copies every pretrained tensor whose shape fits, matches classification-head rows by class name, and logs the count as "Transferred N/M items from pretrained weights". Datasets with different numbers of classes leave the final classification convolutions with different shapes, so task vectors from those runs can be added only in the layers the runs share.

The arithmetic also breaks down when vectors conflict or are scaled too far, because large λ values push the weights out of the region where the fine-tuned models perform well. Negation lowers performance on whatever the fine-tuning data covered, which makes it a lightweight form of machine unlearning. Task Arithmetic in the Tangent Space by Guillermo Ortiz-Jimenez, Alessandro Favero, and Pascal Frossard traced successful arithmetic to weight disentanglement, in which separate directions in weight space change the model's outputs on separate, non-overlapping regions of the input space, and showed that fine-tuning a linearized version of the model strengthens it.

Task Arithmetic vs. Related Concepts#

• Task vectors: The weight differences that task arithmetic operates on. Task arithmetic is the set of operations: addition, negation, and analogies.
• Model soups: Average several fine-tunes of one task to improve that task. Task arithmetic adds fine-tunes of different tasks, and adding two task vectors with λ = 0.5 equals averaging the two fine-tuned models.
• Model editing: Changes one fact or narrow behavior with a targeted, often closed-form weight update. Task arithmetic edits whole tasks at once.
• Multi-task training: Trains one model on all tasks' data together, which avoids interference but needs every dataset at once. Task arithmetic combines models whose data may no longer be available, and it avoids the forgetting that comes from fine-tuning the tasks one after another.