Model Merging

モデルマージでは、トレーニング済みの2つ以上のモデルの重みを、通常はパラメーターを平均または加算して1つのモデルに統合します。これにより、複数のネットワークを実行したり、すべてのデータで再トレーニングしたりせずに、各モデルの能力を引き継いだモデルを作成できます。モデルのアーキテクチャが共通であり、ほとんどの手法では事前トレーニング済みの開始点も共通している場合に有効です。ファインチューニングでは各コピーが重み空間内をわずかに移動するため、モデルの重みはパラメーター単位で依然として対応しています。モデルマージは、オープンな大規模言語モデル (LLMs)の構築や、LoRAアダプターまたはStable Diffusionチェックポイントの統合に広く使われています。また、MetaのLlama 3では、各段階で異なるデータまたはハイパーパラメーターを用いて得られたモデルを、事後トレーニング時に平均しました。

モデルマージの仕組み#

マージ手法は、何を統合するか、モデル間の競合をどのように解決するかによって異なります。

• モデルスープ: 同じベースモデルから異なるハイパーパラメーターでファインチューニングしたモデルの重みを平均します。ファインチューニング済みのViT-Gモデルによる貪欲なスープは、ImageNetでトップ1精度90.94%を達成し、当時の最先端の結果となりました。
• 確率的重み平均化: 1回のトレーニング実行中の異なる時点で保存したチェックポイントを平均します。これにより、損失曲面のより広く平坦な領域に収束する傾向があります。
• タスク算術: 各ファインチューニング済みモデルからベースモデルの重みを引いてタスクベクトルを求め、スケーリングしたタスクベクトルをベースモデルに加えてスキルを統合するか、ベクトルを引いて特定の振る舞いを取り除きます。
• TIES-Merging: 3つの手順でタスクベクトル間の干渉を減らします。まず、ファインチューニング中の変化がわずかなパラメーターをリセットし、各パラメーターの符号を決定して、その符号と一致する値だけをマージします。
• DARE: 各タスクベクトルの値の大半をランダムにドロップし、残りを再スケーリングしてからマージします。これにより、各タスクの能力を維持しながら、冗長な変更を取り除きます。

1回のトレーニング実行中に平均化する処理は、通常のトレーニングにも含まれています。Ultralytics YOLO26のトレーニングでは、最適化対象の重みと並行して、モデルの重みの指数移動平均 (EMA) を保持します。トレーニングが進むにつれて減衰率は0.9999に近づき、保存されるlast.ptおよびbest.ptチェックポイントには、この平均化されたモデルが格納されます。

タスクベクトル手法には共通の構造があります。各ファインチューニング済みモデルをベースモデルとの差分で表し、その差分を調整して統合した後、結果をベースモデルに加えます。

mergekitなどのオープンソースツールは、2つのモデル間の球面線形補間 (SLERP)を含む、こうしたLLM向けの手法を実装しています。

共通の開始点が重要な理由#

ニューラルネットワークの隠れユニットは、ネットワークの計算内容を変えずに並べ替えられます。そのため、異なるランダム初期化からトレーニングした2つのネットワークでは、同じ特徴が同じ位置に格納されることはほとんどなく、平均すると無関係な重みが混ざります。同じベースモデルからファインチューニングしたモデルは、ベースモデルの配置を保ち、低い損失の同じ領域にとどまるため、平均しても機能します。この例では、手書き数字を使って小規模なPyTorch分類器をトレーニングし、残りのトレーニングデータを半分ずつ使って2つのコピーをファインチューニングし、それらのマージと、別々の初期化からトレーニングした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)
order = torch.randperm(len(X))
X, y = torch.tensor(X / 16, dtype=torch.float32)[order], torch.tensor(y)[order]


def mlp():
    """Return a small classifier for 8x8 digit images."""
    return nn.Sequential(nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 10))


def train(model, rows, steps):
    """Train a model with SGD on the given rows and return it."""
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
    for _ in range(steps):
        optimizer.zero_grad()
        nn.functional.cross_entropy(model(X[rows]), y[rows]).backward()
        optimizer.step()
    return model


base = train(mlp(), slice(0, 300), 400)  # shared pretrained starting point
halves = (slice(300, 900), slice(900, 1500))  # two different fine-tuning datasets
for name, (a, b) in [
    ("fine-tuned from one base", [train(copy.deepcopy(base), rows, 100) for rows in halves]),
    ("trained from separate inits", [train(mlp(), rows, 400) for rows in halves]),
]:
    # Merge: average every weight tensor of the two models
    merged = copy.deepcopy(a)
    merged.load_state_dict({k: (a.state_dict()[k] + b.state_dict()[k]) / 2 for k in a.state_dict()})
    accuracy = [(m(X[1500:]).argmax(1) == y[1500:]).float().mean().item() for m in (a, b, merged)]
    print(f"{name}: {accuracy[0]:.3f} and {accuracy[1]:.3f}, merged {accuracy[2]:.3f}")

保留した297個の数字では、1つのベースモデルからファインチューニングした2つのモデルのスコアは0.956と0.946で、平均モデルはどちらよりも高い0.970でした。これは、それぞれがデータの半分から学んだ内容を統合したためです。別々の初期化からトレーニングした2つのモデルのスコアは0.946と0.956でしたが、平均モデルのスコアは0.808に低下しました。Git Re-Basinは、まず一方のモデルの隠れユニットを並べ替えてもう一方のモデルに合わせることで、独立してトレーニングしたモデルをマージします。

モデルマージと関連用語の比較#

• モデルアンサンブル: すべてのモデルを保持して出力を統合するため、モデル数に応じてメモリー使用量と推論レイテンシーが増加します。マージモデルのサイズと速度は単一モデルと同等です。
• 知識蒸留: 1つ以上の教師モデルを模倣する新しい生徒モデルをトレーニングするため、トレーニングデータと計算リソースが必要です。モデルマージは、既存の重みに対する直接的な算術演算です。
• マルチタスクトレーニング: すべてのタスクのデータを一緒に使って1つのモデルをトレーニングします。通常、マージではそのデータを利用せずに到達することを目指す精度の基準になります。
• 連合学習: 各ラウンドで、別々のデバイス上でトレーニングされたモデルの重みを平均します。各ラウンドでは、すべてのクライアントが同じグローバルモデルから開始します。この共通の開始点が、平均化を可能にします。