Stochastic Weight Averaging (SWA) é uma técnica de média de modelos em Deep learning que melhora a generalização ao calcular a média dos pesos de uma Neural network ao longo de múltiplos pontos de sua trajetória de treinamento. Diferentemente do treinamento tradicional, que retém apenas os pesos finais, o SWA coleta pesos em intervalos regulares durante a fase final do treinamento e calcula sua média aritmética. Esse procedimento simples frequentemente leva a soluções mais planas na paisagem de perda, que estão associadas a melhor desempenho em testes e robustez. O SWA foi introduzido por Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov e Andrew Gordon Wilson em seu artigo de 2018 "Averaging Weights Leads to Wider Optima and Better Generalization". O método não requer alterações na arquitetura subjacente do modelo ou na função de perda, tornando-o fácil de integrar em pipelines de treinamento existentes.
A ideia central por trás do SWA é que, durante o treinamento com descida de gradiente estocástica (SGD) ou suas variantes, os parâmetros do modelo oscilam em torno de uma região de baixa perda. Calcular a média desses parâmetros pode suavizar as oscilações e produzir uma solução que se encontra em um mínimo amplo e plano. Tais mínimos planos são hipotetizados para generalizar melhor do que mínimos acentuados, pois são menos sensíveis a perturbações e mudanças entre as distribuições de treinamento e teste. O SWA demonstrou melhorar a precisão e a calibração em várias arquiteturas e tarefas, incluindo classificação de imagens, detecção de objetos e modelagem de linguagem.
Mecanismo e Algoritmo
O SWA opera em duas fases. Na primeira fase, o modelo é treinado com um cronograma padrão de taxa de aprendizado (por exemplo, anelamento de cosseno ou decaimento por etapas) para alcançar uma região próxima a um mínimo. Na segunda fase, o treinamento continua com uma taxa de aprendizado constante ou cíclica, e os pesos são coletados ao final de cada época ou após um número fixo de iterações. Os pesos do SWA são atualizados como uma média móvel: n_swa = n_swa + 1; w_swa = (w_swa * (n_swa - 1) + w) / n_swa, onde w são os pesos atuais e w_swa são os pesos médios. O modelo final usa w_swa.
Um detalhe importante é que as estatísticas de normalização em lote devem ser recalculadas após a média, porque a média móvel e a variância das camadas de normalização em lote não são calculadas em média junto com os pesos. Na prática, após os pesos do SWA serem calculados, uma passagem direta sobre os dados de treinamento é realizada para atualizar as estatísticas de normalização em lote. Este passo é crucial para modelos que usam normalização em lote, como ResNets.
Motivação Teórica
O sucesso do SWA é frequentemente explicado pela geometria da paisagem de perda. Pesquisas de Izmailov et al. e trabalhos subsequentes mostraram que o SGD tende a convergir para pontos próximos à fronteira de uma região de baixa perda, enquanto calcular a média de múltiplos pontos move a solução em direção ao centro dessa região. Esse ponto central geralmente se encontra em uma bacia mais plana, que está associada a melhor generalização. A conexão entre mínimos planos e generalização tem sido estudada extensivamente, com relações com limites PAC-Bayes e a nitidez da função de perda.
Outra perspectiva relaciona o SWA à inferência bayesiana. Sob certas suposições, a trajetória do SGD pode ser vista como uma cadeia de Markov que amostra de uma distribuição posterior. Calcular a média ao longo dessa trajetória aproxima a média posterior, que é a previsão ótima de Bayes para erro quadrático. Essa interpretação liga o SWA a métodos de Monte Carlo de cadeia de Markov com gradiente estocástico, embora o SWA seja mais simples e não exija ajuste cuidadoso das escalas de ruído.
Variantes e Extensões
Várias variantes do SWA foram propostas para melhorar ainda mais o desempenho. Uma extensão notável é o Stochastic Weight Averaging Gaussian (SWAG), introduzido por Wesley Maddox, Timur Garipov, Pavel Izmailov, Dmitry Vetrov e Andrew Gordon Wilson em 2019. O SWAG aproxima a distribuição posterior dos pesos ajustando uma Gaussiana usando os primeiros e segundos momentos dos pesos coletados. Isso permite estimativa de incerteza e melhor calibração, e pode ser usado para aprendizado profundo bayesiano.
Outra variante é o Fast Geometric Ensembling (FGE), que usa um cronograma cíclico de taxa de aprendizado para coletar pesos de diferentes modos da paisagem de perda. O SWA pode ser combinado com o FGE para alcançar desempenho ainda melhor. Além disso, o SWA foi integrado a técnicas como taxas de aprendizado, aumento de dados e poda para melhorar robustez e eficiência.
Aplicações
O SWA foi aplicado com sucesso em uma ampla gama de domínios. Em visão computacional, ele melhora a precisão de classificadores de imagens em conjuntos de dados como CIFAR-10, CIFAR-100 e ImageNet. Por exemplo, o SWA demonstrou aumentar a precisão de teste de um PreAct ResNet-164 no CIFAR-100 de 81,15% para 82,90% sem alterações arquitetônicas. Em detecção de objetos, o SWA ajuda a estabilizar o treinamento e melhorar a precisão média média.
Em processamento de linguagem natural, o SWA tem sido usado para treinar Transformers e grandes modelos de linguagem de forma mais eficaz. Foi demonstrado que ele reduz o sobreajuste e melhora a perplexidade em tarefas de modelagem de linguagem. O SWA também beneficia modelos generativos, como GANs e modelos de difusão, fornecendo atualizações de pesos mais suaves, o que pode levar a treinamento mais estável e amostras de maior qualidade.
Além do aprendizado supervisionado padrão, o SWA foi adaptado para aprendizado semissupervisionado, transferência de aprendizado e adaptação de domínio. Sua simplicidade e generalidade o tornam uma ferramenta valiosa no kit de ferramentas do profissional de Machine learning.
Comparação com Outras Técnicas
O SWA é frequentemente comparado com outros métodos de regularização e ensemble. Diferentemente de ensembles tradicionais, que calculam a média das previsões de múltiplos modelos treinados independentemente, o SWA calcula a média dos pesos, o que é computacionalmente mais barato e requer apenas uma única execução de treinamento. Isso o torna particularmente atraente quando o armazenamento do modelo ou o custo de inferência é uma preocupação.
Em comparação com técnicas como Dropout ou normalização em lote, o SWA não modifica a dinâmica de treinamento, mas sim pós-processa os pesos. Ele pode ser usado em conjunto com esses métodos. O SWA também difere do Adam e de outros otimizadores adaptativos, que mantêm taxas de aprendizado por parâmetro; o SWA é tipicamente aplicado sobre trajetórias de SGD ou Adam.
Uma limitação do SWA é que ele requer um cronograma para determinar quando começar a coletar pesos. Se a coleta começar muito cedo, a média pode ser dominada por pesos subótimos; se começar tarde demais, os benefícios diminuem. No entanto, na prática, uma regra simples é começar a coletar após a taxa de aprendizado ter decaído para um valor baixo.
Considerações Práticas
Implementar o SWA é direto. A maioria dos frameworks de aprendizado profundo, como PyTorch e TensorFlow, fornece utilitários ou exemplos para o SWA. Os principais passos são: (1) treinar o modelo com um cronograma padrão, (2) continuar o treinamento com uma taxa de aprendizado constante ou cíclica, (3) coletar pesos em intervalos, (4) calcular a média móvel e (5) recalcular as estatísticas de normalização em lote.
O SWA demonstrou ser robusto às escolhas de hiperparâmetros. A taxa de aprendizado para a segunda fase é tipicamente definida como o valor mínimo do cronograma inicial ou uma pequena constante. A frequência de coleta de pesos pode ser a cada época ou a cada poucas centenas de iterações. Na prática, coletar pesos por 5 a 20 épocas é frequentemente suficiente.
Impacto e Recepção
A introdução do SWA teve um impacto significativo na comunidade de aprendizado profundo. Ele forneceu um método simples, porém eficaz, para melhorar a generalização sem custo computacional adicional durante a inferência. O artigo foi amplamente citado e inspirou inúmeros trabalhos subsequentes sobre média de pesos e mínimos planos. O SWA é agora uma técnica padrão em muitos pipelines de treinamento, especialmente em pesquisa acadêmica e competições.
O SWA também influenciou o desenvolvimento de outros métodos de média, como a Média Móvel Exponencial (EMA), que é comumente usada no treinamento de modelos generativos e grandes modelos de linguagem. A EMA é uma variante que dá mais peso aos pesos recentes e é frequentemente usada na prática devido à sua simplicidade e eficácia. O SWA e a EMA são complementares, e alguns frameworks permitem alternar entre eles.
Limitações e Direções Futuras
Apesar de seus benefícios, o SWA tem limitações. Ele assume que a paisagem de perda é relativamente suave e que calcular a média dos pesos é significativo, o que pode não valer para todas as arquiteturas ou tarefas. Por exemplo, modelos com normalização em lote exigem manuseio cuidadoso, e modelos com espaços de parâmetros complexos (por exemplo, redes recorrentes) podem não se beneficiar tanto. Além disso, o SWA não fornece estimativas de incerteza, a menos que seja estendido (por exemplo, SWAG).
Direções futuras de pesquisa incluem o desenvolvimento de cronogramas de média adaptativos, a combinação do SWA com métodos bayesianos e a aplicação do SWA a novos domínios, como aprendizado por reforço e aprendizado federado. À medida que o aprendizado profundo continua a evoluir, técnicas de média de pesos como o SWA permanecem uma ferramenta fundamental para melhorar o desempenho e a confiabilidade do modelo.