Task Arithmetic

L’arithmétique des tâches permet de modifier un modèle préentraîné en additionnant et en soustrayant des vecteurs de tâche, c’est-à-dire les variations de poids produites par l’ajustement fin pour chaque tâche. Additionner les vecteurs de plusieurs tâches aux poids du modèle préentraîné donne un modèle capable de les traiter toutes ; soustraire un vecteur rend le modèle moins performant sur la tâche correspondante tout en laissant les autres en grande partie intactes ; combiner des vecteurs par analogie transfère une compétence vers un contexte sans données étiquetées. Chaque modification consiste en une addition élément par élément des poids et ne nécessite donc ni données d’entraînement ni calcul de gradients, au-delà des ajustements fins initiaux. Gabriel Ilharco et ses collègues de l’Université de Washington, de Microsoft Research et de l’Allen Institute for AI ont présenté cette approche dans Editing Models with Task Arithmetic, publié à ICLR 2023.

Les trois opérations#

Les trois opérations produisent de nouveaux poids θ_new = θ_pretrained + λτ_new, où le coefficient d’échelle λ est choisi sur des données de validation mises de côté. Ilharco et al. les ont testées sur des modèles CLIP ajustés finement sur huit tâches de classification d’images, notamment la reconnaissance de panneaux de signalisation et d’images satellite, ainsi que sur les modèles de langage GPT-2 et T5.

Négation utilise τ_new = −τ pour faire oublier une tâche au modèle. La négation des huit vecteurs de tâche CLIP a réduit de 45,8 points de pourcentage la précision moyenne du modèle ViT-L/14 sur ces tâches, tout en conservant au moins 95 % de sa précision sur ImageNet. La soustraction d’un vecteur ajusté finement sur des commentaires toxiques a fait passer la proportion de générations toxiques de GPT-2 Large de 4,8 % à 0,8 %, tandis que la perplexité sur du texte ordinaire restait à moins de 0,5 point du modèle d’origine.

Addition utilise τ_new = τ_1 + τ_2 + … pour créer des modèles multitâches. L’addition des vecteurs de deux tâches CLIP a donné un modèle dont la précision moyenne atteignait 98,9 % de celle des deux modèles spécialisés ; l’addition des huit vecteurs a atteint 91,2 %.

Analogies s’appliquent lorsque des tâches sont liées comme « A est à B ce que C est à D », à l’aide de τ_D ≈ τ_C + (τ_B − τ_A). Ilharco et al. ont créé un modèle d’analyse des sentiments pour les avis Yelp sans données Yelp étiquetées, en ajoutant un vecteur de sentiment d’avis Amazon à la différence entre les vecteurs de modélisation du langage entraînés sur des textes Yelp et Amazon.

Ces opérations fonctionnent parce que les vecteurs de tâche de tâches différentes pointent dans des directions presque orthogonales, avec des similarités cosinus proches de zéro ; en ajouter un modifie donc à peine ce qu’un autre encode. Les interférences augmentent tout de même avec le nombre de tâches, et des méthodes de fusion de modèles comme TIES-Merging les réduisent en supprimant les variations les plus faibles de chaque vecteur et en résolvant les conflits de signe entre paramètres avant l’addition.

Addition et négation de compétences en code#

Cet exemple PyTorch préentraîne brièvement un petit classifieur sur des chiffres manuscrits de 8x8, puis ajuste finement sa première couche pour deux tâches : la tâche A classe les chiffres de 0 à 4 et la tâche B ceux de 5 à 9, chacune utilisant ses propres cinq logits de sortie gelés, comme chaque tâche CLIP utilise sa propre tête de classification gelée dans l’article. Il évalue ensuite les poids préentraînés auxquels sont appliquées différentes combinaisons des deux vecteurs de tâche :

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}")

Sur les 797 chiffres de test, le modèle préentraîné obtient 0.802 pour la tâche A et 0.897 pour la tâche B. L’ajout du seul vecteur de la tâche A fait passer le résultat de la tâche A à 0.925, mais réduit celui de la tâche B à 0.850 : c’est l’oubli catastrophique de l’ajustement fin classique ; l’ajout du seul vecteur de la tâche B produit l’effet inverse. L’ajout des deux donne un modèle qui atteint 0.922 et 0.940, des scores proches de ceux de chaque modèle spécialisé pour sa tâche. La soustraction du vecteur de la tâche A fait chuter son score à 0.616, tandis que celui de la tâche B ne passe que de 0.897 à 0.887.

Conditions et limites#

L’arithmétique des tâches exige que tous les ajustements fins partent des mêmes poids préentraînés et utilisent la même architecture. Lorsqu’Ultralytics YOLO26 s’entraîne à partir de yolo26n.pt sur un nouveau jeu de données, il copie chaque tenseur préentraîné dont la forme convient, fait correspondre les lignes de la tête de classification selon le nom de classe et consigne le nombre dans « Transferred N/M items from pretrained weights ». Les jeux de données comportant des nombres de classes différents donnent des formes différentes aux convolutions finales de classification ; les vecteurs de tâche de ces exécutions ne peuvent donc être additionnés que dans les couches qu’elles ont en commun.

L’arithmétique échoue aussi lorsque les vecteurs sont incompatibles ou trop fortement mis à l’échelle, car des valeurs élevées de λ éloignent les poids de la région où les modèles ajustés finement fonctionnent bien. La négation réduit les performances sur les données couvertes par l’ajustement fin, ce qui en fait une forme légère de désapprentissage automatique. Dans Task Arithmetic in the Tangent Space, Guillermo Ortiz-Jimenez, Alessandro Favero et Pascal Frossard ont relié la réussite de l’arithmétique à la dissociation des poids, où des directions distinctes dans l’espace des poids modifient les sorties du modèle sur des régions disjointes de l’espace d’entrée ; ils ont montré qu’un ajustement fin d’une version linéarisée du modèle renforce ce phénomène.

Arithmétique des tâches et notions apparentées#

• Vecteurs de tâche : Les différences de poids manipulées par l’arithmétique des tâches. L’arithmétique des tâches désigne l’ensemble des opérations : addition, négation et analogies.
• Soups de modèles : Fais la moyenne de plusieurs ajustements fins d'une même tâche pour améliorer cette tâche. L'arithmétique des tâches additionne des ajustements fins de tâches différentes, et l'ajout de deux vecteurs de tâche avec λ = 0.5 équivaut à faire la moyenne des deux modèles ajustés finement.
• Édition de modèles : Modifie un fait ou un comportement précis à l'aide d'une mise à jour ciblée des poids, souvent sous forme fermée. L'arithmétique des tâches modifie des tâches entières à la fois.
• Entraînement multitâche : Entraîne un modèle sur les données de toutes les tâches simultanément, ce qui évite les interférences, mais nécessite tous les jeux de données à la fois. L'arithmétique des tâches combine des modèles dont les données ne sont peut-être plus disponibles, et évite l'oubli causé par l'ajustement fin successif des tâches.