Task Arithmetic

L'aritmetica dei task è un metodo per modificare un modello preaddestrato sommando e sottraendo vettori di task, ovvero le variazioni dei pesi prodotte dal fine-tuning su ciascun task. Sommare ai pesi del modello preaddestrato i vettori di diversi task produce un modello capace di gestirli tutti; sottrarre un vettore di task peggiora le prestazioni del modello su quel task lasciando perlopiù intatte le altre, mentre combinare vettori per analogia trasferisce una competenza a un contesto senza dati etichettati. Ogni modifica è una somma elemento per elemento dei pesi e non richiede dati di addestramento né passaggi di gradiente oltre ai fine-tuning originali. Gabriel Ilharco e colleghi dell'Università di Washington, di Microsoft Research e dell'Allen Institute for AI hanno presentato l'approccio in Editing Models with Task Arithmetic, pubblicato a ICLR 2023.

Le tre operazioni#

Tutte e tre le operazioni producono nuovi pesi θ_new = θ_pretrained + λτ_new, dove il coefficiente di scala λ viene scelto su dati di convalida tenuti da parte. Ilharco et al. le hanno testate su modelli CLIP sottoposti a fine-tuning su otto task di classificazione delle immagini, come il riconoscimento dei segnali stradali e delle immagini satellitari, e su modelli linguistici GPT-2 e T5.

Negazione usa τ_new = −τ per far dimenticare al modello un task. Negare gli otto vettori di task di CLIP ha ridotto di 45,8 punti percentuali l'accuratezza media del modello ViT-L/14 su quei task, mantenendo almeno il 95% dell'accuratezza su ImageNet. Sottrarre un vettore sottoposto a fine-tuning su commenti tossici ha ridotto dal 4,8% allo 0,8% la quota di generazioni tossiche di GPT-2 Large, mantenendo la perplexity sul testo normale entro 0,5 punti rispetto al modello originale.

Addizione usa τ_new = τ_1 + τ_2 + … per creare modelli multi-task. Sommare i vettori di due task di CLIP ha prodotto un modello con un'accuratezza media pari al 98,9% di quella dei due modelli specializzati; sommando tutti e otto i vettori, il valore è arrivato al 91,2%.

Analogie si applicano quando i task sono correlati secondo lo schema «A sta a B come C sta a D», usando τ_D ≈ τ_C + (τ_B − τ_A). Ilharco et al. hanno creato un modello di analisi del sentiment per le recensioni di Yelp senza dati Yelp etichettati, aggiungendo un vettore del sentiment delle recensioni Amazon alla differenza tra i vettori di modellazione del linguaggio addestrati su testi Yelp e Amazon.

Queste operazioni funzionano perché i vettori di task diversi puntano in direzioni quasi ortogonali, con similarità coseno vicine a zero; sommarne uno modifica appena ciò che un altro codifica. L'interferenza cresce comunque con il numero di task, e i metodi di merging dei modelli, come TIES-Merging, la riducono eliminando le variazioni più piccole di ciascun vettore e risolvendo i conflitti tra i segni dei parametri prima di sommarli.

Aggiunta e negazione di competenze nel codice#

Questo esempio PyTorch preaddestra brevemente un piccolo classificatore su cifre scritte a mano di 8x8 pixel, poi esegue il fine-tuning del primo livello per due task: il task A classifica le cifre 0-4 e il task B le cifre 5-9, ciascuno leggendo i propri cinque logit di output congelati, come nell'articolo ogni task CLIP usa la propria testa di classificazione congelata. Infine, valuta i pesi preaddestrati più diverse combinazioni dei due vettori di task:

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

Sulle 797 cifre di test, il modello preaddestrato ottiene 0.802 sul task A e 0.897 sul task B. Aggiungere solo il vettore del task A porta il task A a 0.925, ma abbassa il task B a 0.850: è l'oblio catastrofico del normale fine-tuning; aggiungere solo il vettore del task B produce l'effetto opposto. Sommarli entrambi dà un unico modello con 0.922 e 0.940, valori vicini a quelli di ciascun modello specializzato sul proprio task. Sottrarre il vettore del task A porta il task A a 0.616, mentre il task B passa solo da 0.897 a 0.887.

Requisiti e limiti#

L'aritmetica dei task richiede che ogni fine-tuning parta dagli stessi pesi preaddestrati e dalla stessa architettura. Quando Ultralytics YOLO26 si addestra da yolo26n.pt su un nuovo dataset, copia ogni tensore preaddestrato la cui forma è compatibile, abbina le righe della testa di classificazione in base al nome della classe e registra il numero con «Transferred N/M items from pretrained weights». I dataset con un numero diverso di classi producono convoluzioni finali di classificazione con forme diverse; perciò i vettori di task ottenuti da tali esecuzioni possono essere sommati solo nei livelli condivisi.

L'aritmetica smette inoltre di funzionare quando i vettori sono in conflitto o vengono scalati eccessivamente, perché valori elevati di λ spingono i pesi fuori dalla regione in cui i modelli sottoposti a fine-tuning funzionano bene. La negazione riduce le prestazioni su tutto ciò che era coperto dai dati di fine-tuning, diventando così una forma leggera di machine unlearning. In Task Arithmetic in the Tangent Space, Guillermo Ortiz-Jimenez, Alessandro Favero e Pascal Frossard hanno ricondotto il successo dell'aritmetica alla separazione dei pesi: direzioni distinte nello spazio dei pesi modificano gli output del modello su regioni distinte e non sovrapposte dello spazio degli input. Hanno inoltre mostrato che il fine-tuning di una versione linearizzata del modello rafforza tale proprietà.

Aritmetica dei task e concetti correlati#

• Vettori di task: Le differenze tra i pesi su cui opera l'aritmetica dei task. L'aritmetica dei task è l'insieme delle operazioni: addizione, negazione e analogie.
• Ensemble di modelli: Media diversi fine-tuning dello stesso task per migliorare quel task. L'aritmetica dei task aggiunge fine-tuning di task diversi, e sommare due vettori di task con λ = 0.5 equivale a fare la media dei due modelli sottoposti a fine-tuning.
• Modifica dei modelli: Modifica un fatto o un comportamento circoscritto con un aggiornamento mirato dei pesi, spesso in forma chiusa. L'aritmetica dei task modifica interi task in una sola volta.
• Addestramento multitask: Addestra un unico modello sui dati di tutti i task insieme, evitando interferenze ma richiedendo tutti i dataset contemporaneamente. L'aritmetica dei task combina modelli i cui dati potrebbero non essere più disponibili ed evita l'oblio causato dal fine-tuning dei task uno dopo l'altro.