How to Scale Your EMA
Authors: name: Dan Busbridge; name: Jason Ramapuram; name: Pierre Ablin; name: Tatiana Likhomanenko; name: Eeshan Gunesh Dhekane; name: Xavier Suau; name: Russ Webb
Published: 2023-07-25 20:33:48

Abstract: 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.