Task Arithmetic

Арифметика задач — это способ редактирования предварительно обученной модели путём сложения и вычитания векторов задач, то есть изменений весов, полученных при дообучении для каждой задачи. Если прибавить к весам предварительно обученной модели векторы нескольких задач, получится одна модель, которая решает их все. Вычитание вектора задачи ухудшает работу модели в этой задаче, почти не затрагивая остальные, а комбинация векторов по аналогии переносит навык в условия, для которых нет размеченных данных. Каждое изменение сводится к поэлементному сложению весов, поэтому для него не нужны обучающие данные или шаги градиентного спуска сверх тех, что потребовались для исходного дообучения. Габриэль Ильярко и его коллеги из Вашингтонского университета, Microsoft Research и Института искусственного интеллекта Аллена представили этот подход в работе Editing Models with Task Arithmetic, опубликованной на ICLR 2023.

Три операции#

Все три операции создают новые веса по формуле θ_new = θ_pretrained + λτ_new, где коэффициент масштабирования λ выбирается на отложенных валидационных данных. Ильярко и соавторы проверили их на моделях CLIP, дообученных для восьми задач классификации изображений, таких как распознавание дорожных знаков и спутниковых снимков, а также на языковых моделях GPT-2 и T5.

Отрицание использует τ_new = −τ, чтобы модель забыла задачу. Отрицание восьми векторов задач CLIP снизило среднюю точность модели ViT-L/14 на этих задачах на 45,8 процентного пункта, сохранив не менее 95% точности на ImageNet. Вычитание вектора, полученного при дообучении на токсичных комментариях, снизило долю токсичных ответов GPT-2 Large с 4,8% до 0,8%; перплексия на обычном тексте при этом отличалась от исходной модели не более чем на 0,5 пункта.

Сложение использует τ_new = τ_1 + τ_2 + … для создания многозадачных моделей. Сложение векторов двух задач CLIP позволило создать одну модель, средняя точность которой составила 98,9% от точности двух специализированных моделей, а сложение всех восьми векторов дало результат 91,2%.

Аналогии применяются, когда задачи связаны отношением «A относится к B так же, как C к D», с использованием формулы τ_D ≈ τ_C + (τ_B − τ_A). Ильярко и соавторы создали модель анализа тональности отзывов Yelp без размеченных данных Yelp, прибавив вектор тональности отзывов Amazon к разности между векторами языковых моделей, обученных на текстах Yelp и Amazon.

Эти операции работают, потому что векторы разных задач направлены почти ортогонально и имеют косинусное сходство, близкое к нулю. Поэтому сложение одного вектора почти не меняет то, что кодирует другой. Однако с ростом числа задач усиливается взаимное влияние, и методы объединения моделей, например TIES-Merging, снижают его: они отсекают наименьшие изменения каждого вектора и разрешают конфликт параметров с противоположными знаками перед сложением.

Сложение и отрицание навыков в коде#

В этом примере на PyTorch небольшой классификатор недолго предварительно обучается на рукописных цифрах размером 8x8, после чего его первый слой дообучается для двух задач: задача A — классификация цифр 0–4, задача B — цифр 5–9. Для каждой задачи используются собственные пять замороженных выходных логитов, как в работе для каждой задачи CLIP используется собственная замороженная голова классификации. Затем оцениваются предварительно обученные веса с разными комбинациями двух векторов задач:

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

На 797 тестовых цифрах предварительно обученная модель получает 0.802 в задаче A и 0.897 в задаче B. Добавление только вектора задачи A повышает результат задачи A до 0.925, но снижает результат задачи B до 0.850 — это катастрофическое забывание, характерное для обычного дообучения. Добавление только вектора задачи B даёт обратный эффект. Добавление обоих векторов позволяет создать одну модель с результатами 0.922 и 0.940, близкими к результатам каждой специализированной модели в её собственной задаче. Вычитание вектора задачи A снижает результат задачи A до 0.616, тогда как результат задачи B меняется лишь с 0.897 на 0.887.

Требования и ограничения#

Для арифметики задач необходимо, чтобы каждое дообучение начиналось с одних и тех же предварительно обученных весов и использовало одну архитектуру. Когда Ultralytics YOLO26 обучается с yolo26n.pt на новом наборе данных, модель копирует все предварительно обученные тензоры подходящей формы, сопоставляет строки головы классификации по названиям классов и записывает число в журнал в виде «Transferred N/M items from pretrained weights». Если число классов в наборах данных различается, формы финальных свёрточных слоёв классификации тоже будут различаться, поэтому векторы задач из таких запусков можно складывать только для общих слоёв.

Арифметика также перестаёт работать, когда векторы конфликтуют или масштабируются слишком сильно: большие значения λ выводят веса за пределы области, в которой дообученные модели работают хорошо. Отрицание ухудшает результаты по всему, что охватывали данные дообучения, поэтому это облегчённая форма машинного забывания. В работе «Арифметика задач в касательном пространстве» Гильермо Ортис-Хименес, Алессандро Фаверо и Паскаль Фроссар связали успешное применение арифметики с разделением весов: разные направления в пространстве весов меняют выходы модели на отдельных непересекающихся областях входного пространства. Авторы также показали, что дообучение линеаризованной версии модели усиливает этот эффект.

Арифметика задач и связанные понятия#

• Векторы задач: Разности весов, с которыми работает арифметика задач. Арифметика задач — это набор операций: сложение, отрицание и аналогии.
• Смешивание моделей: Усредни несколько вариантов дообучения для одной задачи, чтобы улучшить результаты на этой задаче. Арифметика задач объединяет варианты дообучения для разных задач: сложение двух векторов задач с λ = 0.5 равносильно усреднению двух дообученных моделей.
• Редактирование моделей: Меняет один факт или узкое поведение с помощью целевого, часто имеющего решение в замкнутой форме, обновления весов. Арифметика задач редактирует целые задачи одновременно.
• Многозадачное обучение: Обучает одну модель одновременно на данных всех задач, что позволяет избежать взаимного влияния, но требует наличия всех наборов данных сразу. Арифметика задач объединяет модели, данные для которых могут быть уже недоступны, и предотвращает забывание, возникающее при последовательном дообучении на разных задачах.