Model Soups

Eine Modellsuppe ist ein einzelnes Modell, das durch Mitteln der Gewichte mehrerer Modelle entsteht, die mit unterschiedlichen Hyperparametern vom selben vortrainierten Modell feinabgestimmt wurden. Nach einer Hyperparametersuche behält man üblicherweise den Lauf mit dem besten Validierungs-Ergebnis und verwirft die übrigen. Eine Suppe nutzt diese verworfenen Läufe erneut: Da jeder mit denselben vortrainierten Modellgewichten begann, sind ihre Parameter weiterhin aufeinander abgestimmt. Ihr Mittelwert erzielt oft ein besseres Ergebnis als jedes einzelne Modell und verursacht beim Ausführen genau denselben Aufwand wie ein Modell. Mitchell Wortsman und Kollegen der University of Washington, Google, Meta AI, der Universität Tel Aviv und der Columbia University stellten das Verfahren in Modellsuppen: Das Mitteln der Gewichte mehrerer feinabgestimmter Modelle verbessert die Genauigkeit ohne längere Inferenzzeit vor, veröffentlicht auf ICML 2022. Eine gierige Suppe aus 58 feinabgestimmten ViT-G/14-Modellen, die mit Googles JFT-3B-Datensatz vortrainiert wurden, erreichte auf ImageNet eine Top-1-Genauigkeit von 90,94 % und damit seinerzeit einen Spitzenwert.

Gleichmäßige und gierige Modellsuppen#

Bei der Feinabstimmung bewegt sich ein großes vortrainiertes Netz im Gewichtsraum nur über eine kurze Strecke. Daher landen Modelle, die davon mit unterschiedlichen Lernraten, Datenerweiterungen, Glättungen der Beschriftungen oder Trainingsdauern feinabgestimmt wurden, meist in derselben breiten Region mit geringem Verlust. Ihr Mittelwert liegt nahe der Mitte dieser Region und verallgemeinert besser als ihre Randbereiche. Das Paper beschreibt zwei Hauptverfahren:

• Gleichmäßige Suppe: Mittelt alle feinabgestimmten Modelle. Das funktioniert, wenn alle Bestandteile gut sind. Bei einer sehr hohen Lernrate kann ein Lauf jedoch den gemeinsamen Bereich verlassen, sodass seine Einbeziehung die Genauigkeit senkt.
• Gierige Suppe: Sortiert die Modelle nach ihrer Genauigkeit auf zurückgehaltenen Validierungsdaten, beginnt mit dem besten Modell und fügt das jeweils nächste nur hinzu, wenn die Validierungsgenauigkeit der gemittelten Suppe dadurch nicht sinkt. Da sie mit dem besten Modell beginnt, ist eine gierige Suppe auf dem Validierungssatz nie schlechter als dieses Modell.

Für CLIP ViT-B/32, das mit 72 zufällig gewählten Hyperparameterkonfigurationen auf ImageNet feinabgestimmt wurde, erzielte das beste Einzelmodell 80,38 %. Die Uniform Soup aus allen 72 Modellen erzielte 79,97 %, weil einige Läufe den gemeinsamen Bereich niedriger Verlustwerte verlassen hatten, während die Greedy Soup fünf Modelle behielt und 81,03 % erzielte. Auf fünf Testdatensätzen mit Verteilungsverschiebung, darunter ImageNet-Skizzen und -Darstellungen, erzielte die Greedy Soup durchschnittlich 50,75 % gegenüber 47,83 % beim besten Einzelmodell. Die Studie berichtet außerdem von Verbesserungen bei ALIGN und kleineren Zugewinnen bei der Feinabstimmung von BERT und T5 für die Textklassifizierung.

Dieses PyTorch-Beispiel trainiert zunächst einen kleinen Klassifikator mit handgeschriebenen Ziffern der Größe 8x8, stimmt neun Kopien mit unterschiedlichen Lernraten und Eingaberauschen fein ab und vergleicht ihre gleichmäßige Suppe mit einem Ensemble aus allen neun Modellen:

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

Die neun feinabgestimmten Modelle erzielen bei den 997 zurückgehaltenen Ziffern Werte zwischen 0.881 und 0.936, im Durchschnitt 0.917. Ihre gleichmäßige Suppe erreicht 0.934 und liegt damit nahe am besten Einzellauf sowie leicht über dem Wert des Ensembles von 0.930. Das Ensemble benötigt neun Vorwärtsdurchläufe pro Bild, die Suppe dagegen nur einen.

Modellsuppen in der Praxis#

Modellsuppen sind besonders nützlich, wenn ohnehin eine Parametersuche durchgeführt wird, etwa beim Transferlernen eines Vision Transformers (ViT) oder eines anderen Basismodells für eine neue Aufgabe der Bildklassifizierung. Die Hyperparameteroptimierung von Ultralytics YOLO ist eine solche Suche: model.tune() trainiert ein Modell pro Satz veränderter Hyperparameter und löscht aus Speicherplatzgründen die Läufe, die nicht gerade die besten sind. Behalten werden nur der Lauf und die Gewichte der besten Iteration – genau der Schritt, bei dem eine Suppe die Auswahl des besten Modells ersetzt.

Die Bestandteile müssen dieselbe Architektur und denselben vortrainierten Ausgangspunkt haben, denn Netze mit unterschiedlichen zufälligen Initialisierungen speichern dieselben Merkmale in unterschiedlichen verborgenen Einheiten. Ihr Mittelwert zerstört daher die Genauigkeit. Die endgültigen Gewichte verursachen bei der Inferenz keine zusätzlichen Kosten, sodass die Suppe dieselbe Inferenzlatenz und denselben Speicherbedarf wie ein einzelnes Modell hat.

Modellsuppen und verwandte Verfahren#

• Modellensemble: Mittelt die Vorhersagen mehrerer Modelle, sodass der Inferenzaufwand mit jedem weiteren Modell steigt. Bei der CLIP-ViT-B/32-Parametersuche erzielte ein Logit-Ensemble auf ImageNet 81,19 %, etwas mehr als die gierige Suppe mit 81,03 %. Dafür musste es bei jedem Bild alle Modelle ausführen.
• Modellzusammenführung: Die umfassendere Familie von Verfahren zur Kombination von Gewichten. Modellsuppen sind der Fall, bei dem alle Modelle für dieselbe Aufgabe feinabgestimmt wurden; bei der Aufgabenarithmetik werden Modelle kombiniert, die für unterschiedliche Aufgaben feinabgestimmt wurden.
• Mitteln von Gewichten innerhalb eines Laufs: Stochastic Weight Averaging (SWA) und der exponentiell gleitende Mittelwert (EMA) der Gewichte mitteln Prüfpunkte entlang eines einzelnen Trainingsverlaufs. Modellsuppen mitteln unabhängige Läufe. Das Paper zeigte, dass das Mitteln von Modellen, die bereits EMA oder SWA verwenden, die Genauigkeit dennoch erhöht. Die beiden Arten des Mittelns ergänzen sich also.
• WiSE-FT: Interpoliert zwischen einem Zero-Shot-Modell und seiner feinabgestimmten Version. Diese Suppe aus zwei Bestandteilen bewahrt mehr von der Robustheit des vortrainierten Modells gegenüber Verteilungsverschiebungen.