Model Soups

Смешанная модель — это одна модель, полученная усреднением весов нескольких моделей, дообученных от одной и той же предварительно обученной модели с разными гиперпараметрами. Обычно после перебора гиперпараметров оставляют запуск с лучшей оценкой на валидации, а остальные удаляют. Смешивание позволяет повторно использовать отброшенные запуски: поскольку каждый из них начинался с одних и тех же предварительно обученных весов модели, их параметры по-прежнему согласованы, а их усреднённая модель часто показывает результат лучше любой из исходных и обходится при запуске ровно столько же, сколько одна модель. Митчелл Вортсман и его коллеги из Вашингтонского университета, Google, Meta AI, Тель-Авивского университета и Колумбийского университета представили этот метод в статье «Смешивание моделей: усреднение весов нескольких дообученных моделей повышает точность без увеличения времени инференса», опубликованной на ICML 2022. Жадное смешивание 58 дообученных моделей ViT-G/14, предварительно обученных на наборе данных JFT-3B от Google, обеспечило точность top-1 90,94% на ImageNet — на тот момент это был результат уровня лучших в отрасли.

Равномерное и жадное смешивание#

Дообучение сдвигает большую предварительно обученную сеть лишь на небольшое расстояние в пространстве весов, поэтому модели, дообученные с разными скоростями обучения, аугментациями данных, сглаживанием меток или длительностью обучения, обычно оказываются в одной широкой области низких значений функции потерь. Усреднение помещает модель ближе к центру этой области, что обеспечивает лучшую обобщающую способность, чем у моделей на её границах. В статье описаны два основных способа:

• Равномерное смешивание: Усредняет все дообученные модели. Метод работает, когда все компоненты хороши, но запуск с очень высокой скоростью обучения может выйти за пределы общей области, и его включение в среднее снижает точность.
• Жадное смешивание: Сортирует модели по точности на отложенных валидационных данных, начинает с лучшей и добавляет следующую модель, только если точность усреднённой модели на валидации не снижается. Поскольку жадное смешивание начинается с лучшей модели, на валидационном наборе оно никогда не уступает этой модели.

Для CLIP ViT-B/32, дообученной на ImageNet с 72 случайными наборами гиперпараметров, лучшая отдельная модель показала 80,38%. Равномерное смешивание всех 72 моделей показало 79,97%, поскольку некоторые запуски вышли за пределы общей области низких потерь; жадное смешивание сохранило пять моделей и показало 81,03%. На пяти тестовых наборах со сдвигом распределения, таких как эскизы и стилизованные изображения ImageNet, средний результат жадного смешивания составил 50,75% против 47,83% у лучшей отдельной модели. В статье также описаны улучшения для ALIGN и менее значительные улучшения при дообучении BERT и T5 для классификации текста.

В этом примере на PyTorch небольшая модель классификатора предварительно обучается на рукописных цифрах размером 8x8, затем девять её копий дообучаются с разными скоростями обучения и аугментацией входных данных шумом; после этого их равномерное смешивание сравнивается с ансамблем всех девяти моделей:

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

На 997 отложенных изображениях цифр девять дообученных моделей показывают результаты от 0.881 до 0.936, в среднем 0.917. Равномерное смешивание показывает 0.934 — почти столько же, сколько лучший отдельный запуск, и немного выше, чем 0.930 у ансамбля, которому для каждого изображения нужны девять прямых проходов, тогда как смешанной модели достаточно одного.

Смешивание на практике#

Смешивание особенно полезно в случаях, когда перебор гиперпараметров всё равно выполняется, например при переносе обучения визуального трансформера (ViT) или другой базовой модели на новую задачу классификации изображений. Так устроен и подбор гиперпараметров Ultralytics YOLO: model.tune() обучает по одной модели на каждом наборе изменённых гиперпараметров и для экономии места удаляет запуски, которые уже не являются лучшими, оставляя только запуск и веса лучшей итерации. Именно этот выбор лучшей модели и заменяет смешивание.

Все компоненты должны иметь одинаковую архитектуру и общую предварительно обученную исходную модель: сети, обученные с разными случайными начальными параметрами, хранят одинаковые признаки в разных скрытых нейронах, а их усреднение уничтожает точность. Итоговые веса не требуют дополнительных затрат при инференсе, поэтому смешанная модель сохраняет задержку инференса и расход памяти одной модели.

Смешивание моделей и родственные методы#

• Ансамбль моделей: Усредняет предсказания нескольких моделей, поэтому стоимость инференса увеличивается с каждым участником. При переборе для CLIP ViT-B/32 ансамбль логитов показал 81,19% на ImageNet — немного выше, чем жадное смешивание с результатом 81,03%, — но для каждого изображения запускались все модели ансамбля.
• Объединение моделей: Более широкое семейство методов комбинирования весов. Смешивание — это случай, когда все модели дообучались на одной задаче, а арифметика задач объединяет модели, дообученные на разных задачах.
• Усреднение весов в рамках одного запуска: Стохастическое усреднение весов (SWA) и экспоненциальное скользящее среднее весов (EMA) усредняют контрольные точки в рамках одной траектории обучения. Смешивание усредняет независимые запуски. В статье показано, что смешивание моделей, уже использующих EMA или SWA, дополнительно повышает точность, поэтому эти два способа усреднения дополняют друг друга.
• WiSE-FT: Интерполирует между моделью zero-shot и её дообученной версией — смешиванием двух компонентов, которое позволяет сохранить больше устойчивости предварительно обученной модели к сдвигу распределения.