Comment mettre à l'échelle votre EMA

How to Scale Your EMA

July 25, 2023

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

cs.AI

Résumé

Préserver la dynamique d'entraînement pour différentes tailles de lot est un outil essentiel en apprentissage automatique pratique, car il permet de faire un compromis entre la taille du lot et le temps d'exécution réel. Ce compromis est généralement rendu possible par une règle de mise à l'échelle. Par exemple, dans la descente de gradient stochastique, il convient de mettre à l'échelle le taux d'apprentissage linéairement avec la taille du lot. Un autre outil important en apprentissage automatique pratique est la moyenne mobile exponentielle (EMA) du modèle, qui est une copie du modèle ne recevant pas d'informations de gradient, mais suivant plutôt son modèle cible avec une certaine inertie. Cette EMA du modèle peut améliorer la robustesse et les propriétés de généralisation en apprentissage supervisé, stabiliser l'étiquetage pseudo-supervisé et fournir un signal d'apprentissage pour l'apprentissage auto-supervisé (SSL). Les travaux antérieurs ont traité la EMA du modèle séparément de l'optimisation, entraînant des dynamiques d'entraînement différentes selon les tailles de lot et une performance moindre du modèle. Dans ce travail, nous proposons une règle de mise à l'échelle pour l'optimisation en présence de EMA de modèle et démontrons sa validité sur une gamme d'architectures, d'optimiseurs et de modalités de données. Nous montrons également la validité de cette règle lorsque la EMA du modèle contribue à l'optimisation du modèle cible, nous permettant d'entraîner des méthodes d'étiquetage pseudo-supervisé et de SSL basées sur la EMA pour des tailles de lot petites et grandes. Pour le SSL, nous permettons l'entraînement de BYOL jusqu'à une taille de lot de 24 576 sans sacrifier les performances, ce qui réduit optimalement le temps d'exécution réel d'un facteur 6.

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.