Model Soups

モデルスープとは、同じ事前学習済みモデルから異なるハイパーパラメーターでファインチューニングした複数のモデルの重みを平均して作る、単一のモデルです。通常、ハイパーパラメーターの探索後は、検証スコアが最も高い実行結果を残し、ほかは破棄します。スープでは、破棄される実行結果を再利用します。各モデルは同じ事前学習済みのモデル重みから始めているため、パラメーターの対応関係が保たれています。その平均は、多くの場合どのモデルよりも高いスコアを達成しながら、実行コストは単一モデルとまったく同じです。ワシントン大学、Google、Meta AI、テルアビブ大学、コロンビア大学のMitchell Wortsmanらは、Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference timeでこの手法を紹介し、ICML 2022で発表しました。GoogleのJFT-3Bデータセットで事前学習したViT-G/14のファインチューニング済みモデル58個を使った貪欲スープは、ImageNetで90.94%のトップ1精度に達し、当時の最先端結果となりました。

一様スープと貪欲スープ#

ファインチューニングでは、大規模な事前学習済みネットワークを重み空間内でわずかに移動させます。そのため、異なる学習率、データ拡張、ラベルスムージング、または学習期間でファインチューニングしたモデルは、低い損失を示す同じ広い領域に位置する傾向があります。それらを平均すると、その領域の中央付近に位置し、端にあるモデルより汎化性能が高くなります。論文では、主に次の2つの手順を定義しています。

• 一様スープ： ファインチューニング済みのすべてのモデルを平均します。すべての構成要素が良好な場合は機能しますが、学習率が非常に高い実行結果は共有領域から外れることがあり、それを平均に含めると精度が低下します。
• 貪欲スープ： 保留した検証データでモデルを精度順に並べ、最良のモデルから始めます。次のモデルを加えるのは、平均したスープの検証精度が低下しない場合に限ります。最良モデルから始めるため、検証セットでは貪欲スープがそのモデルより悪くなることはありません。

ImageNetで72通りのランダムなハイパーパラメーター設定を使ってファインチューニングしたCLIP ViT-B/32では、最良の単一モデルのスコアは80.38%でした。72モデルすべての一様スープは79.97%でした。一部の実行結果が共有する低損失領域から外れていたためです。一方、貪欲スープは5モデルを採用し、81.03%を記録しました。ImageNetスケッチやレンディションなど、分布シフトを含む5つのテストセットでは、貪欲スープの平均は50.75%で、最良の単一モデルの47.83%を上回りました。論文ではALIGNでの改善も報告されており、テキスト分類でBERTとT5をファインチューニングした場合にも、小幅ながら改善が見られました。

このPyTorchの例では、8x8の手書き数字で小型分類器を事前学習し、異なる学習率と入力ノイズのデータ拡張を使って9つのコピーをファインチューニングした後、それらの一様スープと9モデルのアンサンブルを比較します。

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 train(model, first, last, lr, noise, seed):
    """Train a copy of the model on random batches from rows first to last, adding input noise."""
    model, generator = copy.deepcopy(model), torch.Generator().manual_seed(seed)
    optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9)
    for _ in range(300):
        rows = torch.randint(first, last, (32,), generator=generator)
        optimizer.zero_grad()
        inputs = X[rows] + noise * torch.randn(32, 64, generator=generator)
        nn.functional.cross_entropy(model(inputs), y[rows]).backward()
        optimizer.step()
    return model


def accuracy(logits):
    """Return the fraction of the 997 test digits classified correctly."""
    return (logits.argmax(1) == y[800:]).float().mean().item()


base = train(nn.Sequential(nn.Linear(64, 256), nn.ReLU(), nn.Linear(256, 10)), 0, 300, 0.05, 0, 0)  # pretrained
sweep = [(lr, noise) for lr in (0.01, 0.03, 0.1) for noise in (0, 0.2, 0.4)]  # 9 hyperparameter settings
ingredients = [train(base, 300, 500, lr, noise, seed) for seed, (lr, noise) in enumerate(sweep)]

# Uniform soup: average every weight tensor of the nine fine-tuned models
soup = copy.deepcopy(base)
soup.load_state_dict({k: sum(m.state_dict()[k] for m in ingredients) / 9 for k in soup.state_dict()})

with torch.no_grad():
    singles = [accuracy(m(X[800:])) for m in ingredients]
    print(f"Fine-tuned models: {min(singles):.3f} to {max(singles):.3f}, mean {sum(singles) / 9:.3f}")
    print(f"Uniform soup: {accuracy(soup(X[800:])):.3f}")
    print(f"Ensemble of 9: {accuracy(sum(m(X[800:]).softmax(1) for m in ingredients)):.3f}")

9つのファインチューニング済みモデルは、保留した997個の数字で0.881から0.936の範囲のスコアを示し、平均は0.917です。一様スープのスコアは0.934で、最良の単一実行結果に近く、アンサンブルの0.930をわずかに上回ります。アンサンブルでは画像ごとに9回のフォワードパスが必要ですが、スープでは1回で済みます。

実際のモデルスープ活用#

モデルスープは、ビジョントランスフォーマー（ViT）やその他の基盤モデルを新しい画像分類タスクに転移学習する場合など、もともとハイパーパラメーター探索を行う場面で特に有効です。Ultralytics YOLOのハイパーパラメーターチューニングもそのような探索です。model.tune()は、変更したハイパーパラメーターの各セットで1つのモデルを学習し、ストレージを節約するため、現在の最良結果ではない実行結果を削除します。つまり、最良のイテレーションの実行結果と重みだけを残します。スープは、この「最良のものを選ぶ」手順に置き換わる方法です。

構成要素のモデルは、同じアーキテクチャと事前学習済みの初期状態を共有する必要があります。異なるランダム初期値から学習したネットワークは、同じ特徴を異なる隠れユニットに格納するため、平均すると精度が損なわれるからです。最終的な重みは推論時に追加コストがかからないため、スープの推論レイテンシとメモリ使用量は単一モデルと同じです。

モデルスープと関連手法の比較#

• モデルアンサンブル： 複数モデルの予測を平均するため、メンバーが増えるたびに推論コストも増加します。CLIP ViT-B/32の探索では、ロジットアンサンブルがImageNetで81.19%を記録し、貪欲スープの81.03%をわずかに上回りましたが、画像ごとにすべてのメンバーモデルを実行しました。
• モデルマージ： 重みを組み合わせる手法全般を指します。すべてのモデルが同じタスクでファインチューニングされている場合がスープであり、異なるタスクでファインチューニングしたモデルを加算するのがタスク算術です。
• 1回の実行内での重み平均： 確率的重み平均（SWA）と重みの指数移動平均（EMA）は、1本の学習軌跡に沿ったチェックポイントを平均します。スープは独立した実行結果を平均します。論文では、EMAまたはSWAをすでに使っているモデルをスープにしても精度が上がることが示されており、2種類の平均化の効果は加算されます。
• WiSE-FT： ゼロショットモデルとそのファインチューニング版の間を補間する、2つの構成要素からなるスープです。分布シフトに対する事前学習済みモデルのロバスト性をより多く保持します。