Model Soups

Uma sopa de modelos é um único modelo criado pela média dos pesos de vários modelos que foram ajustados finamente a partir do mesmo modelo pré-treinado com diferentes hiperparâmetros. A prática habitual após uma busca de hiperparâmetros é manter a execução com a melhor pontuação de validação e descartar as demais. Uma sopa reaproveita essas execuções descartadas: como cada uma começou com os mesmos pesos do modelo pré-treinado, seus parâmetros continuam alinhados, e a média deles costuma obter uma pontuação superior à de qualquer um deles, com o mesmo custo de execução de um único modelo. Mitchell Wortsman e colegas da Universidade de Washington, do Google, da Meta AI, da Universidade de Tel Aviv e da Universidade Columbia apresentaram o método em Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time, publicado na ICML 2022. Uma sopa gulosa de 58 modelos ViT-G/14 ajustados finamente, pré-treinados no conjunto de dados JFT-3B do Google, alcançou 90,94% de precisão top-1 no ImageNet, um resultado de ponta na época.

Sopas uniformes e gulosas#

O ajuste fino desloca uma grande rede pré-treinada apenas uma curta distância no espaço dos pesos, então modelos ajustados a partir dela com diferentes taxas de aprendizado, aumentos de dados, suavização de rótulos ou durações de treinamento tendem a se situar na mesma ampla região de baixa perda. A média entre eles fica perto do centro dessa região, generalizando melhor do que suas extremidades. O artigo define duas receitas principais:

• Sopa uniforme: Calcula a média de todos os modelos ajustados finamente. Funciona quando todos os ingredientes são bons, mas uma execução com uma taxa de aprendizado muito alta pode sair da região compartilhada, e incluí-la na média reduz a precisão.
• Sopa gulosa: Ordena os modelos pela precisão nos dados de validação reservados, começa pelo melhor e só adiciona cada modelo seguinte se a precisão de validação da sopa média não cair. Como começa pelo melhor modelo, uma sopa gulosa nunca tem desempenho inferior ao desse modelo no conjunto de validação.

Para o CLIP ViT-B/32 ajustado finamente no ImageNet com 72 configurações aleatórias de hiperparâmetros, o melhor modelo individual obteve 80,38%. A sopa uniforme dos 72 modelos obteve 79,97%, pois algumas execuções saíram da região compartilhada de baixa perda, enquanto a sopa gulosa manteve cinco modelos e obteve 81,03%. Em cinco conjuntos de teste com mudança de distribuição, como esboços e representações do ImageNet, a sopa gulosa obteve uma média de 50,75%, contra 47,83% do melhor modelo individual. O artigo também relata ganhos para o ALIGN e ganhos menores ao ajustar finamente o BERT e o T5 para classificação de texto.

Este exemplo de PyTorch pré-treina um classificador pequeno com dígitos manuscritos de 8x8, ajusta finamente nove cópias com diferentes taxas de aprendizado e aumento de ruído na entrada e compara a sopa uniforme delas com um conjunto dos nove 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}")

Os nove modelos ajustados finamente obtêm pontuações entre 0.881 e 0.936 nos 997 dígitos reservados, com média de 0.917. A sopa uniforme obtém 0.934, perto da melhor execução individual e ligeiramente acima dos 0.930 do conjunto, que exige nove passagens de inferência por imagem, enquanto a sopa precisa de apenas uma.

Sopas na prática#

As sopas são mais úteis quando uma busca já seria feita de qualquer forma, por exemplo, ao aplicar aprendizado por transferência a um transformador de visão (ViT) ou outro modelo fundamental para uma nova tarefa de classificação de imagens. O ajuste de hiperparâmetros do Ultralytics YOLO é uma busca desse tipo: model.tune() treina um modelo para cada conjunto de hiperparâmetros modificados e, para economizar armazenamento, exclui as execuções que não são as melhores no momento, mantendo apenas a execução e os pesos da melhor iteração — a etapa de escolher o melhor que uma sopa substitui.

Os componentes precisam compartilhar uma arquitetura e um ponto de partida pré-treinado, pois redes treinadas com diferentes inicializações aleatórias armazenam as mesmas características em unidades ocultas diferentes, e calcular a média entre elas destrói a precisão. Os pesos finais não têm custo adicional na inferência, então a sopa mantém a latência de inferência e o uso de memória de um único modelo.

Sopas de modelos vs. métodos relacionados#

• Conjunto de modelos: Calcula a média das previsões de vários modelos, fazendo o custo de inferência crescer com cada membro. Na busca com CLIP ViT-B/32, um conjunto de logits obteve 81,19% no ImageNet, ligeiramente acima dos 81,03% da sopa gulosa, mas executou todos os modelos do conjunto para cada imagem.
• Mesclagem de modelos: A família mais ampla de métodos que combinam pesos. As sopas são o caso em que todos os modelos foram ajustados finamente para a mesma tarefa, enquanto a aritmética de tarefas combina modelos ajustados finamente para tarefas diferentes.
• Média de pesos em uma única execução: A média estocástica de pesos (SWA) e a média móvel exponencial (EMA) dos pesos calculam a média dos pontos de verificação ao longo de uma única trajetória de treinamento. As sopas calculam a média de execuções independentes, e o artigo constatou que combinar em uma sopa modelos que já usam EMA ou SWA ainda aumenta a precisão, portanto os dois tipos de média se complementam.
• WiSE-FT: Interpola entre um modelo de zero-shot e sua versão ajustada finamente, uma sopa de dois componentes que preserva mais a robustez do modelo pré-treinado à mudança de distribuição.