Task Arithmetic

La aritmética de tareas permite editar un modelo preentrenado sumando y restando vectores de tarea, es decir, los cambios en los pesos producidos al hacer ajuste fino para cada tarea. Sumar a los pesos del modelo preentrenado los vectores de varias tareas produce un modelo que las resuelve todas; restar un vector de tarea empeora el rendimiento del modelo en esa tarea y apenas afecta a las demás; y combinar vectores mediante una analogía transfiere una habilidad a un entorno sin datos etiquetados. Cada edición consiste en sumar elemento a elemento los pesos, por lo que no requiere datos de entrenamiento ni pasos de gradiente aparte de los ajustes finos originales. Gabriel Ilharco y sus colegas de la Universidad de Washington, Microsoft Research y Allen Institute for AI presentaron el método en Editing Models with Task Arithmetic, publicado en ICLR 2023.

Las tres operaciones#

Las tres operaciones producen nuevos pesos θ_new = θ_pretrained + λτ_new, donde el coeficiente de escala λ se elige con datos de validación reservados. Ilharco y sus colaboradores las probaron con modelos CLIP ajustados para ocho tareas de clasificación de imágenes, como señales de tráfico e imágenes por satélite, y con modelos de lenguaje GPT-2 y T5.

Negación utiliza τ_new = −τ para hacer que un modelo olvide una tarea. Negar los ocho vectores de tarea de CLIP redujo en 45,8 puntos porcentuales la precisión media del modelo ViT-L/14 en esas tareas, y mantuvo al menos el 95 % de su precisión en ImageNet. Restar un vector ajustado con comentarios tóxicos redujo del 4,8 % al 0,8 % la proporción de generaciones tóxicas de GPT-2 Large, y la perplejidad en texto normal quedó a menos de 0,5 puntos de la del modelo original.

Suma utiliza τ_new = τ_1 + τ_2 + … para crear modelos multitarea. Sumar los vectores de dos tareas de CLIP produjo un modelo con una precisión media equivalente al 98,9 % de la de los dos modelos especializados; al sumar los ocho vectores, se alcanzó el 91,2 %.

Analogías se aplican cuando las tareas se relacionan como «A es a B lo que C es a D», mediante τ_D ≈ τ_C + (τ_B − τ_A). Ilharco y sus colaboradores crearon un modelo de análisis de sentimientos para reseñas de Yelp sin datos etiquetados de Yelp, sumando un vector de sentimientos de reseñas de Amazon a la diferencia entre vectores de modelado del lenguaje entrenados con texto de Yelp y Amazon.

Estas operaciones funcionan porque los vectores de tareas distintas apuntan en direcciones casi ortogonales, con similitudes coseno cercanas a cero, de modo que sumar uno apenas modifica lo que codifica otro. La interferencia aumenta aun así con el número de tareas, y los métodos de fusión de modelos, como TIES-Merging, la reducen recortando los cambios más pequeños de cada vector y resolviendo los parámetros cuyos signos no coinciden antes de sumarlos.

Suma y negación de habilidades en código#

Este ejemplo de PyTorch preentrena brevemente un clasificador pequeño con dígitos escritos a mano de 8x8 y después ajusta su primera capa para dos tareas: la tarea A clasifica los dígitos del 0 al 4 y la tarea B, del 5 al 9; cada tarea utiliza sus propios cinco logits de salida congelados, igual que cada tarea de CLIP utiliza su propia cabeza de clasificación congelada en el artículo. A continuación, evalúa los pesos preentrenados más distintas combinaciones de los dos vectores de tarea:

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

De los 797 dígitos de prueba, el modelo preentrenado obtiene 0.802 en la tarea A y 0.897 en la tarea B. Sumar solo el vector de la tarea A eleva la tarea A a 0.925, pero reduce la tarea B a 0.850, el olvido catastrófico propio del ajuste fino normal; sumar solo el vector de la tarea B produce el efecto inverso. Sumarlos ambos da un modelo con 0.922 y 0.940, cerca del rendimiento de cada especialista en su propia tarea. Restar el vector de la tarea A reduce la tarea A a 0.616, mientras que la tarea B solo pasa de 0.897 a 0.887.

Requisitos y límites#

La aritmética de tareas requiere que todos los ajustes finos partan de los mismos pesos preentrenados y usen la misma arquitectura. Cuando Ultralytics YOLO26 se entrena desde yolo26n.pt con un conjunto de datos nuevo, copia cada tensor preentrenado cuya forma encaja, empareja las filas de la cabeza de clasificación por nombre de clase y registra el recuento como «Transferred N/M items from pretrained weights». Los conjuntos de datos con distinto número de clases generan convoluciones de clasificación finales con formas diferentes, por lo que los vectores de tarea de esas ejecuciones solo se pueden sumar en las capas que comparten.

La aritmética también falla cuando los vectores entran en conflicto o se escalan demasiado, porque los valores grandes de λ alejan los pesos de la región en la que los modelos ajustados funcionan bien. La negación reduce el rendimiento en lo que cubrían los datos de ajuste fino, por lo que constituye una forma ligera de desaprendizaje automático. Task Arithmetic in the Tangent Space, de Guillermo Ortiz-Jimenez, Alessandro Favero y Pascal Frossard, atribuyó el éxito de la aritmética a la disentrelazación de pesos, en la que direcciones distintas del espacio de pesos cambian las salidas del modelo en regiones distintas y no superpuestas del espacio de entrada, y mostró que ajustar una versión linealizada del modelo la refuerza.

Aritmética de tareas frente a conceptos relacionados#

• Vectores de tarea: Las diferencias entre pesos sobre las que opera la aritmética de tareas. La aritmética de tareas es el conjunto de operaciones: suma, negación y analogías.
• Mezclas de modelos: Promedia varios ajustes finos de una tarea para mejorarla. La aritmética de tareas suma ajustes finos de tareas distintas, y sumar dos vectores de tarea con λ = 0.5 equivale a promediar los dos modelos ajustados.
• Edición de modelos: Cambia un hecho o un comportamiento concreto mediante una actualización de pesos específica, a menudo en forma cerrada. La aritmética de tareas edita tareas completas de una sola vez.
• Entrenamiento multitarea: Entrena un modelo con los datos de todas las tareas a la vez, lo que evita interferencias, pero requiere disponer de todos los conjuntos de datos simultáneamente. La aritmética de tareas combina modelos cuyos datos quizá ya no estén disponibles y evita el olvido que provoca ajustar las tareas una tras otra.