SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Put another way, an optimizer scaling rule and the EMA Scaling Rule each introduce their own discretization errors. In the case where EMA plays a role in 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 a 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 to 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 are close to that SDE, providing we appropriately scale hyperparameters (see Section 2.2). Each of these discretizations shares the total continuous time  $T = \hat{\eta} \times \widehat{N}_{\mathrm{iter}}^{10}$  of the underlying SDE, but each discretization has a different number of iterations  $\widehat{N}_{\mathrm{iter}} = N_{\mathrm{iter}} / \kappa$ .
This perspective is already adopted, perhaps by accident in some domains. For example, in Computer Vision (CV), it is typical to compare model performance after optimization on ImageNet1k after a number of epochs, whilst also specifying 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 a property of the discretization procedure itself. Instead, the discretization procedure should be viewed as a numerical approximation method for the SDE we 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  $\kappa >2$  we wait until the pre-training Word Error Rate (WER) is sufficiently low.
There may be cases where one is working with a 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  $\kappa$ . 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 quantities should be scaled to achieve the desired condition in the same amount of time, i.e. as above  $\widehat{N}_{\mathrm{iter}} = N_{\mathrm{iter}} / \kappa$ , where  $N_{\mathrm{iter}}$  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  $T = \hat{\eta}^2 \times \widehat{N}_{\mathrm{iter}}$  (Malladi et al., 2022).
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  $\hat{B} = \kappa B$ , use a learning rate  $\hat{\eta} = \kappa \eta$  (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  $\kappa B$ .
Next, we give the two scaling rules known for the adaptive 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  $\hat{B} = \kappa B$ , use a learning rate  $\hat{\eta} = \sqrt{\kappa}\eta$ , beta coefficient  $\hat{\beta} = 1 - \kappa \times (1 - \beta)$ , and adaptivity parameter  $\hat{\epsilon} = \frac{\epsilon}{\sqrt{\kappa}}$  (Malladi et al., 2022).
Definition C.3 (Adam Scaling Rule). When running Adam (Kingma & Ba, 2015) with batch size  $\hat{B} = \kappa B$ , use a learning rate  $\hat{\eta} = \sqrt{\kappa}\eta$ , beta coefficients  $\hat{\beta}_1 = 1 - \kappa \times (1 - \beta_1)$ ,  $\hat{\beta}_2 = 1 - \kappa \times (1 - \beta_2)$ , and adaptivity parameter  $\hat{\epsilon} = \frac{\epsilon}{\sqrt{\kappa}}$  (Malladi et al., 2022).
Next, we present a 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  $\hat{B} = \kappa B$ , use a momentum  $\hat{\rho} = \rho^{\kappa}$  and scale other optimizers according to their own scaling rules.
Concretely, if we are using SGD in the presence of a model EMA, Definitions C.1 and C.4 state that we should take  $\hat{\eta} = \kappa \eta$  and  $\hat{\rho} = \rho^{\kappa}$  when scaling by  $\kappa = \hat{B} / B$ .
The final scaling rule is for weight decay, and follows from the scaling logic discussed in Appendix C.1 and Krizhevsky (2014). If we take the weight decay regularization penalty  $\lambda$  defined at batch size  $B$ , what should the weight decay  $\hat{\lambda}$  be for batch size  $\hat{B} = \kappa B$ ? For simplicity, consider  $\kappa$  updates of optimization of parameters  $\theta_{t}$  in the presence of weight decay only
$$
\boldsymbol {\theta} _ {t + \kappa} = \boldsymbol {\theta} _ {t + \kappa - 1} - \eta \lambda \boldsymbol {\theta} _ {t + \kappa - 1} = (1 - \eta \lambda) \boldsymbol {\theta} _ {t + \kappa - 1} = (1 - \eta \lambda) ^ {\kappa} \boldsymbol {\theta} _ {t}. \tag {11}
$$
Therefore, to match the effect of weight decay with a single iteration step, we need to match
$$
1 - \hat {\eta} \hat {\lambda} = (1 - \eta \lambda) ^ {\kappa}. \tag {12}
$$
Solving for  $\hat{\lambda}$  and expanding around  $\eta \approx 0$  gives
$$
\hat {\lambda} = \frac {1 - (1 - \eta \lambda) ^ {\kappa}}{\hat {\eta}} \approx \frac {\eta}{\hat {\eta}} \times \kappa \lambda + O (\eta). \tag {13}
$$
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  $\hat{B} = \kappa B$ , use a penalty term  $\hat{\lambda} = (\kappa \hat{\eta} / \eta) \lambda$ , where  $\hat{\eta}$  and  $\eta$  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 implies that using linear scaling for the learning rate  $\eta$  then the weight decay penalty is automatically scaled, and when using square-root scaling for the learning rate  $\eta$  (e.g. in the case of the Adam Scaling Rule (Definition C.3)) then the weight decay penalty should also be scaled with a square-root as 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  $\theta_{t + 1} = (1 - \lambda)\theta_t$ , then the scaling rule is optimizer-independent, and becomes linear for small weight decay, i.e.  $\hat{\lambda} = \kappa \lambda$ , and for arbitrary  $\lambda$  takes the form  $\hat{\lambda} = 1 - (1 - \lambda)^{\kappa}$ .
Table 2: Scaled learning rates  $\hat{\eta}$  at different batch sizes  $\hat{B} = \kappa B$  given reference learning rates  $\eta$  defined at batch size  $B$ . The reference values of each column are bolded. Note that this is only valid when there is a 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, i.e. total sequence length.
Batch size B	ˆη = κη [SGD]	ˆη = √κη [RMSProp, Adam]
B = 256	B = 512	B = 256	B = 4096
η = 0.1	η = 0.3	η = 0.1	η = 10-3	η = 4.8	η = 10-3
32	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
512	0.2	0.6	0.1	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
8192	3.2	9.6	1.6	0.00566	6.78823	0.00141
16384	6.4	19.2	3.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  $\hat{\rho} = \rho^{\kappa}$  at different batch sizes  $\hat{B} = \kappa B$  given reference momenta  $\rho$  defined at batch size  $B$ . The reference values of each column are boldened. Again in the sequence learning setup, batch size should be appropriately replaced with a notion of sequence length.
Batch size B	B = 256	B = 4096
ρ = 0.9999	ρ = 0.999	ρ = 0.99	ρ = 0.996	ρ = 0.992	ρ = 0.99	ρ = 0.97
32	0.99999	0.99987	0.99874	0.99997	0.99994	0.99992	0.99976
64	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 a base learning rate  $\eta$  defined at 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, it is also common to just present 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 a single reference batch size  $B$  at which all hyperparameters are defined, and then scale from there. In this case, it is easiest to compute using linear scaling the learning rate at the redefined base batch size  $\eta = \tilde{\kappa}\eta_{\mathrm{orig}}$ , where  $\tilde{\kappa} = B / B_{\mathrm{orig}}$ , and then scale this new reference  $\eta$  as  $\hat{\eta} = \kappa \eta$ ,  $\kappa = \hat{B} / B$ , along with e.g. the momentum defined at  $B$ .
As this process can be slightly frustrating, we have provided tables of typical learning rates in Table 2 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 are dynamics that are challenging to replicate at larger batch
Algorithm 1 Stochastic Gradient Descent with Progressive Scaling
Require: Base learning rate  $\eta$  , base momentum  $\rho$  for base batch size  $B$    
Require: Initial target model parameters  $\theta$  and model EMA parameters  $\zeta$    
Require: Epochs  $E$  and schedule of batch sizes  $\mathcal{B} = B_1,B_2,\ldots ,B_E$    
Require: Loss function  $\mathcal{L}$    
for  $e$  in 1,2...,E do  $\begin{array}{rl} & {\hat{B}\gets \mathcal{B}[e]}\\ & {\kappa \gets \hat{B} /B}\\ & {\hat{\eta}\gets \kappa \eta}\\ & {\hat{\rho}\gets \rho^{\kappa}} \end{array}$  Get current batch size  $\kappa \leftarrow \hat{B} /B$  Compute scaling factor Get scaled learning rate  $\hat{\rho}\gets \rho^{\kappa}$  Get scaled momentum   
for  $b$  in 1,2...,floor(E/hat) do Sample a minibatch of  $\hat{B}$  samples  $\mathcal{X} = \{\pmb{x}^{(1)},\dots,\pmb{x}^{(\hat{B})}\}$ $\theta \gets \theta -(\hat{\eta} /\hat{B})\sum_{x\in \chi}\nabla_{\theta}\mathcal{L}(x;\theta ,\zeta)$  SGD Update end for end for
sizes. To remove ambiguity, in Algorithm 1 we provide pseudo-code for how to use Progressive Scaling.
In Algorithm 1, the prefactor of the SGD update could also have been written  $\eta /B$ , although an equivalent use of the base momentum is not possible.
Finally, we outline how to extend Algorithm 1 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  $\theta$ .
2. 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  $T_{\mathrm{fixed}}$ . We did not 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 is trained using SGD. We employ a strategy similar to Malladi et al. (2022), where we 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 1 in (Li et al., 2019)). The set  $G$  is the set of continuous functions  $\mathbb{R}^d \to \mathbb{R}$  with at most polynomial growth, i.e., for  $g \in G$  there exists two scalars  $\kappa_1, \kappa_2 > 0$  such that for all  $x \in \mathbb{R}^d$ , we have  $|g(x)| \leq \kappa_1(1 + \|x\|^{\kappa_2})$ .
For an integer  $\alpha > 0$ ,  $G^{\alpha}$  is the set of functions  $\mathbb{R}^d \to \mathbb{R}$  that are  $\alpha$ -times continuously differentiable and such that all their derivatives up to order  $\alpha$  are in  $G$ .
Similarly to Malladi et al. (2022), we use Noisy Gradient Oracle with Scale Parameter (NGOS) to define the update rules on the parameters.
Definition D.2 (Noisy Gradient Oracle with Scale Parameter (NGOS), adaptation of (Malladi et al., 2022)). A NGOS is a tuple  $\mathcal{G}_{\sigma} = (f, \Sigma, \mathcal{Z}_{\sigma})$ . Given a noise scale parameter  $\sigma > 0$ , the NGOS  $\mathcal{G}_{\sigma}$
takes as input the parameters  $\pmb{\theta}$  and outputs a random vector  $\mathbf{g} = \nabla f(\pmb{\theta},\pmb{\zeta}) + \sigma \in$  where  $\nabla f(\pmb{\theta},\pmb{\zeta})$  is the gradient of  $f$  with respect to  $\pmb{\theta}$  at  $(\pmb{\theta},\pmb{\zeta})$ , and  $\epsilon$  is a random vector drawn from the distribution  $\mathcal{Z}_{\sigma}(\pmb{\theta},\pmb{\zeta})$  with zero mean and covariance  $\Sigma (\pmb {\theta},\pmb {\zeta})$ .
Note that in the above definition, the probability distribution  $\mathcal{Z}_{\sigma}(\theta ,\zeta)$  is allowed to change with the scale  $\sigma$  , but its first two moments  its mean and its covariance  are fixed with  $\sigma$  . We have the following theorem for model EMA under optimization with SGD:
Theorem D.1 (SDE for  $\mathrm{SGD} + \mathrm{EMA}$ ). Consider the couple  $\mathbf{x}_k = (\theta_k, \zeta_k)$  where  $\theta_k$  are the iterates of SGD with a NGOS (Definition D.2) and  $\zeta_k$  is an EMA of  $\theta_k$ , defined, starting from  $\mathbf{x}_0 = \mathbf{x}_0$ , by
$$
\boldsymbol {\theta} _ {k + 1} = \boldsymbol {\theta} _ {k} - \eta \mathbf {g} _ {k}, \text {w i t h} \mathbf {g} _ {k} = \nabla f \left(\boldsymbol {\theta} _ {k}, \zeta_ {k}\right) + \sigma \boldsymbol {\epsilon} _ {k}, \text {a n d} \boldsymbol {\epsilon} _ {k} \sim \mathcal {Z} _ {\sigma} \left(\boldsymbol {\theta} _ {k}, \zeta_ {k}\right), \tag {14}
$$
$$
\zeta_ {k + 1} = \rho \zeta_ {k} + (1 - \rho) \theta_ {k}. \tag {15}
$$
Define  $\beta_0 = (1 - \rho) / \eta$ ,  $\sigma_0 = \sigma \sqrt{\eta}$ , and define the SDE for  $X_{t} = (\Theta_{t},Z_{t})$ , starting from  $X_0 = x_0$ , by
$$
d \Theta_ {t} = - \nabla f \left(\Theta_ {t}, Z _ {t}\right) d t + \sigma_ {0} \Sigma \left(\Theta_ {t}, Z _ {t}\right) ^ {\frac {1}{2}} d W _ {t}, \text {w i t h} W _ {t} \text {a W i e n r p r o c e s s} \tag {16}
$$
$$
d Z _ {t} = \beta_ {0} \left(\Theta_ {t} - Z _ {t}\right) d t. \tag {17}
$$
Assume that  $f$  is continuously differentiable, with  $f \in G^3$  and  $\Sigma^{\frac{1}{2}} \in G^2$  (Definition D.1). Then, for any time horizon  $T > 0$  and test function  $g \in G^2$ , there exists a constant  $c > 0$  such that
$$
\max  _ {k = 0, \dots , \lfloor T / \eta \rfloor} | \mathbb {E} [ g (X _ {\eta k}) ] - \mathbb {E} [ g (\mathbf {x} _ {k}) ] | \leq c \times \eta . \tag {18}
$$
Proof. The proof uses the same tools as in Li et al. (2019). Define  $\Delta(\theta, \zeta) = \eta(-\nabla f(\theta, \zeta) + \sigma \epsilon, \beta_0(\theta - \zeta))$  with  $\epsilon \sim \mathcal{Z}_{\sigma}(\theta, \zeta)$  the one-step update for the SGD + EMA update, such that  $\mathbf{x}_{k+1} = \mathbf{x}_k + \Delta(\mathbf{x}_k)$ . We have the first two moments:
$$
\mathbb {E} [ \Delta (\boldsymbol {\theta}, \zeta) ] = \eta (- \nabla f (\boldsymbol {\theta}, \zeta), \beta_ {0} (\boldsymbol {\theta} - \zeta)) \tag {19}
$$
$$
\mathbb {V} [ \Delta (\boldsymbol {\theta}, \zeta) ] = \eta \sigma_ {0} ^ {2} \left[ \begin{array}{c c} \Sigma (\boldsymbol {\theta}, \zeta) & 0 \\ 0 & 0 \end{array} \right] \tag {20}
$$
and the higher-order moments are  $O(\eta^2)$ . Similarly, let  $\tilde{\Delta} (\pmb {\theta},\pmb {\zeta})$  be the solution at time  $\eta$  of the SDE defined by Equation 6 starting from  $X_0 = (\pmb {\theta},\pmb {\zeta})$ . From Ito's formula, we also obtain
$$
\mathbb {E} [ \tilde {\Delta} (\boldsymbol {\theta}, \zeta) ] = \eta (- \nabla f (\boldsymbol {\theta}), \beta_ {0} (\boldsymbol {\theta} - \zeta)) \tag {21}
$$
$$
\mathbb {V} [ \tilde {\Delta} (\boldsymbol {\theta}, \zeta) ] = \eta \sigma_ {0} ^ {2} \left[ \begin{array}{c c} \Sigma (\boldsymbol {\theta}, \zeta) & 0 \\ 0 & 0 \end{array} \right] \tag {22}
$$
and the higher-order moments are  $O(\eta^2)$ . 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 to the advertised theorem.
This theorem is a simple adaptation of the results of Li et al. (2019). Intuitively, it is expected that  $X_{t}$  and  $\mathbf{x}_k$  are close since  $\mathbf{x}_k$  is the Euler-Maruyama discretization of  $X_{t}$  with learning rate  $\eta$ . We then have the corollary.
Corollary D.1.1 (Validity of the EMA Scaling Rule). Assume that  $f$  is continuously differentiable, with  $f \in G^3$  and  $\Sigma^{\frac{1}{2}} \in G^2$ . Let  $\theta_k^B, \zeta_k^B$  the iterates of the Equation 5 with batch size  $B$  and hyperparameters  $\eta, \rho$ . Let  $\theta_k^{\kappa B}, \zeta_k^{\kappa B}$  be iterates with batch size  $\kappa B$ , learning rate  $\eta$  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 > 0$  and function  $g \in G^2$ , there exists a constant  $d > 0$  such that
$$
\max  _ {k = 0, \dots , \lfloor T / \eta \rfloor} | \mathbb {E} [ g \left(\boldsymbol {\theta} _ {\lfloor k / \kappa \rfloor} ^ {\kappa B}, \boldsymbol {\zeta} _ {\lfloor k / \kappa \rfloor} ^ {\kappa B}\right) ] - \mathbb {E} [ g (\boldsymbol {\theta} _ {k}, \boldsymbol {\zeta} _ {k}) ] | \leq d \times \eta . \tag {23}
$$
Proof. The proof is similar to Malladi et al. (2022). Under the scaling rule, both  $\mathbf{x}_k = (\theta_k,\zeta_k)$  and  $\hat{\mathbf{x}}_{[k / \kappa ]} = (\theta_{[k / \kappa ]}^{\kappa B},\zeta_{[k / \kappa ]}^{\kappa B})$  have the same limiting SDE. Hence we have from the previous theorem that for all test function  $g$ , we can find  $c,c^{\prime}$  such that
$$
\max  _ {k = 0, \dots , \lfloor T / \eta \rfloor} | \mathbb {E} [ g (X _ {\eta k}) ] - \mathbb {E} [ g (\mathbf {x} _ {k}) ] | \leq c \times \eta \text {a n d} \max  _ {k = 0, \dots , \lfloor T / \eta \rfloor} | \mathbb {E} [ g (X _ {\eta k}) ] - \mathbb {E} [ g (\hat {\mathbf {x}} _ {\lfloor k / \kappa \rfloor}) ] | \leq c ^ {\prime} \times \eta . \tag {24}
$$
The triangle inequality then gives
$$
\max  _ {k = 0, \dots , \lfloor T / \eta \rfloor} | \mathbb {E} [ g (\hat {\mathbf {x}} _ {\lfloor k / \kappa \rfloor}) ] - \mathbb {E} [ g (\mathbf {x} _ {k}) ] | \leq (c + c ^ {\prime}) \times \eta . \tag {25}
$$
Hence, taking  $d = c + 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 an added EMA. Like above, we consider that the loss function  $f$  also depends on the EMA tracking parameter  $\zeta_{k}$ . We begin with RMSProp with EMA, which iterates:
$$
\mathbf {v} _ {k + 1} = \gamma \mathbf {v} _ {k} + (1 - \gamma) \mathbf {g} _ {k} ^ {2}, \text {w i t h} \mathbf {g} _ {k} = \nabla f (\boldsymbol {\theta} _ {k}, \zeta_ {k}) + \sigma \epsilon_ {k}, \text {a n d} \epsilon_ {k} \sim \mathcal {Z} _ {\sigma} (\boldsymbol {\theta} _ {k}, \zeta_ {k}), \tag {26}
$$
$$
\boldsymbol {\theta} _ {k + 1} = \boldsymbol {\theta} _ {k} - \eta \left(\sqrt {\mathbf {v} _ {k}} + \varepsilon\right) ^ {- 1} \times \mathbf {g} _ {k} \tag {27}
$$
$$
\zeta_ {k + 1} = \rho \zeta_ {k} + (1 - \rho) \theta_ {k}. \tag {28}
$$
Like in Malladi et al. (2022), we place ourselves in the high noise regime, in which the term  $\mathbf{g}_k^2$  in Equation 26 is approximated by  $\mathbf{g}_k^2 \simeq \sigma^2 \mathrm{diag}(\Sigma (\pmb{\theta}_k, \pmb{\zeta}_k))$ . We use the same scaling rules, with an additional one for  $\rho$ :
$$
\gamma_ {0} = (1 - \gamma) / \eta^ {2}, \sigma_ {0} = \sigma \eta , \varepsilon_ {0} = \varepsilon \eta , \text {a n d} \beta_ {0} = (1 - \rho) / \eta^ {2}, \tag {29}
$$
and we let  $\mathbf{u}_k = \mathbf{v}_k / \sigma^2$ . The equations for RMSProp with EMA then become, using only these new variables and  $\eta$ :
$$
\mathbf {u} _ {k + 1} - \mathbf {u} _ {k} = \eta^ {2} \gamma_ {0} (\operatorname {d i a g} (\Sigma (\boldsymbol {\theta} _ {k}, \zeta_ {k})) - \mathbf {u} _ {k}), \tag {30}
$$
$$
\boldsymbol {\theta} _ {k + 1} - \boldsymbol {\theta} _ {k} = - \left(\sqrt {\mathbf {u} _ {k}} + \varepsilon_ {0}\right) ^ {- 1} \left(\eta^ {2} \nabla f \left(\boldsymbol {\theta} _ {k}, \zeta_ {k}\right) + \eta \epsilon_ {k}\right) \tag {31}
$$
$$
\zeta_ {k + 1} - \zeta_ {k} = \eta^ {2} \beta_ {0} \left(\theta_ {k} - \zeta_ {k}\right). \tag {32}
$$
This formulation makes it clear that these iterations can be seen as the discretization of the SDE
$$
d U _ {t} = \gamma_ {0} (\operatorname {d i a g} (\Sigma (\Theta_ {t}, Z _ {t})) - U _ {t}) d t, \tag {33}
$$