Como Dimensionar Sua EMA

How to Scale Your EMA

July 25, 2023

Autores: Dan Busbridge, Jason Ramapuram, Pierre Ablin, Tatiana Likhomanenko, Eeshan Gunesh Dhekane, Xavier Suau, Russ Webb

cs.AI

Resumo

Preservar a dinâmica de treinamento em diferentes tamanhos de lotes é uma ferramenta importante para o aprendizado de máquina prático, pois permite a troca entre o tamanho do lote e o tempo de execução. Essa troca é tipicamente habilitada por uma regra de escalonamento; por exemplo, no gradiente descendente estocástico, deve-se escalar a taxa de aprendizado linearmente com o tamanho do lote. Outra ferramenta importante para o aprendizado de máquina prático é a Média Móvel Exponencial (EMA, do inglês Exponential Moving Average) do modelo, que é uma cópia do modelo que não recebe informações de gradiente, mas segue seu modelo alvo com algum momento. Essa EMA do modelo pode melhorar a robustez e as propriedades de generalização do aprendizado supervisionado, estabilizar a pseudo-rotulagem e fornecer um sinal de aprendizado para o Aprendizado Auto-Supervisionado (SSL, do inglês Self-Supervised Learning). Trabalhos anteriores trataram a EMA do modelo separadamente da otimização, levando a diferentes dinâmicas de treinamento em diferentes tamanhos de lotes e a um desempenho inferior do modelo. Neste trabalho, fornecemos uma regra de escalonamento para otimização na presença de EMAs de modelo e demonstramos sua validade em uma variedade de arquiteturas, otimizadores e modalidades de dados. Também mostramos a validade da regra quando a EMA do modelo contribui para a otimização do modelo alvo, permitindo-nos treinar métodos de pseudo-rotulagem e SSL baseados em EMA em tamanhos de lote pequenos e grandes. Para SSL, permitimos o treinamento de BYOL até o tamanho de lote de 24.576 sem sacrificar o desempenho, resultando em uma redução ideal de 6 vezes no tempo de execução.

English

Preserving training dynamics across batch sizes is an important tool for practical machine learning as it enables the trade-off between batch size and wall-clock time. This trade-off is typically enabled by a scaling rule, for example, in stochastic gradient descent, one should scale the learning rate linearly with the batch size. Another important tool for practical machine learning is the model Exponential Moving Average (EMA), which is a model copy that does not receive gradient information, but instead follows its target model with some momentum. This model EMA can improve the robustness and generalization properties of supervised learning, stabilize pseudo-labeling, and provide a learning signal for Self-Supervised Learning (SSL). Prior works have treated the model EMA separately from optimization, leading to different training dynamics across batch sizes and lower model performance. In this work, we provide a scaling rule for optimization in the presence of model EMAs and demonstrate its validity across a range of architectures, optimizers, and data modalities. We also show the rule's validity where the model EMA contributes to the optimization of the target model, enabling us to train EMA-based pseudo-labeling and SSL methods at small and large batch sizes. For SSL, we enable training of BYOL up to batch size 24,576 without sacrificing performance, optimally a 6times wall-clock time reduction.