title: How to Scale Your EMA
summary: 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.
paper_authors._id: 64c36d0b2941512c4c584516
paper_authors.name: Dan Busbridge
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2023-07-28T07:51:53.750Z
paper_authors._id: 64c3726f2a5eaefd000cdedd
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/64c3726f2a5eaefd000cdedd/iwFifH1sWQy7agW3eTmNQ.png
paper_authors.fullname: Dan Busbridge
paper_authors.type: user
paper_authors.user: dbusbridge
paper_authors._id: 64c36d0b2941512c4c584517
paper_authors.name: Jason Ramapuram
paper_authors._id: 64c36d0b2941512c4c584518
paper_authors.name: Pierre Ablin
paper_authors._id: 64c36d0b2941512c4c584519
paper_authors.name: Tatiana Likhomanenko
paper_authors._id: 64c36d0b2941512c4c58451a
paper_authors.name: Eeshan Gunesh Dhekane
paper_authors._id: 64c36d0b2941512c4c58451b
paper_authors.name: Xavier Suau
paper_authors._id: 64c36d0b2941512c4c58451c
paper_authors.name: Russ Webb
paper_title: How to Scale Your EMA
paper_summary: 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.
paper_ai_summary: 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.
paper_ai_keywords: stochastic gradient descent Exponential Moving Average (EMA) robustness generalization supervised learning pseudo-labeling Self-Supervised Learning (SSL) BYOL wall-clock time