Model Soups

A model soup is a single model made by averaging the weights of several models that were fine-tuned from the same pretrained model with different hyperparameters. The usual practice after a hyperparameter sweep is to keep the run with the best validation score and throw the rest away. A soup reuses those discarded runs: because each one started from the same pretrained model weights, their parameters still line up, and their average often scores higher than any of them while costing exactly as much to run as one model. Mitchell Wortsman and colleagues at the University of Washington, Google, Meta AI, Tel Aviv University, and Columbia introduced the method in Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time, published at ICML 2022. A greedy soup of 58 fine-tuned ViT-G/14 models, pretrained on Google's JFT-3B dataset, reached 90.94% top-1 accuracy on ImageNet, a state-of-the-art result at the time.

Uniform and Greedy Soups#

Fine-tuning moves a large pretrained network only a short distance in weight space, so models fine-tuned from it with different learning rates, data augmentation, label smoothing, or training lengths tend to land in the same wide region of low loss. Averaging them lands near the middle of that region, which generalizes better than its edges. The paper defines two main recipes:

• Uniform soup: Averages every fine-tuned model. It works when all ingredients are good, but a run with a very high learning rate can leave the shared region, and averaging it in lowers accuracy.
• Greedy soup: Sorts the models by accuracy on held-out validation data, starts from the best one, and adds each next model only if the averaged soup's validation accuracy does not drop. Because it starts from the best model, a greedy soup is never worse than that model on the validation set.

For CLIP ViT-B/32 fine-tuned on ImageNet with 72 random hyperparameter settings, the best single model scored 80.38%. The uniform soup of all 72 scored 79.97%, because some runs had left the shared low-loss region, while the greedy soup kept five models and scored 81.03%. On five distribution-shift test sets, such as ImageNet sketches and renditions, the greedy soup averaged 50.75% against 47.83% for the best single model. The paper also reports gains for ALIGN, and smaller ones when fine-tuning BERT and T5 on text classification.

This PyTorch example pretrains a small classifier on 8x8 handwritten digits, fine-tunes nine copies with different learning rates and input-noise augmentation, and compares their uniform soup with an ensemble of all nine:

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

The nine fine-tuned models score between 0.881 and 0.936 on the 997 held-out digits, 0.917 on average. Their uniform soup scores 0.934, close to the best single run, and slightly above the ensemble's 0.930, which needs nine forward passes per image where the soup needs one.

Soups in Practice#

Soups are most useful where a sweep is run anyway, for example when transfer learning a vision transformer (ViT) or other foundation model to a new image classification task. Ultralytics YOLO hyperparameter tuning is such a sweep: model.tune() trains one model per set of mutated hyperparameters and, to save storage, deletes the runs that are not the current best, keeping only the best iteration's run and weights, the pick-the-best step that a soup replaces.

The ingredients must share an architecture and a pretrained starting point, because networks trained from different random initializations store the same features in different hidden units, and averaging them destroys accuracy. The final weights cost nothing extra at inference, so the soup keeps the inference latency and memory of one model.

Model Soups vs. Related Methods#

• Model ensemble: Averages the predictions of several models, so inference cost grows with every member. In the CLIP ViT-B/32 sweep, a logit ensemble scored 81.19% on ImageNet, slightly above the greedy soup's 81.03%, while running every member model for each image.
• Model merging: The broader family of weight-combining methods. Soups are the case where all models were fine-tuned on the same task, while task arithmetic adds models fine-tuned on different tasks.
• Weight averaging within one run: Stochastic weight averaging (SWA) and the exponential moving average (EMA) of weights average checkpoints along a single training trajectory. Soups average independent runs, and the paper found that souping models that already use EMA or SWA still raises accuracy, so the two kinds of averaging add up.
• WiSE-FT: Interpolates between a zero-shot model and its fine-tuned version, a two-ingredient soup that keeps more of the pretrained model's robustness to distribution shift.