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 6$\times$ wall-clock time reduction.

翻译：跨批次大小保持训练动态是实用机器学习的重要工具，它能够在批次大小与挂钟时间之间进行权衡。这种权衡通常通过缩放规则实现，例如在随机梯度下降中，学习率应与批次大小线性缩放。实用机器学习中另一重要工具是模型指数移动平均（EMA），这是一种不接收梯度信息、而是以一定动量跟随其目标模型的模型副本。模型EMA能够提升监督学习的鲁棒性与泛化能力、稳定伪标签生成，并为自监督学习（SSL）提供学习信号。以往研究将模型EMA与优化过程分开处理，导致不同批次大小下的训练动态差异和模型性能下降。本文提出了含模型EMA的优化缩放规则，并验证了其在多种架构、优化器和数据模态下的有效性。我们还展示了该规则在模型EMA参与目标模型优化时的有效性，从而能够在小批量和大批量场景下训练基于EMA的伪标签方法和SSL方法。对于SSL，我们实现了BYOL在批次大小高达24,576时的训练且不损失性能，最佳情况下挂钟时间减少6倍。