  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 machine learning tool is the
model EMA, a functional copy of a target model, whose parameters move towards
those of its target model according to an Exponential Moving Average (EMA) at a
rate parameterized by a momentum hyperparameter. This model EMA can improve the
robustness and generalization of supervised learning, stabilize
pseudo-labeling, and provide a learning signal for Self-Supervised Learning
(SSL). Prior works have not considered the optimization of the model EMA when
performing scaling, 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 a model EMA and demonstrate the rule's 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, a 6$\times$ wall-clock time
reduction under idealized hardware settings.
