Model Soups

Una mezcla de modelos es un único modelo creado al promediar los pesos de varios modelos que se han ajustado a partir del mismo modelo preentrenado con distintos hiperparámetros. La práctica habitual tras una búsqueda de hiperparámetros consiste en conservar la ejecución con la mejor puntuación de validación y descartar las demás. Una mezcla reutiliza esas ejecuciones descartadas: como todas partieron de los mismos pesos del modelo preentrenado, sus parámetros siguen estando alineados, y su promedio suele obtener una puntuación superior a la de cualquiera de ellas, con exactamente el mismo coste de ejecución que un solo modelo. Mitchell Wortsman y sus colaboradores de la Universidad de Washington, Google, Meta AI, la Universidad de Tel Aviv y la Universidad de Columbia presentaron el método en Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time, publicado en ICML 2022. Una mezcla codiciosa de 58 modelos ViT-G/14 ajustados, preentrenados con el conjunto de datos JFT-3B de Google, alcanzó una precisión top-1 del 90,94 % en ImageNet, un resultado de vanguardia en aquel momento.

Mezclas uniformes y codiciosas#

El ajuste fino desplaza una gran red preentrenada solo una corta distancia en el espacio de pesos, por lo que los modelos ajustados a partir de ella con distintos índices de aprendizaje, aumentos de datos, suavizado de etiquetas o duraciones de entrenamiento suelen acabar en la misma amplia región de baja pérdida. Al promediarlos, el resultado queda cerca del centro de esa región y generaliza mejor que sus extremos. El artículo define dos recetas principales:

• Mezcla uniforme: Promedia todos los modelos ajustados. Funciona cuando todos los componentes son buenos, pero una ejecución con un índice de aprendizaje muy alto puede salirse de la región compartida, y promediarla reduce la precisión.
• Mezcla codiciosa: Ordena los modelos según su precisión en datos de validación reservados, empieza por el mejor y solo añade cada modelo siguiente si la precisión de validación de la mezcla promediada no disminuye. Como parte del mejor modelo, una mezcla codiciosa nunca obtiene peores resultados que ese modelo en el conjunto de validación.

Para CLIP ViT-B/32 ajustado con ImageNet mediante 72 configuraciones aleatorias de hiperparámetros, el mejor modelo individual obtuvo una puntuación del 80,38 %. La mezcla uniforme de los 72 obtuvo un 79,97 %, porque algunas ejecuciones habían salido de la región compartida de baja pérdida, mientras que la mezcla codiciosa conservó cinco modelos y obtuvo un 81,03 %. En cinco conjuntos de prueba con cambio de distribución, como los bocetos y las representaciones de ImageNet, la mezcla codiciosa obtuvo un promedio del 50,75 %, frente al 47,83 % del mejor modelo individual. El artículo también informa de mejoras para ALIGN y de otras más pequeñas al ajustar BERT y T5 para la clasificación de texto.

Este ejemplo de PyTorch preentrena un clasificador pequeño con dígitos manuscritos de 8x8, ajusta nueve copias con distintos índices de aprendizaje y aumentos de datos con ruido de entrada, y compara su mezcla uniforme con un conjunto de los nueve modelos:

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

Los nueve modelos ajustados obtienen puntuaciones entre 0.881 y 0.936 en los 997 dígitos reservados, con una media de 0.917. Su mezcla uniforme obtiene 0.934, cerca del mejor resultado individual y ligeramente por encima de 0.930 del conjunto, que necesita nueve pasadas hacia delante por imagen, mientras que la mezcla solo necesita una.

Mezclas en la práctica#

Las mezclas son más útiles cuando ya se realiza una búsqueda, por ejemplo, al aplicar aprendizaje por transferencia a un transformador de visión (ViT) u otro modelo fundacional para una nueva tarea de clasificación de imágenes. El ajuste de hiperparámetros de Ultralytics YOLO es una de esas búsquedas: model.tune() entrena un modelo por cada conjunto de hiperparámetros modificados y, para ahorrar espacio de almacenamiento, elimina las ejecuciones que no son las mejores en ese momento y conserva solo la ejecución y los pesos de la mejor iteración; una mezcla sustituye este paso de elegir el mejor.

Los componentes deben compartir una arquitectura y un punto de partida preentrenado, porque las redes entrenadas con distintas inicializaciones aleatorias almacenan las mismas características en unidades ocultas diferentes, y promediarlas destruye la precisión. Los pesos finales no suponen ningún coste adicional durante la inferencia, así que la mezcla mantiene la latencia de inferencia y la memoria de un solo modelo.

Mezclas de modelos frente a métodos relacionados#

• Conjunto de modelos: Promedia las predicciones de varios modelos, por lo que el coste de inferencia aumenta con cada componente. En la búsqueda de CLIP ViT-B/32, un conjunto de logits obtuvo un 81,19 % en ImageNet, ligeramente por encima del 81,03 % de la mezcla codiciosa, pero ejecutó todos los modelos del conjunto para cada imagen.
• Fusión de modelos: La familia más amplia de métodos para combinar pesos. Las mezclas son el caso en que todos los modelos se ajustaron a la misma tarea, mientras que la aritmética de tareas suma modelos ajustados a tareas distintas.
• Promedio de pesos en una misma ejecución: El promedio estocástico de pesos (SWA) y la media móvil exponencial (EMA) de pesos promedian puntos de control a lo largo de una única trayectoria de entrenamiento. Las mezclas promedian ejecuciones independientes, y el artículo observó que mezclar modelos que ya utilizan EMA o SWA también aumenta la precisión; por tanto, los dos tipos de promedio se complementan.
• WiSE-FT: Interpola entre un modelo de cero ejemplos y su versión ajustada, una mezcla de dos componentes que conserva mejor la robustez del modelo preentrenado frente a cambios de distribución.