How to Scale Your EMA
Authors: Dan Busbridge
Jason Ramapuram
Pierre Ablin
Tatiana Likhomanenko
Eeshan Gunesh Dhekane
Xavier Suau
Russell Webb
Venue: NeurIPS 2023 spotlight (2023)
Keywords: Optimization
scaling rules
EMA
exponential moving average
self-supervised learning
pseudo-labelling
semi-supervised learning
BYOL
distillation
speech
vision
TL;DR: We derive and verify a scaling law for model exponential moving averages, allowing the systematic scaling of model classes depending on EMA, like distillation self-supervised methods.

Abstract: 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 machine learning tool is the model EMA, a functional copy of a target model, whose parameters move towards those of its target model according to an Exponential Moving Average (EMA) at a rate parameterized by a momentum hyperparameter. This model EMA can improve the robustness and generalization of supervised learning, stabilize pseudo-labeling, and provide a learning signal for Self-Supervised Learning (SSL). Prior works have not considered the optimization of the model EMA when performing scaling, 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 a model EMA and demonstrate the rule's 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, a 6$\times$ wall-clock time reduction under idealized hardware settings.

How to Scale Your EMA
Dan Busbridge"
Jason Ramapuram"
Pierre Ablin"
Tatiana Likhomanenko
Eeshan Gunesh Dhekane
Xavier Suau
Russ Webb
Apple
(dbusbridge, jramapuram, p_ablin, antares, eeshan, xsuaucuadros, rwebb) eapple. com
Abstract
Preserving training dynamics across batch sizes is 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 scaling rule, for example, in stochas- tic gradient descent, one should scale the learning rate linearly with the batch size. Another important machine learning tool is the model EMA, functional copy of target model, whose parameters move towards those of its target model according an Exponential Moving Average (EMA) at rate parameterized by a momentum hyperparameter. This model EMA can improve the robustness and generalization supervised learning. stabilize pseudo-labeling, and provide a learning signal for Self- Self-Supervised Learning (SSL). Prior works have not considered the optimization of the model EMA when performing scaling, leading different training dynam- ics across batch sizes and lower model performance. In this work, we provide a scaling rule for optimization in the presence of a model EMA and demonstrate the rule's validity across range of architectures, optimizers, and data modali- ties. We also show the rule's s validity where the model EMA contributes the optimization of the target model, enabling us train EMA-based pscudo-labeling and SSL methods at small and large batch sizes. For SSL, we enable training BYOL up to batch size 24,576 without sacrificing performance, a 6x wall-clock time reduction under idealized hardware settings.
Introduction
With data and models becoming progressively larger (Chen et al., 2020; Kaplan et al., 2020; Bom- masani et al., 2021; Srivastava al., 2022), the ability reduce training wall-clock time is a re- quirement for practical Machine Learning (ML) at scale. Optimizer scaling rules allow us to find faster learning procedures that produce similar results. For example, the linear sealing rule for Stochastic Gradient Descent (SGD) (Krizhevsky, 2014; Goyal et al., 2017), states that the learning rate should be scaled linearly with the batch size. This optimizer scaling works both ways. Access to larger computational resources means one can train equivalent models in reduced wall-clock time. Alternatively, with access to limited computational resources, larger distributed computations can be replicated at increased wall-clock time.
Primary contributor. For detailed breakdown author contributions see Appendix
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Many ML algorithms rely on model EMA, a functional copy of target model, whose parame- ters move towards those of its target model according to an Exponential Moving Average (EMA) (Definition 1.1) at a rate parameterized by a momentum hyperparameter p.
Definition 1.1 (EMA Update). The EMA update for the model EMA parameters 3e following target model parameters O, at iteration with momentum p = 1 Pp is
S = p + (1 p) p) 0, = (1 -1- Pp) 3, + pp 8
(I)
The model EMA has a number of desirable properties: i) the model EMA inhabits wider minima than the target model, reducing overfitting and improving generalization (Ruppert, 1988; Polyak & Juditsky, 1992; Huang et al.. 2017; Izmailov et al., 2018; He et al., 2022); ii) compared to the tar- get model, the model EMA moves slowly, making it useful as stabilizer for networks governing Bellman updates in reinforcement learning, (Lillicrap et al., 2016); and iii) the model EMA is rela- tively cheap to compute, whilst providing valid model but different to the target model. This third property has made the model EMA a common choice for the teacher in many distillation setups, from semi-supervised learning (Tarvainen & Valpola, 2017; Sohn et al., 2020; Manohar al., 2021; Higuchi et al., 2022), to Self-Super Learning (SSL) methods like Bootstrap Your Own Latent (BYOL) (Grill et al., 2020), DINO (Caron et al., 2021), and data2vee (Baevski et al., 2022b,a).
Despite its significant role in optimization, a recipe for adapting the EMA Update (Definition 1.1) when changing batch size has, to the best of our knowledge, been absent. To address this, we derive an EMA Scaling Rule (Definition 1. 1.2) which states how the EMA momentum p hyperparameter should be modified
Definition 1.2 (EMA Scaling Rule). When computing the EMA update (Definition 1.1) of a model undergoing stochastic optimization with batch size B = kB, use momentum p = p" and scale other optimizers according to their own scaling rules.
In Definition 1.2, the momentum p. which defined batch size B. typically corresponds "good hyperparameter choice", although this does not need to be the case in general. In this paper, we make the following contributions.
1. With the assumptions of Goyal et al. (2017), we derive an EMA Scaling Rule: the EMA update momentum should be scaled exponentially with the batch size (Definition 1.2).
2. To validate this EMA Scaling Rule theoretically, we propose Stochastic Differential Equation (SDE) approximations of optimization in the presence of a model EMA (Section 2.2). This model EMA contributes to the loss, covering cemi-supervieed learning and SSL. We prove that these approximations are first order weak approximations, and that our EMA Scaling Rule is correct in the SDE limit under realistic gradient assumptions (Corollary 2.1.1).
3. We empirically validate the EMA Scaling Rule in synthetic settings (Section 3.1) and real- world settings where the model EMA plays an increasingly significant role in optimization: i) where the model EMA is used during inference instead of the target model (Section 3.2); ii) pseudo-labeling. where the model EMA (teacher) follows the target model (student), and the student is optimized on mixture of a) labeled data and b) data without labels, whose pseudo- labels are produced by the teacher (Section 3.3); and iii) self-supervised learning, which is the same as the semi-supervised case, except there is no labeled data (Section 3.4).
4. We observe that pseudo-labeling and SSL training dynamies during optimizer warm-up uurm-up are not always able to be replicated at large batch sizes using only the EMA Scaling Rule. We propose and verify practical methods to overcome this limitation, enabling us to scale to a batch size of 24.576 with BYL Vision Transformers (ViTs), reducing wall-clock training by 6x under idealized hardware scenarios while maintaining performance of the batch size 4096 baseline.
Finally, to aid practitioners looking to scale, in Appendix C we provide a Scaling Toolbox, which gives practical advice on how to scale systematically, collecting known scaling rules, and explaining how to think about the SDE perspective of optimization.
The target model usually undergoes gradient-hased optimization, but this does not have he the case. 3 We stress that the study of momentum in gradient-basee optimizers is not the focus of this work. We refer to Smith & Le (2018); Li al. (2019) for discussion scaling rules for these methods.

---

