How to Scale Your EMA

TL;DR

This paper introduces a new scaling rule for optimizing Exponential Moving Average (EMA) models in machine learning, enabling consistent training dynamics across different batch sizes and improving efficiency in self-supervised learning.

Contribution

The authors propose a novel scaling rule for EMA optimization that maintains training dynamics across batch sizes and validate it across various architectures and modalities.

Findings

Validates the scaling rule across multiple architectures and data types.

Enables training of BYOL with batch sizes up to 24,576 without performance loss.

Achieves a 6× reduction in wall-clock time for SSL training.

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