How to Scale Your EMA
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.
64c36d0b2941512c4c584516
Dan Busbridge
claimed_verified
2023-07-28T07:51:53.750Z
64c3726f2a5eaefd000cdedd
https://cdn-avatars.huggingface.co/v1/production/uploads/64c3726f2a5eaefd000cdedd/iwFifH1sWQy7agW3eTmNQ.png
Dan Busbridge
user
dbusbridge
64c36d0b2941512c4c584517
Jason Ramapuram
64c36d0b2941512c4c584518
Pierre Ablin
64c36d0b2941512c4c584519
Tatiana Likhomanenko
64c36d0b2941512c4c58451a
Eeshan Gunesh Dhekane
64c36d0b2941512c4c58451b
Xavier Suau
64c36d0b2941512c4c58451c
Russ Webb
How to Scale Your EMA
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.
A scaling rule for optimization with Exponential Moving Average (EMA) improves performance and enables efficient training of pseudo-labeling and SSL methods, including BYOL, at various batch sizes.
stochastic gradient descent Exponential Moving Average (EMA) robustness generalization supervised learning pseudo-labeling Self-Supervised Learning (SSL) BYOL wall-clock time