SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Additional results In Figure 12 we present a more detailed view of the results in Section 3.2. First, we see that for all train metrics, model trajectories match, and that a learning rate step schedule after warmup is present. As discussed in Figure 12, a gap in EMA Test Top-1 trajectories begins at scaling  $\kappa = 4$ , with a more pronounced effect visible at  $\kappa = 8$ . From Figure 12 it is clear that the (non-EMA)
Test Top-1 performance trajectory is no longer matching at these scalings, demonstrating that the problem is not due to a breakdown of the EMA Scaling Rule, but rather, that the model is overfitting at larger batch sizes due to batch normalization (Ioffe & Szegedy, 2015).
Figure 12: ResNetv2-50 Polyak-Ruppert averaging on ImageNet1k for different scalings  $\kappa$ . The baseline model ( $\kappa = 1$ , black dashed) uses batch size 1024 and momentum  $\rho_{B} = 0.9999$ , is scaled down to a batch size of 512 (left), and up to a batch size of 4096 (right) with (blue,  $\rho = \rho_{B}^{\kappa}$ ) and without (red,  $\rho = \rho_{B}$ ) the EMA Scaling Rule (Definition 1.2). Bands indicate the mean and standard deviation across three runs.
F.3 Applying the EMA Scaling Rule to Batch Normalization
In Section 3.2 and Appendix F.2, we investigated a range of scalings  $\kappa$ , with and without applying the EMA Scaling Rule to the Polyak momentum. In those experiments, we maintained Batch Normalization (Ioffe & Szegedy, 2015) coefficients of  $\rho_{\mathrm{BN}} = 0.9$  throughout<sup>11</sup>, i.e. the EMA Scaling Rule is not applied. The running statistics of Batch Normalization are an EMA with values determined by  $\rho_{\mathrm{BN}}$  and so it is reasonable to suspect we should apply the EMA Scaling Rule to  $\rho_{\mathrm{BN}}$  also.
In Figure 13 we investigate the effect of applying the EMA Scaling Rule to Batch Normalization coefficients, using  $\hat{\rho}_{\mathrm{BN}} = \rho_{\mathrm{BN}}^{\kappa}$ . We observe that the Test Top-1 trajectories with the EMA Scaling Rule applied are slightly closer to the reference trajectories for scalings  $\kappa \geq 2$  than those trajectories without the EMA Scaling Rule. As the effect is not particularly large, at least in this setup, we do pursue further abulating applications of the EMA Scaling Rule to batch normalization coefficients, and always use  $\rho_{\mathrm{BN}} = 0.1$  for Batch Normalization, independent of  $\kappa$ .
In many ML frameworks, this value is defined using  $\beta_{\rho} = 1 - \rho$ , i.e. the default is 0.1 and corresponds to  $\beta_{\mathrm{BN}}$  rather than 0.9 corresponding to  $\rho_{\mathrm{BN}}$ . We use  $\rho_{\mathrm{BN}}$  to maintain consistency across this work.
Figure 13: ResNetv2-50 Polyak-Ruppert averaging on ImageNet1k for different scalings  $\kappa$ . The baseline model ( $\kappa = 1$ , black dashed) uses batch size 1024 and momentum  $\rho_{B} = 0.9999$ , is scaled down to a batch size of 512 (left), and up to a batch size of 4096 (right) with the EMA Scaling Rule applied to only model parameters (blue,  $\rho = \rho_{B}^{\kappa}$ ), and model parameters and buffers (orange,  $\rho = \rho_{B}^{\kappa}(\dagger)$ ). Bands indicate the mean and standard deviation across three runs.
G Additional details and results for Automatic Speech Recognition (ASR)
In this section we provide additional details for the speech recognition experiments in both the supervised and semi-supervised case.
Data We use the LibriSpeech dataset (Panayotov et al., 2015) which is a dataset of audio-transcription pairs. For supervised Polyak-Ruppert averaging experiments, we use train-clean-100 as training data, and for semi-supervised pseudo-labeling experiments, we use train-clean-100 as the labeled and train-clean-360 and train-other-500 as the unlabeled data. The standard LibriSpeech validation sets (dev-clean and dev-other) are used to tune all hyperparameters, as well as to select the best models. Test sets (test-clean and test-other) are only used for reporting final model performance, measured in WER without an external language model. We maintain the original  $16\mathrm{kHz}$  sampling rate, and compute log-mel filterbanks with 80 coefficients for a 25ms sliding window, stripped by 10ms, later normalized to zero mean and unit variance for each input sequence.
Acoustic model We employ a vanilla encoder-based only transformer model trained with the Connectionist Temporal Classification (CTC) loss (Graves et al., 2006). We use the training configuration from Likhomanenko et al. (2021a), which has three stages: i) 1D convolutions to perform striding (kernel of 7 with stride of 3); ii) a Transformer encoder with 36 layers, post-LayerNorm, four attention heads, an embedding dimension of 768, an MLP dimension of 3072, a dropout frequency of 0.3, and a layer drop frequency of 0.3; and iii) a linear layer to map to the target vocabulary<sup>12</sup>. To reduce model training time by a factor of approximately  $2 - 3 \times$ , and to reduce memory footprint,
12The token set of this vocabulary consists of the 26 English alphabet letters augmented with the apostrophe and a word boundary token.
Table 6: Hyperparameters summary for speech recognition task for supervised (left) and semi-supervised pseudo-labeling (right) training with a vanilla transformer. The  $0.3 \rightarrow 0.1$  in the dropout and layer drop rates indicates that a rate of 0.3 is used during pre-training, and a rate of 0.1 is used during pseudo-labeling.
Supervised	Pseudo-Labeling
Librispeech test-clean / test-other WER	7.8/19.1	4.8/11.5
Optimizer	Adam	Adam
Optimizer scaling rule	Adam	Adam
Base (β1,β2)	(0.995, 0.999)	(0.995, 0.999)
Base learning rate	0.0001	0.0001
Base learning rate warmup (steps)	64k	64k
Learning rate schedule	Fixed (no decay)	Fixed (no decay)
Learning rate minimum value	0	0
Base training duration (steps)	400k	500k
Base batch size (dynamic)	8 × 290s	8 × 290s
Base teacher momentum	0.99995	0.9999
Weight decay	None	None
Numerical precision	bf16	bf16
Augmentation stack	SpecAug	SpecAug
Dropout	0.3	0.3 → 0.1
Layer drop	0.3	0.3 → 0.1
Gradient clipping	1	1
Labeled:unlabeled data ratio	N/A	1:3
Base pre-training steps	N/A	20k
Base start of EMA accumulation (steps)	N/A	19k
we use CAPE positional embeddings (Likhomanenko et al., 2021b) instead of relative positional embeddings (Shaw et al., 2018): both models perform similarly.
Training Here we discuss our training procedure for base batch size  $B = 8 \times 290s$ , which is adapted from Likhomanenko et al. (2021a), and is summarized in Table 6. We use SpecAugment (Park et al., 2019) activated after 5k steps of training: two frequency masks with frequency mask parameter  $F = 30$ , ten time masks with maximum time-mask ratio  $p = 0.1$  and time mask parameter  $T = 50$  are used; time warping is not used.
One difference in setup is we use the Adam optimizer, whereas Likhomanenko et al. (2021a) used Adagrad (Duchi et al., 2010). Even though Adagrad can be viewed as a particular limit ( $\beta_{1} = 0$  and  $\beta_{2} \to 1$ ) of Adam (Kingma & Ba, 2015), we were unable to produce reasonable optimization in practice when applying the Adam Scaling Rule of Malladi et al. (2022) in this limit. As a consequence, we chose to work with the Adam optimizer, where its scaling rule has been shown to work (Malladi et al., 2022), and we take  $\beta_{1} = 0.995$ ,  $\beta_{2} = 0.999$ , and  $\epsilon = 10^{-8}$ . We obtained similar results for  $\beta_{1} = 0.99$ . Finally, we use a linear learning rate warmup (64k steps) after which the learning rate is kept constant until convergence. This performance can be improved further by using a step decay schedule as shown in prior work. We also apply gradient clipping of 1, and do not use weight decay.
Pseudo-Labeling The pseudo-labeling process comprises of two stages: i) The pre-training phase, where we train model on labeled data for 20k steps with model EMA accumulation starting after 19k steps; and ii) the pseudo-labeling phase, where we involve unlabeled data by generating pseudolabels from the model EMA (teacher) and provide them to the model (student) as if they were ground-truth labels. Pseudo-labels are generated without any dropout applied to the teacher, and no data augmentation is applied for the corresponding inputs. To produce the pseudo-label, we use hard transcription (Definition G.1)
Definition G.1 (Hard Transcription). For a sequence of frames, select the most probable token per frame, removing repetitions and the CTC blank token. For example, "h##eell##ll####oo" is transformed into "hello", where "#" is the CTC blank token.
These hard transcriptions are then used as transcription for student optimization. We use a 1:3 proportion of labeled to unlabeled data as this was found to be optimal in Likhomanenko et al. (2021a), and we decrease model dropout and layer drop rates to 0.1 after pre-training phase. As we have access to the ground-truth labels on the data being treated as unlabeled, we can track
pseudo-label quality by computing pseudo-labels on this data, and compute the WER against their ground-truth. Pseudo-label quality is the primary metric to evaluate progress on unlabeled data, as loss on pseudo-labeled data is unreliable when a teacher model and pseudo-labels are evolving with each time step.
Scaling of batch size Sequential data is typically processed using dynamic batching as it is more computationally efficient than using a fixed number of sequences (Ott et al., 2019). In our work, we use dynamic batching of  $\sim 290$ s audio per GPU. Moreover, for CTC we do not apply any additional sequence normalization. We experimented with fixed batching, but did not observe any significant differences in conclusions compared with the dynamic batching.
We note that dynamic batching is a more challenging setting for achieving systematic scaling, as the number of independent sequences in any given batch may change, and the i.i.d. assumption does not hold at the frame level. Despite these violations of the assumptions of Section 2.2, our results demonstrate that the Adam Scaling Rule (Definition C.3, Malladi et al. (2022)) holds in the case of dynamic batches, as does our EMA Scaling Rule (Definition 1.2).
The base batch size is set to  $B = 8 \times 290s$ , and in our experiments we scale down to batch size of  $B = 2 \times 290s$  and up to batch size of  $B = 128 \times 290s$ . The number of warmup and pre-training steps, steps before SpecAugment is turn on and model EMA is accumulated are scaled according to Appendix C.1.
Compute All experiments are done using A100 80GB 8GPU nodes with bfloat16 precision training. While for supervised training evaluation of different EMA decay values is done in parallel during a single run, for pseudo-labeling every EMA decay value needs separate training. Final models training compute is detailed in Tables 7 and 8. Total compute, including e.g. code development, runs with errors, and debugging, is 61k GPUh.
Table 7: Compute usage for supervised model for speech recognition task in Figure 14. Values include node allocation times (typically a small % of corresponding total runtime), giving a practical estimate of reproduction cost. All experiments conducted are using 80Gb A100s with fast interconnect.
Batch Size	GPUs	Time (h)	Compute/Run (GPUh)	Runs	Compute (GPUh)
2 × 290s	2	222	444	1	222
4 × 290s	4	108	432	1	432
8 × 290s	8	64	512	1	512
16 × 290s	16	54	864	1	896
32 × 290s	32	37	1,184	1	1,184
Total	3,436
Table 8: Compute usage for continuous pseudo-labeling for the speech recognition task in Figure 16. Values include node allocation times (typically a small % of corresponding total runtime), giving a practical estimate of reproduction cost. All experiments conducted are using 80Gb A100s with fast interconnect.
Batch Size	GPUs	Time (h)	Compute/Run (GPUh)	Runs	Compute (GPUh)
2 × 290s	2	225	550	2	1,110
4 × 290s	4	120	480	2	960
8 × 290s	8	72	576	1	576
16 × 290s	16	45	720	2	1,440
32 × 290s	32	33	1,056	4	4,224
64 × 290s	64	25	1,600	2	3,200
Total	11,510
G.1 Additional experimental settings and detailed metrics
We present detailed comparison between models trained with and without EMA Scaling Rule in Figures 14 and 15 for supervised training and in Figures 16 and 17 for semi-supervised training.
First, we observe that if the Adam Scaling Rule does not hold perfectly[13] (there is a mismatch between trajectories for the model before pseudo-labels are involved) the EMA Scaling Rule also
<sup>13</sup>See Malladi et al. (2022) for a discussion on scenarios that lead to a breakdown of the Adam Scaling Rule.
Figure 14: 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 (model non-EMA is marked by orange). The Adam Scaling Rule (Malladi et al. (2022), Definition C.3) is used throughout. For momentum  $\rho_{B} = 0.9999$  we observe similar trajectories for all models.
Figure 15: 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.999$  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 (model non-EMA is marked by orange). The Adam Scaling Rule (Malladi et al. (2022), Definition C.3) is used throughout. If momentum  $\rho_{B}$  is small and accumulation history is short we observe no any significant difference between models which all are matching the reference trajectory despite scaling  $\kappa$ .
gives discrepancies with the reference trajectory, however they are negligible compared to models trained without EMA Scaling Rule. For the semi-supervised training, to alleviate the difficulties with a breakdown of the Adam Scaling Rule for large  $\kappa$  we postpone the pseudo-labeling process until the model reaches similar WER as the baseline. This allows us to align the initial model conditions for pseudo-labeling. In this scenario we are able to match the reference trajectory up to  $\kappa = 8$ .
Figure 16: Transformer pseudo-labeling on LibriSpeech (trained on train-clean-100 as labeled and the rest of LibriSpeech as unlabeled) 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 = 2 \times 290s$  (left) and up to  $B = 64 \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. For  $\kappa \leq 2$ , we start pseudo-labeling after  $20k / \kappa$  training steps; while for  $\kappa > 2$ , we start when pre-training WER matches the baseline WER ( $24k / \kappa$  for  $\kappa = 4$  and  $29k / \kappa$  for  $\kappa = 8$ ). For  $\kappa = 4$  we experimented with both variants: we start pseudo-labeling after  $20k / \kappa$  (dashed) and when pre-training WER matches the baseline WER (solid,  $24k / \kappa$ ).
Figure 17: Transformer pseudo-labeling on LibriSpeech (using train-clean-100 as labeled) 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.999$ . We investigate dynamic batch sizes down to  $B = 2 \times 290s$  (left) and up to  $B = 16 \times 290s$  (right), with (blue,  $\rho = \rho_{B}^{\kappa}$ ) and without (red,  $\rho = \rho_{B}$ ) the EMA Scaling Rule. The Adam Scaling Rule is used throughout.
We note that this result reveals that errors for the Adam Scaling Rule and the EMA Scaling Rule are contributing, although the way in which they contribute is different, and one can dominate the other. We observe in Figure 16 that if the initial conditions of the models are similar (attained by using the same WER as a condition to begin pseudo-labeling) then the error from the EMA Scaling Rule dominates over that of the Adam Scaling Rule, causing a divergence in training dynamics.
Second, we observe in practice that the EMA Scaling Rule holds for both fixed batching (a sequence length in the batch can vary significantly) and for dynamic batching (when total number of frames in the batch is fixed, though padding still is accounted to the this amount). This shows that EMA Scaling Rule is applicable to sequential data too.
Third, we observe in Figures 15 and 17 that for smaller values of  $\rho_{B}$ , scaling with or without EMA Scaling Rule behave similarly, and reference trajectories match in the supervised and semi-supervised cases. However, if the momentum is too large, the teacher moves slowly and is uninformative, whereas if the momentum is too low, the teacher and the student are effectively be the same model, implying: i) the student will self-predict with high confidence, removing any benefits of distillation<sup>14</sup>; and ii) training instability or model divergence will happen in the low-resource settings (Likhomanenko et al., 2021a; Higuchi et al., 2022).
G.2 Scaling to  $\kappa = 16$  with Progressive Scaling
Finally, we aim to scale semi-supervised pseudo-labeling further to  $\kappa = 16$ . In this case we observe that Adam Scaling Rule does not hold in the pre-training phase and there is no model convergence. To overcome this, we apply Progressive Scaling (Definition 3.2). We pre-train models on supervised data with  $\kappa = 8$  for  $29\mathrm{k}$  of reference steps (model EMA accumulation starts at  $28\mathrm{k}$  steps). We then scale to  $\kappa = 16$  and begin pseudo-labeling. We see in Figure 18 that Progressive Scaling enables us to scale pseudo-labeling to  $\kappa = 16$  with (middle) and without (left) the EMA Scaling Rule. Second, models with the EMA Scaling Rule track the baseline much closer than models without the EMA Scaling Rule, although a small gap is present. We further experimented with Progressive Scaling, postponed the transition condition to the  $\kappa = 16$  until  $75\mathrm{k}$  reference steps. In Figure 18 (right), we see this scaled model tracks the reference trajectory, and so using a combination of the EMA Scaling Rule and Progressive Scaling, we are able to scale pseudo-labeling to  $\kappa = 16$ , corresponding to a dynamic batch size of  $128\times 290s$ .
Figure 18: Transformer pseudo-labeling on LibriSpeech (trained on train-clean-100 as labeled and the rest of LibriSpeech as unlabeled) with different Progressive Scaling from  $\kappa = 8$  to  $\kappa = 16$  ( $\kappa = 8 \rightarrow 16$ ). 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$ . The scaling with  $\kappa = 8$  is shown with lighter color for reference from Figure 16. We investigate dynamic batch sizes progressively from  $B = 64 \times 290$ s to  $B = 128 \times 290$ s, with (blue,  $\rho = \rho_{B}^{\kappa}$ ) and without (red,  $\rho = \rho_{B}$ ) the EMA Scaling Rule. For reference (top) we show the learning rate schedule with Progressive Scaling. The Adam Scaling Rule (Malladi et al. (2022), Definition C.3) is used throughout. Left and middle correspond to Progressive Scaling with scale from  $\kappa = 8$  to  $\kappa = 16$  at 29k steps, while right corresponds to 75k steps.
H Additional details and results for self-supervised image representation learning
Organization This appendix is structured into three sections. We first give an overview of our chosen SSL method BYOL (Appendix H.1), our recipe for training BYOL using ResNet 18s (Appendix H.2), our recipe for training BYOL using Vision Transformers (ViTs) (Appendix H.3), ablations of normalization approaches that lead to the development of this recipe (Appendix H.4), and additional results corresponding to longer training duration (Appendix H.5) and further understanding the impact of Progressive Scaling (Appendix H.6).
Second, we demonstrate that the EMA Scaling Rule combined with Progressive Scaling can scale a ResNet-50 BYOL model trained with LARS to batch size 32,768 without performance drop, demonstrating the empirical utility of the tools we provide outside of their theoretical validity (Appendix H.10).
Finally, we show that it is possible to systematically scale DINO (Caron et al., 2021) using a combination of Progressive Scaling and the EMA Scaling Rule, providing a solution for researchers and practitioners wanting to train DINO at scale.
H.1 Components of self-supervised learning
First, a key component of many SSL methods is the stop-gradient or StopGrad (Definition H.1).
Definition H.1 (Stop Gradient/StopGrad(\cdot)). The stop-gradient operator StopGrad(\cdot) prevents the flow of gradient information
$$
\frac {d f (\operatorname {S t o p G r a d} (h (x ; \omega)) ; \theta)}{d \omega} \equiv 0 \tag {72}
$$
for all parametric functions  $h$  and  $f$  and for all parameters  $\theta$  and  $\omega$ .
Applying a stop-gradient is sometimes called detaching in the literature. Now, we introduce the update rule of our representative SSL method BYOL in Definition H.2.
Definition H.2 (BYOL Update). BYOL learns unsupervised features by minimizing the cosine distance between the predictions of a student backbone  $f(\cdot ;\theta)$  (typically a ResNet or Vision Transformer), projected through  $h(\cdot ;\omega)$  (typically a Multi-Layer Perceptron (MLP)), and the predictions of an EMA teacher  $f(\cdot ;\zeta)$  (Grill et al., 2020). The update for the parameters of BYOL is then
$$
\left(\boldsymbol {\theta} _ {t + 1}, \boldsymbol {\omega} _ {t + 1}\right) = \left(\boldsymbol {\theta} _ {t}, \boldsymbol {\omega} _ {t}\right) - \eta \times \frac {1}{B} \sum_ {x \in \mathbb {B}} \nabla_ {\left(\boldsymbol {\theta}, \boldsymbol {\omega}\right)} \mathcal {L} \left(x; \boldsymbol {\theta} _ {t}, \boldsymbol {\omega} _ {t}, \zeta_ {t}\right) \tag {73}
$$
$$
\zeta_ {t + 1} = \rho \zeta_ {t} + (1 - \rho) \theta_ {t + 1} \tag {74}
$$
$$
w i t h \mathcal {L} \left(x; \boldsymbol {\theta} _ {t}, \omega_ {t}, \zeta_ {t}\right) = \frac {1}{2} \cos \left[ h \left(f \left(x _ {1}; \boldsymbol {\theta} _ {t}\right); \omega_ {t}\right), S t o p G r a d \left(f \left(x _ {2}; \zeta_ {t}\right)\right) \right] + \left(x _ {1} \leftrightarrow x _ {2}\right), \tag {75}
$$
where  $\cos (\pmb {a},\pmb {b})\equiv 1 - \pmb {a}\cdot \pmb {b} / (||\pmb {a}||||\pmb {b}||)$  is the cosine distance, and  $x_{1}$  and  $x_{2}$  are two views of a single.   
variate  $x$  , often produced by augmentations, and  $x_{1}\leftrightarrow x_{2}$  denotes symmetrization over  $x_{1}$  and  $x_{2}$
As noted in Section 3.4, 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).
One more piece of technology typically employed in SSL is a tracking probe (Definition H.3) which we will use to evaluate the performance of BYOL on downstream tasks of interest, for example, image classification.