SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 29: DINO ViT-B/16 on CIFAR-10 for different scalings  $\kappa$  and base teacher momentum  $\rho_{B} = 0.992$ . We show how the hyperparameters  $\rho, B$  and learning rate change with the Progressive Scaling in Definition 3.2. These hyperparameters correspond to the training runs in Figure 28. Those for Figure 27 are identical, with the exception of  $\rho$  that starts at 0.996 instead of 0.992.
In Figure 29 we show how the hyperparameters  $\rho, B$  and learning rate change with the progressive scaling in Definition 3.2.
We also attempted to use a sharp batch size transition (Figures 30 and 31), which leads to the collapse pheonomena observed in prior work. This collapse happens with and without the EMA Scaling Rule. We suspect this is due to dynamics specific to DINO's early phase that are even more challenging to replicate under discretization than those of BYOL.
Figure 30: DINO ViT-B/16 on CIFAR-10 for different scalings  $\kappa$  and base teacher momentum  $\rho_{B} = 0.992$ . The baseline model ( $\kappa = 1$ , black dashed) uses batch size 1024 and center momentum  $\rho_{c} = 0.9$ , and is scaled up from batch size 2048 (left) to 4096 (right) with (blue,  $\rho = \rho_{B}^{\kappa}$ ) and without (red,  $\rho = \rho_{B}$ ) the EMA Scaling Rule. Progressive Scaling is employed with a sharp transition at epoch 100, leading to a collapse phenomenon.
Figure 31: DINO ViT-B/16 on CIFAR-10 with  $\rho_{B} = 0.992$  and a sharp transition in batch size at epoch 100. We show how the hyperparameters  $\rho, B$  and learning rate change with sudden scaling. These hyperparameters correspond to the training runs in Figure 30.
Compute The compute usage for the DINO investigations is detailed in Table 17.
Table 17: Compute usage for DINO investigations. 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.
Batch Size	GPUs	Time (h)	Compute/Run (GPUh)	Runs	Compute (GPUh)
1,024	24	6.8	163.5	2	327.0
1,024 → 2,048	40	4.6	182.4	1	182.4
1,024 → 3,072	48	4.0	189.9	1	189.9
1,024 → 4,096	64	3.3	212.3	1	212.3
1,024 → 2,048 (100 Epochs)	40	4.8	190.6	4	762.3
1,024 → 3,072 (100 Epochs)	48	4.0	192.5	4	769.9
1,024 → 4,096 (100 Epochs)	64	3.6	232.1	4	928.2
All other compute, e.g. code development, runs with errors, and debugging	38239.2
Total	41,611.1
Our results in this section show it is possible to scale DINO to large batch sizes without sacrificing performance by using both the EMA Scaling Rule and Progressive Scaling, providing the batch size schedule of Progressive Scaling is not sudden.
I Additional details on numerical stability
A general analysis of overflow and underflow of the EMA Update (Definition 1.1) or EMA Scaling Rule (Definition 1.2) for different momenta  $\rho$ , particularly for IEE-754 floating point values, is beyond the scope of this work due to non-linearity from mechanisms like gradual underflow (IEEE, 2019).
In our setting, do not suffer from practical overflow or underflow issues through exponentiation when applying the EMA Scaling Rule, as FP32 precision allows a maximum  $\rho = 1 - \varepsilon$ , or minimum  $\rho = \varepsilon$  with  $\varepsilon \approx 1.2 \times 10^{-7}$ . Take self-supervised image representation learning (Section 3.4) as a baseline, with  $\kappa = 1$  corresponding to batch size  $B = 4096$  with momentum  $\rho_{B} = 0.996$ . The maximum value of  $\rho$  corresponds to scaling  $\kappa = \log (\rho_{B}) / \log (\varepsilon) \approx 1 / (32K)$ , give a batch size less than one, while the minimum value of  $\rho$  corresponds to scaling  $\kappa = \log (\rho_{B}) / \log (1 - \varepsilon) \approx 4K$ , giving a batch size  $B \approx 8M$  which is beyond current hardware feasibility, and beyond the breakdown of known optimizer scaling rules (Li et al., 2021).
To examine how momentum may induce numerical errors in practice during training, we train a linear regression model with a Polyak-Ruppert average Definition 3.1, and track the difference between FP32 model weights and weights in i) BF16; ii) FP16; and iii) a second FP32 run, which act as a proxy for overflow and underflow. In Figure 32 we plot these differences using
Figure 32: Numerical precision of target and EMA networks compared to an FP32 reference on a regression task for a range of momenta.
the maximum absolute difference between model parameters, where the maximum is taken over individual weights
$$
\operatorname {M a x A b s D i f f} \left(\mathrm {d t y p e}\right) = \max  _ {i = 1} ^ {P} \left| \theta_ {i} ^ {\mathrm {F P 3 2}} - \theta_ {i} ^ {\mathrm {d t y p e}} \right|, \tag {82}
$$
where  $P$  is the number of parameters in the model. We observe that when model weights and EMA weights are FP16 (never done in practice), an increasing variance happens for FP16 as the value of the momentum  $\rho$  approaches 0.99999, whereas BF16 and FP32 are stable. We stress that all experiments presented in the paper store weights for target model and EMA in FP32 and use automatic-mixed precision to cast them to BF16 during training, and so do not encounter momentum-induced overflow or underflow.
J Contributions
All authors contributed to writing this paper, designing the experiments, discussing results at each stage of the project.
Preliminary work Derivation of the EMA Scaling Rule with learning rate  $\eta = 0$ , initial synthetic and self-supervised ImageNet1k experiments done by Dan Busbridge.
EMA scaling rules for constant gradients Original proof of Equation 4 and the form of  $\delta (\eta ,\rho ,\kappa)$  in Equation 54 done by Eeshan Gunesh Dhekane. Final proof presented in Appendix E.1 done by Dan Busbridge, verified by Eeshan Gunesh Dhekane and Pierre Ablin.
EMA approximation theorems with SDEs Proofs of validity of EMA Scaling Rule in the SDE limit presented in Section 2.2 and Appendix D done by Pierre Ablin.
Polyak-Ruppert averaging in a simple setting Design of noisy parabola setting of Section 3.1 and initial experiments done by Russ Webb. Design of  $\rho^*$  -optimality search (Equation 10), final experiments and analysis of Section 3.1 and Appendix F.1 done by Dan Busbridge.
Polyak-Ruppert averaging on image classification ResNetv2-50 reproduction (Table 4) and baseline momentum identification done by Jason Ramapuram. Final ImageNet1k experiments and analysis of Section 3.2 and Appendices F.2 and F.3 done by Dan Busbridge.
Automatic speech recognition Experiments and analysis of automatic speech recognition using Polyak-Ruppert averaging (Section 3.2) and continuous pseudo-labeling (Section 3.3 and Appendix G), as well as design choice of a seed model to start pseudo-labeling (aligning quality of the seed models for different batch size settings before pseudo-labeling process) done by Tatiana Likhomanenko.
Self-supervised image representation learning BYOL ResNet-18 recipe (Table 9) and experiments on CIFAR10 using SGD (Figure 5), and BYOL ResNet-50 experiments using LARS (Appendix H.10) done by Dan Busbridge. BYOL ResNet 50 baseline implementation and BYOL ViT
recipe (Table 10) done by Jason Ramapuram. BYOL ViT exploratory ablations done by Eeshan Gunesh Dhekane and Jason Ramapuram. All final BYOL ViT experiments and analysis (Figure 6 and Appendices H.4 to H.6) done by Jason Ramapuram. Baseline DINO reproduction done by Dan Busbridge. DINO experiments and analysis (Appendix H.11) done by Xavier Suau Cuadros.
Progressive Scaling Progressive Scaling (Definition 3.2 and Algorithm 1) is proposed by Dan Busbridge based on discussions with Xavier Suau Cuadros, Tatiana Likhomanenko, Jason Ramapuram, Russ Webb, and the authors of Malladi et al. (2022). Adaptation of progressive scaling to semi-supervised learning in automatic speech recognition (Appendix G.2) done by Tatiana Likhomanenko, and to self-supervised learning in vision done by Dan Busbridge and Jason Ramapuram for BYOL (Figures 5 and 6 and Appendices H.4, H.5 and H.10) and Xavier Suau Cuadros for DINO (Appendix H.11).
Limiting behavior of Polyak-Ruppert averaging Original proof of limiting behavior of Polyak-Ruppert averaging done by Eeshan Gunesh Dhekane. Final proof presented in Appendix E.2 done by Dan Busbridge, verified by Eeshan Gunesh Dhekane.
Numerical stability analysis Polyak-Ruppert experiment (Figure 32) using linear regression for various floating point precisions done by Jason Ramapuram.
Implementation details Investigations carried out in two distributed, scalable frameworks: Jax for automatic speech recognition experiments, done by Tatiana Likhomanenko; and PyTorch for all remaining investigations, done by Dan Busbridge, Xavier Suau Cuadros, Eeshan Gunesh Dhekane, Jason Ramapuram and Russ Webb. Initial implementation of progressive scaling experiments for incremental-style strategies (e.g. Appendix H.5) showing feasibility done by Jason Ramapuram, and subsequent progressive scaling implementations for smooth strategies (e.g. Appendices H.10 and H.11) done by Dan Busbridge and Xavier Suau Cuadros.