Task Arithmetic

タスク算術は、タスクごとのファインチューニングで生じた重みの変化であるタスクベクトルを加算・減算して、事前学習済みモデルを編集する手法です。複数のタスクのタスクベクトルを事前学習済みのモデル重みに加えると、それらすべてを処理する1つのモデルが得られます。タスクベクトルを減算すると、ほかのタスクの性能をほぼ保ちながら、そのタスクの性能を低下させられます。また、ベクトルを比喩的に組み合わせると、ラベル付きデータのない環境にスキルを転移できます。編集はすべて重みへの要素単位の加算であり、元のファインチューニング以外の学習データや勾配ステップは必要ありません。University of Washington、Microsoft Research、Allen Institute for AIのGabriel Ilharcoらは、ICLR 2023で発表した「Editing Models with Task Arithmetic」でこの手法を紹介しました。

3つの操作#

3つの操作はいずれも、重みθ_new = θ_pretrained + λτ_newを生成します。スケーリング係数λは、ホールドアウト検証データを使って選択します。Ilharcoらは、交通標識や衛星画像など8つの画像分類タスクでファインチューニングしたCLIPモデルと、GPT-2およびT5の言語モデルでこれらの操作を検証しました。

反転ではτ_new = −τを使い、モデルにタスクを忘れさせます。8つのCLIPタスクベクトルを反転すると、ViT-L/14モデルのそれらのタスクに対する平均精度は45.8ポイント低下しましたが、ImageNet精度は少なくとも95%保たれました。有害なコメントでファインチューニングしたベクトルを減算すると、GPT-2 Largeの有害な生成の割合は4.8%から0.8%に下がり、通常のテキストに対するパープレキシティは元のモデルとの差が0.5ポイント以内でした。

加算ではτ_new = τ_1 + τ_2 + …を使い、マルチタスクモデルを構築します。CLIPの2つのタスクベクトルを加えると、平均で2つの特化モデルの精度の98.9%を達成する1つのモデルが得られ、8つすべてを加えると91.2%に達しました。

類推は、タスク間の関係が「Aに対するBは、Cに対するDのようなもの」という場合に使い、τ_D ≈ τ_C + (τ_B − τ_A)を適用します。Ilharcoらは、Yelpのラベル付きデータを使わず、Amazonレビューの感情ベクトルを、YelpとAmazonのテキストで学習した言語モデルベクトルの差に加えることで、Yelpレビューの感情分析モデルを構築しました。

異なるタスクのタスクベクトルはほぼ直交する方向を向き、コサイン類似度が0に近いため、これらの操作が機能します。そのため、あるベクトルを加えても別のベクトルが表す内容はほとんど変化しません。それでもタスク数が増えると干渉が大きくなります。モデルマージ手法のTIES-Mergingなどは、各ベクトルの変化量が小さい部分を削り、加算前に符号が一致しないパラメーターを調整して干渉を減らします。

コードでスキルを加算・反転する#

このPyTorchの例では、8x8の手書き数字を使って小さな分類器を短時間事前学習し、次に最初の層を2つのタスク向けにファインチューニングします。タスクAは数字0～4、タスクBは数字5～9を分類し、それぞれ独立した5つの凍結済み出力ロジットから読み取ります。論文における各CLIPタスクが独自の凍結済み分類ヘッドを使うのと同様です。その後、事前学習済み重みと2つのタスクベクトルのさまざまな組み合わせを評価します：

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桁では、事前学習済みモデルのタスクAのスコアは0.802、タスクBのスコアは0.897です。タスクAのベクトルだけを加えると、タスクAは0.925に上がりますが、タスクBは0.850に下がります。これは通常のファインチューニングにおける破滅的忘却であり、タスクBのベクトルだけを加えると逆の結果になります。両方を加えると、タスクAは0.922、タスクBは0.940となる1つのモデルが得られ、それぞれのタスクに特化したモデルに近い精度を達成します。タスクAのベクトルを減算すると、タスクAは0.616に下がりますが、タスクBは0.897から0.887への変化にとどまります。

要件と制限#

タスク算術を適用するには、すべてのファインチューニングで同じアーキテクチャの同じ事前学習済み重みを使う必要があります。Ultralytics YOLO26を新しいデータセットで学習する際、yolo26n.ptから開始すると、形状が合う事前学習済みテンソルをすべてコピーし、クラス名に基づいて分類ヘッドの行を対応付け、「Transferred N/M items from pretrained weights」とログに出力します。クラス数の異なるデータセットでは、最終的な分類畳み込みの形状が異なるため、これらの実行結果から得たタスクベクトルを加算できるのは、共通する層だけです。

ベクトルが競合したり、スケールを上げすぎたりすると、算術処理はうまく機能しません。λが大きいと、ファインチューニング済みモデルの性能が保たれる領域から重みが外れるためです。反転すると、ファインチューニングデータが対象とした内容の性能が下がるため、軽量な機械アンラーニングの一種になります。Guillermo Ortiz-Jimenez、Alessandro Favero、Pascal Frossardによる「Task Arithmetic in the Tangent Space」では、算術処理がうまく機能する理由を重みの分離性に結び付けました。重み空間の別々の方向が、入力空間の重ならない領域に対するモデル出力を個別に変化させるというものです。また、モデルを線形化したバージョンをファインチューニングすると、この性質が強まることも示しました。

タスク算術と関連概念#

• タスクベクトル：タスク算術で操作する重みの差分です。タスク算術とは、加算、反転、類推という一連の操作を指します。
• モデルスープ： 1つのタスクで複数のファインチューニング済みモデルを平均し、そのタスクの性能を向上させます。タスク算術では異なるタスクのファインチューニング結果を加算します。2つのタスクベクトルを λ = 0.5 で加算すると、2つのファインチューニング済みモデルの平均になります。
• モデル編集： 狙いを定めた、多くの場合は閉形式の重み更新によって、1つの事実または限定的な動作を変更します。タスク算術では、タスク全体を一度に編集します。
• マルチタスク学習： すべてのタスクのデータをまとめて1つのモデルで学習するため、干渉を避けられますが、すべてのデータセットを同時に用意する必要があります。タスク算術は、データがすでに利用できない可能性のあるモデルを組み合わせられ、タスクを順番にファインチューニングすることで生じる忘却も回避できます。