Model Soups

Une soupe de modèles est un modèle unique obtenu en faisant la moyenne des poids de plusieurs modèles ajustés finement à partir du même modèle préentraîné, avec différents hyperparamètres. Après un balayage d'hyperparamètres, la pratique habituelle consiste à conserver l'exécution ayant obtenu le meilleur score de validation et à supprimer les autres. Une soupe réutilise ces exécutions écartées : comme chacune est partie des mêmes poids de modèle préentraînés, leurs paramètres restent alignés, et leur moyenne obtient souvent un score supérieur à chacune d'entre elles, tout en coûtant exactement autant à exécuter qu'un seul modèle. Mitchell Wortsman et ses collègues de l'Université de Washington, de Google, de Meta AI, de l'Université de Tel-Aviv et de Columbia ont présenté cette méthode dans Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time, publié à ICML 2022. Une soupe gloutonne de 58 modèles ViT-G/14 ajustés finement, préentraînés sur le jeu de données JFT-3B de Google, a atteint une précision top-1 de 90.94 % sur ImageNet, un résultat de pointe à l'époque.

Soupes uniformes et gloutonnes#

L'ajustement fin ne déplace qu'un peu un grand réseau préentraîné dans l'espace des poids ; les modèles ajustés finement à partir de celui-ci avec différents taux d'apprentissage, augmentations de données, lissages d'étiquettes ou durées d'entraînement tendent donc à se retrouver dans la même vaste région de faible perte. Leur moyenne se situe près du centre de cette région, où la généralisation est meilleure qu'à ses bords. L'article définit deux recettes principales :

• Soupe uniforme : Fait la moyenne de tous les modèles ajustés finement. Cette méthode fonctionne lorsque tous les ingrédients sont bons, mais une exécution avec un taux d'apprentissage très élevé peut sortir de la région commune ; l'inclure dans la moyenne réduit alors la précision.
• Soupe gloutonne : Trie les modèles selon leur précision sur les données de validation mises de côté, commence par le meilleur et n'ajoute chaque modèle suivant que si la précision de validation de la soupe moyennée ne baisse pas. Comme elle part du meilleur modèle, une soupe gloutonne n'est jamais moins performante que celui-ci sur l'ensemble de validation.

Pour le modèle CLIP ViT-B/32 affiné sur ImageNet à l’aide de 72 configurations aléatoires d’hyperparamètres, le meilleur modèle individuel a obtenu 80,38 %. La soupe uniforme des 72 modèles a obtenu 79,97 %, car certaines exécutions avaient quitté la région commune de faible perte, tandis que la soupe gloutonne a conservé cinq modèles et obtenu 81,03 %. Sur cinq jeux de test soumis à un décalage de distribution, notamment les croquis et les rendus d’ImageNet, la soupe gloutonne a obtenu une moyenne de 50,75 %, contre 47,83 % pour le meilleur modèle individuel. L’article fait également état de gains pour ALIGN, ainsi que de gains plus modestes lors de l’ajustement fin de BERT et T5 pour la classification de texte.

Cet exemple PyTorch préentraîne un petit classifieur sur des chiffres manuscrits de 8x8 pixels, ajuste finement neuf copies avec différents taux d'apprentissage et différentes augmentations par bruit d'entrée, puis compare leur soupe uniforme à un ensemble des neuf modèles :

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

Les neuf modèles ajustés finement obtiennent des scores compris entre 0.881 et 0.936 sur les 997 chiffres mis de côté, avec une moyenne de 0.917. Leur soupe uniforme obtient 0.934, un résultat proche de celui de la meilleure exécution seule et légèrement supérieur aux 0.930 de l'ensemble, qui nécessite neuf passes avant par image contre une seule pour la soupe.

Les soupes en pratique#

Les soupes sont surtout utiles lorsqu'un balayage est de toute façon prévu, par exemple lors de l'apprentissage par transfert d'un transformeur de vision (ViT) ou d'un autre modèle de fondation pour une nouvelle tâche de classification d'images. Le réglage des hyperparamètres d'Ultralytics YOLO est un tel balayage : model.tune() entraîne un modèle par ensemble d'hyperparamètres mutés et, pour économiser de l'espace de stockage, supprime les exécutions qui ne sont pas les meilleures du moment, en ne conservant que l'exécution et les poids de la meilleure itération ; une soupe remplace cette étape de sélection du meilleur modèle.

Les modèles composants doivent partager la même architecture et le même point de départ préentraîné, car les réseaux entraînés à partir d'initialisations aléatoires différentes répartissent les mêmes caractéristiques dans des unités cachées différentes, et leur moyenne détruit la précision. Les poids finaux n'entraînent aucun coût supplémentaire à l'inférence ; la soupe conserve donc la latence d'inférence et la mémoire d'un seul modèle.

Soupes de modèles et méthodes connexes#

• Ensemble de modèles : Fait la moyenne des prédictions de plusieurs modèles ; le coût d'inférence augmente donc avec chaque modèle ajouté. Lors du balayage CLIP ViT-B/32, un ensemble de logits a obtenu 81.19 % sur ImageNet, légèrement au-dessus des 81.03 % de la soupe gloutonne, tout en exécutant chaque modèle de l'ensemble pour chaque image.
• Fusion de modèles : La famille plus large des méthodes de combinaison des poids. Les soupes correspondent au cas où tous les modèles ont été ajustés finement sur la même tâche, tandis que l'arithmétique des tâches combine des modèles ajustés finement sur des tâches différentes.
• Moyenne des poids au cours d'une seule exécution : La moyenne stochastique des poids (SWA) et la moyenne mobile exponentielle (EMA) des poids calculent la moyenne des points de contrôle le long d'une seule trajectoire d'entraînement. Les soupes calculent la moyenne d'exécutions indépendantes, et l'article a montré que souper des modèles qui utilisent déjà EMA ou SWA augmente encore la précision ; les deux types de moyennes se cumulent donc.
• WiSE-FT : Interpole entre un modèle zero-shot et sa version ajustée finement, soit une soupe à deux ingrédients qui préserve davantage la robustesse du modèle préentraîné face au décalage de distribution.