2 The EMA Scaling Rule
We begin with an informal discussion of scaling rules and motivate the existence of an exponential scaling rule for the momentum parameter controlling the update of the model EMA.
2.1 Background and an informal discussion of scaling rules
Consider model with parameters 0, at iteration r updated with SGD (Definition 2.1). Definition 2.1 (SGD Update). The SGD update for model with parameters 9, at iteration t given a minibatch {x(h) ~ Px b- 1,2, ........, _B} of B |B] samples with learning rate n is
0+1 = 0, - 17 x Σ VxL(x;0,,,
where L is the loss function. Vx 8,) is the parameter gradient for the sample x at iteration t, and the B are Independent and Identically Distributed (t.i..)) from Px.
Iterating over sequence of independent minibatches Bo, B,.... Bx-1 produces model parameters
1
01+x 0, -,7 B
VoL(x;
(3)
xEB,
If gradients vary slowly Vol((x;; v VoL(x; (x O.... k- one SGD step with кス batch B U,B, of size B xB results in 0.+1 = 0.4:: yielding the SGD Scaling Rule (Definition 2.2).
Definition 2.2 (SGD Scaling Rule). When running SGD (Definition 2.1) with batch size B kB, use a learning rate i = kn (Krizhevsky, 2014; Goyal et al., 2017).
For clarity in this work, we adopt the naming convention [Algorithm Name] Scaling Rule, which means all parameters of those algorithms are appropriately scaled from batch size B to xB.
As discussed in Goyal et al. (2017), although the assumption of slowly changing gradients is strong, if it is true, then 0,+ 0.41 only if ktj. The validity of the SGD Scaling Rule has been formally studied. In particular, there was ambiguity regarding whether the scaling should be a square-root or linear (Krizhevsky, 2014). SDE approaches have resolved this ambiguity, and have been used to estimate the scaling when the SGD Scaling Rule is no longer guaranteed to hold (Li et al., 2021) 1).
To address model parameter EMAs, first restate the EMA Update (Definition 1.1).
Definition 1.1 (EMA Update). The EMA update for the model EMA parameters &, following target model parameters O, at iteration with momentum = 1 Bp is
S = p3, + (1- 0, = (1 Bp) 2, + 0,-
(I)
The model EMA parameters I do not typically receive gradient information, we take the convention that p is close to one, and the Bp subscript will be omitted where it is clear from the context.
Assuming again that gradients change slowly VxL(xx; 0r+j, Sra,) VeL 0, 3,) g. for gradient iterating over x independent minibatches produces model states (see Appendix E. for derivation)
1
0
0, 3,
- kg 0, 0 O (n Bp)
=
(1
p)
p
=
pm 3, + (1- (1-p")
(4)
g
0
0
g
g
The first row is the SGD Scaling Rule (Definition 2.2). The third row implements the slowly chang- ing gradients assumption for the first row. The second row is equivalent to a single EMA up- date (Definition 1.1) with momentum ò = p": we can take single SGD update with batch size B- xB and learning rate n kr,, and a single EMA update with momentum p², and we get (Ô,+1, 3(m1) (0 Sexx) up to terms O11 Bp). This yields the EMA Scaling Rule (Definition 1.2).
3

---

Definition 1.2 (EMA Scaling Rule). When computing the EMA update (Definition 1.1) of a model undergoing stochastic optimization with batch size B = kB, use momentum P = px and scale other optimizers according to their own scaling rules.
The EMA Sealing Rule was derived for SGI,, and is extended to other optimizers in the following way. An optimizer scaling rule ensures 0/+1 = Oz-xe satisfying identification for the first row. Next, the zeroth order term in 17 x Bp in the second row in Equation 4 is cedimizer-independee therefore unchanged. Finally, the first order terms in x fp in the second row, corresponding to the scaling rule error, are an EMA accumulation of target model updates under optimization, which is still (n x Bp), although its functional form may be different for different optimizers.
and
The above discussion is intended to give intuition for why the EMA momentum should be scaled exponentially, As we have used the same slow-moving gradient assumption as the original SGD Scaling Rule, this may cast doubt on whether our rule correct. To remove this ambiguity, we will follow Smith & Le (2018); Li et al. (2021); Malladi et al. (2022), and show that the EMA Scaling Rule (Definition 1.2) is correct in the SDE limit under more realistic gradient assumptions.
2.2 The EMA Scaling Rule through the lens of stochastic differential equations
SDEs are a tool typically used to obtain scaling rules from first principles (Li et al., 2021; Malladi et al., 2022). In the following, we use SDEs to obtain strong theoretical guarantees for the EMA Scaling Rule found in Section 2. We consider the following discrete dynamics for EMA:
x++1 0x Bk- with Ba = ff (n) + ogg and E ~ -, Sn),
(5)
Sx+1 Sx (1-p) p) Oz.
where o > is the noise scale, (0,, (x) is the gradient noise distribution, assumed to be zero-mean and variance Σ(Ox. (x) independent of o, and (f (x) (of 5t). We posit a dependency of the loss f on the EMA& ī in order to cover semi-supervised (Section 3.3) and SSL (Section 3.4). The case of Tolyakk-Ruppen averaging (Section 3.2), is covered by letting f be independent of Ç.
We aim to obtain an SDE approximation of Equation 5 as 17 goes to zero. The scaling rule for iterations of 0 is well known (Li et al., 2021): we let σ = - The analysis of Section 2.1 gives the scaling rule i = nę and و = p* Linearizing this rule near r = gives ý = 1 - x (1 (1-p), which is a linear relationship between 1 p and 17. We therefore let fo = (1 p)/n and consider the SDE
d0, f-V Z.) dt oa Z.)+ dW, with W. Wiener process,
(6)
dZ, = dz)
where 9, and 2, are SDE variables relating to model and EMA parameters respectively. The SDE in Equation 6 approximates the discrete iterations of Equation 5 when the learning rate 7) goes to zero. One way to see this is that an Euler-Maruyama discretization of the SDE with learning rate exactly recovers the discrete iterations. More formally, we have Theorem 2.1, which is in the same spirit as those found in Li et al. (2021); Malladi et al. (2022). In the theorem, G is the set of functions with derivatives up order a that have at most polynomial growth (see Definition D.I).
Theorem 2.1 (SDE for SGD + EMA; informal see Theorem D.1), Assume that f is continuously
differentiable, with f E G². Let 0,, be solutions of Equation 6, and 0x, Se iterationn of Equation
with Σ E G². Then, for any time horizon > and function g E G², there exists a constant e > 0 independent of n such that
[ Znk)] ([g 50]1 S cc n7 max k= [1)n]
Theorem 2.1 formalizes the intuition that the SDE is an accurate approximation of the discrete iterations. In turn, it allows validating the scaling rule in the same spirit as in Malladi et al. (2022). Corollary 2.1.1 (Validity of the EMA Scaling Rule). Assume that f is continuously differentiable, with f f G and E E G2. Let "55",((" be iterations of Equation 5 with batch size and hyperpa- rameters 7, p. Let ((KB) x(xB) be iterates with batch size kB, and n determined by the SGD Scaling Rule (Definition 2.2) and p determined by the EMA Scaling Rule ( (Defmition 1.2). Then, for any time horizon T > 0 and function g E G?, there exists a constant c > 0 independent of n such that
[B[g(((KB)
(KB) [KjK] <(E9))1 ) Blg(0 (B) ..5 _(8)) I1 S cx
max
k=
T/nl
4

---

Table I: The role of the model EMA in the optimization of (0, 5) given target model 0 for different tech- niques, ordered by increasing influence of the EMA model. All statements assume a momentum s p and that the target model 0 is subject to stochastic optimization at a batch size B.
TECHNIQUE
ROLE OF MODEL. EMA
SUPPPRR AV-
O undergoes optimization and is tracked by 5. which does not affect 0. is an estimate of 0 with time horizon and variance determined by B and p.
ERAGING, SEC. 3.2
CONTINUOUS PSEUDO-LABELING, SEC. 3.3
Averaging. After Pre- Pre-Training 3 (teacher) produces largets for 0 (student) from unlabeled data, which is com- bined with labeled data. The optimization endpoint dependent on and p.
Pre-Trainngg is as above in
SELF-SUPERVISED
As above in After Pre-Training, except there is labeled data. The optimiza- tion endpoint is dependent on and p.
LEARNING, SEC. 3.4
Corollary 2.1.1 shows that two trajectories with different batch sizes are close in the limit of small learning rate, demonstrating the validity of Definition 1.2. A natural follow-up question is what happens when an adaptive optimizer is used instead of SGD? Malladi et al. (2022) study this without an EMA and characterize how hyperparameters change with the noise scale. In particular, they show that under a high gradient noise hypothesis, there exists a limiting SDE. In Appendix D, we derive the limiting SDEs for RMSProp and Adam with an EMA. Although a formal proof of closeness between the iterations and these SDEs is beyond the scope of this work, these SDEs indicate that the EMA Scaling Rule holds for adaptive algorithms. We demonstrate this empirically in Section 3.
Experiments
Now that we have derived and shown the validity of the EMA Scaling Rule, we verify empirically. The experiments validate the EMA Scaling Rule for variety of uses of EMA and are ordered by
increasing influence of the role of EMA on the optimization procedure (see Table 1). The baseline in all of our experiments is without the EMA Scaling Rule, which applies all known relevant scaling rules except the EMA Scaling Rule, and represents previous best practice.
3.1 Polyak-Ruppert averaging in a simple setting
At inference, it is typical to use a model EMA, known as
Averaging (Definition 3.1).
Definition 3.1 Dolyak--uppert Average). When optimizing model parameters 0, compute their EMA 3 (Definition 1.1). Use instead of 0 at inference (Polyak & Juditsky, 1992; Ruppert, 1988).
We begin by showing the EMA Scaling Rule is required to match parameter trajectories in a simple setting. Consider the optimization of 0 in a noisy parabola whose loss L(0) is parameterized by coefficients for curvature a > 0, scaled additive noise b 2 0, and additive noise c 2 0:
L(() 0²
0k+1 0g- 1) gg
g a 0g + EK.
EK N (0.
The scaling factor x in the covariance denominator implements gradient noise reduction as scaling (i.e. batch size) increases (Jastrzebski et al., 2017). Let 0 E R be optimized with SGD (Defini- tion 2.1) and E R be a Ruprert average (Definition 3.1) for with momentum p = B. At scaling x 1, we use fR= // 10 and An 10t iterations, to yield total time T = fB x To keep gradients ((1) and gradient noise non-negligible, we take a = 1, b b= 0.5, and c = 0.
First, observe the effect of scaling on single run (Figure la) by tracking the position of the model EMA. We see that at scaling x 8 or к = 256, the runs using the EMA Scaling Rule match the bascline trajectory, whereas the runs using the baseline momentum do not, with a greater deviation induced by greater scaling x. Even at x = 8, there is a significant difference between scaled and unscaled trajectories, despite the seemingly small numerical difference of their momenta".
Second, we consider whether the EMA Scaling Rule is optimal. To do this, inspired by the SDE analysis (Section 2.2), we define the approximation error, Err(p, K, g, of test function y for a given
Momentum enters optimization exponentially; small changes can lead to very different updates.
5

---

Scaling == 8
Scaling X 256
EMA Position 4
1.0
Approximation Error
10%
Momentum
10-1
0.8
EMA.
1
= Pa
Pe
0.90
PMA.p=P
10
0.6
8
Pa
256,
256,
P
EMA, p Pj
Modal
0.85
0.0
0.5
0.0
0.5
1.0
16 256 1024
16 256 1024
Continuous Time
Continuous Time
Scaling Factor
Scaling Factor
(a) Trajectory of the model EMA 3 under different scalings with 1 PB 7B= = 10-
(b) Choices for momentum (left) with corresponding ap-
proximation errors (Equation 10) (right).
Figure I: (a) We show the effect of scaling by comparing model EMA trajectories of the baseline (K 1, black dashed) x = 8 (left) and x == 256 (right), with (p = P.. blue) and without (p PB. red) the EMA Scaling Rule. (b, left) The momentum according for different scaling rules and the empirically optimal o² (Equation 10), (b, right) The approximation error (Bquation 10) of trajectorics in (b, left) and the target model (orange).
EMA Test Top-1 (%)
Scaling 1/2
Scaling ==
Scaling
Scaling
75
-
512.p
2048, P PG
4096,
8192 p Pg
512,
2018,
1095,
8192.
PM
1024, Ph
0
15
30
45
60
75
90
0
15
30
45
60
75
90
0
15
30
45
60
75
15
30
45
60
75
Train Epochs
Train Epochs
Train Epuchs
Train Epochs
Figure 2: RexNetv2-50 Polyak- Mlppe-Ruppert averaging on ImageNetlk for different scalings The baseline model
(x = 1, black dashed) uses batch size 1024 and momentum pB = 0.9999, scaled down a batch size of 512 (left), and up to batch size of 4096 (right) with (blue, P = p and without (red. = PB) the EMA Scaling Rule (Definition 1.2). Bands indicate the mean and standard deviation across three runs.
scaling k using momentum p. and the value of the momentum p (k,g) that minimizes this error:
BB((t) - Eg(c(( (52)) k=0....Tijn
p' (x. g) = arg min Eri(p, K. g).
Enr(p,x,g) = max
(10)
P
For scalings x E (1. 4,.., 1024), we determine the optimal momentum p and compare it to the EMA Scaling Rule (Figure 1b, left). The scaling rule tracks the p" until K = 256, when the pa become systematically higher. We see target model eтror increase at = 256 (Figure 1b, right). As the target model error is EMA-independent, this indicates that the SGD Scaling Rule is breaking. At the lower scaling k = 64, there is an inflection point in the EMA Scaling Rule approximation error, before the model error grows. This difference indicates the O Bp) terms of Equation 4 are beginning to influence the EMA update. Finally, these observations are true in D = 100 dimensions, (Appendix F.I), and we stress that not changing the momentum at every scaling x induces large approximation error, indicating there is merit to using the EMA Scaling Rule.
3.2 Supervised learning real data with Polyak-Rupper averaging
We now tum to real-world classification where the target model optimizes a parametric log- 0) with inputs and labels (x, y) drawn from a joint distribution p(y, x).
likelihood max, log pog (y|x;
Image Classification We consider a variant of the original SGD Scaling Rule result (Goyal al., 2017) and train a ResNetv2 (He et al., 201 201 on ImageNet1 k (Russakovsky et al., 2014) (Figure 2) using three step learning rate schedule. The base momentum PB 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 (loffe & Szegedy, 2015) statistics". We observe that without the EMA Scaling Rule, there is significant drop in model EMA test performance, whereas with the EMA Scaling Rule, we
$ 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 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 к (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 x 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, trajecto- ries match, 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 con- vergence is due to the longer training time compared to the EMA horizon discussed in Table (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 necessary 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 growing gap between the baseline and EMA-scaled trajec- tories as we increase K. Inspecting the train loss and non-EMA test WER, which do not depend on the EMA update (see Figure 14, Appendix G.I), indicates this due a breakdown of the Adam Scaling Rule. In summary, evaluation ASR shows that the EMA Scaling Rule holds in practice for sequential data with dynamic batch sizes, as well as when using adaptive optimization.
EMA test- other WER
Scaling K == 1/4
Scaling K = 1/2
Scaling K 2
Scaling K = 4
100
29050. -Po
3 290r. Ph
8 2900, - Ph
2x 200.
4x 200
16 2905,
2900,
B 2x 20xx200.p-1 -
4x 24x.9 P8
- 290,, Pb
D- 290.0 -
50
2 Train Steps
2 Train Steps
2 Train Steps
2 Train Steps
x10
x10
x x10
100
Figure 3: Transformer Polyak-Ruppert averaging on LibriSpeeck (trained on train-clean-100) with different
scalings x. The baseline (x = 1, black dashed) is trained with Adam and momentum PB 0.99995 at dynamic batch size B = 8 x 290s, which corresponds to a single train step on the x-axis. We investigate dynamic batch sizes down to B = x 290s (left) and up to B = 32 x 290s (right). with (blue, p p$) and without (red, p PB) 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- -upersueee 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 in- volved in the overall optimization (Likhomanenko al., 2021a, la, 2022; Manohar et al., 2021; Higuchi et al., 2022). We first pre-train target model (student) on a limited labeled set for short period (e.g- 20k steps of B = 8 x 290s²). Concurrently, the student updates a model EMA (teucher). After pre- training, 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 ground-trui truth transcriptions, 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 scaled batch sizes. When the EMA Scaling Rule is not used, Figure reveals a significant difference in PL quality trajectory, leading to a higher test WER.
For x > 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 PI. quality at the start of pseudo- labeling (20k steps of B x 290s), which affects the training dynamics (Berrebbi et al., 2023). To
NNot that number of steps is batch size dependent and should be scaled by 1/x (see Appendix C).
7

---

Scaling k= 1/2
Scaling k=
Scaling k = 4
Scaling X =
test-other oot-otthe WER PL WER
50-
100
290 Pa
8x 290,
4= 190r
8 16 200, Po
A= 32x 200
8= 64 200k
50
290,. PM
16x 290.,
250b, PE
64% 290k,
2 Train Steps
0
0
2 Train Steps
x <10
Train Steps
10
Train Steps
10
x 110
Figure 4: Transformer sreeeing on LibriSpeech with different scalings The baseline (k 1. black dashed) is trained with Adam a dynamic batch size of 8x290 seconds, which corresponds a single train step on the x-axis. The model EMA (teacher) updated with momentum Pa = 0.9999. We investigate dynamic batch sizes down to B = 4 x 290s (left) and up to B = 64 x 290s (right), with (blue, p = P5) and without (red, p PB) the EMA Scaling Rule. The Adam Scaling Rule (Malladi et al. (2022), Definition C.3) is used throughout. For pre-training WER matches the baseline WER.
start pseudo-labeling after 20k/x training steps: while for 2, start when
alleviate the Adam Scaling Rule mismatch effect for k > 2, we postpone the pseudo-labeling until pre-training on labeled data gives similar validation WER, see Appendix G. With this heuristic, can match the baseline trajectory with the EMA Scaling Rule up to K = 8 (Figure 4).
In summary, (a) model EMA affects the optimization process of psendo-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- tre-traannng 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 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)² for investigation into scaling 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 al., 2020; Koppula et al., 2022). Since BYO learns through self-referential distillation, momentum plays a significant role in optimization. We analyze: i) ResNet- 18 (He et al., 20) CIFARIO (Krizhevsky et al., 2014) (Figure 5) using SGD (Definition 2. 2.1); and ii) a ViT-B/16 (Dosovitskiy et al., 2021) on ImageNet using Adam (Loshchilov & Hutter, 2019). A recipe for BYOL using ViTs is provided in Appendix H.3.
ResNet-18 on R-FAR-10 We begin with a ResNet-18 18 model and short training duration to enable quick iteration, and an SGD optimizer it has as known scaling rule. This allows to probe the EMA Scaling Rule without potential confounders like poor gradient- based optimizer scaling".
We observe that without the EMA Scaling Rule, there is a drop in test top-1 linear probe (Defini- tion 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 range of base learning rates !/B and momenta PB in Appendix H.8. 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 hyperpa- rameters at B, slowly increase the batch size to the desired largest batch size during training. At any intermediate batch size B xB, all hyperparameters are scaled according to their sealing rules.
"The BYOL FMA update (Equation 74) uses 0,+1 instead of our analyzed i (Equation 4). The effect upon the overall EMA update is O x fp) and so is captured by the EMA Scaling Rule (Definition 1.2). FFor competitive performance with the reference BYOL (Grill et al., 2020) using a ResNet-50, adaptive optimization, and longer training duration, Appendix H. and Figure 26.
8

---

Scaling
Scaling 4
Scaling 8
Progressive Scaling
Test Top-1 (56
40
2048, p P
4046 P PE
8 8192, P P5
P =
0 2018, p Pa
D 4196.
B 8192.
0 31920 p
8 1024.
B 1024,
B 1024,
1024.
0.7
Train Loss
0.6
0.5
0.3
75
Train Epochs
Train Epochs
Train Epochs
Train Epochs
Figure 5: ResNet-18 BYOL on CIFARIO for different k. The haseline (x = 1, black dashed) uses batch size 1024 and momentum PB 0.992, and is scaled from batch size 2048 (left) to 8192 (third) with (blue, f and without (red. & PB) the EMA Scaling Rule. At K 8 also run progressive scaling (right), with transitions 10 (green) and (orange) epochs. Bands indicate mean and standard deviation across three runs.
Progressive Scaling K 2 Progressive Scaling k= 4 Progressive Scaling K x=6 Progressive Scaling K
Test Top-1 (%)
23
0-
D=81 p PF
18384, 18384, Pu
P5
32768, p -P
B=8192, p Pa
24576, /
Ph
32768. p Pe
4096. p Pa
B 4096,
24576 P PS
B 9096, Pe
10a
24576. Ph
Train Loss
= 4096.
10"1
200
2000
Train Epochs
Train Epochs
Train Epochs
Train Epochs
Figure 6: BYOL ViT-B/16 on ImageNetlk for different scalings x. The baseline model (x = 1, black dashed) uses batch size 4096 and teacher momentum р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, p PH), with the EMA Scaling Rule but without progressive scaling (blue, P$). without the EMA Scaling Rule but with progressive scaling (purple. p PB), and without either (red, p PB). Progressive scaling transitions from the reference model at epoch 60. See Appendix H.6 for discussion on BYOL progressive scaling.
We see that transitioning to the higher batch size during the warmup period results in model optimization trajectory that diverges from the baseline, whereas transitioning after warmup results in matching final trajectories of the scaled and baseline models. 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 BYO dynamics during warmup.
Vision Transformers on ImageNet1k Progressive Scaling coupled with the EMA Scaling Rule 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 al., 2017; Huo al., 2021). Despite the breakdown, there is only small drop 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 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" e 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 hatch 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.
4 Related work
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 P1, and parameters transform. The P12 transformation 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) 8) investigated the benefits of scheduling the batch size at 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 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 < 1. Our work provides recipe to scale down and up in x. 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 al., 2022). In contrast, work presents practical tools to scale large batch sizes in the presence of an EMA, enabling practical training of these SSL. methods on large scale data.
5 Conclusion
when changing the batch size by factor of k, exponentiate the momentum of the EMA update to the power of K. 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 the choice of EMA momentum.
We provide an EMA Scaling Rule:
We prove the validity of the EMA Scaling Rule by deriving first-order SDE approximations of dis- crete model optimization when 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.c. Ruppert averaging) in speech and vision domains, pscudo-labeling in speech, and self-supervised image rep- resentation 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 differ- ences in optimization trajectories. For example, we can scale the BYOL self- supervisd 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 scale the EMA momentum correctly, in the same way that learning rates and other optimizer hyperparameters are scaled.
9 Instead of single large batch transition as in Figure 6 we perform a sequential transition in Appendix H.5. We find that slow increase in batch size minimizes the magnitude of the perturbation and leads to final model with higher effective linear probe top-1 than the reference by approximately 1.17%.
10

---

6 Acknowledgements
We thank Miguel Sarabia del Castillo, Adam Golinski, Pau Rodriguez Lopez, Skyler Seto, Ami- tis 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.
References
Lei Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization. CoRR, abs/1607.06450, 2016. URL http::tttaaaae.ee aaaiv.iiiv....... barg/abs/nue /100/160...00 06450.
Alexei Baevski, Arun Babu, Wei-Ning Hsu, and Michael Auli. Efficient self-supervised learn- ing with contextualized target representations for vision, speech and language. CoRR, abs/2212.07525,, 2022a. doi: 10448550/r2xi.22222.7722 URL https: sestts::/doot org/10. //00//////... 48550 aariiv..iivvv 2212. 07525.
Alexei Baevski, Wei-Ning Hsu, Qiantong Xu, Arun Babu, Jiatao Gu, and Michael Auli. Data2vec: A general framework for self- supervieeee learning in speech, vision and language. In International Conference on Machine Learning, pp. 1298-1291 1312. PMLR, 2022b.
Dan Berrebbi, Ronan Collobert, Samy Bengio, Navdeep Jaitly, and Tatiana Likhomanenko. Contin- uous pseudo-labeling from the start. The Eleventh International Conference on Learning Rep- resentations, 2023. URL https: : /opnreeieee net / forum? id=mmtdGr2bbug
Rishi Bommasani, Drew A. Hudson, Ehsan Adeli, Russ B. Altman, Simran Arora, Sydney Arx, Michael S. Bernstein, Jeannette Bohg. Antoine Bosselut, Emma Brunskill, Erik Brynjolfsson, Shyamal Buch, Dallas Card, Rodrigo Castellon, Niladri S. Chatterji, Annie S. Chen, Kathleen Creel, Jared Quiney Davis, Dorottya Demszky, Chris Donahue, Moussa Doumbouya, Esin Dur- mus, Stefano Ermon, John Etchemendy, Kawin Ethayarajh, Li Fei-Fei, Chelsea Finn, Trevor Gale, Lauren Gillespie, Karan Goel, Noah D. Goodman, Shelby Grossman, Neel Guha, Tatsunori Hashimoto, Peter Henderson, John Hewitt, Daniel E. Ho, Jenny Hong, Kyle Hsu, Jing Huang, Thomas Icard, Saahil Jain, Dan Jurafsky, Pratyusha Kalluri, Siddharth Karamcheti, Geoff Keel- ing, Fereshte Khani, Omar Khattab, Pang Wei Koh, Mark S. Krass, Ranjay Krishna, Rohith Ku- ditipudi, and et al. On the opportunities and risks of foundation models. CoRR, abs/2108.072 072558, 2021. URL https: / arxiv. org/ abs/ / 2208 07258.
Andy Brock, Soham De, Samuel L. Smith, and Karen Simonyan. High-performance large-scale image recognition without normalization. In Marina Meila and Tong Zhang (eds.). Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 1059-1071, PMLR, 2021. URL http::h//// http:///proceeinggg. mlr. www.ssrsssvv....... v19 bttpkboooaaaaaa/ html.
Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan- Tien Lin (eds.), Advances Neural Information Processing Systems 33: Annual Confer- ence Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020. URL https: httpp///////ccceeeeeeeeeee neurips. cccpaperr/220///ee 7D0e50cc0f11602287744112222222 htw.://fffaab22ccc-a---ttaattat Abstract html.
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Ar- mand Joulin. Emerging properties in self-supervised vision transformers. CoRR, abs/2104.14294 2021. URL https: arxiv. org/ hagg/bbs/22 2 2004. 10004.142000 14294.
Ting Chen, Simon Kornblith, Kevin Swersky, Mohammad Norouzi, and Geoffrey E. Hin- ton. Big self-supervised models are strong semi-supervised learners. In Hugo Larochelle,
11

---

Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Aisuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neu- ral Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, vir- tual, 2020. URL https:/ http:///oooceeeeeeeeeeeee/e / neurips cc/ hccppaprr/220// / fcbe95ccdd551da1 (11.5511d181207000ooo
81208 wwwco0000100000500000 aat.sss-bbttrtttttcc html.
Xinlei Chen, Saining Xie, and Kaiming He. An empirical study of training self-supervised vi- sion transformers. In 2021 IEEE/CVF International Conference on Computer Vision, ICCV 2021, Montreal, QC, Canada, October 10-17, 2021, pp. 9620-9629. IEEE, 2021. doi: N0.1109/CCV44922.200..0000 00050..
URL https: doi org/ 10. 1109/ ICCV48922.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. Utthttps: /owp:///eeeeeee.eeee/mmeme. net/ forum? id=YicbFdNTTy.
John C. Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learn- ing and stochastic optimization. In Adam Tauman Kalai and Mehryar Mohri (eds.), COLT 2010 The 23rd Conference on Learning Theory, Haifa, Israel, June 27-29, 2010, pp. 257- 269. Omnipress, 2010. URL COLT201 CoLT2010proceddngs .pd.engge...f
haifa. il. ibm. com www.ccommppppppe /apprs///s///s////
Abe and Aug 2020.
Fetterman and Josh Albrecht. contrastive learning with "bootstrap your own latent" (byol), wtt.ttelllgeen.../ Lererraattve--Laninng html.
Understanding self-supervised
URL
understanding-selI htt.sele-ssuuppeeeeeeooootttttte
Priya Goyal, Piotr Dollár, Ross B. Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangging Jia, and Kaiming He. Accurate, large minibatch SGD: training im- agenet in hour. CoRR, abs/1706.02677, 2017. URL http: /ttp:///axivvooo org/ abss/706. 02677.
Alex Graves, Santiago Fernández, Faustino Gomez, and Jürgen Schmidhuber. Connectionist tem- poral classification: labelling unsegmented sequence data with recurrent neural networks. In Proceedings of the 23rd international conference on Machine learning, pp. (89-376,200 2006.
Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre H. Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Ávila Pires, Zhaohan Guo, Mohammad Ghesh- laghi Azar, Bilal Piot, Koray Kavukcuoglu, Rémi Munos, and Michal Valko. Boot- strap your latent A new approach self-supervised learning. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan- Tien Lin (eds.), Advances Neural Information Processing Systems 33: Annual Conference on Neu- ral Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, vir- tual, 2020. URL https: httpp:://rocceeet http://eooeeeeeeeeeeeeeee/eessee neurips. cc/paper/2020 20/22/0000hhh///// /ash/ Abst html.
Junxian He, Jiatao Gu, Jiajun Shen, and Mare' Aurelio Ranzato. Revisiting self-training for neural sequence generation. In International Conference Learning Representations, 2020. URL
net foorummdd-SJ-e.. -.-rum?d--S.ggdnAVK..... gdnAVKDHI
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In 2015 IEEE International Conference on Computer Vision, ICCV 2015, Santiago, Chile, December 7-13, 2015, pp. 1026-1034 1034. IEEE Com- puter Society, 2015. doi: 10. 00.1109/1CCV.2055.222 URL https: // doi. org/ 10.1001. 1109/ ICCV ECCV.2013 2015.005100 123.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In 2016 IEEE Conference Computer Vision and Pattern Recognition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016, pp. 770-778, IEEE Computer Society, 2016a. doi: URL ht tttpps:/ /oiii org/ 10. 1
2016. 90.
12

---

Kaiming He, Xiangyu Zhang. Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In Bastian Leibe, Jiri Matas, Nicu Sebe, and Max Welling (eds.), Computer Vi- sion ECCV 2016 14th European Conference, Amsterdam, The Netherlands, October 1 11-14, 2016, Proceedings, Part IV, volume 9908 of Lecture Notes in Computer Science, pp. 630-645. Springer, 2016b. doi: 10.1007/978-3-319 -644399944----
URL https: / /tw. /
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross B. Girshick. Masked autoencoders are scalable vision learners. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022, pp. 15979-15988. IEEE, 2022. doi: 10. 10.110//VPR32268.2200 2022..055.. URL URLhttps: htt./ooooooooooooooo org 10. 1109/ CVPR52688 2022. 01553.
Yosuke Higuchi, Niko Moritz, Jonathan Le Roux, and Takaaki Hori. Momentum pseudo-labeling: Semi-supervised asr with continuously improving pseudo-labels. IEEE Journal of Selected Topics in Signal Processing, (666):1424-1433 2022.
Gao Huang, Yixuan Li, Geoff Pleiss, Zhuang Liu, John E. Hopcroft, and Kilian Q Weinberger. Snapshot ensembles: Train I, get M for free. In 5th International Conference on Learning Repre- sentations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. Open- Review.net, 2017. URL https: /pwnneveeeeeeeeeeeeeeeeeeee
Zhouyuan Huo. Bin Gu, and Heng Huang. Large batch optimization for deep learning using new complete layer-wise adaptive rate scaling. In Thirty- AAAI Conference on Artificial Intelligence, AAAI 2021, Thirty-Third Conference on Innovative Applications of Artificial In- telligence, IAAI 2021, The Eleventh Symposium on Educational Advances in Artificial Intelli- gence, EAA/ 2021, Virtual Event, February 2-9, 2021, pp. 7883-7890. AAAI Press, 2021. URL https: /ojs. aaai. aaai.org/ iidex. php/ hAAT/article//vew// wtt.e///vewww/66....11 16962.
IEFF. leee standard for floating-point arithmetic.. IEEE Std 754-2019 (Revision of IEEE 754-2008), pp. 1-84, 2019. doi: 10. N0.1109/EEEESTTD201993335
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In Francis R. Bach and David M. Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning, ICML 2015, Lille, France, 6-11 July 2015, volume 37 of JMLR Workshop and Conference Proceedings, pp. 448-4564666 JMLR.org. 2015, URL http: / http:/poceeednng// mlr. press/ v37 1offe15 html.
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry P. Vetrov, and Andrew Gordon Wil- son. Averaging weights leads to wider optima and better generalization. In Amir Globerson and Ricardo Silva (eds.). Proceedings of the Thirty- Fourth Conference on Uncertainty in Artificial In- telligence, UAI 2018, Monterey, California, USA, August 6-10, 2018 pp. 2018. URL http:///// / awai.org//aaa..a.a........... www.oorrgg......... hut//aaa11////// wwwaa/////proeeeinnnee https//apppps///eeeeeee /313 pdf.
AUAI Press,
Stanislaw Jastrzebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Ben- gio, and Amos J. Storkey. Three factors influencing minima in SGD. CoRR, abs/171 abs/1711.04623, 2017. URL http: //axxiv org/ 1711. 04623.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. CORR, abs/2001. .08361, 2020. URL https: // /ttp/xxxiv..////////////////aava www.ii.or///..oo..oo
08361.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Truck Proceedinggs 2015. URL http: //arxiv. wwgg/bs/////oooooooooo.ooo..oeat 1 4 12. 6980.
Lestie Kish. Survey Sampling, volume 59. Cambridge University Press, 1965. doi: 10.1017/ $00033555001332133.
Skanda Koppula, Yazhe Li, Evan Shelhamer, Andrew Jaegle, Nikhil Parthasarathy, Relja Arand- jelovic, João Carreira, and Olivier J. Hénaff. Where should I spend my flops? efficiency evalu- ations of visual pre-training methods. CoRR, abs/2209. 15589, 2022. doi: 10.445500/axxvv2229. 15589. URL https://///. // //ooi. .agi.org//e org/10. 48550/ /atpvaaiivv./... 2209. 15589.
13

---

Alex Krizhevsky. One weird trick for parallelizing convolutional neural networks CoRR, (bs/1404.5997, 2014. URL pattp://arxer wwepxiiv.rrg..oo org/abs 1404. 5997.
Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Cifar-10 (canadian institute for advanced re- search). 2014. URL http: /wwww cs. toronto. edu/-kiiz/cifaan html.
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape of neural nets. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 31: Annual Conference Neural Information Process- ing Systems 2018, NeurlPS 2018, December 3-8, 2018, Montréal, Canada, pp. 6391 ppaper/2018/harl
6401, 2018. URL https: httt::::/prrceeeeee ptwpeeeeeoeeeeseeeeooe neurips. (1l1bbb3e66b50bbb666666667777777 (66779166696-- -oers--bbsrrctct html.
Qianxiao Li, Cheng Tai, and E Weinan. Stochastic modified equations and dynamics of stochastic gradient algorithms i: Mathematical foundations. The Journal of Machine Learning Research, 20 (1):1474-152) 2019.
Zhiyuan Li, Sadhika Malladi, and Sanjeev Arora. On the validity of modeling SGD with stochastic differential equations (sdes). In Marc Aurelio Ranzato, Alina Beygelz- imer, Yann Dauphin, Percy Liang. and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neural Informa- tion Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pp. 12712 12725, 2021. URL https: / http://e/oocc//o/sseeeoe///. www.gii..nriip......m cc/ wttps/ccpppppppa/ee/eeeeee hww.sectit.-.bttattt....... html.
1996296622665009221111166666-169
Tatiana Likhomanenko, Qiantong Xu, Jacob Kahn, Gabriel Synnaeve, and Ronan Collobert. slimipl: Language-monn model-fooee iterative pseudo-labeling. Proc. Interspeech, 2021a.
Tatiana Likhomanenko, Qiantong Xu, Gabriel Synnaeve, Ronan Collobert, and Alex Rogozhnikov. Cape: Encoding relative positions with continuous augmented positional embeddings. Advances in Neural Information Processing Systems, 34, 2021b.
Tatiana Likhomanenko, Ronan Collobert, Navdeep Jaitly, and Samy Bengio. Continuous soft pseudo-labeling in ASR. In I Can't Believe It's Not Better Workshop: Understanding Deep Learn- ing Through Empirical Falsification, 2022. URL https: /ppprrevieeereeeee///.mee nett/ forum?
Timothy P. Lillicrap. Jonathan J. Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. In Yoshua Bengio and Yann LeCun (eds.), 4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings, 2016. URL http: //arxiv. org/ /ass///////e 509. 02971.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In 7th International Con- ference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019, Open- Review.net, 2019. URL httpss://ppenreeiiew
Sadhika Malladi, Kaifeng Lyu, Abhishek Panigrahi, and Sanjeev Arora. On the sdes and scaling rules for adaptive gradient algorithms. In NeurIPS, 2022. htte//apprr22000 html.
URL
cctpspcccpppppre http:://apepee 2at.a00002...00244444----------- 32ac710102f0620d6
Vimal Manohar, Tatiana Likhomanenko, Qiantong Xu, Wei-Ning Hsu, Ronan Collobert, Yatharth Saraf. Geoffrey Zweig, and AbdeIrahman Mohamed. Kaizen: Continuously improving teacher using exponential moving average for semi-supervised speech recognition. In 2021 IEEE Auto- matic Speech Recognition and Understanding Workshop (ASRU). pp. 518-525. IEEE, 2021.
Daisuke Niizumi, Daiki Takeuchi, Yasunori Ohishi, Noboru Harada, and Kunio Kashino. BYOL for audio: Exploring pre-trained general-purpose audio representations. IEEE ACM Trans. Audio Speech Lang. Process., 31:137-151, 2023. doi: 10. 109/TASLL //doi. org/10../00 11 1 09 /TASLP 2022. 3221 007.
URL ht https:
14

---

Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Mare Szafraniec, Vasil Khali- dov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, Mahmoud As- sran, Nicolas Ballas, Wojciech Galuba, Russell Howes, Po- Yao Huang, Shang-Wers Weng-Weenn Li, Ishan Misra, Michael G. Rabhat, Vasu Sharma, Gabriel Synnaeve, Hu Xu, Hervé Jégou, Julien Mairal, Patrick Labatut, Armand Joulin, and Piotr Bojanowski. Dinov2: Learning robust visual fea- tures without supervision. CoRR, abs/2304.071933 2023, doi: 10448550aaxxv.2344007999 URL ht https: //oii orooorg/////oo 10. 48550/ arxiv 2304. 07193.
Myle Ott, Sergey Edunov, Alexci Bacvski, Angela Fan, Sam Gross, Nathan Ng. David Grangier, and Michael Auli. fairseq: A fast, extensible toolkit for sequence modeling. In Proceedings of Demonstrations, 2019.
NAACL-HLT 2019:
Vassil Panayotov, Guoguo Chen, Daniel Povey, and Sanjeev Khudanpur. Librispeech: an corpus based on public domain audio books. In 2015 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 5206-5210. IEEE, 2015.
Daniel S Park, William Chan, Yu Zhang, Boung-Cheng Chiu, Barret Zoph, Ekin D Cubuk, and Quoe VLe Specaugment: A simple data augmentation method for automatic speech recognition. Proc. Interspeech 2019, pp. 2613-2617, 2019
B. T. Polyak and A. B. Juditsky. Acceleration of stochastic approximation by averaging. SIAM Journal Control and Optimization, 30(4):838 -85, 1992. doi: 10.1137/0330046. URL https: /tt.
org 10.
Siyuan Qiao, Huiyu Wang, Chenxi Liu, Wei Shen, and Alan L. Yuille. Weight standardization. CoRR, 0t0//9033115200 2019. URL ht //axxiv////vv.oo////eo org/ aaaaabaso900 19900 10520.
Pierre H. Richemond, Jean-Bastien Grill, Florent Altché, Corentin Tallec, Florian Strub, Andrew Brock, Samuel L. Smith, Soham De, Razvan Pascanu, Bilal Piot, and Michal Valko. BYOL works even without batch statistics. CoRR, abs/2010.10241, 2020. URL httts://arxii: org/
10241.
Pierre H. Richemond, Allison C. Tam, Yunhao Tang, Florian Strub, Bilal Piot, and Felix Hill. The edge of orthogonality: A simple view of what makes BYOL tick. CoRR, abs/2302.04817, 2023. doi: 10.48500/2Xiv.2322..444 2302 04817.
URL https: org/10..//... 48550/ arxiv.
David Ruppert. Efficient estimations from a slowly convergent robbins-monro process. 1988.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang. Andrej Karpathy, Aditya Khosla, Michael S. Bernstein, and Li Fei-Fei. Imagenet large scale visual recognition challenge. CoRR, Abs/1440.0055, 2014. URL http: /tap::///axvvi..//aaaeaaeeaaaaar org/ abs / 1409. 0575.
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representa- tions. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 2 (Short Pupers), pp-
464-468. 2018.
Samuel L. Smith and Quoc Le. A bayesian perspective generalization and stochastic gradient descent. In 6th International Conference on Learning Representations, ICLR 2018. Vancouver, BC, Canada, April 30- May 2018, Conference Track Proceedings. OpenReview.net, 2018. URL httphttpps... :tttlp:://ppeeeeee //wppppeeeie.....oom/// net/ forum? id= 4yg0z.
Samuel L. Smith, Pieter-Jan Kindermans, Chris Ying, and Quoe V. Le. Don't decay the learning rate, increase the batch size. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30- -0-33 MMy 3, 2018, Conference Track Proceedings. OpenReview.net, 2018. URL https:// wpw.reviee..eeee/////p/.. net forum? -i=:iiddiiiiiee. 1BXCZ.
Kihyuk Sohn, David Berthelot, Nicholas Carlini, Zizhao Zhang, Han Zhang, Colin A Raffel, Ekin Dogus Cubuk, Alexey Kurakin, and Chun-Liang Li. Fixmatch: Simplifying semi-supervised learning with consistency and confidence. Advances in neural information processing systems, 33:596-608, 2020.
15

---

Aarohi Srivastava, Abhinav Rastogi, Abhishek Rao, Abu Awal Md Shoeb, Abubakar Abid, Adam Fisch, Adam R. Brown, Adam Santoro, Aditya Gupta, Adrià Garriga-Alonso, Agnieszka Kluska, Aitor Lewkowycz, Akshat Agarwal, Alethea Power, Alex Ray, Alex Warstadt, Alexander W. Kocurek, Ali Safaya, Ali Tazarv, Alice Xiang, Alicia Parrish, Allen Nie, Aman Hussain, Amanda Askell, Amanda Dsouza, Ameet Rahane, Anantharaman S. Iyer, Anders Andreassen, Andrea Santilli, Andreas Stuhlmüller, Andrew M. Dai, Andrew La, Andrew K. Lampinen, Andy Zou, Angela Jiang, Angelica Chen, Anh Vuong, Animesh Gupta, Anna Gottardi, Antonio Norelli, Anu Venkatesh, Arash Gholamidavoodi, Arfa Tabassum, Arul Menezes, Arun Kirubarajan, Asher Mullokandov, Ashish Sabharwal, Austin Herrick, Avia Efrat, Aykut Erdem, Ayla Karakas, and et al. Beyond the imitation game: Quantifying and extrapolating the capabilities of language models. CoRR, (ww/2206.041155 2022. doi:
URL https:/
doi org/ 48550/4555/a00 arxiv. 2206 04615.
Antti Tarvainen and Harri Valpola. Mean teachers are better role models: Weight-averaged con- sisteney targets improve semi-supervised deep learning results. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA. USA, neurips copss.cc/p wtt.cc/cppppp //ppp//201////// pp. 1195-1204, 2017. URL https::/ppooceediiiiii Abstrefff--srrrtttcccccc. html.
Yuandong Tian, Xinlei Chen, and Surya Ganguli. Understanding self-supervised learning dynam- ics without contrastive pairs. In Marina Mcila and Tong Zhang (eds.), Proceedings of the 38th International Conference Machine Learning, ICML 2021. 18-24 July 2021, Virtual Event, vol- ume 139 of Proceedings of Machine Learning Research, pp. 19228-10278... PMLR, 2021. URL ht http:/ / proceedings..... mlr. press/v139) /tt//iian///// 1a. html.
Tijmen Tieleman, Geoffrey Hinton, et al. Lecture 6.5-rmsprop: Divide the gradient by running average of its recent magnitude. COURSERA: Neural networks for machine learning, 4(2):26-31, 2012.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Ro- man Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, neurips. ssppperr/201/ Abstracc html.
pp. 5998 6008, 2017. URL / http:///e///cceeee///eeee hashh/355ee225554cccc1bbbooo
Yuxin Wu and Kaiming He. Group normalization. In Vittorio Ferrari, Martial Hebert, Cristian Sminchisescu, and Yair Weiss (eds.), Computer Vision - ECCV 2018 - 15th European Conference, Munich, Germany, September 8-14, 2018, Proceedings, Part XIII, volume 11217 of Lecture Notes doi:
in Computer Science pp. 39 Springer, 2018.
URL
ht https:: //do:////oi.// ooog/o///// 10. 1007 002100777788 0010778---0000011----1
Qiantong Xu, Tatiana Likhomanenko, Jacob Kahn, Awni Hannun, Gabriel Synnaeve, and Ronan Collobert. Iterative pseudo-labeling for speech recognition. Proc. Interspeech 2020, pp. 1006- 1010, 2020.
Yang You, Igor Gitman, and Boris Ginsburg. Scaling SGD batch size to 32k for imagenet training. CORR, 208/1708.10 038888888 2017. URL http: /ttp//aavvvvv//////////ee/./.avv org/ abs////ae 03888.
Yang You, Jing Li, Sashank J. Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiao- dan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsich. Large batch optimization for deep learning: Training BERT in 76 minutes. In 8th International Conference on Learning Repre- sentations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL www.:.ee/oorrooeee formmoormm wiw.mmddddSSxxxttWWWW..
16

---

Appendices
Broader impact
18
Limitations
18
The scaling toolbox: practical methods for enabling systematic scaling
19
C.I The continuous time/SDE perspective
19
C.2 Scaling rules for optimization
20
C.3 Commonly used values of hyperparameters different batch sizes
21
C.4 Progressive scaling
21
D EMA approximation theorems with SDEs
22
D.I SGD with model EMA
22
D.2
Adaptive gradient methods with model EMA
24
Additional proofs
25
E.1
Iterations of SGD + EMA
25
E.2
Limiting behavior of Polyak- Polyak-Ruppert Ruppert averaging
26
Additional details and results for Polyak-Ruppert Ruppert averaging
27
F.I
Noisy parabola
28
F.2 Image Classification
30
F3
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 x = 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
H.9 ResNet-18 additional scaling analysis
45
H.10 Scaling Re BYOL using LARS and Progressive Scaling
46
H.1l Preventing collapse phenomena in DINO at scale
47
I Additional details on numerical stability
51
Contributions
52
17

---

A Broader impact
This work shows how to adapt Machine Learning (ML) optimization in the presence of 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.
Our EMA Scaling Rule lowers compute usage as 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 pn at some batch size (for example, the momentum that gives the best transfer performance), then the optimal value at another batch size ß is exactly the one given by the EMA Scaling Rule ý = p for scaling x B/B.
3. Our EMA Scaling Rule enables researchers 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:
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 greater num- ber 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 an extra hyperpa- rameter 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.
B Limitations
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):
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 model EMA, the model optimization trajectories at the batch sizes and xB, 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 x due to discretization error (Li et al., 2021). Second, we know that if the gradient noise not sufficiently large, the SDE approximation for Adam does not exist (Malladi 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 optimiza- tion methods RMSProp and Adam (Malladi et al., 2022). Empirically, a square-root scaling law for LAMB (You 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 ievoolution-basee archi- tectures. Despite this, we are able to demonstrate in Appendix H. that combination of the EMA Scaling Rule and progressive scaling can match, or surpass baseline BYOL performance 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.
18

---

3. It may be the case that the optimal performance attainable by a given model setup exists at level of discretization/graieen 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 at large k can be considerable. In cases where the model EMA plays role in the overall optimization, the error introduced by the EMA Scaling Rule can break the preservation of model dynamics.
Put another way, an optimizer scaling rule and the EMA Scaling Rule each introduce their own dis- cretization errors. the case where EMA plays a role optimization, as soon as the discretization error of either the optimizer scaling rule or the EMA Scaling Rule is large, the error for the joint optimization procedure is large. This is at least as bad as cases that do not use a model EMA during the optimization process.
C The scaling toolbox: practical methods for enabling systematic scaling
There are many different components involved in preserving optimization dynamics at different batch sizes. In this appendix we collect into single place the different concepts and values that we found useful in practice, in an attempt to make the practice of scaling as accessible as possible.
C.1 The continuous time/SDE perspective
Here we discuss the mindset difference required when trying preserve training dynamics. In ML we typically use stochastic optimization, leading us to think of the optimization in terms of performing updates, or stepping the optimizer. This notion has become more common in the era of large datasets, where it may be the case that we only see a fraction of the dataset during optimization.
For dynamics preservation under scaling, we suggest that it is simpler to consider the amount of data seen by the training process, or alternatively, the amount of continuous time in the discretization of SDEs view. The reason is the following. The SDE scaling rule results (Definition 1.2, Li et al. (2019, 2021); Malladi et al. (2022)) follow from showing that different discretizations of the SDE close to that SDE, providing we appropriately scale hyperparameters (see Section 2.2). Each of these discretizations shares the total continuous time T fjx Nier of the underlying SDE, but each discretization has different number of iterations Nuer = Naer/K.
This perspective is already adopted, perhaps by accident in some domains. For example, in Com- puter Vision (CV), it is typical to compare model performance after optimization on ImageNetlk after a number of epochs, whilst also specifing a learning rate warmup after a number of epochs. This transforms the schedule into the form wait until the process meets [condition], where here [condition] is when the process has seen sufficiently many samples.
More generally, we can specify any condition that is not property of the discretization procedure itself. Instead, the discretization procedure should be viewed as a numerical approximation method for the SDE are evolving, and the properties of that discretization process (like number of steps) are not of specific interest in the world view where we do decouple optimization from the batch size. A specific example of this more general case is present in Section 3.3, where for scaling x > 2 we wait until the pre- pre-training Word Error Rate (WER) is sufficienlly low.
There may be cases where one is working with setup that is explicitly defined in terms of quantities related to the discretization process. Indeed, the optimizer hyperparameters are examples of these, and need to be scaled accordingly with x. The other typical example of this is conditions based on the number of optimizer steps, rather than the number of epochs. In this case, these quantitics should be scaled to achieve the desired condition in the same amount of time, i.e. as above Nu = NNwe/x x, where Nuer is the number of iterations specified at the base batch size B. Concretely, if training is specified in a number of steps, then doubling the batch size implies you should train for half the number of steps.
This is in the case of SGD, for RMSProp and Adam one should use I = if x Niter (Malladi al., 2022).
19

---

C.2 Scaling rules for optimization
For ease of reference, we collect all the scaling rules related to batch size modification we are aware of. We begin with the most well-known, the SGD Scaling Rule (Definitions 2.2 and C.1).
Definition C.1 (SGD Scaling Rule). When running SGD (Definition 2.1) with batch size B kB, use learning rate îj = Kπ) (Krizhevsky, 2014; Goyal et al., 2017).
The SGD Scaling Rule is also known as the Linear Scaling Rule (LSR), although for clarity, this work adopts the naming convention [Algorithm Name] Scaling Rule, which means all parameters of those algorithms are appropriately scaled from batch size B to xB.
Next, we give the two scaling rules known for the adapative optimizers RMSProp (Tieleman et al., 2012) and Adam (Kingma Ba, 2015) in Definition C.2 and Definition C.3 respectively.
Definition C.2 (RMSProp Scaling Rule). When running RMSProp (Tieleman et al., 2012) with batch size B = kB, use a learning rate i vkn, beta coefficient B = x (1 -1- B8), and adaptivity (Malladi et al., 2022).
parameter = <K
Definition C.3 (Adam Scaling Rule). When running Adam (Kingma & Ba, 2015) with batch size = kB, use learning rate vkn, beta coefficients p. = x (1 f.), P2 = x (1 BB2), and adaptivity parameter vkc
(Malladi et al., 2022).
Next, we present contribution of this work, the EMA Scaling Rule (Definitions 1.2 and C.4), which extends the above scaling rules to allow the presence of a model EMA which is able to contribute to the overall optimization (see Appendices D and E.1 for derivations).
Definition C.4 (EMA Scaling Rule). When computing the EMA update (Definition 1.1) of a model undergoing stochastic optimization with batch size B = kB, use momentum ý = pp and scale other optimizers according to their own scaling rules.
Concretely, if are using SGD in the presence of model EMA, Definitions C.1 and C.4 state that should take Kmj and P p" when scaling by B/B.
The final scaling rule for weight decay, and follows from the scaling logic discussed in Ap- pendix C.1 and Krizhevsky (2014). If we take the weight decay regularization penalty 2. defined at batch size B, what should the weight decay be for batch size B xB? For simplicity, consider updates optimization of parameters 9, in the presence of weight decay only
++x 0 17 00 = (1 - 2) 0, (1 - 17 2)x 0,
(11)
Therefore, to match the effect of weight decay with a single iteration step, we need to match
1- q Å (1 2)*
(12)
Solving for î and expanding around ~ 0 gives
i 1 - (1 7 2)*
x "
xxx++0())
(13)
"
i
This leads to the Weight Decay Scaling Rule (Definition C.5).
Definition C.5 (Weight Decay Scaling Rule). When using weight decay with batch size B = kB, use a penalty term å == (Kń/n) A, where n and n represent the scaled and unscaled learning rates of the corresponding optimizer (Krizhevsky, 2014; Li et al., 2018; Loshchilov & Hutter, 2019).
The Weight Decay Scaling Rule implics that using linear scaling for the learning rate 1j then the weight decay penalty is automatically scaled, and when using square-root scaling for the learning rate 7 (c.g. in the case the Adam Scaling Rule (Definition C.3)) then the weight decay penalty should also be scaled with square-root is proposed in Loshchilov & Hutter (2019).
Finally, we see that if the implementation of weight decay does not have an update scaled by the learning rate, i.e. the update is 0,+1 (1 A) 6,, then the scaling rule is
and
becomes linear for small weight decay, i.c. = xd, and for arbitrary 2 takes the form Á = 1- 1-(1-Ai (1-1)*.
20

---

Table 2: Scaled learning rates ij at different batch sizes B xB given reference learning rates defined at batch size B. The reference values of cach column are boldened. Note that this is only valid when there is notion of single sample. In the sequence learning setup (for example, in Section 3.3), the notion of batch size should be appropriately replaced with the dynamic batch size, total sequence length.
ij = Km) [SGD]
= vkn IRMSProp, Adam]
B 256
B 512
B 256
B 4096
Batch size B
0.1
0.3
0.1
10-3
4.8
10-3
0.0125
0.0375
0.00625
0.00035
0.42426
0.00009
64
0.025
0.075
0.0125
0.0005
0.6
0.00013
128
0.05
0.15
0.025
0.00071
0.84853
0.00018
256
0.1
0.3
0.05
0.001
1.2
0.00025
0.00141
1.69706
0.00035
1024
0.4
1.2
0.2
0.002
2.4
0.0005
2048
0.8
2.4
0.4
0.00283
3.39411
0.00071
4096
1.6
4.8
0.8
0.004
4.8
0.001
0.00566
6.78823
0.00141
16384
6.4
19.2
0.008
9.6
0.002
32768
12.8
38.4
6.4
0.01131
13.57645
0.00283
65536
25.6
76.8
12.8
0.016
19.2
0.004
Table 3: Scaled EMA momenta f pK at different batch sizes B = given reference momenta p defined batch size B. The reference values of each column are boldened. Again in the sequence learning setup, batch size should be appropriately replaced with notion of sequence length.
256
4096
Batch size B
0.9999
p 0.999
P 0.99
p 0,996
p 0.992
p 0.99
P = 0.97
0.99999
0.99987
0.99874
0,99997
0.99994
0.99992
0.99976
0.99997
0.99975
0.99749
0,99994
0,99987
0.99984
0.99952
128
0.99995
0.9995
0.99499
0,99987
0.99975
0.99969
0.99905
256
0.9999
0,999
0.99
0,99975
0,9995
0.99937
0.9981
512
0.9998
0.998
0.9801
0,9995
0,999
0.99874
0.9962
1024
0.9996
0.99601
0.9606
0,999
0,99799
0.99749
0.99241
2048
0.9992
0,99203
0.92274
0,998
0,99599
0.99499
0,98489
4096
0.9984
0.98412
0.85146
0.996
0.992
0.99
0.97
8192
0.9968
0.96849
0.72498
0,99202
0,98406
0.9801
0,9409
16384
0.99362
0.93798
0.5256
0,9841
0.96838
0.9606
0.88529
32768
0.98728
0.8798
0.27625
0,96844
0.93776
0.92274
0.78374
65536
0.97472
0.77405
0.07632
0.93788
0.8794
0.85146
0.61425
C.3
Commonly used values of hyperparameters at different batch sizes
In the literature it is common to give base learning rate n defined batch size 256, implicitly using the SGD Scaling Rule, even when using the Adam optimizer. Because the scaling of other optimization hyperparameters was not understood until recently, these for the experiment, e.g. the Adam betas and epsilon, and the EMA momentum, implicitly defined at the scale of the experiment, for example at batch size 4096. One way to deal with this in practice is to define single reference batch size B at which all hyperparameters are defined, and then scale from there. this case, is easiest to compute using linear scaling the learning rate at the redefined base batch size 7 R Morig. where ă B/ge and then scale this reference ! as f = Kπ,, B/B, along with e.g. the momentum defined at B.
also common to just present
As this process can be slightly frustrating, we have provided tables of typical learning rates in Table and momenta in Table 3.
C.4 Progressive scaling
In Section 3.4 we introduced Progressive Scaling (Definition 3.2) to test our hypothesis that early in the BYOL training procedure, there dynamics that are challenging to replicate at larger batch
21

---

Algorithm Stochastic Gradient Descent with Progressive Scaling
Require: Base learning rate base momentum for base batch size B Require: Initial target model parameters 0 and model EMA parameters 3 Require: Epochs E and schedule of batch sizes B = B₁, B2. Bz..., Bg
Require: Loss function L
for e in 1, do
[[e]
Get current batch size
B/B
Compute scaling factor
i - kmj
Get scaled learning rate
p*
Get scaled momentum
for b in 1, (borrEEB)) do
Sample minibatch of 8 samples X (x x(.),
0 0 (f/BB Exex vg[(x;0,g)
SGD Update
3 p (1 0)
EMA Update
end for
end for
sizes. To remove ambiguity, in Algorithm we provide pseudo-code for how Progressive Scaling.
In Algorithm I, the prefactor of the SGD update could also have been written n/B, although an equivalent use of the base momentum is not possible.
Finally, we outline how to extend Algorithm to more complex setups, like those presented in Section 3.4:
1. Optimizer scaling rules are used appropriately, for example the Adam scaling rule in case of using the Adam optimizer to update parameters
Schedules for hyperparameters are computed using the base hyperparameters, and are then modified by application of the scaling law in epoch (outer) loop.
3. Schedules for hyperparameters at the step rather than epoch level can be achieved in practice through recomputing the schedule and updating the notion of minibatch index appropriately throughout training.
All of the above techniques are used in Section 3.4. In addition, scheduling batch sizes within epoch is possible, providing one maintains a notion of computation within some fixed continuous time Tixed We did investigate this scenario.
D
EMA approximation theorems with SDEs
D.1 SGD with model EMA
We will now derive the EMA scaling rule when tracking model parameters and the model trained using SGD. We employ strategy similar to Malladi et al. (2022), where associate to each iterative process a Stochastic Differential Equation (SDE). In order to control the distance between the SDE and the discrete process, we use the tools from Li et al. (2019).
Definition D.1 (Polynomial growth, Definition in (Li et al., 2019)). The set G is the set of con- tinuous functions I - R with most polynomial growth, i.e., for g e G there exists two scalars K1, K2 > 0 such that for all x Rd,, we have ig(x)) S (1 [[=))
For an integer 0, G is the set of functions Bdd - R that are a-times continuously differentiable and such that all their derivatives up to order a are in G.
Similarly Malladi et al. (2022), define the update rules on the parameters.
Noisy Gradient Oracle with Scale Parameter (NGOS) to
Definition D.2 (Noisy Gradient Oracle with Scale Parameter (NGOS), adaptation of (Malladi al., 2022)). A NGOS is tuple Go Zo)). Given noise scale parameter 0, the NGOS Go
22

---

takes as input the parameters O and outputs a random vector g Vf(0,)) + de where Vf(0, () is the gradient of f with respect to 0 at (0, (), and is a random vector drawn from the distribution (0, () with zero mean and covariance Σ(0, Ç).
Note that in the above definition, the probability distribution (0, is allowed to change with the scale G, but its first two moments - its mean and its covariance are fixed with o. We have the following theorem for model EMA under optimization with SGD:
Theorem D.1 (SDE for SGD + EMA). Consider the couple Xx = (0x, where Og are the iterates of SGD with a NGOS (Defimition D.2) and SI is an EMA of 0x. defined, starting from X, = wy, by
0k+1 0x ggk. with 8k Vf(0%, 51) dek, and EX (0k, St),
(14)
(K-1 = pīa + (1 11-17 p)0g
(15)
Define Bo (1 - 1) = a Vī and define the SDE X, (0,, 7T), starting from Xo = 2o, by
d0; (V, Z,) + o ( Z,) with W, a Wiener process
(16)
dZ, f. Z,)dt
(17)
Assume that f is continuously differentiable, with f E G and Σ E G (Definition D.I). Then, for any time horizon T and test test function 9 E G² there exists constant > 0 such that
max k=0.....7711
[B(g(x,)) B[g(x))] S cxn
(18)
Proof. The proof uses the same tools in Li al. (2019). Define A(O, () n(-vf(O,C) σe, (0 ()) with e (., () the one-step update for the SGD + EMA update, such that Xk+1 = Xg A(xg). We have the first moments:
B[A(O, [)] = 7- ((0..))0000-)))
(19)
VIA(O. ()] = go x(. 5)
0
0
(20)
and the higher-order moments are O Similarly, let A(0, () be the solution at time of the SDE defined by Equation 6 starting from X, (0, (). From Ito's formula, we also obtain
BJÂ(O, ()] (- ff( (o ())
(21)
x(0 0
VIÅ ()] Oao 0
(22)
and the higher-order moments are Onn²). Hence, the moments of the discrete iteration and of the SDE match up to second order. Following the same proof technique as in Li et al. (2019) then leads the advertized theorem.
This theorem is simple adaptation of the results of Li et al. (2019). Intuitively, it is expected that X, and X are close since x; is the Suler-Maruyame discretization of X, with learning rate 1). We then have the corollary.
Corollary D.1.1 (Validity of the EMA Scaling Rule). Assume that f is continuously differentiable, with f E G and Σ E G². Let 0% 2h the iterates of the Equation 5 with batch size B and hyper- parameters n), p. Let e% xx.B be iterates with batch size xB, learning rate determined by the SGD Scaling Rule (Definition 2.2) and momentum determined by the EMA Scaling Rule, linear version (Definition 1.2). Then, for any time horizon T and function g G², there exists a constant d such that
max 1T/n)
[B [R/K) (T/x)) ) (lg 31)1 s dxn
(23)
Proof. The proof is similar to Malladi et al. (2022). Under the scaling rule, both x = (0g, 5x) and
{ [k/x] 2xB 5(K/x)) have the same limiting SDE. Hence we have from the previous theorem
that for all test function g. find c' such that
max ke/0.....T/m
S c exn and
max (T//
S c'xn. (24)
23

---

The triangle inequality then gives
[R/<(( B[g(x)]] (c+c))n
(25)
k=0,...77m
Hence, taking d == c gives the expected result.
D.2 Adaptive gradient methods with model EMA
We now turn to the case where one uses an adaptive gradient method rather than SGD to train the model. We follow derivations similar to those of Malladi et al. (2022), with added EMA. Like above, we consider that the loss function f also depends on the EMA tracking parameter 32- We begin with RMSProp with EMA, which iterates:
Vk+1 vv
B) with ge (00kk, 5x) dee, and EK (0ke Sz),
(26)
0k+1 0g 77 Vv + 2) x gi
(27)
Sk+1 pSk + (1 - p)
(28)
Like in Malladi et al. (2022), we place ourselves in the high noise regime, in which the term ge in Equation 26 is approximated by g d'diag( E(0g, 5x)). We the same scaling rules, with an additional one for p:
Yo (1 (1-r)/n ao on, En, and An (1 p)/n²
(29)
and we let uy v2/o². The equations for RMSProp with EMA then become, using only these new variables and
Ug+1 Uy 17 ('yo(dia (x), uk),
(30)
0k+1 02 -vvu + Eo) (n) ff (x) + nE2)
(31)
Ik+1 LA = = Bo - {x)
(32)
This formulation makes it clear that these iterations can be seen as the discretization of the SDE
dU, yo diag [E, - U,)dt,
(33)
d0,- (,, eq) (2)
(34)
dZ,- fo d,)
(35)
with step size 7². Of course, we recover the SDE of Malladi al. (2022) in the case where fo = 0. A formal proof of closeness between the iterates and the SDE trajectory is out of the scope of the present paper since it would imply redoing much of the theoretical work developed in Malladi al. (2022). Still, the previous informal analysis hints that for RMSProp. the scaling rule in Equation 29 should be used. In other words, given a certain set of hyperparameters y. and p. if the batch size goes from to B x x B, the noise level becomes à = a/ Vic and keeping the quantities in Equation 29 constant means that we should use new hyperparameters
y = (1 - x 1= X Vx, and X = (1 (1-p) xx.
The linear rule ý = - (1 (1- x is at the first order equivalent to the exponential scaling rule = p². Hence, even though the limiting SDE differs greatly from that of SGD, and even though the scaling rule regarding the learning rate differs, we recover for the momentum term p the exact same scaling rule for SGD.
We finish the discussion with the case of Adam, which leads once again the same rule for SGD. Adam with EMA tracking of the network parameters iterates
m = fim (1 B.) with ga = (00g 3x) oeg. and Ex (0, 3z).
(36)
Vk+1 = B2V (1 Be)
(37)
m = mk+1/( ptttl)
(38)
Vk+1 = Vk-1/(1 (1 Bk+t))
(39)
Oc11 = O₁ 7(( Vi 8) 1 m.
(40)
Se.1 = pb: (1 p)0
(41)
24

---

Here, we use the same minor modification of the iterations as in Malladi et al. (2022), where we use V₁ instead of Vk+: in the denominator of the 0₂ update.
We consider the following scaling for the hyperparameters
(1 B.)n''
-2) on, en, and Bo (1 p)/n"
(42)
and y: (t)
1- (111-xxp(-c1)) (2(1) 1 exp((x2=))) and Hn vE/o". The SDE for Adam EMA
given by
dM, = ci (V) - - (,, 2,) 2,)
(43)
dU; c2 (diag(((O,, 2,)) U,)dt
(44)
(t)
d0, (1(1) (an U, + 60 vo M,di
(45)
dz, = fo dt.
(46)
This is once again the same SDE as in Malladi et al. (2022) with the added EMA term. Like previously, this SDE hints at the fact that scaling rule eq. 42 should be used. In other words, given set of hyperparameters $1. f2. n). and p, if the batch size goes from B to x B, then the noise level becomes = o/ vi and keeping quantities in eq. 42 constant means that we should use as new hyperparameters
8 = (1 BB)xx =1 - (1 (e) xx, - VR, and 1 (1 p) x We once again recover linear rule for p which is equivalent the exponential scaling rule = p* in the limit p 0.
E
Additional proofs
E.1 Iterations of SGD EMA
Here we derive a critical component the EMA Scaling Rule, the matrix equation of Equation from which the EMA Scaling Rule (Definition 1.2) follows.
Theorem E.1 (Iterations of SGD EMA). Assuming that gradients change slowly over iterations of SGD (Defimition 2) and EMA (Deftnition 1.1): VoL(x; Orej. L(tz) = VoL(x; 0,, 5,) = g. for j 1, ..., and representative gradient g iterating over k independent minibatches produces
model states
0/,x]]
1
-7
0,
O--nKg
S1+K
1
p
p 3. + (1 0²) + (1 x Bp)
(47)
g
0
Proof. First note that for matrices of the form
ag.2
A
1,1
a1,1
0
(48)
their multiplication follows
0
0g,2
1
0
bo.2
A B
=
41,1
0111
1 b
b
0
0
0
1
0
ag.2 bo.2
1
41,
b
a1, bT,1
(1 a1,1) bo.2
(49)
0
0
and
1
0
do,2 boj
1
0
Co.2
ABC =
- 0i, bin
b
(1
an,) bo.2
1- C1,1
C1,1
0
0
0
0
0
1
0
an,2 bo.2 + Ca.2
a1,
1i,
C11
an
(1 a1,)
boz + a,, by,)) Co22
(50)
0
1
25

---

By induction
0
K 00,2
A
1
a 1.1
1.1
aa $1,1, K)
(51)
where
K-1
1 1 di,1
8( a022 u1,1,
up,2
d1,1)
20,2
k
for a1, # 1.
(52)
1
01,1
follows that
0
11
1
0
K 7
1 1-p
p
0
p*
p
8(n. p. x) 1
(53)
0
1
0
0
where the EMA Scaling Rule error
8(n.p.k) (-n) K 1 1-p" pm
1-p = (-r) (-) = O(n Bp).
(54)
where Bp =1-p and the approximation is around p = 1.
E.2 Limiting behavior of Polyak-Rupper averaging
Here we sketch the asymptotic behavior of target model 0 and its EMA 3. Let us assume that converges the stationary distribution lim 0 0", 0 (0). We are interested in statistical properties of ? = lim S. as this will formalize the notion of how the EMA depends on the time-horizon defined by its momentum p as discussed in Table
As a warm-up, for n independent random variables X,....,Xp X2, we know that the sample mean & (x,,22,.. X,, Xn) has the statistical properties
0²
x μ4
Var[*]
(55)
where μ and the population mean and variance. This gives an idea of what to expect. As will now show, the expectation of ( should have no time-horizon dependence, whereas the variance of 5 will depend on its time horizon (i.e. the number of samples it integrates over) which is defined by p.
In the case of weighted sum
((w)_
W; Xis
(56)
then if the x; are Independent and Identically Distributed (i.i.d.), then
([x(w)] =
n w; E[x;] n 14
W
1 n
Wi.
(57)
and for the variance (Kish, 1965)
[ar[x((]] = w o²
0² = Var[x]]
(58)
We can verify that we reproduce the well-known result in Equation 55 in the case where all weights are equal to 1 follows
: Wi
2
Var[x(w)] =
1
(59)
n?
n² n
26

---

In the case of an exponential moving average we have
p',
(60)
Let's consider the specific case where at iteration k which is sufficiently large that and 0 have converged to their stationary distributions. From k, the iterations unfold as
Sen --:==pttt- p SK + (1 p)
O--
(61)
E00
We rearrange for terms in ?
t-k
21+1 - + St (1 p)
p
0:-b
(62)
and before proceeding to the final result, using n
k, we compute the convenient quantities
R-1
1-p"
p
X
(63)
N
n
1-p
R-1
1
2n
0²
X
(64)
n
n
p²
Taking expectation of Equation 62 and setting statistics their stationary values, we have
(1 (1-p") B[3"] (1 p) n [[0*] = (1 - (1-p") B[0"),
(65)
where we have used the result in Equation 57. follows that for p we have
B [[O"),
(66)
independent of p. Finally, can take the variance of Equation 62. First the left hand side
Var [3++1 - p" 3x] = Var [[++1] + p²" Var [& = (1 + 2² Var [3"].
(67)
Next the right hand side
1 pmn = (1 (1- p)2
Var (1 - p) Σ
0,-1
= =1 (1-1)) p)² Var
Σ 0,1
1- p- Var[0"]
(68)
i=0
i=0
Finally, equating left and right hand sizes and rearranging for Varl(") gives
1 1 p- 1 Var [0*] 1+ p2 1 p
Var [3"]
(69)
In the limit L + the momentum-dependen prefactor becomes
1 2n
1-
1-p
1 1-P. Var [0"].
lim 1400
lim Var [3')
(70)
1+ p2n
1+
1+p
1+
Equations 69 and 70 validate our intuition. When p - 0, then ğ behaves like independent T, 0, the momentum-dependent prefactor
with their variance and expectation matching. When serves as an aggregator over the history when t is sufficiently large compared to k, reducing the variance Var[5"] ] but preserving its expectation. This formalizes the notion of time horizon discussed Table 1.
F Additional details and results for Polyak-Ruppert averaging
Additional background Polyak-Ruppers averaging (Definition 3.1) is a simplification of Stochas- tie Weight Averaging (SWA) (Izmailov et al., 2018) which uses a more complex multi-cycle sched- ule based weighting of the model parameters. Both Definition 3.1 and SWA present similar favor- able properties like wider minima and better generalization (Izmailov et al., 2018). For example, He et al. (2022) observed that a supervised ViT-H/14 overfits on ImageNet (Russakovsky et al., 2014) without a model EMA, achieving an accuracy of 80.9%. Equipping a Ruppert average (p 0.9999) alleviated overfitting and gave 83, 1 % accuracy.
27

---

Organization In this appendix, we look at additional momenta for one-dimensional noisy parabola, as well as extensions D-dimensions (Appendix F.I), provide a more detailed view of the results of Section 3.2 (Appendix F.2), and investigate the scenario where the EMA Scal- ing Rule (Definition 1.2) is applied to batch normalization (loffe Szegedy, 2015) coefficients (Appendix F.3).
F.1
Noisy parabola
Additional one-dimensional examples First we consider additional one-dimensional examples, investigating the effect of modifying the base momentum PB- We present PB 0.99 in Figure 7, and
PB 0.999 Figure 8. The results for PB = 0.9999 are presented in main text in Figure
Scaling X
Scaling X 256
Approximation Error
EMA Position L
1.0
Momentum P
1.0
0.8
Pa
Pa
0.5
EMA, P P-p
P = Phi
0.6
Pb
256
Ps
EMA.
10-
42
Model
0.4
0.0
0.0
0.5
0.5
1.0
64
256 1024
16
64
256 1024
Continuous Time
Continuous Time
Scaling Factor K
Scaling Factor K
(a) Trajectory of model EMA under different scalings k. with PB = 0.99, AB = 10 4
(b) Choices for momentum (left) with corresponding ap- proximation errors (Equation 10) (right).
Figure 7: (a) We show the effect of scaling by comparing model EMA trajectories of the baseline (k= black dashed) to x (left) and 256 (right), with (p = p blue) and without (p PB, red) EMA Scaling Rule. (b, left) The momentum according for different scaling rules and the empirically optimal p² (Equation 10). (b, right) The approximation error (Equation 10) of trajectories in (b, left) and the target model (orange). Error for p² is computed using hold-out mitigate overfitting.
Scaling 8
Scaling X 256
Approximation Error
EMA Position 5
1.0
10
1.00
Momentum
0.8
Pe
1.p
Pb
0.75
EMA,
EMA, P
P-
2t P- PE
0.50
EMA.
Ph
==
Ph
Model
0.1
0.25
0.0
0.5
0.0
0.5
1.0
16
256
1024
16
256 1024
Continuous Time
Continuous Time
Scaling Factor x
Scaling Factor k
(a) Trajectory of the model EMA under different scalings k, with Pn = 0.999, A/8 = 10
(b) Choices for momentum (left) with corresponding ap- proximation errors (Equation 10) (right).
Figure 8: (a) We show the effect of scaling by comparing model EMA trajectories of the baseline (K black dashed) to
(left) and 256 (right), with (p = P.. blue) and without (p PB, red) the EMA Scaling Rule. (b. left) The momentum according for different scaling rules and the empirically optimal p² (Equation 10). (b, right) The approximation error (Equation 10) of trajectories in (b, left) and the target model (orange). Error for is computed using hold-out mitigate overfitting.
As described by the scaling error term in Equation 54, the approximation error at given is higher for lower momenta p. For large range of scalings x, the EMA Scaling Rule and the optimal momenta p' are consistent. In summary, we see the synthetic experiments validate the results of Section 3.1 for a range of momenta p.
Our final of the synthetic noisy parabola will consider an extension to D dimensions. Consider the optimization of 0 E RD in noisy parabola at the origin:
Examples in higher dimensions
(() a eTe, 2
c₁ N Dgi
(0.
0 0 Bk.
= a 0 + ck.
(71)
for curvature 0, scaled additive b 0, and additive > 0 noise coefficients. The scaling factor x in the covariance denominator implements the reduction in gradient noise as the scaling (i.e., the batch size) increases (Jastrzebski et al., 2017). Let 0 E RD be optimized with SGD (Definition 2.1)
28

---

and let there be Polyak-Ruppert average (Definition 3.1) 3 E R with momentum p = 1- B for 0 We consider dimensionalities
(Figure 9), 16 (Figure 10), and D 100 (Figure 11). We observe no significant differences in the EMA scaling behavior as we vary dimensions.
Scaling K
Scaling K 256
Approximation Errur
EMA Norm ||
1.0
Momentum P
10"
Approximatonn
0.9
K= Pa
K 1.IE Pa
0.95
EMA, P Pe
EMA,
0.8
== P-
K 256, P5
EMA,
PE
k= Pn
256, PV
0.90
Model
0.3
1.0
16
256
16 256 1024
Continuous Time
Continuous Time
Scaling Factor
Scaling Factor
(a) Norm of the model EMA 3 under different scalings K. with PB = 0.9999, n 10 1 2.
(b) Choices for momentum (left) with corresponding ap-
proximation errors (Equation 10) (right).
Figure 9: (a) We show the effect of scaling by comparing model EMA trajectories of the baseline (x black dashed) to
(left) and = 256 (right), with (o = p.. blue) and without (p = РB. red) EMA Scaling Rule. (b, left) The momentum according for different scaling rules and the empirically optimal p² (Fquation 10). (h, right) The approximation error (Equation 10) of trajectories in (b, left) and the target model (orange). Error for p" is computed using hold-out mitigate overfitting.
Scaling X
Scaling X 256
approximation Error
EMA Norm ||G|
1.0
Momentum Q. 1.00
10P
0.9
Pa
1.PE Pa
0.95
EMA, P -p'
EMA, P =P
pia.
EMA, Pa
10-2
p
Pa
p
PM
0:90
0.0
0.5
1.0
0.0
0.5 Continuous Time
1.0
16
64
250
1074
16
64
256 1024
Continuous Time
Scaling Factor
Scaling Factor K
(a) Norm of the model EMA 3 under different scalings k. with PB 0.9999, nB = 10
(b) Choices for momentum (left) with corresponding ap- proximation errors (Equation 10) (right).
Figure 10: (a) We show the effect of scaling by comparing model EMA trajectories of the baseline (k = 1, 8 (left) and 256 (right), with (p A blue) and without Scaling Rule. (b, left) The momentum according for different scaling rules and the empirically optimal p (Equation 10). (b, right) The approximation eтror (Equation 10) of trajectories in (b, left) and the target model (orange). Error for p" computed using hold-out to mitigate overfitting.
black dashed)
Pp. red) the EMA
Scaling K ==8
Scaling X = 256
Ioproximation Error
EMA Norm (
1.0
Momentum R. $1000
10
0.9
PP
1.
EMA, Pl
EMA, P
0.8
256, pir
EMA, PE
10-2
0.90
Model
0.0
0.5
0.5
1.0
16
61
256
1024
16
64
256 1024
Continuous Time
Continuous Time
Scaling Factor
Scaling Factor
(a) Norm of the model EMA 3 under different scalings with PB 0.9999, n 10 4 D 100.
(b) Choices for momentum (left) with corresponding ap- proximation errors (Equation 10) (right).
Figure 11: (a) We show the effect of scaling by comparing model EMA trajectories of the baseline (k = 1, black dashed) to 8 (left) and = 256 (right), with (p P blue) and without (p = PB. red) the EMA Scaling Rule. (b, left) The momentum according for different scaling rules and the empirically optimal p (Equation 10). (b, right) The approximation error (Equation 10) of trajectories in (b, left) and the target model (orange). Error for p" is computed using a hold-out to mitigate overfitting.
Compute The compute usage for the noisy parabola experiments is relatively small, with each run taking less than one minute single CPU, and so we do not detail this compute usage as do in the other experimental sections.
29

---

Table Supervised ResNett22-0 hyperprmette used in Polyak-Ruppert Averaging Averaggn experiments.
Supervised ResNetv2-50
ImageNetIk Test Top-1
76.27 = 0.10%
ImageNetIk EMA Test Top-1
76.55 = 0.07%
Weight initialization
mimngnn meeiiinnnnonooeeee normal (relu) BatchNorm
Backbone normalization
Synchronized BatchNorm over replicas Learning rate schedule
No
Multi step:
x0.1 at [30, 60, 60] epochs
Learning rate warmup (epochs)
Learning rate minimum value Training duration (epochs)
1 x 10
90
Optimizer
SGD Momentum
SGD momentum
0.9
Optimizer scaling rule
Linear 0.4
Base learning rate
Base batch size
1024
Base Polyak momentum
0,9999
Weight decay
1 X 10
Weight decay scaling rule
None Yes bf16
Weight decay skip bias
Numerical precision
Augmentation stack
ImageNet
Label smoothing rate
0.1
F.2 Image Classification
Hyperparameters
We present the base hyperparameters for image experiments in Table 4.
Data For large scale vision evaluation, we use the ImageNetlk dataset (Russakovsky et al., 2014), a widely used dataset containing approximately 1.2 million labeled images, distributed almost uni- formly across 1000 different object classes, like animals, plants, and vehicles.
The images in ImageNetlk cropped to a standard size (in our case, 224 x 224), before further processing. This is part of the standard ImageNet augmentation stack for convolutional networks mentioned in Table 4.
not consistent in resolution. To handle this, they are resized and
Compute usage rized in Table 5,
The compute usage image classification Polyak-Rupperi averaging is summa-
Table 5: Compute usage for image classification Polyak-Ruppert averaging in Figures 2 and 13. The three runs for the batch size 1.024 baseline correspond three seeds, and the nine runs for all other batch sizes correspond to using and not using the EMA Scaling Rule shown in Figure 2, and its application to Batch Normalization shown in Figure 13. All experiments conducted are using 80Gb A100s.
Batch Size
GPUs
Time (h)
Compute/Run (GPUh)
Runs
Compute (GPUh)
512
35.3
282.4
9
2,541.6
1,024
17.1
137.0
410.9
2,048
13.3
106.7
960.6
4,096
4.2
33.5
9
301.9 19
8,192
16
2.8
44.8
9
403.6
All other compute, e.g. code development, runs with errors, and debugging
25,768.3
Total
30386.8
Additional results Figure 12 we present more detailed view of the results Section 3.2. First, we see that for all train metrics, model trajectories match, and that learning rate step schedule after warmup present. discussed in Figure 12, gap in EMA Test Top-1 trajectories begins scaling = 4, with more pronounced effect visible at x 8. From Figure 12 it is clear that the (non-EMA)
30

---

Test Top-1 performance trajectory is no longer matching at these scalings, demonstrating that the problem is not due a breakdown of the EMA Scaling Rule, but rather, that the model is overfitting at larger batch sizes due to batch normalization (loffe & Szegedy, 2015).
Scaling
Scaling 2
Scaling
Scaling 8
Train Model Loss
8-
512,
2048, PC
409
K192 /
2114,
8192. Ps
B
P
1024.
Train Top- (9)
60-
40
0-
Test Top-1 (96)
60
20
EMA Test Top-1 (%)
80-
40-
20-
15
30
45
63
75
90
0
15
30
45
60
75
90
15
30
4.
80
75
90
01
15
30
25
60
75
90
Train Epochs
Train Epochs
Train Epochs
Train Epochs
Figure 12: ResNerv2-50 Polyak-Ruppert averaging ImageNetlk for different scalings The baseline model (x black dashed) uses batch size 1024 and momentum pB 0.9999. scaled down batch size of 512 (left), and up to batch size 4096 (right) with (blue, P p$) and without (red, p = pB) 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 k, with and without applying the EMA Scaling Rule to the Polyak momentum. In those experiments, we maintained Batch Normal- ization (Loffe Szegedy. 2015) coefficients of pbn = 0.9 throughout" i.e. the EMA Scaling Rule is not applied. The running statistics of Batch Normalization are an EMA with values determined by PBN and so is reasonable to suspect we should apply the EMA Scaling Rule to PBN also
In Figure 13 we investigate the effect of applying the EMA Scaling Rule to Batch Normalization coefficients, using PBN PÊN- We observe that the Test Top-1 trajectories with the EMA Scaling Rule applied are slightly closer to the reference trajectories for scalings K 2 2 than those trajectories without the EMA Scaling Rule. As the effect is not particularly large, at least in this setup, do pursue further ablating applications of the EMA Scaling Rule to batch normalization coefficients, and always use PBN 0.1 for Batch Normalization, independent of
In many ML frameworks, this value is defined using BB 11 i.e. the default is 0.1 and corresponds to BBn rather than 0.9 corresponding to PBN. We use PBN maintain consistency across this work.
31

---

Sculing K 1/2
Scaling 2
Scaling
Scaling K
Train Model Loss
512, p p-mp (f)
B 2048,
4096, p p
8192,
512.p
2043,
B 1024, p- Pa
B 1024, p Pir
1024, p Py
1024, p
Train Tog-1 (%)
60
20
80
Test Top-1 (96)
40-
EMA Test Top-1 (96)
60-
40
0--
Train Epochs
Train Epochs
Train Epochs
Train Epochs
Figure 13: ResNerv2-50 Polyak- Polyak-Ruppert averaging on ImageNerlk for different scalings x. The baseline model
(k = 1. black dashed) uses batch size 1024 and momentum PB = 0.9999, is scaled down to a batch size of 512 (left). and up a batch size of 4096 (right) with the EMA Scaling Rule applied to only model parameters (blue, p pG) and model parameters and buffers (orange, p P (f)). Bands indicate the 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-10 as training data, and for semi-supervised pseudo-labeling experiments, we use train-clean- as the labeled and train-clean-360 and train-other-500 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 per- formance, measured in WER without an external language model. We maintain the original 16kHz sampling rate, and compute log-mel filterbanks with 80 coefficients for 25ms sliding window, strided by 10ms, later normalized to zero mean and unit variance for each input sequence.
Acoustic model We employ vanilla encoder-based only transformer model trained with the Con- nectionist Temporal Classification (CTC) loss (Graves et al., 2006). We use the training configura- tion from Likhomanenko al. (2021a), which has three stages: 1D convolutions perform strid- ing (kernel of 7 with stride of 3); ii) Transformer encoder with 36 layers, post-LayerNorm four attention heads, an embedding dimension of 768, MLP dimension of 3072, a dropout frequency of 0.3, and layer drop frequency of 0.3; and iii) a linear layer to map to the target vocabulary' 12 To reduce model training time by factor of approximately 2 3x, and reduce memory footprint,
12ehe token set of this vocabulary consists of the 26 English alphabet letters augmented with the apostrophe and a word boundary token.
32

---

Table 6: Hyperparameters summary for speech recognition task for supervised (leff) and bemi-supervised pseudo-labeling (right) training with vanilla transformer. The 0.3 -> 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-Labeligg
Librispeech test-clean/ / test-other WER
7.8/19.1
4.8/11.5
Optimizer
Adam
Adam
Optimizer scaling rule
Adam
Adam
Base (f1, f2)
(0.995, 0.999)
(0.995, 0.999)
Base learning rate
0.0001
0.0001
Basc lcarning rate warmup (steps)
64k
64k
Learning rate schedule
Fixed (no decay) 0
Fixed (no decay) 0
Learning rate minimum value
Base training duration (steps)
400k
500k
Base batch size (dynamic)
x 290s
8 290s
Base teacher momentum Weight decay
0,99995
0,9999
None bf16
None bf16
Numerical precision
Augmentation stack
SpecAug
SpecAug
Dropout
0.3
0.3 0.1
Layer drop
0.3 1
0.3 - 0.1 1
Gradient clipping
Labeledzunlabeled data ratio
N/A
1:3
Base pre-training steps
N/A
20k
Basc start of EMA accumulation (stops)
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 = 8x290s, which adapted from Likhomanenko et al. (2021a), and is summarized in Table 6. We use SpecAugment (Park et al., 2019) activated after steps of training: two frequency masks with frequency mask parameter F= = 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 (Bi 0 and Bg - 1) of Adam (Kingma & Ba, 2015), we were unable to produce reasonable optimization in practice when applying the Adam Scaling Rule of Malladi al. (2022) this limit. As conse- quence, we chose to work with the Adam optimizer, where its scaling rule has been shown to work (Malladi al., 2022), and we take fi = 0.995, f2 = 0.999, and = 10-8 We obtained similar results for &1 0.99. Finally, we use linear learning rate warmup (64k steps) after which the learning rate kept constant until convergence. This performance can be improved further by using 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 pseudo- labels 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.
Definition G.1 (Hard Transcription). For sequence of frames, select the most probable token per frame, removing repetitions and the CTC blank token. For example, "h#fteell#fl##Hoo"" 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 have access to the ground-trutt labels on the data being treated as unlabeled, we track
33

---

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 evaluate progress on unlabeled data, as loss on pseudo-labeled data is unreliable when 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 fixed number of sequences (Ott et al., 2019). In our work, use dynamic batching of ~290s 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 more challenging setting for achieving systematic scaling, the number of independent sequences in any given batch may change, and the i.i.d. assumption does not hold 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 the case of dynamic batches, as does our EMA Scaling Rule (Definition 1.2).
The base batch size is set to 8 x 290s, and in our experiments we scale down to batch size of B 2 x 290s and up to batch size of B 128 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.I.
t 16 precision training. While for supervised training evaluation of different EMA decay values is done in parallel during single run, for pseudo-labeling every EMA decay value needs separate training. Final mod- els training compute is detailed in Tables 7 and 8. Total compute, including e.g. code development, runs with errors, and debugging, is 61k GPUh.
Compute All experiments are done using A 100 80GB 8GPU nodes with
Table 7: Compute usage for supervised model for speech recognition task in Figure 14. Values include node allocation times (typically small % of corresponding total runtime), giving practical estimate of reproduction cost. All experiments conducted using 80Gb A100s with fast interconnect.
Batch Size
GPUs
Time (h)
Compute/Run (GPUh)
Runs
Compute (GPUh)
2 290s
2
222
4 290s
4
108
432
432
8 x 290s
8
64
512
512
16 x 290s
16
54
864
896
32 290s
32
37
1,184
1,184
Total
3,436
Table 8: Compute usage for continuous pscudo-labeling for the specch recognition task in Figure 16. Values include node allocation times (typically small % of corresponding total runtime), giving practical estimate reproduction cost. All experiments conducted are using 80Gb A100s with fast intercounect.
Batch Size
GPUs
Time (h)
Compute/Run (GPUh)
Runs
Compute (GPUh)
x 290s
225
550
1,110
x 290s
120
480
2
960
8 x 290s
72
576
I
576
16 290s
16
45
720
2
1,440
290s
33
1,056
4.224
64 x 290s
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 Fig- ures 14 and 15 for supervised training and in Figures 16 and 17 for semi-supervised training.
First, observe that if the Adam Scaling Rule does not hold perfectly (there is mismatch between trajectories for the model before pseudo-labels are involved) the EMA Scaling Rule also
See Malladi al. (2022) for discussion scenarios that lead breakdown of the Adam Scaling Rule.
34

---

Scaling 1/4
Scaling K 1/2
Scaling 2
Scaling
Train Loss
400
200
EMA lest-other WER test-other WER
100
80
60
40
2000
290
p Ps
200
2906,
4x 290%
8 16 290,,
8- 2900
AB
290,
40
20
Train Steps
Train Steps
x10
Train Steps
%10
Train Steps
Figure 14: Transformer Polyak-Ruppert averaging LibriSpeech (trained train-clean-100) with different
scalings x. The baseline (x = 1, black dashed) is trained with Adam and momentum pg 0,99995 at dynamic batch size B = 8 x 290s, which corresponds to a single train step on the x-axis. We investigate dynamic batch sizes down to B = x 290s (left) and up to B = 32 x 290s (right). with (blue, p= Pf) and without (red, p PB) the EMA Scaling Rule (model Mon-EMA marked by orange). The Adam Scaling Rule (Malladi al. (2022), Definition C.3) used throughout. For momentum PB 0.9999 we observe similar trajectories for all models.
Scaling 1/4
Scaling k = 1/2
Scaling x=
Scaling K
Train Loss
400
200
EMA test test-olhin other WER test test-other other WER
100
80
60
20-
R= 8x 700
8xx 100s. Ax
100s
Nm 700r P
B- 200
4. 290
B 16 200,,
8- 9- 200
2=02200.
16x 290., P8
250.
60
40
20
Train Steps
x10
Train Steps
x10
Train Steps
x10
Train Steps
Figure 15: Transformer Polyak-Rupper averaging on LibriSpeech (trained on train-cleant
with different
scalings K. The baseline (k = 1, black dashed) trained with Adam and momentum Pn = 0.999 a dynamic batch size B = 8 x 290s, which corresponds to a single train step on the x-axis. We investigate dynamic batch sizes down to B = x 290s (left) and up to B = 32 x 290s (right), with (blue, p p.) and without (red, p PB) the EMA Scaling Rule (model -onAEMA marked by orange). The Adam Scaling Rule (Malladi al. (2022), Definition C.3) is used throughout. If momentum PB small and accumulation history short observe no any significant difference between models which all are matching the reference trajectory despite scaling k.
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 breakdown of the Adam Scaling Rule for large 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
able match the reference trajectory up to k 8.
35

---

Scaling
Scaling
Scaling
Scaling
Scaling
600
P P
2000,0
250h, B 16 250 Af
32x 200, Po die 32x 200,
3526 P B-6d 290m
600
B 2: 200, 0 pf
8- 200,
of
Sepervisol
200
100
1.WERR
40-
20
100
WEi
73
Steps
X300
Steps
Trin Steps
1000
Train Staps
100P
Steps
X10
Figure 16: Transformer pseudo-labeling on LibriSpeech (Irained on train-clean- 100 as labeled and the rest of LibriSpeech as unlabeled) with different scalings The baseline (x 1, black dashed) is trained with Adam a dynamic butch size of 290 seconds, which corresponds model EMA (teacher) is updated with momentum PB 0.9999, We investigate dynamic batch sizes down to = 2 x 290s (left) and up to B = 64 x 290s (right), with (blue, P = ff) and without (red, P PB) the EMA Scaling Rule. The Adam Scaling Rule (Malladi et al. (2022), Definition C.3) is used throughout. For K S 2, we start pseudo-labeling after 20k/K training steps; while for > 2, we start when pre-training WER matches the baseline WER (24K/K for and 29k/x for 8). For = experimented with both variants: we start pseudo-labeling after 20k/x (dashed) and when pre-training WER matches the baseline WER (solid, 24k/x).
single train step on the x-axis. The
Supervised Train Loss
Scaling k= 1/4
Scaling = 1/2
Scaling
600
B-8> 250. - Ps
B 250 Pa 4x 290r
B- 8ix 290 p Ao # 250.,
#=2x2200,
400
2000,p
200. -
200
80
PL WER
60
100
test-other WER
60
0
Train Steps
100
Train Steps
X10
Train Steps
110
Figure 17: Transformer pseudo-labeling on LibriSpeech (using train-clean- 100 as labeled) with different scal- ings K. The baseline (x = 1. black dashed) is trained with Adam at a dynamic batch size of x 290 seconds, which corresponds to single train step on the x-axis. The model EMA (teacher) updated with momentum PB 0.999. We investigate dynamic batch sizes down to B 2 290s (left) and up to x6x2900 (right). with (blue, p = pT) and without (red, p = PB) the FMA 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 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 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.
36

---

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 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 PB. scaling with or without EMA Scaling Rule behave similarly, and reference trajectories match the supervised and semi- supervised cases. However, if the momentum is too large, the teacher moves slowly and is uninfor- mative, whereas if the momentum is too low, the teacher and the student are effectively be the same model, implying: i) the student will self-pree predict with high confidence, removing any benefits of dis- tillation 4; and ii) training instability or model divergence will happen in the low-resource settings (Likhomanenko et al., 202121 Higuchi et al., 2022).
G.2 Scaling to x 16 with Progressive Scaling
Finally, we aim scale semi-supervised pseudo-labeling further to K 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 к = 8 for 29k of reference steps (model EMA accumulation starts at 28k steps). We then scale to к 16 and begin pseudo-labeling We in Figure 18 that Progressive Scaling enables us to scale pseudo-labeling to 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 x 16 until 75k 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 x = 16, corresponding to dynamic batch size of 128 290s.
110- Scaling 8
Scaling == 8 16
Scaling 16
Leannng Rac
60
WER
40-
0
lost-othr WER
100
B Ax 2M05,
8 Nx 7M0
8 8x 750
2000,
2000,p
128 2005 2
R= 64 - 138% 7906
8 61 178 200
Train Stops
x 10
Train Steps
x10'
Train Steps
X 10
Figure 18: Transformer pseudo-labeling on LibriSpeech (trained on train-clean- 1000 as labeled and the rest of LibriSpeech as unlabeled) with different Progressive Scaling from k = 8 to k = 16 (x = 8 - 16). The baseline (x = 1. black dashed) is trained with Adam at dynamic butch size of 8 x 290 seconds, which corresponds single train step the x-axis. The model EMA (teacher) updated with momentum Pa 0.9999. The scaling with
shown with lighter color for reference from Figure 16. We investigate dynamic batch sizes progressively from B = 64 x 290s to 128 x 290s, with (bluc, p = p8) and without (red, p = pB) the EMA Scaling Rule. For reference (top) we show the lcarning rate schedule with Progressive Scaling. The Adam Scaling Rule (Malladi et al. (2022), Definition C.3) is used throughout. Left and middle correspon to Progressive Scaling with scale from
16 at 29k steps, while right corresponds to 75k steps.
1e et al. (2020) alleviated the problem with the proper amount of noise during student model training, whilst Xu et al. (2020) used beam-search decoding with language model.
37

---

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. our recipe for training BYOL using ResNet 18s (Ap- pendix H.2), our recipe for training BYOL using Vision Transformers (ViTs) (Appendix H.3), abla- tions 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 understand- ing the impact of Progressive Scaling (Appendix H.6).
Second, we demonstrate that the EMA Scaling Rule combined with Progressive Scaling can scale ResNet-50 BYL model trained with LARS batch size 32,768 without performance drop, demonstrating the empirical utility of the tools we provide outside of their theoretical validity (Ap- pendix H.
Finally, we show that it is possible to systematically scale DINO (Caron et al., 2021) using a com- bination of Progressive Scaling and the EMA Scaling Rule, providing 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 lop-ggraient or StopGrad (Definition H.1). Definition H.1 (Stop Gradien/SStoGrrad( )). The stop-gradient operator StopGrud( prevents the flow of gradient information
df (tottpCCratt((x: =
(72)
deo
for all parametric functions and f and all parameters 0 and w.
Applying stop-gradient is sometimes called detaching in the literature. Now, we introduce the update rule of our representative SSL method BYOL Definition H.2.
Definition H.2 (BYOL Update). BYOL learns unsupervised features by minimizing the cosine dis- tance between the predictions of a student backbone f(:00) (typically a ResNet or Vision Trans- former), projected through h(::w)) (typically Multi-Layer Perceptron (MLP)), and the predictions of an EMA teacher f(:;5) (Grill et al., 2020). The update for the parameters of BYOL is then
(0/+1, (0+11 = (0,, (0)) - -
1
Σ
V (o...)) L(x; c0 5,)
(73)
B
x
S++1 = p 3, + (1 - p) 0/+1
(74)
1
with L(x; 0,, w. {,) cos [hf (x]f(x;:))) w) StopGrad(f(ss;;3.) + (x1 X2).
(75)
where cos(a, a, b) = (((aal/(((aa])))) is the cosine distance, and xı and x2 are two views of a single variate x, often produced by augmentations, and x1 + X2 denotes symmetrization over X1 and x2.
As noted in Section 3.e BYOL EMA update (Equation 74) uses 0,41 instead of our analyzed 0, (Equation 4). The effect upon the overall EMA update is O Bp) and is captured by the EMA Scaling Rule (Definition 1.2).
One more piece of technology typically employed in SSL is tracking probe (Definition H.3) which we will use to evaluate the performance of BYOL on downstream tasks of interest, for example, image classification.
Definition H.3 (Tracking Probe/Linear Probe). When optimizing model parameters or of an SSL method, simultaneously optimize the parameters & of a probe model r( :)))) under a downstream objective [(d), For example, in classification, with data x and samples
L (d) (x. O,, E,) = - log P(y (Str(((oorrrd ())));
(76)
(total) (x, 0 wi T.. 5,) [x; 0, G {,) ((d) (x, y, 60, E.).
(77)
The is a probe for the teacher, which typically the better choice due to Supperr averaging effects (see Section 3.2). When the r is linear model, the tracking probe is called linear probe.
38

---

Table 9: BYOL ResNet-18 hyperparameters for CIFARIO
ResNet-18
Weight initialization
kaiming_um uniform (He et al., 2015)
Backbone normalization
BatchNorm
Head normalization
BatchNorm
Synchronized BatchNorm over replicas
Yes
Learning rate schedule
Single Cycle Cosine
Learning rate warmup (epochs)
20
Learning rate minimum value
0
Training duration (epochs)
100
Optimizer
SGD
Optimizer scaling rule
SGD
Optimizer momentum
0.9
Gradient clipping
0.1
Base learning rate
0.02
Base hatch
1024
Base teacher momentum
0.992
Weight decay
1 X 10-6
Weight decay scaling rule
None
Weight decay skip bias
Yes
Numerical precision
Lf32
Augmentation stack
BYOL CIFAR10
It is also typical to use Batch Normalization layer without trainable affine terms before this linear layer as in He et al. (2022) to stabilize probe training. this case, the running statistics can be absorbed into a definition of the linear layer weights and biases, and so this is still a linear probe, although we will call this pre-bn linear probe to remove ambiguity.
H.2 A ResNet-18 recipe for BYOL
Hyperparameters We present the base hyperparameters for training BYOL with a ResNet-18 backbone using SGD in Table This recipe was developed by starting from a well-known BYOL ResNet-50 recipe (Grill et al., 2020), adapting the input augmentations for CIFAR10, and performing a search over learning rate choices for SGD optimizer.
H.3 A Vision Transformer recipe for BYOL
Hyperparameters We present the base hyperparameters for training BYOL with a ViT-B/16 back- bone in Table 10. This recipe was developed by starting from well-known supervised ViT-B/16 recipe (He et al., 2022) and performing a search over weight decay and learning rate hyperparameter choices. We find that BYOL performs well with heavy weight decay (2 0.3) and low learning rate (1 = 3) ) at a base batch size B = 4096. The AdamW optimizer is used, and so for scaling to other batch sizes B- = xB we use the Adam Scaling Rule (Definition C.3) We use a pre-bn linear probe as discussed in Appendix H.1. Finally, the performance of BYOL can be further improved by employing multicrop (Caron et al., 2020) by = +2% in absolute test top-1 performance on Im- ageNetlk compared to without multicrop, however, as this is not our focus, we omit this from the presented recipe.
Additional background Achieving large scale SSL training with ViTs to large scale SSL train- ing has been a long standing goal in the community. MoCo-v3 (Chen et al., 2021) enables the use of ViTs with contrastive learning, but achieves this through modifications of the ViT training procedures, including gradient freezing on the image patching layer, and re-introducing Batch Nor- malization to post-attention MLP layers. Despite these modifications, MoCo-v3 was only trained up to batch size of 6144, where model performance begins to suffer (Chen et al., 2021). In Figure 6 we demonstrate that combining dynamic batch scaling (Appendix C.4) with the EMA Scaling Rule
We note that Adam (Kingma & Ba, 2015) and AdamW (Loshchilov & Hutter, 2019) are equivalent in the limit of zero weight decay, and that the Adam Scaling Rule (Definition C.3) was derived with zero weight decay (Malladi et al., 2022).
39

---

Table 10: BYOL ViT-B/16 hyperparameters.
BYOL VIT-B/16
ImageNetIk Linear Probe Test Top-1
74.47% (Figure 19)
Weight initialization
trunc normal (.02)
Backbone normalization
LayerNorm BatchNorm
Head normalization
Synchronized BatchNorm over replicas
No
Learning schedule
Single Cycle Cosine
Learning rate warmup (epochs)
40
Learning rate minimum value
1 x 10
Training duration (epochs)
480
Optimizer
AdamW
Optimizer scaling rule
Adam (0.9, 0.95)
Base (f1.82)
Base learning
1 x 10-3
Base batch size
4096
Base teacher momentum
0.99
Weight decay
0.3
Weight decay scaling rule
None Yes bf16
Weight decay skip bias
Numerical precision
Augmentation stack
BYOL (Grill et al., 2020)
Stochastic depth
0.1
(Definition 1.2) enables BYL to be trained using ViTs batch sizes of 24,576 without any drop in performance compared to the reference batch size of 4096. We emphasize that the piecewise transitions in the schedules are important for preserving training dynamics.
H.4
The role of Batch Normalization and Layer Normalization in BYOL with ViTs
x
2.5
BN. B-3072
70-
BN. B=24576
60
I R-9072
Teacher Momentum
0.998
2.0
LN, B-24576
Test Top-1 (96)
Train Loss
0.6
0.996
Learning Rate
1.5
30
0.4
0.994
1.0
20
0.2
10-
0.900
0.0
200
400
0
200
400
0
200
400
0
200
400
Train Epochs
Train Epochs
Train Epochs
Train Epochs
Figure 19: BYOL ViT-B/16 on ImageNetlk for different scalings x. present runs comparing LayerNorm (bluc) to BatchNorm (red) in the projection and prediction heads of BYOL ViT models for batch size 3072 (dashed) and 24,576 (solid) without the EMA Scaling Rule. corresponds to B 4096. In all scenarios the transformer backbone only LayerNorm. We truneate the training of the large batch size LayerNorm variant to preserve compute (indicated by x).
Here compare the roles of Batch Normalization (BatchNorm, Ioffe & Szegedy (2015)) and Layer Normalization (LayerNorm, Ba et al. (2016)) in the projection and prediction heads of B BOL (Grill al., 2020) using ViTs.
been observed that BatchNorm plays a critical role BYOL predictor and projector dynam- ics (Fetterman & Albrecht, 2020), and using either LayerNorm normulization significantly decreases model performance. Subsequently, was demonstrated (Richemond al., 2020) that competitive BYOL performance could be achieved through combination of Group Normaliza-
40

---

tion (GroupNorm, Wu & He (2018)) and Weight Standardization (Qiao et al., 2019). Additionally, Richemond et al. (2020) showed that if BatchNorm used the backbone, one can use LayerNorm or normalization in the predictor and projector without any performance drop.
In this work, bone, projector and predictor (see Figure 19), decoupling BYOL's reliance on batch statistics, desirable trait for a representation learning algorithm (Brock et al., 2021). At batch size 3072, using LayerNorm in the predictor and projector achieves competítive performance (74.10%) performing slightly worse than using BatchNorm (74.47%). At the larger batch size of 24,576, runs perform significantly worse as the EMA Scaling Rule was not applied.
show it is possible to train BYOL ViT using only LayerNorm across the back-
H.5
Longer training duration with incremental Progressive Scaling
x10
2.5
70-
0.7
0.99
60-
0.6
2.0-
0.98
Test T6p-1 (15)
50
3072
Train Loss
0.5
0.97
Learnnng Rate
10-
B=4144
0.4
0.96
30-
== 9216
0.3
Teacher
1.0-
12288
0.95
B 15360
0.2
8 = 18432
0.94
0.5
10-
B = 21501
0.1
74576
0.0
0.92
0.0
100
400 Train Epochs
200
300
500
100
400 Train Epochs
200
300
500
100
300 Train Epochs
200
400
500
100 200 300 400 500
Train Epochs
Figure 20: BYOL ViT-B/16 on ImageNetk for different scalings k. The baseline model (k = 0.75, black dashed) uses batch size 3072 and teacher momentum = 0.99. We increment the batch size by 3072 every 60 epochs to a final batch size of 24,576 using Progressive Scaling (Definition 3.2).
Here the same base hyperparameters Table 10, except that train for 480 instead of 300 epochs. To mitigate the student impulse phenomena discussed in Section 3.4, in Figure 20 we investigate increasing the batch size every 60 epochs using Progressive Scaling (Definition 3.2). We observe that this more gradual procedure enables closer tracking of the baseline train loss trajec- tory. Additionally, this procedure results in a scaled linear probe performance that outperforms the baseline (75.64% compared to the baseline performance of 74.47%). ). The same procedure can be ap- plied to the LayerNorm variant discussed in Appendix H.4, which produces similar result (75.09% compared to the baseline performance of 74.10%).
H.6 Building intuition around Progressive Scaling and momentum sensitivity
Our final BYOL VIT results to help build intuition around Progressive Scaling (Definition 3.2), as well as when the EMA Scaling Rule is most important. In Figure 21 we explore transition- ing from the baseline batch size 4096 model to batch size 24.576 in single transition after 60 epochs. After this transition, we continue training for 240 epochs for a range of momenta: p E (0.8, 0.9, 0.95, 0.97, 0.9867, 0.994, 0.999) without the EMA Scaling Rule.
We observe that after the transition, any 0.9 S p S 0.994 produces a linear probe performance that matches or outperforms the baseline at the end of training. This indicates that after the initial training period, BYOL becomes less sensitive to the choice of teacher momentum. Note that without the initial 60 epochs of training with batch size 4096, all models, including those employing the EMA Scaling Rule diverge (see B = 24, 576 in Figure 6).
We present an illustration for why this might happen in Figure 22. First, we see that using the EMA Scaling Rule always keeps the model within the acceptable momentum region. We also wee that not using the EMA Scaling Rule can keep the model within the acceptable momentum region for range of batch sizes, depending on how large wide in momenta the acceptable region at the base batch size. Finally, see that the momentum value matters much more at low values of momenta
41

---

x10
1.000
70-
0.7
0.975
60
0.950
Teacher Momentum
Test Top-1 (56)
50
Learning Rate
Train Loss
0,925
10
0.900
0.4
80
0.800
0,900
0.3
0.875
20
0,95
0.850
0.970
0.2
0.5
10
0.987
0.825
0.994
0.999
0.800
Train Epochs
Train Epochs
Train Epochs
Train Epochs
Figure 21: BYOL VIT-B/16 on ImageNetIk for different momenta p. The baseline model (p 0.99, black dashed) batch size 4096. At the 60th epoch we apply Progressive Scaling (Definition 3.2) and transition batch size 24576. We train for a further 240 epochs without EMA scaling for range momenta: p {0.9, 0.95, 0.97, 0. 98987, 0.994).
(the acceptable momentum region shrinks), whereas large momenta, this region of acceptability widens.
1.00
0.95
Momentum
0.90
Pmnin
Pmax
0.85
0.80
128
256
1024 2048
4096
Batch Size
Figure 22: A hypothetical scenario where there upper and lower limit
qualitatively leading
the same result. We assume at batch size 1024 there is an upper (Pmax, black dashdot) and lower (pmin. black dashed) limit for valid momenta. We show what happens if we start with PB = 0.95 at batch size of 4096, and scale with (p P blue) and without (p PB. red) the EMA Scaling Rule.
H.7 Compute usage for VIT BYOL investigation
We now summarize the compute usage for the BYOL ViT experiments. First we detail cost of reproducing Figure 6 in Table 11, Figure 19 in Table 12, and Figure 20 in Table 13.
Table 11: Compute usage for baseline VIT BYOL investigation in Figure 6. Values include node allocation times (typically small % af corresponding total runtime), giving practical estimate of reproduction cost. All experiments conducted are using 80Gb A100s, experiments indicated by (f) have faster interconnect.
Batch Size
GPUs
Time (h) Compute/Run (GPUh)
Runs
Compute (GPUh)
4,096
32
16.6
531.4
1,062.7
8,192
48
14.1
678.0
2
1,356.1
16,384
96
8.3
800.3
1,600.6
24,576
128
850.1
3,400.4
32,768
176
4.1
721.7
2
1,443.4
Total
8,863.1
42

---

Table 12: Compute usage for ViT BYOL investigation into BatchNorm and LayerNorm variants in Figure 19. Values include node allocation times (typically small corresponding total runtime), giving a practical estimate of reproduction cost. All experiments conducted are using 80Gb A A000s, and were run for 480 epochs, except those indicated by (+) were truncated early (see Figure 19 for more details).
Batch Size
Normalization
GPUs
Time (h)
Compute (GPUh)
3,072
BatchNorm
16
47.9
766.0
3,072
LayerNorm
48.0
768.7
24,576
BatchNorm
128
14.8
1900.4
24,576
LayerNorm
128
3.5
451.1
Total
3,886.2
Table 13: Compute usage for VIT BYOL investigation into incremental scaling in Figure 20. Values include node allocation times (typically small % of corresponding total runtime), giving practical estimate of production cost. experiments conducted are using 80Gb A100s 60 epochs. Stage 0 corresponding to the baseline Figure 20 is the run detailed the first row of Table 12, using batch size of 3,072, Batch Normalization, and 16 GPUs. Computing only the first 60 epochs of stage 0 corresponds to approximately 127.7 GPUh, which would bring the total cost of Figure 20 to 1,432.9 GPUh.
Stage
Batch Size
GPUs
Time (h)
Compute (GPUh)
6,144
32
3.5
113.0
9,216
48
149.8
3
12,288
64
2.8
176.0
4
15,360
80
2.3
186.5
5
18,432
96
2.1
202.9
21,504
112
2.1
235.8
24,576
128
1.9
241.3
Total
1,305.2
Next, the cost of single momentum ablation presented in Figure 21 is 240 epochs at batch size 24,576, which is 240/480 x 1900.4 GPUh 950.2 GPUh, giving total cost over seven runs of 6651.4 GPUh.
Finally, providing full view of the investigations carried out for the ViT BYOL given in Table 14.
Table 14: Total compute usage for ViT BYOL investigations.
Compute (GPUh)
Baselines (Figure 6 and Table 11)
8,863.1
BatchNorm and LayerNorm (Figure 19 and Table 12)
3,886.2
Incremental scaling (Figure 20 and Table 13)
1,305.2
Momentum ablations (Figure 21)
6,651.4
All other compute, e.g. code development. runs with errors, and debugging
84,984.1
Total
105,690.0
H.8
ResNet-18 hyperparameter sensitivity analysis
To demonstrate that the EMA Scaling Rule works for broad range of optimization hyperparameters (i.e. beyond those presented in Figure 5 and Section 3.4), we provide sensitivity analysis for base teacher momentum PB and base learning rate TJ8 in the challenging setting of BYOL.
Base teacher momentum In Figure 23 we show the effect of changing the base teacher momen- tum PB. defined at batch size 1024. The EMA Scaling Rule is robust to modifications of momentum down to PB 0.946 this particular setting. Below PB 0,946, matching poor, although the smallest momentum in this setting corresponds to 0.841 æ 0.5, which is particularly small teacher momentum, and is unlike to provide utility over the using the target model (see Appendix E.2).
43

---

Scaling K
Scaling
Scaling K
Scaling K
75
75
Test Top- 1 (%)
Test Top- I (%)
50-
90-
f 2018,
B 2096,
Pi
25-
D- 2018, -P
4096.
2048,
B
1024, p
Pb
1024. Ps
1024, Pa
P2
Thain Loss
Thain Loss
0.4
0.2
0.2
25
50
75
100
25
50
73
100
25
50
75
100
25
50
75
100
Train Epochs
Train Epochs
Train Epochs
Train Epochs
(a) PB = 0.841
(b) 0.946
Scaling K 2
Scaling K
Scaling ==
Scaling K
75
Test Top-1 (%)
Test Top-1 (%)
50
50
25-
2048,
40% P 4006 Pe
25
P
2048,
2048,
-v
0-
R
1024.
B 1024, A
1074,
1074
Train Loss
0.6-
Train Loss
0.6
75 100
100
25
75
100
Train Epochs
Train Epochs
Train Epochs
Train Epochs
(c) PB 0.974
(d) PB 0.987
Scaling
Scaling
Scaling ==2
Scaling X
Test Top-1 (%)
Test Top-1 (%)
25
2018, p P5
6096 - Ps 4090,
2018,
P5
4096
Ps
2048,
8
2048,
4000
0-
R = 1074,
1974
1024
1074
Thain Loss
0.4-
Train
0.4
0.2
0.2
50
75
100
100
25
75
75
100
Train Epochs
Train Epochs
Train Epechs
Train Epochs
(e) PB = 0.992
(f) PB 0.997
Figure 23: R8 BYOL CIFARIO teacher momentum sensitivity (188 = 0.08) for scalings x (2.4) and base teacher momenta рB (0.841. 0.946, 0.974, 0.987, 0.992, 0.997) defined dashed) uses batch size 1024, and scaled from batch size 2048 (left) to 4096 (right) with (blue, p pB) and without (red, p PB) the EMA Scaling Rule. Bands indicate mean and standard deviation across three runs.
The baseline (k black
Base learning rate In Figure 24 we show the effect of changing the base learning rate nJs defined at batch size 1024. The EMA Scaling Rule is robust over a wide range of learning rates. At the largest learning rate na 0.5 matching starts to become poor scaling 4.
44

---

Scaling K
Scaling
Scaling K
Scaling K
75
75
Test Top- 1 (%)
Test Top- 1 (%)
50-
50-
f 2018,
2096, Pi
25-
D- 2018, -P
4096. 1924, Pb
4EN% Pt
P Pb
1024. Pr
1024, Pa
Thain Loss
Thain Loss
0.4
0.2
0.2
25
50
75
100 0
25
50
75
100
25
50
75
100
25
50
75
100
Train Epochs
Train Epochs
Train Epochs
Train Epochs
(a) = 0.01
(b) 1B = 0.02
Scaling K 2
Scaling K
Scaling ==
Scaling X
75
Test Top-1 (%)
Test Top-1 (%)
50
50
25-
2048.
40% ps 4006 Pv
25
P
40%
2048, AN
2048,
4006
0-
R
1024.
1024, A
1024,
1034
Train Loss
0.6-
Train Loss
0.6
75 100
100
25
100
Train Epochs
Train Epochs
Train Epochs
Train Epochs
(c) n/ 0.04
(d) 1B 0.15
Scaling
Scaling
Scaling K == =
Scaling X
75
Test Top-1 (%)
Test Top-1 (%)
25
2018, p P5
4096
- Ps 1974
== 2018, P5
4096
2048,
8
4090
2048,
4096
0.
R = 1074,
1074
1074
Thain Loss
Train Loss
0.6
0.4-
0.4-
0.2
0.2
75
100
100
25
50 75
100
Train Epochs
Train Epochs
Train Epochs
Train Epochs
(e) y 0.2
(f) BB 0.5
Figure 24: R8 BYOL on CIFARIO learning rate sensitivity (pp = 0.992) for scalings x (2, 4) and base learning rates пв e [0.01, 0.02, 0.04, 0.15, 0.20, 0.50) defined batch size 1024, and is scaled from batch size 2048 (left) 4096 (right) with (blue, p = pB) and without (red, p= = PB) the EMA Scaling Rule. Bands indicate mean and standard deviation across three runs.
The baseline
black dashed) uses
H.9 ResNet-18 additional scaling analysis
To demonstrate that the EMA Scaling Rule works for a broad range of scalings k (i.e. beyond those presented in Figure and Section 3.4), investigate scaling down to 1/8 in Figure 25. We see that the EMA Scaling Rule works well down to the small batch size of 128, although matching is not perfect. We suspect this is due the presence of Batch Normalization layers in the ResNet- 18 architecture, which underperform at small batch sizes (loffe & Szegedy, 2015). The synthetic analysis of Section 3.1 instead demonstrated the EMA Scaling Rule holding for scalings spanning factors of x that differ by 1024, with scaling error insensitive to the value of x for sufficiently low (see Figure lb).
45

---

Scaling x= 1/8
Scaling K 1/4
Scaling == 1/2
Test Top-1 1 (%)
50
25
128. p Ph
p Pb
-512, P
B 128.
B
256,
1024,
PB
1024, Pb
Train Loss
0.6
0.2
25
50
75
100 0
25
50
75
100
25
50
75
100
Train Epochs
Train Epochs
Train Epochs
Figure 25: ResNet-18 BYOL on CIFARIO lower sealing analysis (nh = 0.08, Ps = 0.992) for scalings x (1/8, 1/4, 1/2). The baseline (x = 1. black dashed) uses batch size 1024. and is scaled from batch size 128 (left) to 512 (right) with (blue, p P and without (red, PB) the EMA Scaling Rule. Bands indicate mean and standard deviation across three runs.
H.10 Scaling a ResNet-50 BYOL using LARS and Progressive Scaling
Here investigate whether Progressive Scaling and the EMA Scaling Rule can be used in practice where there is no known optimizer SDE approximation. We use the default 300 epoch configuration for BYOL (Grill et al., 2020) in Figure 26. We see that although trajectories during training do not match, we are able to match or surpass the linear probe performance of the BYOL baseline at the larger batch size if 32.768. This indicates that the contributions of our work have practical utility beyond the theoretical constraints.
The compute usage for the BYOL ResNet using LARS is detailed in Table 15.
Table 15:
Compute usage for ResNet LARS investigation Figure 26. Values include node allocation
times (typically a small % corresponding total runtime), giving practical estimate of reproduction cost. All experiments conducted are using 80Gb 100s.
Batch Size
GPUs
Time (h)
Compute (GPUh)
4, 0999
32, 768 (120 Epochs)
128
14.1
1809.8
090
32, 768 (60 Epochs)
128
12.9
1655.9
4, 090
32.8
524.9
All other compute, e.g. code development, with errors, and debugging
17,654.6
Total
21645.2
46

---

Progressive Scaling 60 Epochs
Progressive Scaling for Epochs
Test Top-1 (%)
60
40-
B
4096
32, 768 (60 Epochs)
8 4096
32, 768 (120 Epochs)
4296
Parum Norm
300
250
0.6
Teacher STD
0.4
0.2
Teacher Momentumn Train Loss
0.4
0.2
1.00
0.98
0.96
Learning Rate
10
32768
Batch Size
16384
8192
0
50
100
150
200
250
300
50
100
150
200
250
300
Train Epochs
Train Epachs
Figure ResNer50 BYOL on ImageNetIk using LARS for different configurations progressive scaling. The baseline (black dashed) uses batch size 4096 and momentum = 0.99. We consider progressive scaling (blue) smoothly from epoch 60 for 60 epochs (left) and 120 epochs (right) up until batch size 32,768, scaling the learning rate lincarly, and applying the EMA Scaling Rulc.
H.11 Preventing collapse phenomena in DINO at scale
Until now, our representatives SSL method has been BYOL for reasons discussed in Section 3.4. Here, we will turn our attention to Distillation with NO labels (DINO) (Caron et al., 2021), which has the update rule presented in Definition H.4.
Definition H.4 (DINO Update). DINO learns unsupervised features by matching predictions over emergent pseudo-labels of a student buckbone and head f(::0) to those of an EMA teacher f(:;3) through a cross-entropy guided distillation procedure. DINO has a additional centering procedure, which is a form of batch normalization with momentum Pe 0.9 which we do not scale using the
47

---

Table 16: DINO ViT-B/16 Training hyperparameters.
DINO ViT-B/16
CIFARIO Linear Probe Top-1 (pB 0.996) CIFARIO Linear Probe Top-1 (pB = 0.992)
85.38%
86.96%
Weight initialization
trunc norma (- 02)
Normalization
Layer Norm Single Cycle Cosine
Learning rate schedule
Learning warmup (epochs)
50
Learning rate minimum value
1 x 10-6
Training duration (epochs)
280
Optimizer
Adam
Optimizer scaling rule
Adam (0.9, 0.95)
Base (fI, (2)
Base learning rate
3 x 10-4
Base batch size (B)
1024
Base teacher momentum (Pg)
0.992 0.996
Base weight decay
0.04
Weight decay scaling rule
Linear
Weight decay skip hias
Yes
Center Momentum
0.9
Center Momentum Scaling Rule
None
Precision
bf16
Augmentation stack
DINO multi-cua (Caron et al., 2020)
EMA Scaling Rule. The update for the parameters of DINO is
1
Σ
0z11 17- B
(78)
xCB
S++1 p5. + (1 0
(79)
C+11 Pc C, + (1 Pc) Ex
(80)
with L(x; S, Hff(x f.). 3i) - ct) (x1 + X2).
(81)
where H(a.b) = MM Pm a) lm is the cross-entropy between categorical distributions over M (emergent pseudo-Jclasses Jclasses given logits a. b e RMM x1 and X2 are two views of single variate often produced by augmentations, and x + X2 denotes symmetrization over X1 and x2-
In practice, DINO employs multi-crop (Caron et al., 2021). We omit this detail for clarity of presen- tation, although we do use multi-crop in the experiments that follow.
Our interest DINO is due the difficulty in its optimization and in particular, preventing collapse phenomena in DINO at batch sizes above 1024, which is an open research problem. In this section, we will show that combination the EMA Scaling Rule (Definition 1.2) and Progressive Scaling (Definition 3.2) enable training of DINO beyond batch size 1024 without sacrificing performance.
Hyperparameters
Base hyperparameters are presented in Table 16.
Results In Figures 27 and 28 we show the results obtained training DINO on CIFAR-10 with PB = 0.996 and PB = 0.992 respectively at the reference batch size of 1024. We employ smooth Progressive Scaling (Definition 3.2) between epochs 120 and 180.
At batch size 2048, the training loss matches the reference only when the EMA Scaling Rule is applied, whereas the run without the scaling rule diverges from the reference. The impact of this divergence is emphasized as we consider the larger batch size of 4096. Here. there is also a gap with the FMA Scaling Rule, however approximately three times smaller than the gap without the EMA Scaling Rule.
16 an example, see https: / thub. wwt. issuecomment
http:////iiiii..oooo/acoooooooeo http:///mmccbbotteeeecc///////// watpnooisssss//ae 1tsseccoment--811535555555 881453515.
48

---

Additionally, we observe that using PB = 0.992 yields higher Top-1 accuracy over PB = 0.996, and our experiments, using the EMA Scaling Rule always performs better in terms of linear probe performance than not using the scaling rule.
Progressive Scaling K = 2
Progressive Scaling K = Progressive Scaling K
Test Top-1 (%
75
50
25
B 1024,p PB
1024, DO
1024.p P PB
2048, P PB
3072. PR
4096. P PB
B P P PA
3072, PM
4096. P P8
11
Train Loss
10
100
200
100
200
Train Epochs
Train Epochs
Train Epochs
Figure 27: DINO VIT-B/16 CIFAR-10 10 for different scalings and base teacher momentum Ps 0.996. The baseline model black dashed) uses batch size 1024 and center momentum Pe 0.9, and scaled up from batch size 2048 (left) to 4096 (right) with (blue, P = of) and without (red, p = pB) the EMA Scaling Rule. Between epochs 100 and 180 we scale the batch size using progressive scaling (Definition 3.2).
Progressive Scaling K 22 Progressive Scaling K =3 Progressive Scaling ==
Test Top-1 (%)
75
50
25-
B= p024,
1024.
8 1024,
B= 2 Pb
== 3072,
4096, p
B= P Pb
= 3072, Pb
4096,
Train Loss
100
200
100
200
100
200
Train Epochs
Train Epochs
Train Epochs
Figure 28: DINO VIT-B/16 CIFAR-10 for different scalings x and base teacher momentum PB 0.992. The scaled up
baseline model (x black dashed) uses batch size 1024 center momentum Pe 0.9, from batch size 2048 (left) to 4096 (right) with (blue, ) and without (red, = =ps)) the EMA Scaling Rule. Between epochs 100 and 180 we scale the batch size using progressive scaling (Definition 3.2).
49

---

In Figure 29 we show how the hyperparameters B and learning rate change with the progressive scaling Definition 3.2.
x x10*
Teacher Momentum (p)
1.000
4000
B 1024, P#
B 2048, Pa
0.995
Batch Size (B)
Learning Rate
2048.
PE
3000
B 3072, pa
0.990
B= 3072, PM
B 4096,
0.985
2000
B 4 ps
0.980
1000
Train Epochs
Train Epochs
Train Epochs
Figure 29: DINO ViT-B/16 on CIFAR-JO 10 for different scalings and base teacher momentum P3 = 0.992. We show how the hyperparameters p. 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 that starts at 0.996 instead of 0.992.
We also attempted to use 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 due to dynamics specific to DINO's early phase that are even more challenging replicate under discretization than those of BYOL.
Progressive Scaling K = 2 Progressive Scaling k = Progressive Scaling k= 4
Test Top-1 (%)
75
50
25
B= 1024,
8= 1024.
1024, 8 = 4096.
B B p Pa
= 3072.
B 2048, Pb
= 3072, P3
8 = 4096.
Train Loss
10
100
200
100
200
100
200
Train Epochs
Train Epochs
Train Epochs
Figure 30: DINO VIT-B/16 CIFAR-10 10 for different scalings x and base teacher momentum PB 0.992. The scaled up
baseline model (x black dashed) uses batch size 1024 and center momentum Pe 0.9, from batch size 2048 (left) to 4096 (right) with (blue, pg) and without (red, = PB) the EMA Scaling Rule. Progressive Scaling is employed with sharp transition at epoch 100, leading to collapse phenomenon.
50

---

x10
Teacher Momentum (p)
1.000
4000
B 1024, pø
2048. p P8
0.998
Batch Size (B)
Learning Rate
B 2048, p PM
0.996
3000
B 3072, P#
== 3072, P
0.994
4096. P
2000
4096, p Pa
0.992
0.990
1000
0
200
200 Train Epochs
200
Train Epochs
Train Epochs
Figure 31: DINO VIT-B/16 CIFAR-JO 10 with PB 0.992 and sharp transition in hatch size epoch 100. We show how the hyperparameters p. and learning change with sudden scaling. These hyperparameters correspond the training rums 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 small % of corresponding total runtime), giving 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, 1,0 - 2, 048
40
4.6
182.4
182.4
1, 1,024 - 3, 072
48
4.0
189.9
189.9
1,024 - 4, 096
64
3.3
212.3
212.3
1,024 - 2, 048 (100 Epochs)
40
190.6
762.3
1,024 - 3. 072 (100 Epochs)
48
4.0
192.5
769.9
1,024 - 4. 096 (100 Epochs)
64
3.6
232.1
4
928.2
All other compute, code development, runs with errors, and debugging
38239.2
Total
41,611.1
Our results in this section show it is possible 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
general analysis overflow and underflow of the EMA Update (Definition 1.1) or EMA Scaling Rule (Definition 1.2) for different momenta p. 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).
our setting, do not suffer from practical overflow or underflow issues through exponentiation when applying the EMA Scaling Rule, as FP32 precision allows maximum p= 1-c, minimum with = 1.2 x 10-7, Take self-supervised image representation learning (Section 3.4) as baseline, with corresponding
batch size B = 4096 with momentum pB 0.996. The maximum value of corresponds to scaling x = (og( /log(e) 1/(32K), give a batch size less than one, while the minimum value of corresponds scaling x log(Ps) /log(lllg) v 4K, giving batch size B x 8.M 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 practice during training, we train linear regression model with a Polyak-Rupper average Definition 3.1, and and track the difference between FP32 model weights and weights in BP16 6; ii) FP1 6; and iii) a second FP32 run, which act as proxy for overflow and underflow. In Figure 32 we plot these differences using
51

---

2 Scaling K =
Scaling K =
dttyps
max, 0.P32 _0
Targer
EMA at b$16
Target at fp16
EMA at fp16
1p32
EMA at Ip32
0.99
0.999
0.9999
0.99999
0.99
0.999
0.9999
0.99999
Momentum p
Momentum p
Figure 32: Numerical precision of target and EMA networks compared to an PP32 reference on a regression task for range of momenta.
the maximum absolute difference between model parameters, where the maximum taken over individual weights
MaxAbsDiff(dtypp) max 0
- @ltypel
(82)
where 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 6 the value of the momentum, p approaches 0.99999, whereas BF16 and FP32 are stable. We stress that all exper- iments presented in the paper store weights for target model and EMA in FP32 and use automatic- mixed precision to cast them to BF1 during training, and so do not encounter momentum-induce overflow or underflow.
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 17 initial synthetic and self-supervised ImageNetlk experiments done by Dan Busbridge.
EMA scaling rules for constant gradients Original proof of Equation 4 and the form of 8(n P. x) in Equation 54 done by Feshan 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 limit presented in Section 2.2 and Appendix D done by Pierre Ablin.
Proofs of validity of EMA Scaling Rule in the SDE
Polyak-Rupper averaging in simple setting Design of noisy parabola setting of Section 3.1 and initial experiments done by Russ Webb. Design of p" p"-optimality search (Equation 10), final experiments and analysis of Section 3.1 and Appendix F.I done by Dan Busbridge.
Tolyak-Ruppers averaging on image classification ResNetv2-50 reproduction (Table 4) and baseline momentum identification done by Jason Ramapuram. Final ImageNetlk 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 ing Polyak-Ruppert averaging (Section 3.2) and continuous pscudo-labeling (Section 3.3 and Ap- pendix G). as well design choice of seed model 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 18 recipe (Table 9) and experi- ments on CIFARIO using SGD (Figure 5), and BYOL ResNet-50 experiments using LARS (Ap- pendix H0) done by Dan Busbridge. B BOL ResNet 50 baseline implementation and BYOL ViT
52

---

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 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 Ramapu- ram, Russ Webb, and the authors of Malladi et al. (2022). Adaptation of progressive scaling to semi- supervised learning automatic speech recognition (Appendix G.2) done by Tatiana Likhoma- nenko, and to self-supervised learning in vision done by Dan Busbridge and Jason Ramapuram for BYOL (Figures and 6 and Appendices H.4, H.5 and H.10) and Xavier Suau Cuadros for DINO (Appendix H.1l).
Limiting behavior Polyak-Ruppe 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-Rupper 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. done by Dan Busbridge and Xavier Suau Cuadros.
53


Reviews:
[{"review_id": "JNzl98bHcM", "reviewer": "Reviewer_bHsu", "summary": "This paper talks about the scaling rule for model EMA, and proposes that the momentum coefficient of EMA update should be scaled exponentially with the batch size instead of linearly. The authors first informally derive such rule (Equation 4) under the assumption that the gradients change slowly for a few steps. Then the authors theoretically validate such rule by considering the SDE limit of the optimization process (Equation 6) and prove that the optimization trajectory of different batch sizes could still have small difference after we exponentiate by the batch size ratio $\\kappa$. \n\nFor the experiments, the authors first use a noisy parabola and verify that exponentiating the model momentum term would approximate well with the empirical optimal momentum, as illustrated by Figure 1.b. The authors also consider ResNet and transformer for a supervised task, and the results are the same. However, there is a growing gap between the exponentiated momentum with large batch versus the original momentum with small batch (Figure 2), which is attributed to the breakdown of Adam scaling rule. In the final self-supervised experiments, the authors found a performance gap between the scaled model with EMA and the baseline, and thereby introduce progressive scaling which slowly increases batch size during the early stage of training to match the training dynamics at larger batch size.", "questions": "The questions are raised above.", "limitations": "The authors acknowledge that there is still a $O(\\eta * \\beta_\\rho)$ difference for the scaled model with the unscaled one, but didn't have an ablation of the impacts from the learning rate and the impacts from the base momentum coefficient. It would be nice to see such study somewhere.", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well-written and easily understandable, and the experiment section is well-designed that contain both simple convex task and a broad spectrum of deep learning tasks. The idea is built on top of Malladi et al. (2022)'s work but the contribution is clear. The application of scaling model EMA is broad, and I believe this paper make a nice contribution behind it.", "weaknesses": "In figure 6 with subfigure progressive scaling kappa = 6, after applying the progressive scaling, both purple (without EMA scaling) and orange line (with EMA scaling) performs roughly the same. However, the other subfigures (that varies $\\kappa$) in figure 6 show that EMA scaling is helpful besides the progressive scaling. I might misunderstand but could you provide an explanation why $\\kappa$=6 would be quite different from $\\kappa$=4 and $\\kappa$=8?\n\nIn line 210, section 3.2, could you elaborate more on the intuiton behind the sentence: *Inspecting the train loss leads us to conclude this is due to a breakdown of the Adam Scaling Rule*? Such clarification would be non-trivial as the gap is growing when we increase the scaling factor.\n\nAn additional ablation study on both momentum coefficient + learning rate:\nThe model EMA momentum coefficient used in this paper is too close to 1. It would be nice to have an ablation study on the EMA scaling rule vs. the optimal scaling factor when we vary the EMA momentum coefficient away from 1. \nSimilarly, we would expect when the learning rate increases, the EMA scaling rule will break down, but it would be nice to quantify such phase change even with a convex task.\n\nMinor:\nIn footnote 5, the eqn 68 might actually refer to eqn 74.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to read our work. We are happy you found the paper easy to read, acknowledge the contributions made, and see the broad applicability of an EMA Scaling Rule.\n\n### Weaknesses\n\n#### *Figure 6 - ViT scaling*\nWe also note that the use of an EMA Scaling Rule is more important in some circumstances than others. Intuitively, the EMA Scaling Rule is less important when the value of the EMA momentum is less important (such as Polyak Ruppert averaging). \n\nTo investigate in detail, we conducted the sensitivity analysis in Appendix H.5. We found that for good performance, it is sufficient to use a momentum within an acceptable region $\\rho_{\\text{min}}\\leq \\rho \\leq \\rho_{\\text{max}}$, where $\\rho_{\\text{min}}$ and $\\rho_{\\text{max}}$ each achieve a minimum performance. Note that this is different to choosing a momentum whose training dynamics are the same as one from a reference batch size $B$ (the central goal of the EMA Scaling Rule). Application of the EMA Scaling Rule shows that the range of acceptable momenta actually increases with scaling (see Fig. 22). \n    \nThis explains the observation for Figure 6, $\\kappa=6$: after 60 epochs at $B=4096$, the reference momentum $\\rho_B$ enters the acceptable region for $B=24576$; it is not necessary to apply the EMA Scaling Rule in this specific setting to get good performance. We stress however, this strategy is not guaranteed to work, and the momentum given by EMA Scaling Rule will always be at least as good as this choice.\n    \nWe will clarify text around Figure 6 and more strongly highlight Appendix H.5.\n\n#### *How do you know the Adam Scaling Rule is broken?*\n\nIn this experiment, we have used both optimizer scaling (Adam in this case), as well as EMA Scaling. This particular section concerns Polyak-Ruppert averaging, which lets us check independently whether the Adam Scaling Rule and the EMA Scaling Rule are holding. In Appendix G we present Figure 14 which includes Train Loss and test-other WER. Both quantities depend *only* on the target model, and are *not* related to the model parameter EMA. For $\\kappa>2$ (i.e. $\\kappa=4$ shown), the loss and WER curves deviate from the baseline. This is a breakdown of the Adam Scaling Rule as only the Adam optimizer is involved in the training dynamics leading to these quantities.\n\n#### *Additional momentum and learning rate analysis* \n\nWe agree studies of the validity of the EMA Scaling Rule for a wider range of hyperparameters are very useful. In the quadratic setting, analysis of $\\rho_B\\in\\{0.99, 0.999\\}$ is presented in Appendix F.1 alongside higher-dimensional counterparts.\n    \nIn the rebuttal pdf, we present a learning rate $0.01\\leq \\eta_B \\leq 0.5$ and momentum $0.841\\leq \\rho_B \\leq 0.997$ ($B=1024$) sensitivity analysis for the challenging setup of BYOL + SGD in Figures 1 and 2. Note that the lower bound corresponds to  $\\rho=0.5$ at $\\kappa=4$ ($B=4096$), and is a very low momentum value. We see that the EMA Scaling Rule is robust to modifications of the learning rate, and has good scaling behaviour down to momenta of $\\rho_B=0.946$ at $B=1024$ ($\\rho=0.8$ at $B=4096$) in this particular setting.\n    \nThese results and related additional analysis will be included in a new version of the paper.\n\n#### *Minor: In footnote 5, ...* \nThank you for finding this, it will be corrected in a new version of the paper.\n\n### Limitations\n\n#### *Additional momentum and learning rate analysis* \nSee above.\n\n### Summary\n\nAgain, thank you for the time you took reading our work and providing your valuable questions. We will integrate our responses, along with the requested sensitivity analysis into the paper, which we hope will address the raised limitation. These enhancements will contribute to the overall readability and provide a complete full picture for future readers."}, {"review_id": "jYPQcLmxmM", "reviewer": "Reviewer_dfpr", "summary": "The exponential moving average (EMA) of model parameters is used in many contexts, including to improve robustness in supervised learning, and as an important component of self-supervised learning. Given model parameters $\\theta_t$ at iteration $t$, the exponential moving average $\\zeta_t$ is updated as $\\zeta_{t+1} = \\rho \\zeta_t + (1 - \\rho) \\theta_t$, where $\\rho$ is the EMA coefficient. This paper proposes a scaling rule that dictates how $\\rho$ should be modified to preserve training dynamics when the batch size is modified. In particular, when the batch size $B$ is scaled by a factor of $\\kappa$ yielding $\\kappa B$, the EMA coefficient should be exponentiated to the power of $\\kappa$, yielding $\\rho^\\kappa$.\n\nThe authors justify this rule by analogy to the learning rate vs batch size rule from (Goyal et al., 2017), under the assumption that the gradients change extremely slowly between successive steps in training. They also give a justification using an SDE interpretation.\n\nEmpirically, they evaluate the EMA scaling rule on several real-world tasks, including supervised learning for image classification and speech recognition, semi-supervised learning for speech recognition, and self-supervised image representation learning. In each case, they compare the loss curves of: 1) a baseline model trained with batch size $B$ and EMA coefficient $\\rho$; 2) a model trained with batch size $\\kappa B$ and the same EMA coefficient $\\rho$; and 3) a model trained with batch size $\\kappa B$ and modified EMA coefficient $\\rho^\\kappa$. They show that using the scaling rule $\\rho^\\kappa$ produces more similar loss curves to the baseline model than the un-modified $\\rho$.", "questions": "* In many of the experiments, $\\rho$ is very large: in Figure 1(a), $\\rho=0.9999$, in Figure 2 $\\rho = 0.9999$, in Figure 3 $\\rho = 0.9995$, and in Figure 4 $\\rho=0.9999$. Why are these EMA coefficients so large? Are these optimal values at the base batch size, and how were they chosen? How do these experiments behave when $\\rho$ is smaller, like $\\rho=0.9$? As a more general point, because $\\rho$ is exponentiated, it seems that there could be numerical overflow or underflow issues depending on the initial value of $\\rho$ and the scaling $\\kappa$.\n\n* In Eq. 5, what form does the gradient noise distribution $\\mathcal{E}\\_{\\sigma}(\\theta_k, \\zeta_k)$ take? Is it Gaussian? If so, it would be better to use standard notation $\\mathcal{N}$. What does it mean that the gradient noise distribution is \"assumed to be zero-mean and variance $\\Sigma(\\theta_k, \\zeta_k)$ independent of $\\sigma$\"? If the variance is independent of $\\sigma$, why use the notation $\\mathcal{E}_{\\sigma}$?\n\n* See also questions from the weaknesses box.\n", "limitations": "\n* For the real-world experiments, it would be important to report scaling behavior beyond the factors that are shown in the paper. Why is the maximum scaling factor $\\kappa = 8$? Why not continue the scaling to see if the scaling rule continues to hold, and determine where it breaks down? Does it work if $\\kappa \\in \\{ 16, 32, 64, 128, \\dots \\}$? If the reason for not using larger factors is that the base model already used a large minibatch size ($B=1024$), then why not train a base model with a small batch size $B=16$ such that large $\\kappa$ can be evaluated? In addition, because the scaling rule is meant to hold for both increasing and decreasing the batch size, why not have an example reproducing the results of a model trained with $B= 2048$ using a much smaller batch size like $B=16$? This would be a factor of $128$ decrease in batch size.\n\n* This also relates to the question about numerical overflow/underflow, because if the original EMA coefficient is $\\rho = 0.9999$ and $\\kappa = \\frac{1}{128}$ then the adapted coefficient would be $\\rho^{\\frac{1}{128}} = 0.99999922$. What are the limits to the approach, and what are its limitations?", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* The authors present a toy example where they apply the EMA to a noisy quadratic problem, where the \"batch size\" is simulated by adjusting the gradient variance. On this task, they compute the optimal EMA coefficient $\\rho^*$ for various scaling factors $\\kappa$ up to 1024, and show that the $\\rho^\\kappa$ coefficient nearly matches $\\rho^*$ with respect to the approximation error. I appreciate the inclusion of such a toy example.\n\n* The empirical evaluation considers diverse tasks, including both vision (ImageNet1K classification with a ResNet) and language (speech recognition with a Transformer) tasks. These experiments also consider different optimizers (SGD and Adam).\n\n* The authors also evaluate progressive scaling, in which they increase the batch size during training and simultaneously scale the LR and EMA coefficient. They found that in one case (training a BYOC Vision Transformer), this was necessary to obtain decent performance.\n\n* The plots are generally well-made and clear, with good labels.", "weaknesses": "* One of the main drawbacks of the paper is that it is not clearly written. The core idea is very simple, and should be simply explained. However, most of the paper is needlessly obfuscated, making it hard to read. The majority of the definitions (2.1, 2.2, 2.3, and 3.1) are not novel, and are basic concepts that would not typically be presented as formal definitions. This reads as though the authors are trying to add math and formality to the paper, when it is not necessary.\n\n* The plots in Figure 1(a) are not clearly presented, as the legend overlaps with more than half the figure.\n\n* While the experiments use several architectures and datasets, they all measure the same thing, comparing the loss values of models trained with either $\\{ (\\kappa B, \\rho^\\kappa), (\\kappa B, \\rho), (B, \\rho) \\}$. There is very little diversity in this respect. Are there other things that can be done with the ability to match scaling behaviors, like doing hyperparameter optimization in the small-batch regime (which could be run in parallel on multiple devices rather than requiring a large amount of compute for a single experiment) to find the optimal $\\rho$, and then scale it to the large-batch setting using $\\rho^\\kappa$?\n\n* For the Transformer experiment in Figure 3, it is hard to see the long-term behavior of the loss curves. It would help if the y-axis was log-scaled. Here, the red curve ($\\rho$ unmodified) seems to do just as well as $\\rho^\\kappa$, especially in terms of final performance.\n\n* Throughout the paper, the notation $\\times$ is used for simple scalar multiplication. This is non-standard in ML literature, and confusing to the reader, as $\\times$ looks like a cross product. There is no reason to write $\\eta \\times \\frac{1}{B}$ rather than simply $\\frac{\\eta}{B}$.\n\n* In Eq. 6, it is unclear what $\\Theta$ and $Z_t$ are, as they are not defined.\n\n* Probable typo in L163: \"scaled additive $b > 0$\" --> \"scaling $b > 0$\"\n\n* The base results (black dashed lines) in the real-world experiments do not seem to have bands for the mean and std over three runs (only the red and blue curves have bands).\n\n**Minor**\n\n* The math font looks non-standard.\n\n* I do not think that the abstract does a good job describing the exponential moving average.\n\n* The term \"model EMA\" should be more clearly defined at the beginning, as \"the EMA of the model parameters.\"\n\n* In the abstract, it is not clear what is meant by \"Prior works have treated the model EMA separately from optimization.\"\n\n* In the abstract, what is meant by \"optimally a $6\\times$ wall-clock time reduction\"? Is this the actual wall-clock reduction achieved in practice, or under some idealized scenarios? This should be clarified.\n\n* L32: \"These weights are updated through by a momentum hyperparameter.\" --> Bad grammar \"through by,\" and it does not make sense to say that the weights are updated \"by a momentum hyperparameter.\"\n\n* To improve clarity, the introduction should introduce the mathematical form of the EMA, $\\theta' = \\alpha \\theta' + (1 - \\alpha) \\theta'$. The expression for the EMA needs to be introduced before Definition 1.1 because otherwise it is not clear what $\\rho$ corresponds to and how it is used.\n\n* Definition 1.1 should simply be labeled Definition 1, because there is no Definition 1.2. In Def. 1.1, it is confusing to say \"scale other optimizers.\" Should this be \"scale other optimization hyperparameters\"?\n\n* On Line 64, write out \"self-supervised learning\" instead of saying \"SSL\" because in this context it could be misinterpreted as standing for \"semi-supervised learning.\"\n\n* What is the purpose of including the third row in Eq. 4, which keeps the gradient $g$ unchanged? Is this just to highlight the assumption that the gradient changes very slowly, so that the first two rows hold?\n\n* Typo: L115, \"typically use to\" --> \"typically used to\"\n\n* L278 Typo: \"a model EMA alters\" --> \"a model EMA that alters\"\n\n* L307 Typo: \"only when by the EMA\" --> \"only by using the EMA\"\n\n* L210 Typo: \"In summary, ASR\" --> \"In summary, this ASR\"\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to read our work. We are happy you found the plots clear, appreciate the mixture of controlled toy and diverse real world investigations, and find utility in enabling large batch training of BYOL with Vision Transformers. \n\n### Weaknesses\n#### *Obfuscation*\nThe purpose of numbered definitions is to avoid ambiguity for readers as there are many concepts, algorithms and scaling rules in the text. Each definition is accompanied with appropriate context, and relevant citations where it was introduced by prior works.\n#### *Plots in main text all the same thing*\nThe choice of reported metrics is three-fold: 1) they are used in prior works studying scaling rules, e.g. Maladi et al. 2021; 2) a unified comparison across models and domains makes it clear the EMA Scaling Rule is agnostic to them; and 3) These metrics are sufficient to show training dynamics preservation - validating the main claims of the paper. Additional domain-specific metrics are reported in the Appendices.\n#### *Other applications of the EMA Scaling Rule*\nWhere hyperparameters are not from prior works (cited when this is the case), the $\\rho$-optimization procedure you describe is exactly what we did. E.g. in supervised vision, take 8 GPUs and optimise $\\rho_B$ at low batch size. Then scale to larger batch sizes (e.g. across 32 GPUs, with the same per-gpu batch size), and verify the scaling law by comparing $\\rho_B$ to the predicted momentum $\\rho_B^\\kappa$ ($\\kappa=4$ in this example). Total compute budget (GPUh) is the same, whereas training wall-clock time is reduced at higher batch sizes.\n#### *Fig 3*\nWe note that final performance is not the main property we are interested in (we agree that in Fig 3, with/without the EMA Scaling Rule have similar end performance). Our goal is to match the whole trajectory, which does not happen without the EMA Scaling Rule. See also App. H.5 and Fig 22 for a detailed discussion.\n#### *Eq 6*\n$\\Theta$ and $Z_t$ are SDE variables.\n#### *L163*\n$b$ is the coefficient of noise scaled by the gradient norm. Noise is additive so $b$ is the “scaled additive noise” coefficient.\n#### *Bands for baseline*\nWe will show baseline bands, see rebuttal pdf for examples.\n\n### Minor\n#### *Purpose of 3rd row in Eq. 4*\nThis is the constant gradient within-time interval assumption of Goyal et al. 2017. This row contributes to gradient updates of the tracking model on the top row through repeated applications of the stepping matrix. We present a full breakdown of this in App. E.1.\n\n### Questions\n#### *Choice of momenta*\nMomenta are chosen for performance and are model/domain dependent. A popular choice for SSL methods is $\\rho_B=0.996$ at $B=4096$ (or equivalently $\\rho_B=0.9997$ at $B=256$) (Grill et al. 2020). We stress batch size must be considered when determining if a momentum is large (App. C.1). Momenta in our paper are taken from reference recipes and cited, or determined by max-performance hyperparameter search. We also note that the toy quadratic experiments do cover a large momentum range, with the smallest momentum evaluated $0.99^{1024}\\approx 0.00003$ in Fig. 7, far from 1 (EMA Scaling Rule breaks here). Finally, we present a momentum sensitivity analysis in the rebuttal pdf (Figure 1).\n#### *Overflow/underflow*\nAnalysis of overflow/underflow, particularly for [IEE-754 floating point values](https://standards.ieee.org/ieee/754/6210/), is non-linear due mechanisms like gradual underflow.\n    \nIn our setting, $\\rho$ will not suffer from practical overflow or underflow issues through exponentiation, as FP32 precision allows a maximum $\\rho=1-\\epsilon$, or minimum $\\rho=\\epsilon$, $\\epsilon\\approx 1.2\\times 10^{-7}$. Taking SSL as a baseline, this corresponds to $\\kappa\\approx 32K$ or $\\kappa\\approx 4K$ respectively, the former being below $B=1$ and the latter giving $B\\approx 9M$ which is beyond current hardware feasibility, and beyond the breakdown of known optimizer scaling rules.\n   \nTo examine how momentum may induce numerical errors in practice, we present Figure 4 in the rebuttal pdf. This shows a linear regression model with a Polyak-Ruppert average. We study $\\rho\\in[0.99, 0.99999]$ and track the maximum absolute difference between FP32 model weights and weights in {BF16, FP16, a second FP32 run}. This acts as a proxy for overflow and underflow. When model weights and EMA weights are FP16 (never done in practice) we observe an increasing variance only for FP16 as the value of $\\rho$ becomes closer to 0.99999 — BF16 and FP32 are stable here. We stress all experiments presented in the paper store weights (network and EMA) in FP32 and use automatic-mixed precision to cast them to BF16 during training, so never encounter overflow/underflow induced by momentum.\n\n#### *Scaling behavior* \nWe reported real-world experiments for batch sizes up until matching breaks in all cases (Fig 2, 5, 16, 18) while for toy problems — scaling ranges up to $\\kappa=1024$ (see App. F.1). We now provide a scaling sensitivity analysis in the rebuttal pdf (Fig. 3) where scaling is matched down to $\\kappa=1/8$ (batch normalization used in the ResNet 18 results in slight mismatch at low batch size). Moreover, $\\kappa$ is only notational; the batch size or discretization level matters for SDE approximation: e.g. scaling from batch 1024 to 4096 ($\\kappa=4$) means that we can scale also from batch 1 to 4096 ($\\kappa=4096$) as this limit is guaranteed by the properties of the SDE, providing there are no confounding effects like batch normalization. Experiments using $B=1$ will be $4096\\times$ slower to run in practice, thus impractical to perform for real-world problems.\n\n### Summary\n\nThank you again for taking the time to read our work, as well as for the many pointers, suggestions, and typo highlighting.  We will incorporate these changes in the paper revision to improve readability, as well as the additional content prompted by your feedback.\n"}, {"review_id": "473HhiaLQ1", "reviewer": "Reviewer_6nqM", "summary": "This work proposes a scaling rule for EMA in optimization and empirically evaluates its effectiveness across a range of tasks and optimizers. Authors find that proposed method enables large batch training using 24k batch sizes with a 6x wall-clock time reduction.", "questions": "- My main question is regarding the motivation behind such a scaling rule for EMA. Have the authors verified that such a scaling rule is infact needed since for most practical problems, hyper-parameters are just fixed and no such rule is used, still giving a good performance for most models. Is there a fundamental gap (that has potential to improve model quality or speed substantially) the authors are trying to address here?\n- It looks like there is still some tuning required for the scaling rule to work? Can the authors clarify how much and sensitivity to larger models and different architectures?", "limitations": "- Key questions regarding broad applicability of this method are not answered (such as generality of the solution to varying model types, model scales, etc). Therefore it is not clear how impactful this method is.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 1, "strengths": "- Idea of formulating a scaling rule for EMA is novel and not one I have seen much in the recent deep learning / optimization literature\n- Theoretical motivations are clear to understand\n- Strong empirical study", "weaknesses": "- Applicability to modern optimizers is still not clear. Despite the authors discussing some connections to Adam, I think an in-depth discussion of how the theory/scaling rule relates to Adam is not fully clear.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to read our manuscript. We are are happy you find our formulation novel, the theoretical motivations clear, and the study design strong.\n\nWe begin by clarifying a potential misunderstanding. Our work is devoted to the behaviour of exponential moving averages (EMAs) of model parameters, and *not* EMAs used within gradient-based optimizers such as Adam. Investigations for this direction already exist (see e.g. Malladi et al. (2022)). Understanding the dynamics of Adam is *not* a purpose of this work. We do conduct some experiments where we take an EMA of the sequence of parameters produced by the Adam optimizer (e.g. Fig. 3, 4), and discuss some related theoretical aspects in Appendix D.\n\nWe note that the review contains concerns regarding the broad applicability of an EMA Scaling Rule (ESR), and contains a number of questions we will now address, beginning with a restatement of our contribution.\n\nAn optimizer scaling rule allows the same training dynamics to be produced at different experimental scales (batch sizes). A common use is producing similar models in less wall-clock time. Our work derives an ESR, extending this training speedup to methods which use model parameters exponential moving averages, which include (but are not limited to) state-of-the-art semi-supervised methods, like continuous pseudo-labeling, and self-supervised methods like BYOL and DINO.\n\n### Weaknesses\n\n#### *Applicability to modern optimizers*\nOur work does not deal with the momentum in gradient optimizers (Footnote 1), and scaling rules for Adam are already well-studied (Section 4). Our focus is the momentum of the model parameter EMA, whose scaling is independent to an optimizer choice (Line 104 & Appendix D), and is guaranteed, providing the gradient optimizer scaling is known (Appendix B, Limitation 2). \n    \nWe empirically verified that the EMA Scaling Rule works with SGD (with/without momentum), Adam, AdamW, and LARS optimizers, which are popular modern ML choices. Moreover, we empirically focus on different domains (toy data, vision, speech) and different models (ResNet and variants of Transformers). Although some of these optimizers do contain an EMA (such as Adam), these are not model parameter EMAs, and are not the focus of our work.\n\n### Questions\n\n#### *What is the practical requirement for an EMA Scaling Rule?*\n\nFor example, Fig. 2 shows EMA tracking performance goes to zero under batch scaling when going up only $4\\times$ in batch size if the ESR is not used. One could find the ESR-predicted value through a hyperparameter search at the larger batch size, but this comes with large computational cost compared to already knowing the value. We note that our method still requires a *good momentum value* to be known at some computational scale - so at least one search, or one reference recipe is required. After this, the good value at a new scale can be determined by the ESR.\n\n#### *Is there tuning required for the EMA Scaling Rule to work?* \n\nWe discuss in Limitations (App. B) scenarios when the ESR will and will not work. In summary, when the corresponding optimizer scaling rule is known, and discretization error is not large, then the EMA Scaling Rule will work, independent of any other hyperparameter choices. We have presented additional hyperparameter ablations in the rebuttal pdf highlighting this (Fig. 1, 2, 3).\n\n### Limitations\n\n#### *Broad applicability of method* \nWe have restated that our objective is to study the behavior of the model parameter EMA, which is independent of optimizer choice, model class, and problem domain. We also note that, beyond the theoretical guarantees, we empirically validated the EMA Scaling Rule for a large array of models, domains and modern optimizer choices. We hope that the combination of these two answers this point.\n\n### Summary\n\nThank you again for taking the time to review our work. We hope a restatement of our overall goal, and the reemphasis of the wide empirical validation we have performed  addresses the limitation you have raised regarding broad applicability.\n"}, {"review_id": "RCHrxtcAn2", "reviewer": "Reviewer_nnfC", "summary": "In duality to previous work on learning rate scalings, the authors present a novel scaling strategy for the momentum parameter in teacher-student settings, i.e. how to adapt the momentum in the learning dynamics when the batch-size is changed. The work puts an emphasis on the large batch-size setting, which is of particular interest due to its computational efficiency in practice. The obtained scaling is validated through a theoretical analysis and numerous experiments in different settings demonstrate that (a slight modification of) the scaling indeed preserves performance at large batch-sizes. \n", "questions": "1. The CIFAR10 results in Fig. 5 seem surprisingly bad (around 55%-58%).\n2. Are you using a linear or a K-NN probe for the SSL experiments?\n3. After the warm-up phase with progressive scaling, how well is the model already performing at this point? I.e. how much is still gained in the large (or maximal) batch-size regime?\n4. I find it very surprising that the model is so sensitive to the momentum parameter in Polyak-Ruppert averaging. How well is the student performing at the end of training and how much is actually gained by averaging? I would be very surprised if the non-ideally averaged teacher would perform inferior to the student at the end of training.\n", "limitations": "The authors have addressed the limitations of their work.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The paper﻿ is very well-written and easy to follow. The problem of large batch-size training is very relevant as it enables more efficient training (especially when distributing across multiple GPUs) and the resulting scaling rule is very simple. I believe that such insights can be of tremendous use to practicioners, especially if hyperparameter tuning is too expensive, especially in the setting of SSL where the setups are usually already very brittle to hyperparameter changes.\n2. The experimental setup is quite broad, including various student-teacher settings such as Polyak-Ruppert averaging, pseudo-labeling and self-supervised learning. \n3. The technique is nicely motivated through a theoretical analysis, strengthening the particular choice of scaling further. ﻿\n", "weaknesses": "Some important aspects of the experimental setup are not well explained:\n\n1. In the real data experiments (i.e. Fig 2, 3, 4, 5, 6), what does “No EMA Scaling Rule” entail? Does it still apply the “Learning-rate Scaling Rule”? If it does not, then it would be very hard to draw any conclusions regarding \"EMA Scaling” as it is not clear how much of the “heavy-lifting” the learning rate scaling is doing. I hope the authors can clarify this.\n2. One of the main motivation of the works is to make methods like SSL more computationally effective by ensuring that training them with large batch-sizes does not deterioate performance. In order to make it work in practice, the authors need to resort to “progressive scaling”, where the batch-size is only gradually increased after the warm-up phase. How long is the warm-up period exactly (during which a small batch size is used) and how quickly is the batch-size scaled up afterwards? How much time is lost compared to the “EMA Scaling without Progression”? \n3. Another aspect which remained a bit mysterious is the effect of the baseline momentum. For Polyak-Ruppert averaging for instance, its role seems very crucial. If e.g. rho_base were set to 1, the teacher would always be equal to the student and at inference time I would not expect any sort of batch size effect, given that the learning rate scaling correctly works. I would appreciate if the authors could elaborate more on this initial choice of momentum.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to read our work. We are happy you found the paper easy to read, appreciate the empirical and theoretical validation of the EMA Scaling Rule, and recognize the broad applicability of the result to many modern ML settings.\n\n### Weaknesses\n\n#### *1. No EMA Scaling Rule* \nWe will clarify that this means the optimizer is still scaled (i.e. SGD Scaling Rule for learning rate is applied), while no scaling rule is used for the EMA, i.e. the same momentum is used for different batch sizes. This means that the effect we see is due only to the EMA parameter, and not an unscaled optimizer.\n\n#### *2. Computational effectiveness of Progressive Scaling* \nIf one needs to use Progressive Scaling, a greater wall-clock benefit is realized an earlier transition, but note that one does not need to begin scaling from $\\kappa=1$. E.g. for pseudo-labeling experiments in the case of scaling $\\kappa=16$, 15% of training time is done first at $\\kappa=8$ followed by 85% of training time at $\\kappa=16$ (see Appendix, Figure 18), maintaining most of the wall-clock time benefit. We also note that for self-supervised learning, one of our main use cases, performance benefits are given by long training duration; the time cost of Progressive Scaling can be made negligible by training for longer.\n    \nThe warm-up phase should happen once the optimizer SDE approximation is valid. Malladi et al. 2021 showed the Adam Scaling Rule is only valid after finite training and a gradient noise criterion is met. A practical option could be to track this criterion, and begin Progressive Scaling when it is met. In this work, our choices are empirical, and found that the beginning Progressive Scaling after the learning rate warmup phase has completed works well.\n#### *3. Effect of baseline momentum*\nThe baseline momentum at a given batch size affects model performance. In all our experiments, the baseline momentum is taken from reference recipes (and cited accordingly), or are found using max-performance hyperparameter search. We, now also present a baseline momentum sensitivity analysis in the rebuttal pdf (Figure 1) which we will include in the paper.\n    \nOn pathological choices of base momentum, there are two. The first is $\\rho_B = 1$, which means that the EMA model (teacher) is always equal to its initial state, independent of scaling. This is the same as saying the EMA model is infinitely slow, which no change in discretization modification can alter. The second pathological is $\\rho_B = 0$, which means that the EMA model is always equal to the target model (student). This is the same as saying the EMA model is infinitely fast, again which no change in discretization modification can alter. No momentum $rho_B$ can be scaled onto either of these values exactly, but only asymptotically, i.e. $\\lim_{\\kappa\\to\\infty}\\rho_B^\\kappa = 0$ and $\\lim_{\\kappa\\to0}\\rho_B^\\kappa = 1$ for $0<\\rho_B<1$, with the latter limit corresponding to the SDE limit, providing other quantities are scaled appropriately. Each pathological choice does not have the properties of a moving average, and produces trivial results; there is only initial-state dependence, or there is no EMA respectively.\n\n### Questions\n#### *1. CIFAR10 results are not competitive*\nWe agree that these results are not competitive. The small model (ResNet 18) and short training duration were chosen to produce an experimental setup that enabled quick iteration, and the SGD optimizer was chosen as it has as known optimizer scaling rule. This allowed us to precisely probe the claims of the EMA Scaling Rule in the challenging setting of BYOL, where the EMA network provides the learning signal for the whole optimization process.\n    \nCompetitive results can be achieved by: using a larger model (e.g. ResNet 50), training for longer, or using an optimizer like LARS. Performing these changes gives us a model consistent with the reference BYOL implementation (App. H.6, Fig. 23).\n    \nWe also note that the EMA Scaling Rule should also work for sub-optimal hyperparameter settings, which this experiment confirms. We will clarify in the text why we chose this particular setting.\n#### *2. Which probe do you use in SSL?* \nWe will clarify we are using a linear problem in all SSL experiments.\n#### *3. In Progressive Scaling, how much is gained in the large batch regime?*\nIf Progressive Scaling is successful, the model will reach a final performance that is the same as the unscaled model; there is no particular gain one can expect after Progressive Scaling, as one could progressively scale at any point (scaling at the start would give 100% of the benefit, and scaling at the end would give 0%). In our pseudo-labeling settings, we see WER improve from 20% to 10% (Fig. 18) after warmup, and in SSL we see BYOL ViT accuracy increase from 64% to 75% (Fig. 23) after warmup for the last Progressive Scale change.\n#### *4. Sensitivity to momentum in Polyak-Ruppert averaging*\nIn Polyak-Ruppert averaging, we do see gains compared to the student. In vision we see a gain of $\\sim0.3$% (Table 4) and in ASR we improve 1-2% in WER (see App. Fig 14, 15), which is significant. The dramatic fall-off in Fig. 2 happens as the training duration of 90 epochs is not large, and is taken from the reference (Goyal et al. 2017). When scaling happens, with each increase in $\\kappa$ the EMA network should to move exponentially faster in order to adapt with the new discretization. For scaling $\\kappa\\geq 4$, the EMA network without the EMA Scaling Rule moves too slowly to take advantage of any learning the target model has done.\n\n### Summary\n\nAgain, we appreciate the time you took reading our work and sharing your valuable comments and questions. We will integrate our responses, along with additional clarifications, into the paper. These enhancements will contribute to the overall readability, clarity and completeness.\n"}]

Author rebuttal:
We thank all our reviewers for taking the time to review our work. We will adapt the paper based on their insightful comments, feedback, and questions.

There were a number of questions prompting additional investigations. We provide here a pdf containing the results of those investigations. Its contents are:
- Figure 1: A momentum sensitivity analysis
- Figure 2: A learning rate sensitivity analysis
- Figure 3: A scaling analysis to low batch sizes
- Figure 4: A numerical precision analysis.

Specific context and discussions are presented in the corresponding review rebuttals.

Many thanks,

The authors

Meta-review: Reviewers were unanimously enthusiastic about this quality of this study of neural iterate averaging (a smoothing tool that shows up usefully in many parts of deep learning, and is not fully understood for neural nets), deriving an exponential scaling rule via an analysis in the SDE limit, and performing an extensive empirical validation of this scaling. The authors addressed reviewers' questions comprehensively. The insights and results in this work are crisp, and have the potential for broad algorithmic impact; thus, I recommend a spotlight distinction.
Decision: Accept (spotlight)