Task Arithmetic

A aritmética de tarefas é uma forma de editar um modelo pré-treinado somando e subtraindo vetores de tarefa, isto é, as alterações nos pesos produzidas pelo ajuste fino para cada tarefa. Somar aos pesos do modelo pré-treinado os vetores de tarefa de várias tarefas produz um único modelo capaz de lidar com todas elas; subtrair um vetor de tarefa piora o desempenho do modelo nessa tarefa, preservando em grande parte as demais; e combinar vetores por analogia transfere uma habilidade para uma configuração sem dados rotulados. Cada edição é uma soma elemento a elemento dos pesos e, por isso, não exige dados de treinamento nem etapas de gradiente além dos ajustes finos originais. Gabriel Ilharco e colegas da University of Washington, da Microsoft Research e do Allen Institute for AI apresentaram a abordagem em Editing Models with Task Arithmetic, publicado na ICLR 2023.

As três operações#

As três operações produzem novos pesos θ_new = θ_pretrained + λτ_new, em que o coeficiente de escala λ é escolhido com dados de validação reservados. Ilharco et al. testaram as operações em modelos CLIP ajustados em oito tarefas de classificação de imagens, como reconhecimento de sinais de trânsito e imagens de satélite, e em modelos de linguagem GPT-2 e T5.

Negação usa τ_new = −τ para fazer um modelo esquecer uma tarefa. Negar os oito vetores de tarefa do CLIP reduziu em 45,8 pontos percentuais a precisão média do modelo ViT-L/14 nessas tarefas, mantendo pelo menos 95% da precisão no ImageNet. Subtrair um vetor ajustado com comentários tóxicos reduziu de 4,8% para 0,8% a proporção de gerações tóxicas do GPT-2 Large, com perplexidade em textos comuns a no máximo 0,5 ponto do modelo original.

Adição usa τ_new = τ_1 + τ_2 + … para criar modelos multitarefa. Somar os vetores de duas tarefas do CLIP gerou um modelo com 98,9% da precisão média dos dois modelos especializados; ao somar os oito vetores, o resultado foi 91,2%.

Analogias aplicam-se quando as tarefas se relacionam como «A está para B assim como C está para D», usando τ_D ≈ τ_C + (τ_B − τ_A). Ilharco et al. criaram um modelo de análise de sentimentos para avaliações do Yelp sem dados rotulados do Yelp, somando um vetor de sentimentos de avaliações da Amazon à diferença entre os vetores de modelagem de linguagem treinados com textos do Yelp e da Amazon.

Essas operações funcionam porque os vetores de tarefa de tarefas diferentes apontam em direções quase ortogonais, com similaridades do cosseno próximas de zero. Assim, somar um vetor quase não altera o que outro codifica. A interferência ainda aumenta com o número de tarefas, e métodos de fusão de modelos, como TIES-Merging, reduzem esse problema ao eliminar as menores alterações de cada vetor e resolver as divergências de sinal entre parâmetros antes da soma.

Adição e negação de habilidades em código#

Este exemplo de PyTorch pré-treina brevemente um classificador pequeno com dígitos manuscritos de 8x8 pixels e, em seguida, ajusta a primeira camada para duas tarefas: a tarefa A classifica os dígitos de 0 a 4, e a tarefa B, os dígitos de 5 a 9. Cada tarefa lê os próprios cinco logits de saída congelados, assim como cada tarefa do CLIP usa a própria cabeça de classificação congelada no artigo. Depois, o exemplo avalia os pesos pré-treinados mais diferentes combinações dos dois vetores de tarefa:

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

Nos 797 dígitos de teste, o modelo pré-treinado obtém 0.802 na tarefa A e 0.897 na tarefa B. Somar apenas o vetor da tarefa A aumenta a pontuação da tarefa A para 0.925, mas reduz a da tarefa B para 0.850, o esquecimento catastrófico do ajuste fino convencional; somar apenas o vetor da tarefa B produz o efeito inverso. Somar ambos gera um modelo com pontuações de 0.922 e 0.940, próximas das de cada especialista na respectiva tarefa. Subtrair o vetor da tarefa A reduz a pontuação da tarefa A para 0.616, enquanto a da tarefa B muda apenas de 0.897 para 0.887.

Requisitos e limitações#

A aritmética de tarefas exige que cada ajuste fino comece com os mesmos pesos pré-treinados e use a mesma arquitetura. Quando o Ultralytics YOLO26 é treinado a partir de yolo26n.pt com um novo conjunto de dados, copia todos os tensores pré-treinados com dimensões compatíveis, associa as linhas da cabeça de classificação pelos nomes das classes e registra a quantidade como «Transferred N/M items from pretrained weights». Conjuntos de dados com números diferentes de classes resultam em convoluções finais de classificação com dimensões diferentes; assim, só é possível somar os vetores de tarefa das execuções nas camadas que elas têm em comum.

A aritmética também falha quando os vetores entram em conflito ou são escalados excessivamente, pois valores grandes de λ levam os pesos para fora da região em que os modelos ajustados têm bom desempenho. A negação reduz o desempenho em tudo o que os dados de ajuste fino abrangiam, tornando-a uma forma simples de desaprendizado de máquina. Em Task Arithmetic in the Tangent Space, Guillermo Ortiz-Jimenez, Alessandro Favero e Pascal Frossard atribuíram o sucesso da aritmética ao desacoplamento dos pesos, no qual direções distintas no espaço dos pesos alteram as saídas do modelo em regiões separadas e não sobrepostas do espaço de entrada. Também mostraram que ajustar uma versão linearizada do modelo reforça esse efeito.

Aritmética de tarefas e conceitos relacionados#

• Vetores de tarefa: As diferenças entre pesos usadas pela aritmética de tarefas. A aritmética de tarefas é o conjunto de operações: adição, negação e analogias.
• Sopas de modelos: Faz a média de vários ajustes finos para uma mesma tarefa, melhorando essa tarefa. A aritmética de tarefas adiciona ajustes finos de tarefas diferentes, e adicionar dois vetores de tarefa com λ = 0.5 equivale a calcular a média dos dois modelos ajustados.
• Edição de modelos: Altera um fato ou um comportamento específico com uma atualização direcionada dos pesos, muitas vezes em forma fechada. A aritmética de tarefas edita tarefas inteiras de uma só vez.
• Treinamento multitarefa: Treina um modelo com os dados de todas as tarefas em conjunto, evitando interferências, mas exige todos os conjuntos de dados ao mesmo tempo. A aritmética de tarefas combina modelos cujos dados talvez já não estejam disponíveis e evita o esquecimento causado pelo ajuste fino das tarefas uma após a outra.