SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Image Classification We consider a variant of the original SGD Scaling Rule result (Goyal et al., 2017) and train a ResNetv2 (He et al., 2016b) on ImageNet1k (Russakovsky et al., 2014) (Figure 2) using a three step learning rate schedule. The base momentum  $\rho_{B} = 0.9999$  at batch size 1024 was found by hyperparameter optimizing for EMA test performance, and we seek to achieve this optimized performance at different batch sizes. We do not apply the EMA Scaling Rule on the Batch Normalization (Ioffe & Szegedy, 2015) statistics<sup>5</sup>. We observe that without the EMA Scaling Rule, there is a significant drop in model EMA test performance, whereas with the EMA Scaling Rule, we
<sup>5</sup>Since Batch Normalization statistics use an EMA update, it is reasonable to ask whether the EMA Scaling Rule should be applied. We investigate this in Appendix F.3. We find one should apply the scaling rule, however, the effect is less significant than the application of the EMA Scaling Rule to model parameters.
can approximate the baseline model EMA test top-1 performance across all batch sizes. We match baseline EMA statistics across the full trajectory batch size 2048, where the test EMA performance diverges. This is due to non-EMA test performance dropping for high  $\kappa$  (see Appendix F.2). We observe that model EMA top-1 is approximately  $0.2\%$  to  $0.3\%$  higher than the target model.
Automatic Speech Recognition (ASR) We train a transformer (Vaswani et al., 2017) using the Connectionist Temporal Classification (CTC) loss (Graves et al., 2006) and Adam optimizer on the train-clean-100 subset (100h) of LibriSpeech (Panayotov et al., 2015) (for details see Appendix G). We apply the Adam Scaling Rule (Malladi et al. (2022), Definition C.3) and use dynamic batching (minibatch size  $\times$  sequence length  $=$  const  $= 290s$ , and s indicates audio duration in seconds).
Without the EMA Scaling Rule, there is a significant difference in model EMA test Word Error Rate (WER) trajectories compared to the baseline, whereas with the EMA Scaling Rule, trajectories match, as is shown in Figure 3. We note that compared to image classification, in ASR, the model EMA converges to similar final performance irrespective of use of the scaling rule. This convergence is due to the longer training time compared to the EMA horizon as discussed in Table 1 (see Appendix E.2 for a proof sketch). Although in this specific case one can achieve similar final performance without the EMA Scaling Rule, it is necessary to use the EMA Scaling Rule in order to replicate the full training trajectory, which gives guarantees on properties like final performance (see Corollary 2.1.1). We also observe a growing gap between the baseline and EMA-scaled trajectories as we increase  $\kappa$ . Inspecting the train loss and non-EMA test WER, which do not depend on the EMA update (see Figure 14, Appendix G.1), indicates this is due to a breakdown of the Adam Scaling Rule. In summary, evaluation on ASR shows that the EMA Scaling Rule holds in practice for sequential data with dynamic batch sizes, as well as when using adaptive optimization.
Figure 3: Transformer Polyak-Ruppert averaging on LibriSpeech (trained on train-clean-100) with different scalings  $\kappa$ . The baseline ( $\kappa = 1$ , black dashed) is trained with Adam and momentum  $\rho_{B} = 0.99995$  at a dynamic batch size  $B = 8 \times 290s$ , which corresponds to a single train step on the  $x$ -axis. We investigate dynamic batch sizes down to  $B = 2 \times 290s$  (left) and up to  $B = 32 \times 290s$  (right), with (blue,  $\rho = \rho_{B}^{\kappa}$ ), and without (red,  $\rho = \rho_{B}$ ) the EMA Scaling Rule. The Adam Scaling Rule (Malladi et al. (2022), Definition C.3) is used throughout.
3.3 Semi-supervised speech recognition via pseudo-labeling
We continue using the same ASR model and training pipeline of Section 3.2. However, we consider semi-supervised learning via continuous pseudo-labeling where labeled (train-clean-100, 100h) and unlabeled (the rest of LibriSpeech, 860h) data are given during training, and the model EMA is involved in the overall optimization (Likhomanenko et al., 2021a, 2022; Manohar et al., 2021; Higuchi et al., 2022). We first pre-train a target model (student) on a limited labeled set for a short period (e.g. 20k steps of  $B = 8 \times 290s^6$ ). Concurrently, the student updates a model EMA (teacher). After pretraining, we continue training the student with both labeled and unlabeled data, with the teacher first transcribing unlabeled data from the batch producing Pseudo-Labels (PLs). These PLs are treated by the student as ground-truth transactions, and standard supervised optimization is performed.
Compared to Polyak-Ruppert Averaging (Section 3.2), where the model EMA plays no role in the joint optimization, we observe that in PL it is essential to employ the EMA Scaling Rule in order to match the model trajectories at scaled batch sizes. When the EMA Scaling Rule is not used, Figure 4 reveals a significant difference in PL quality trajectory, leading to a higher test WER.
For  $\kappa > 2$ , we found the Adam Scaling Rule does not perfectly match the reference trajectory in the pre-training phase. This results in a significantly different PL quality at the start of pseudolabeling (20k steps of  $B = 8 \times 290$ s), which affects the training dynamics (Berrebbi et al., 2023). To
Note that number of steps is batch size dependent and should be scaled by  $1 / \kappa$  (see Appendix C).
Figure 4: Transformer pseudo-labeling on LibriSpeech with different scalings  $\kappa$ . The baseline ( $\kappa = 1$ , black dashed) is trained with Adam at a dynamic batch size of  $8 \times 290$  seconds, which corresponds to a single train step on the  $x$ -axis. The model EMA (teacher) is updated with momentum  $\rho_{B} = 0.9999$ . We investigate dynamic batch sizes down to  $B = 4 \times 290$ s (left) and up to  $B = 64 \times 290$ s (right), with (blue,  $\rho = \rho_{B}^{\kappa}$ ) and without (red,  $\rho = \rho_{B}$ ) the EMA Scaling Rule. The Adam Scaling Rule (Malladi et al. (2022), Definition C.3) is used throughout. For  $\kappa \leq 2$ , we start pseudo-labeling after  $20\mathrm{k} / \kappa$  training steps; while for  $\kappa > 2$ , we start when pre-training WER matches the baseline WER.
alleviate the Adam Scaling Rule mismatch effect for  $\kappa > 2$ , we postpone the pseudo-labeling until pre-training on labeled data gives similar validation WER, see Appendix G. With this heuristic, we can match the baseline trajectory with the EMA Scaling Rule up to  $\kappa = 8$  (Figure 4).
In summary, (a) model EMA affects the optimization process of pseudo-labeling in ASR resulting in the necessity of EMA Scaling Rule to be applied while scaling the batch size; (b) an optimizer scaling rule breakdown results in the EMA Scaling Rule breakdown but this effect can be alleviated by longer pre-training on labeled data having similar PLs quality at the start across different scalings.
3.4 Self-supervised image representation learning
Finally, we turn our attention to distillation based Self-Supervised Learning (SSL). where the model EMA is the teacher (Grill et al., 2020; Niizumi et al., 2023; Caron et al., 2021; Oquab et al., 2023).
We will use BYOL (Grill et al. (2020), Definition 1.1)<sup>7</sup> for our investigation into scaling as it is well-studied (Tian et al., 2021; Richemond et al., 2023), relatively simple to implement due to minimal hyper-parameters, and obtains competitive results (Grill et al., 2020; Koppula et al., 2022). Since BYOL learns through self-referential distillation, momentum plays a significant role in optimization. We analyze: i) a ResNet-18 (He et al., 2016a) on CIFAR10 (Krizhevsky et al., 2014) (Figure 5) using SGD (Definition 2.1); and ii) a ViT-B/16 (Dosovitskiy et al., 2021) on ImageNet1k using AdamW (Loshchilov & Hutter, 2019). A recipe for BYOL using ViTs is provided in Appendix H.3.
ResNet-18 on CIFAR-10 We begin with a ResNet-18 model and short training duration to enable quick iteration, and an SGD optimizer as it has as known scaling rule. This allows us to probe the EMA Scaling Rule without potential confounders like poor gradient-based optimizer scaling<sup>8</sup>.
We observe that without the EMA Scaling Rule, there is a drop in test top-1 linear probe (Definition H.3) performance compared to the baseline, whereas with the EMA Scaling Rule, we closely match the baseline model until batch size 4096. We show that this result is consistent for a range of base learning rates  $\eta_{B}$  and momenta  $\rho_{B}$  in Appendix H.8. At batch size 8192, we see a performance gap between the scaled model using the EMA Scaling Rule and the baseline. We speculate that this is due to dynamics early in the BYOL training process that are challenging to replicate at larger batch sizes. To test, and potentially circumvent this, we introduce Progressive Scaling (Definition 3.2).
Definition 3.2 (Progressive Scaling, informal; see Appendix C.4). Given batch size  $B$  and hyperparameters at  $B$ , slowly increase the batch size to the desired largest batch size during training. At any intermediate batch size  $\hat{B} = \kappa B$ , all hyperparameters are scaled according to their scaling rules.
The BYOL EMA update (Equation 74) uses  $\theta_{t + 1}$  instead of our analyzed  $\theta_t$  (Equation 4). The effect upon the overall EMA update is  $O(\eta \times \beta_{\rho})$  and so is captured by the EMA Scaling Rule (Definition 1.2).
For competitive performance with the reference BYOL (Grill et al., 2020) using a ResNet-50, adaptive optimization, and longer training duration, see Appendix H.10 and Figure 26.
Figure 5: ResNet-18 BYOL on CIFAR10 for different  $\kappa$ . The baseline ( $\kappa = 1$ , black dashed) uses batch size 1024 and momentum  $\rho_{B} = 0.992$ , and is scaled from batch size 2048 (left) to 8192 (third) with (blue,  $\rho = \rho_{B}^{\kappa}$ ) and without (red,  $\rho = \rho_{B}$ ) the EMA Scaling Rule. At  $\kappa = 8$  we also run progressive scaling (right), with transitions at 10 (green) and 30 (orange) epochs. Bands indicate mean and standard deviation across three runs.
Figure 6: BYOL ViT-B/16 on ImageNet1k for different scalings  $\kappa$ . The baseline model ( $\kappa = 1$ , black dashed) uses batch size 4096 and teacher momentum  $\rho_{B} = 0.99$ , and is scaled from batch size 8192 (left) to 32768 (right) with progressive scaling and the EMA Scaling Rule (Definition 3.2) (orange,  $\rho = \rho_{B}^{\kappa}$ ), with the EMA Scaling Rule but without progressive scaling (blue,  $\rho = \rho_{B}^{\kappa}$ ), without the EMA Scaling Rule but with progressive scaling (purple,  $\rho = \rho_{B}$ ), and without either (red,  $\rho = \rho_{B}$ ). Progressive scaling transitions from the reference model at epoch 60. See Appendix H.6 for a discussion on BYOL progressive scaling.
We see that transitioning to the higher batch size during the warmup period results in a model optimization trajectory that diverges from the baseline, whereas transitioning after warmup results in matching final trajectories of the scaled and baseline models. In summary, progressive scaling allows us to match BYOL dynamics at large batch sizes, provided we transition after the warmup period. This observation is consistent with our hypothesis regarding BYOL dynamics during warmup.
Vision Transformers on ImageNet1k Progressive Scaling coupled with the EMA Scaling Rule is required when scaling BYOL ViTs (Figure 6), enabling baseline loss tracking to a batch size of 24,576. Perfect scaling fails at batch size 32,768, consistent with observations in supervised learning (Goyal et al., 2017; Huo et al., 2021). Despite the breakdown, there is only a small drop in  $1.6\%$  probe performance when using the EMA Scaling Rule, compared to as  $44.56\%$  drop without it. We also observe that it is sometimes possible to match test model performance using only Progressive Scaling and not the EMA Scaling Rule, although this still induces a training loss mismatch. We stress that such an approach is not guaranteed to work and discuss when this approach succeeds and fails in Appendix H.6 and Figure 22.
At the transition point between batch sizes, an impulse perturbation<sup>9</sup> is measured at the student, visible from the training loss. This is recovered from by the learning process, and the new model matches the reference batch size. This perturbation happens in both the AdamW and SGD settings, leading us to suspect this is due to the BYOL learning process, rather than an artifact of optimizer or momentum scaling. However, since this is not directly related to the EMA Scaling Rule proposed in this work, we defer this analysis to future investigation.
Optimizer scaling rules from SDEs The SDE perspective has uncovered optimizer scaling rules and allowed an understanding of their limitations. Smith & Le (2018) used SDEs to uncover the SGD Scaling Rule, while (Li et al., 2021) used SDEs to explain that rule's breakdown in terms of discretization error. The SDE analysis was extended to adaptive optimization by (Malladi et al., 2022), producing an Adam Scaling Rule (Definition C.3), indicating that along with the learning rate, the  $\beta_{1,2}$  and  $\epsilon$  parameters transform. The  $\beta_{1,2}$  transformation is consistent with the EMA Scaling Rule in the SDE limit. Our work differs as it considers a model EMA that alters the objective.
Varying the batch size during training Smith et al. (2018) investigated the benefits of scheduling the batch size at a fixed learning rate as an alternative to scheduling the learning rate at a fixed batch size. These two are equivalent through the SGD Scaling Rule. The authors do not scale the optimizer hyperparameters during this procedure, as they are intentionally replicating the training dynamics of a learning rate schedule. This is in contrast with Progressive Scaling (Definition 3.2) which scales the hyperparameters to maintain the optimization process at different levels of discretization.
Large batch training of SSL distillation methods SSL methods learn representations without labels, meaning they can take advantage of web-scale data. Large batch optimization is required to make use of this data in a reasonable amount of time. Grill et al. (2020) demonstrated algorithmic robustness when reducing the batch size through gradient accumulation and EMA update skipping, which implements an approximation of our EMA Scaling Rule for  $\kappa < 1$ . Our work provides a recipe to scale down and up in  $\kappa$ . MoCo-v3 (Chen et al., 2021) enables contrastively distilled ViTs up to a batch size of 6144, where the model drops in performance. More recently, methods like DINO (Caron et al., 2020) present a worse scenario, and are unable to scale beyond batch size 1024 (Koppula et al., 2022). In contrast, our work presents practical tools to scale to large batch sizes in the presence of an EMA, enabling practical training of these SSL methods on large scale data.
We provide an EMA Scaling Rule: when changing the batch size by a factor of  $\kappa$ , exponentiate the momentum of the EMA update to the power of  $\kappa$ . This scaling rule should be applied in addition to optimizer scaling rules (for example, linearly scaling the SGD learning rate), and enables the scaling of methods which rely on EMA and are sensitive to the choice of EMA momentum.
We prove the validity of the EMA Scaling Rule by deriving first-order SDE approximations of discrete model optimization when a model EMA is present and can contribute to the model objective. We demonstrate empirical support for a variety of uses of EMA, ordered by increasing influence of the role of EMA on the optimization procedure: supervised model tracking (i.e. Polyak-Ruppert averaging) in speech and vision domains, pseudo-labeling in speech, and self-supervised image representation learning. In almost all scenarios, using the EMA Scaling Rule enables matching of training dynamics under batch size modification, whereas not using it results in significant differences in optimization trajectories. For example, we can scale the BYOL self-supervised method to a batch size of 24,576 without any performance loss only when using the EMA Scaling Rule.
While learning rate scaling rules are relatively commonplace in ML, the role of EMA has been overlooked. With this work, we highlight the importance of scaling the EMA momentum, and hope that future works will use the EMA Scaling Rule to scale the EMA momentum correctly, in the same way that learning rates and other optimizer hyperparameters are scaled.
Instead of a single large batch transition as in Figure 6 we perform a sequential transition in Appendix H.5. We find that a slow increase in batch size minimizes the magnitude of the perturbation and leads to a final model with higher effective linear probe top-1 than the reference by approximately  $1.17\%$ .
We thank Miguel Sarabia del Castillo, Adam Golinski, Pau Rodriguez Lopez, Skyler Seto, Amitis Shidani, Barry Theobald, Vimal Thilak, Floris Weers, Luca Zappella, and Shaungfei Zhai for their helpful feedback and critical discussions throughout the process of writing this paper; Okan Akalin, Hassan Babaie, Denise Hui, Mubarak Seyed Ibrahim, Li Li, Cindy Liu, Rajat Phull, Evan Samanas, Guillaume Seguin, and the wider Apple infrastructure team for assistance with developing and running scalable, fault tolerant code; and Kaifeng Lyu and Abhishek Panigrahi for discussion and details regarding scaling rules for adaptive optimizers. Names are in alphabetical order by last name within group.
A Broader impact 18
B Limitations 18
C The scaling toolbox: practical methods for enabling systematic scaling 19
C.1 The continuous time/SDE perspective 19
C.2 Scaling rules for optimization 20
C.3 Commonly used values of hyperparameters at different batch sizes 21
C.4 Progressive scaling 21
D EMA approximation theorems with SDEs 22
D.1 SGD with model EMA 22
D.2 Adaptive gradient methods with model EMA 24
E Additional proofs 25
E.1 Iterations of  $\mathrm{SGD} + \mathrm{EMA}$  25
E.2 Limiting behavior of Polyak-Ruppert averaging 26
F Additional details and results for Polyak-Ruppert averaging 27
F.1 Noisy parabola 28
F.2 Image Classification 30
F.3 Applying the EMA Scaling Rule to Batch Normalization 31
G Additional details and results for Automatic Speech Recognition (ASR) 32
G.1 Additional experimental settings and detailed metrics 34
G.2 Scaling to  $\kappa = 16$  with Progressive Scaling 37
H Additional details and results for self-supervised image representation learning 38
H.1 Components of self-supervised learning 38
H.2 A ResNet-18 recipe for BYOL 39
H.3 A Vision Transformer recipe for BYOL 39
H.4 The role of Batch Normalization and Layer Normalization in BYOL with ViTs 40
H.5 Longer training duration with incremental Progressive Scaling 41
H.6 Building intuition around Progressive Scaling and momentum sensitivity 41
H.7 Compute usage for ViT BYOL investigation 42
H.8 ResNet-18 hyperparameter sensitivity analysis 43
H.9 ResNet-18 additional scaling analysis 45
H.10 Scaling a ResNet-50 BYOL using LARS and Progressive Scaling 46
H.11 Preventing collapse phenomena in DINO at scale 47
I Additional details on numerical stability 51
J Contributions 52
This work shows how to adapt Machine Learning (ML) optimization in the presence of a model Exponential Moving Average (EMA). There are a number of benefits to this:
1. Scaling rules democratize the training of ML models: they give ML researchers the ability to replicate the optimization of large scale systems, even if those researchers do not have access to i) significant parallel computational resources or ii) the technical tooling to do so.
2. Our EMA Scaling Rule lowers compute usage as it removes the necessity for a hyperparameter search over momenta; in the case where our scaling assumptions hold, if we know the value of the optimal momentum  $\rho_{B}$  at some batch size  $B$  (for example, the momentum that gives the best transfer performance), then the optimal value at another batch size  $\hat{B}$  is exactly the one given by the EMA Scaling Rule  $\hat{\rho} = \rho_{B}^{\kappa}$ , for scaling  $\kappa = \hat{B} / B$ .
3. Our EMA Scaling Rule enables researchers to more quickly iterate through experimental ideas, and opens up access to large-scale training (for example, larger models and larger datasets) for Pseudo-Labeling and Self-Supervised Learning (SSL) techniques.
These points have potential negative consequences:
1. As our EMA Scaling Rule enables researchers to iterate the same experiments more quickly, and perform large-scale training with EMA-based methods, this may encourage a greater number of experiments, or the training of larger models. Either of these possibilities leads to greater energy consumption.
2. As the need to determine momentum hyperparameters has now been removed, researchers who were previously discouraged from attempting to scale these methods due to an extra hyperparameter to tune may begin to perform such experiments, leading, once more, to greater energy consumption.
The environmental impact of each of these two points may be significant.
The EMA Scaling Rule provides a recipe for producing training dynamics independent of the batch size used in stochastic optimization. The technology underpinning it will not always give the desired behavior, however.
The first issue occurs with the wording present in the EMA Scaling Rule: [...] and scale other optimizers according to their own scaling rules (Definition 1.2):
1. This statement requires that the given Stochastic Differential Equation (SDE) approximation we are using for the model optimizer is itself providing well-behaved scaling, that is, that in the absence of a model EMA, the model optimization trajectories at the batch sizes  $B$  and  $\kappa B$ , with optimizer hyperparameters appropriately scaled, are close. In general we know this is not true. First, we know that the SDE approximation for Stochastic Gradient Descent (SGD) breaks at a given  $\kappa$  due to discretization error (Li et al., 2021). Second, we know that if the gradient noise is not sufficiently large, the SDE approximation for Adam does not exist (Malladi et al., 2022), i.e. an SDE motivated scaling rule has no meaning.
2. This statement requires knowledge of how to scale the corresponding model optimizer. We have principled ways to achieve this for SGD (Li et al., 2021), and for the adaptive optimization methods RMSProp and Adam (Malladi et al., 2022). Empirically, a square-root scaling law for LAMB (You et al., 2020) has been observed, however, it has not been derived formally. Problematically, there is no known hyperparameter scaling law or SDE approximation known for LARS (You et al., 2017), which has been used in Bootstrap Your Own Latent (BYOL) (Grill et al., 2020) and many other large-scale training procedures for convolution-based architectures. Despite this, we are able to demonstrate in Appendix H.10 that a combination of the EMA Scaling Rule and progressive scaling can match, or surpass baseline BYOL performance at a batch size of 32,768 using LARS, indicating that although the theoretical guarantees may not hold, there is still practical utility in the tools we provide in this work.
3. It may be the case that the optimal performance attainable by a given model setup exists at a level of discretization/gradient noise where no SDE exists. In this case, SDE-derived scaling rules can never be valid, and no scaling of this dynamics can be achieved with known tools.
The second issue is related to the case when the optimizer scaling rule is valid. In this case, the error for the EMA Scaling Rule at finite learning rate  $\eta$  at large  $\kappa$  can be considerable. In cases where the model EMA plays a role in the overall optimization, the error introduced by the EMA Scaling Rule can break the preservation of model dynamics.