SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
d \Theta_ {t} = - \left(\sigma_ {0} \sqrt {U _ {t}} + \varepsilon_ {0}\right) ^ {- 1} (\nabla f (\Theta_ {t}, Z _ {t}) d t + \sigma_ {0} \Sigma (\Theta_ {t}, Z _ {t}) ^ {1 / 2} d W t) \tag {34}
$$
$$
d Z _ {t} = \beta_ {0} \left(\Theta_ {t} - Z _ {t}\right) d t, \tag {35}
$$
with step size  $\eta^2$ . Of course, we recover the SDE of Malladi et al. (2022) in the case where  $\beta_0 = 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 et 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  $\gamma, \eta$  and  $\rho$ , if the batch size goes from  $B$  to  $\hat{B} = \kappa \times B$ , the noise level becomes  $\hat{\sigma} = \sigma / \sqrt{\kappa}$ , and keeping the quantities in Equation 29 constant means that we should use as new hyperparameters
$$
\hat {\gamma} = 1 - (1 - \gamma) \times \kappa , \quad \hat {\eta} = \eta \times \sqrt {\kappa}, \text {a n d} \hat {\rho} = 1 - (1 - \rho) \times \kappa .
$$
The linear rule  $\hat{\rho} = 1 - (1 - \rho) \times \kappa$  is at the first order equivalent to the exponential scaling rule  $\hat{\rho} = \rho^{\kappa}$ . 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  $\rho$  the exact same scaling rule as for SGD.
We finish the discussion with the case of Adam, which leads once again to the same rule as for SGD. Adam with EMA tracking of the network parameters iterates
$$
\mathbf {m} _ {k + 1} = \beta_ {1} \mathbf {m} _ {k} + (1 - \beta_ {1}) \mathbf {g} _ {k}, \text {w i t h} \mathbf {g} _ {k} = \nabla f \left(\boldsymbol {\theta} _ {k}, \zeta_ {k}\right) + \sigma \epsilon_ {k}, \text {a n d} \epsilon_ {k} \sim \mathcal {Z} _ {\sigma} \left(\boldsymbol {\theta} _ {k}, \zeta_ {k}\right), \tag {36}
$$
$$
\mathbf {v} _ {k + 1} = \beta_ {2} \mathbf {v} _ {k} + (1 - \beta_ {2}) \mathbf {g} _ {k} ^ {2} \tag {37}
$$
$$
\tilde {\mathbf {m}} _ {k + 1} = \mathbf {m} _ {k + 1} / \left(1 - \beta_ {1} ^ {k + 1}\right) \tag {38}
$$
$$
\tilde {\mathbf {v}} _ {k + 1} = \mathbf {v} _ {k + 1} / \left(1 - \beta_ {2} ^ {k + 1}\right) \tag {39}
$$
$$
\boldsymbol {\theta} _ {k + 1} = \boldsymbol {\theta} _ {k} - \eta \left(\sqrt {\tilde {\mathbf {v}} _ {k}} + \varepsilon\right) ^ {- 1} \times \tilde {\mathbf {m}} _ {k + 1} \tag {40}
$$
$$
\zeta_ {k + 1} = \rho \zeta_ {k} + (1 - \rho) \theta_ {k}. \tag {41}
$$
Here, we use the same minor modification of the iterations as in Malladi et al. (2022), where we use  $\mathbf{v}_k$  instead of  $\mathbf{v}_{k + 1}$  in the denominator of the  $\theta_{k}$  update.
We consider the following scaling for the hyperparameters
$$
c _ {1} = \left(1 - \beta_ {1}\right) / \eta^ {2}, c _ {2} = \left(1 - \beta_ {2}\right) / \eta^ {2}, \sigma_ {0} = \sigma \eta , \varepsilon_ {0} = \varepsilon \eta , \text {a n d} \beta_ {0} = (1 - \rho) / \eta^ {2}, \tag {42}
$$
and  $\gamma_{1}(t) = 1 - \exp (-c_{1}t)$ ,  $\gamma_{2}(t) = 1 - \exp (-c_{2}t)$ , and  $\mathbf{u}_k = \mathbf{v}_k / \sigma^2$ . The SDE for Adam + EMA is given by
$$
d M _ {t} = c _ {1} \left(\left(\nabla f \left(\Theta_ {t}, Z _ {t}\right) - M _ {t}\right) d t + \sigma_ {0} \Sigma \left(\Theta_ {t}, Z _ {t}\right) ^ {1 / 2} d W _ {t}\right) \tag {43}
$$
$$
d U _ {t} = c _ {2} \left(\operatorname {d i a g} \left(\Sigma \left(\Theta_ {t}, Z _ {t}\right)\right) - U _ {t}\right) d t \tag {44}
$$
$$
d \Theta_ {t} = - \frac {\sqrt {\gamma_ {2} (t)}}{\gamma_ {1} (t)} \left(\sigma_ {0} \sqrt {U _ {t}} + \varepsilon_ {0} \sqrt {\gamma_ {2} (t)}\right) ^ {- 1} \times M _ {t} d t \tag {45}
$$
$$
d Z _ {t} = \beta_ {0} \left(\Theta_ {t} - Z _ {t}\right) d t. \tag {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 the scaling rule in eq. 42 should be used. In other words, given a set of hyperparameters  $\beta_{1},\beta_{2},\eta$  and  $\rho$ , if the batch size goes from  $B$  to  $\kappa \times B$ , then the noise level becomes  $\hat{\sigma} = \sigma /\sqrt{\kappa}$  and keeping quantities in eq. 42 constant means that we should use as new hyperparameters
$$
\hat {\beta} _ {1} = 1 - \left(1 - \beta_ {1}\right) \times \kappa , \quad \hat {\beta} _ {2} = 1 - \left(1 - \beta_ {2}\right) \times \kappa , \quad \hat {\eta} = \eta \times \sqrt {\kappa}, \text {a n d} \hat {\rho} = 1 - (1 - \rho) \times \kappa .
$$
We once again recover a linear rule for  $1 - \rho$  which is equivalent to the exponential scaling rule  $\hat{\rho} = \rho^{\kappa}$  in the limit  $\rho \to 0$ .
E Additional proofs
E.1 Iterations of  $\mathbf{SGD} + \mathbf{EMA}$
Here we derive a critical component of the EMA Scaling Rule, the matrix equation of Equation 4 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 (Definition 2.1) and EMA (Definition 1.1):  $\nabla_{\theta}\mathcal{L}(x;\theta_{t+j},\zeta_{t+j}) \approx \nabla_{\theta}\mathcal{L}(x;\theta_t,\zeta_t) \approx \mathbf{g}$ , for  $j = 1,2,\ldots,\kappa$  and representative gradient  $\mathbf{g}$ , iterating over  $\kappa$  independent minibatches produces model states
$$
\left[ \begin{array}{c} \boldsymbol {\theta} _ {t + \kappa} \\ \boldsymbol {\zeta} _ {t + \kappa} \\ \mathbf {g} \end{array} \right] = \left[ \begin{array}{c c c} 1 & 0 & - \eta \\ 1 - \rho & \rho & 0 \\ 0 & 0 & 1 \end{array} \right] ^ {\kappa} \cdot \left[ \begin{array}{c} \boldsymbol {\theta} _ {t} \\ \boldsymbol {\zeta} _ {t} \\ \mathbf {g} \end{array} \right] = \left[ \begin{array}{c} \boldsymbol {\theta} _ {t} - \eta \kappa \mathbf {g} \\ \rho^ {\kappa} \boldsymbol {\zeta} _ {t} + (1 - \rho^ {\kappa}) \boldsymbol {\theta} _ {t} + O (\eta \times \beta_ {\rho}) \\ \mathbf {g} \end{array} \right]. \tag {47}
$$
Proof. First note that for matrices of the form
$$
\boldsymbol {A} = \left[ \begin{array}{c c c} 1 & 0 & a _ {0, 2} \\ 1 - a _ {1, 1} & a _ {1, 1} & 0 \\ 0 & 0 & 1 \end{array} \right], \tag {48}
$$
their multiplication follows
$$
\begin{array}{l} \boldsymbol {A} \boldsymbol {B} = \left[ \begin{array}{c c c} 1 & 0 & a _ {0, 2} \\ 1 - a _ {1, 1} & a _ {1, 1} & 0 \\ 0 & 0 & 1 \end{array} \right] \left[ \begin{array}{c c c} 1 & 0 & b _ {0, 2} \\ 1 - b _ {1, 1} & b _ {1, 1} & 0 \\ 0 & 0 & 1 \end{array} \right] \\ = \left[ \begin{array}{c c c} 1 & 0 & a _ {0, 2} + b _ {0, 2} \\ 1 - a _ {1, 1} b _ {1, 1} & a _ {1, 1} b _ {1, 1} & (1 - a _ {1, 1}) b _ {0, 2} \\ 0 & 0 & 1 \end{array} \right], \tag {49} \\ \end{array}
$$
$$
\begin{array}{l} \textbf {\textit {A B C}} = \left[ \begin{array}{c c c} 1 & 0 & a _ {0, 2} + b _ {0, 2} \\ 1 - a _ {1, 1} b _ {1, 1} & a _ {1, 1} b _ {1, 1} & (1 - a _ {1, 1}) b _ {0, 2} \\ 0 & 0 & 1 \end{array} \right] \left[ \begin{array}{c c c} 1 & 0 & c _ {0, 2} \\ 1 - c _ {1, 1} & c _ {1, 1} & 0 \\ 0 & 0 & 1 \end{array} \right] \\ = \left[ \begin{array}{c c c} 1 & 0 & a _ {0, 2} + b _ {0, 2} + c _ {0, 2} \\ 1 - a _ {1, 1} b _ {1, 1} c _ {1, 1} & a _ {1, 1} b _ {1, 1} c _ {1, 1} & (1 - a _ {1, 1}) b _ {0, 2} + (1 - a _ {1, 1} b _ {1, 1}) c _ {0, 2} \\ 0 & 0 & 1 \end{array} \right]. \tag {50} \\ \end{array}
$$
By induction
$$
\boldsymbol {A} ^ {\kappa} = \left[ \begin{array}{c c c} 1 & 0 & \kappa \times a _ {0, 2} \\ 1 - a _ {1, 1} ^ {\kappa} & a _ {1, 1} ^ {\kappa} & \delta \left(a _ {0, 2}, a _ {1, 1}, \kappa\right) \\ 0 & 0 & 1 \end{array} \right], \tag {51}
$$
$$
\delta \left(a _ {0, 2}, a _ {1, 1}, \kappa\right) = a _ {0, 2} \sum_ {i = 1} ^ {\kappa - 1} \left(1 - a _ {1, 1} ^ {i}\right) = a _ {0, 2} \left(\kappa - \frac {1 - a _ {1 , 1} ^ {\kappa}}{1 - a _ {1 , 1}}\right), \quad \text {f o r} a _ {1, 1} \neq 1. \tag {52}
$$
$$
\left[ \begin{array}{c c c} 1 & 0 & - \eta \\ 1 - \rho & \rho & 0 \\ 0 & 0 & 1 \end{array} \right] ^ {\kappa} = \left[ \begin{array}{c c c} 1 & 0 & - \kappa \eta \\ 1 - \rho^ {\kappa} & \rho^ {\kappa} & \delta (\eta , \rho , \kappa) \\ 0 & 0 & 1 \end{array} \right] \tag {53}
$$
where the EMA Scaling Rule error
$$
\delta (\eta , \rho , \kappa) = (- \eta) \left(\kappa - \frac {1 - \rho^ {\kappa}}{1 - \rho}\right) \approx (- \eta) \left(\kappa - \kappa + \mathcal {O} \left(\beta_ {\rho}\right)\right) = 0 + \mathcal {O} (\eta \times \beta_ {\rho}), \tag {54}
$$
where  $\beta_{\rho}\equiv 1 - \rho$  and the approximation is around  $\rho = 1$
E.2 Limiting behavior of Polyak-Ruppert averaging
Here we sketch the asymptotic behavior of a target model  $\theta$  and its EMA  $\zeta$ . Let us assume that  $\theta$  converges to the stationary distribution  $\lim_{t\to \infty}\theta_t = \theta^*$ ,  $\theta^{*}\sim p_{\infty}(\theta)$ . We are interested in statistical properties of  $\zeta^{*} = \lim_{t\to \infty}\zeta_{t}$ , as this will formalize the notion of how the EMA depends on the a time-horizon defined by its momentum  $\rho$  as discussed in Table 1.
As a warm-up, for  $n$  independent random variables  $x_1, \ldots, x_2$ , we know that the sample mean  $\bar{x} = \frac{1}{n} (x_1, x_2, \ldots, x_n)$  has the statistical properties
$$
\mathbb {E} [ \bar {x} ] = \mu , \quad \operatorname {V a r} [ \bar {x} ] = \frac {\sigma^ {2}}{n}, \tag {55}
$$
where  $\mu$  and  $\sigma$  are the population mean and variance. This gives us an idea of what to expect. As we will now show, the expectation of  $\zeta^{*}$  should have no time-horizon dependence, whereas the variance of  $\zeta^{*}$  will depend on its time horizon (i.e. the number of samples it integrates over) which is defined by  $\rho$ .
In the case of a weighted sum
$$
\bar {x} ^ {(w)} = \sum_ {i = 1} ^ {n} w _ {i} x _ {i}, \tag {56}
$$
then if the  $x_{i}$  are Independent and Identically Distributed (i.i.d.), then
$$
\mathbb {E} \left[ \bar {x} ^ {(w)} \right] = \sum_ {i = 1} ^ {n} w _ {i} \mathbb {E} \left[ x _ {i} \right] = n \bar {w} \mu , \quad \bar {w} = \frac {1}{n} \sum_ {i = 1} ^ {n} w _ {i}, \tag {57}
$$
and for the variance (Kish, 1965)
$$
\operatorname {V a r} \left[ \bar {x} ^ {(w)} \right] = n \cdot \overline {{w ^ {2}}} \cdot \sigma^ {2} \quad \overline {{w ^ {2}}} = \frac {1}{n} \sum_ {i = 1} ^ {n} w _ {i} ^ {2}, \quad \sigma^ {2} = \operatorname {V a r} \left[ x _ {i} \right]. \tag {58}
$$
We can verify that we reproduce the well-known result in Equation 55 in the case where all weights are equal to  $\frac{1}{n}$  as follows
$$
\forall i: w _ {i} = \frac {1}{n} \Longrightarrow \overline {{w ^ {2}}} = \frac {1}{n} \cdot \sum_ {i = 1} ^ {n} \left(\frac {1}{n}\right) ^ {2} = \frac {1}{n ^ {2}} \Longrightarrow \operatorname {V a r} \left[ \bar {x} ^ {(w)} \right] = n \cdot \frac {1}{n ^ {2}} \cdot \sigma^ {2} = \frac {\sigma^ {2}}{n}. \tag {59}
$$
In the case of an exponential moving average we have
$$
\zeta_ {t + 1} = \rho \zeta_ {t} + (1 - \rho) \theta_ {t} = \rho^ {t} \zeta_ {1} + (1 - \rho) \sum_ {i = 0} ^ {t - 1} \rho^ {i} \theta_ {t - i}. \tag {60}
$$
Let's consider the specific case where we are at iteration  $k$  which is sufficiently large that  $\zeta$  and  $\theta$  have converged to their stationary distributions. From  $k$ , the iterations unfold as
$$
\zeta_ {t + 1} = \rho^ {t + 1 - k} \zeta_ {k} + (1 - \rho) \sum_ {i = 0} ^ {t - k} \rho^ {i} \theta_ {t - i}. \tag {61}
$$
We rearrange for terms in  $\zeta$
$$
\zeta_ {t + 1} - \rho^ {t + 1 - k} \zeta_ {k} = (1 - \rho) \sum_ {i = 0} ^ {t - k} \rho^ {i} \theta_ {t - i}, \tag {62}
$$
and before proceeding to the final result, using  $n = t + 1 - k$ , we compute the convenient quantities
$$
\bar {\rho} = \frac {1}{n} \sum_ {i = 0} ^ {n - 1} \rho^ {i} = \frac {1}{n} \times \frac {1 - \rho^ {n}}{1 - \rho} \tag {63}
$$
$$
\overline {{\rho^ {2}}} = \frac {1}{n} \sum_ {i = 0} ^ {n - 1} \rho^ {2 i} = \frac {1}{n} \times \frac {1 - \rho^ {2 n}}{1 - \rho^ {2}}. \tag {64}
$$
Taking expectation of Equation 62 and setting statistics to their stationary values, we have
$$
\left(1 - \rho^ {n}\right) \mathbb {E} \left[ \zeta^ {*} \right] = (1 - \rho) n \bar {\rho} \mathbb {E} \left[ \theta^ {*} \right] = \left(1 - \rho^ {n}\right) \mathbb {E} \left[ \theta^ {*} \right], \tag {65}
$$
where we have used the result in Equation 57. It follows that for  $\rho \neq 1$  we have
$$
\mathbb {E} \left[ \zeta^ {*} \right] = \mathbb {E} \left[ \theta^ {*} \right], \tag {66}
$$
independent of  $\rho$ . Finally, we can take the variance of Equation 62. First the left hand side
$$
\operatorname {V a r} \left[ \zeta_ {t + 1} - \rho^ {n} \zeta_ {k} \right] = \operatorname {V a r} \left[ \zeta_ {t + 1} \right] + \rho^ {2 n} \operatorname {V a r} \left[ \zeta_ {k} \right] = \left(1 + \rho^ {2 n}\right) \operatorname {V a r} \left[ \zeta^ {*} \right]. \tag {67}
$$
Next the right hand side
$$
\operatorname {V a r} \left[ (1 - \rho) \sum_ {i = 0} ^ {n - 1} \rho^ {i} \theta_ {t - i} \right] = (1 - \rho) ^ {2} \operatorname {V a r} \left[ \sum_ {i = 0} ^ {n - 1} \rho^ {i} \theta_ {t - i} \right] = (1 - \rho) ^ {2} \cdot \left(\frac {1 - \rho^ {2 n}}{1 - \rho^ {2}}\right) \cdot \operatorname {V a r} [ \theta^ {*} ]. \tag {68}
$$
Finally, equating left and right hand sizes and rearranging for  $\mathrm{Var}[\zeta^{*}]$  gives
$$
\operatorname {V a r} \left[ \zeta^ {*} \right] = \frac {1 - \rho^ {2 n}}{1 + \rho^ {2 n}} \cdot \frac {1 - \rho}{1 + \rho} \cdot \operatorname {V a r} \left[ \theta^ {*} \right] \tag {69}
$$
In the limit  $t\to \infty$  , the momentum-dependent prefactor becomes
$$
\lim  _ {t \rightarrow \infty} \left(\frac {1 - \rho^ {2 n}}{1 + \rho^ {2 n}} \cdot \frac {1 - \rho}{1 + \rho}\right) = \frac {1 - \rho}{1 + \rho} \Rightarrow \lim  _ {t \rightarrow \infty} \operatorname {V a r} \left[ \zeta^ {*} \right] = \frac {1 - \rho}{1 + \rho} \cdot \operatorname {V a r} \left[ \theta^ {*} \right]. \tag {70}
$$
Equations 69 and 70 validate our intuition. When  $\rho \rightarrow 0$ , then  $\zeta$  behaves like  $\theta$  independent of  $T$ , with their variance and expectation matching. When  $\rho > 0$ , the momentum-dependent prefactor serves as an aggregator over the history when  $t$  is sufficiently large compared to  $k$ , reducing the variance  $\mathrm{Var}[\zeta^*]$  but preserving its expectation. This formalizes the notion of time horizon discussed in Table 1.
F Additional details and results for Polyak-Ruppert averaging
Additional background Polyak-Ruppert averaging (Definition 3.1) is a simplification of Stochastic Weight Averaging (SWA) (Izmailov et al., 2018) which uses a more complex multi-cycle schedule based weighting of the model parameters. Both Definition 3.1 and SWA present similar favorable 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 ImageNet1k (Russakovsky et al., 2014) without a model EMA, achieving an accuracy of  $80.9\%$ . Equipping a Polyak-Ruppert average  $(\rho = 0.9999)$  alleviated overfitting and gave a  $83.1\%$  accuracy.
Organization In this appendix, we look at additional momenta for one-dimensional noisy parabola, as well as extensions to  $D$ -dimensions (Appendix F.1), provide a more detailed view of the results of Section 3.2 (Appendix F.2), and investigate the scenario where the EMA Scaling Rule (Definition 1.2) is applied to batch normalization (Ioffe & 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  $\rho_{B}$ . We present  $\rho_{B} = 0.99$  in Figure 7, and  $\rho_{B} = 0.999$  in Figure 8. The results for  $\rho_{B} = 0.9999$  are presented in main text in Figure 1.
(a) Trajectory of the model EMA  $\zeta$  under different scalings  $\kappa$ , with  $\rho_{B} = 0.99$ ,  $\eta_{B} = 10^{-4}$ .
(b) Choices for momentum (left) with corresponding approximation errors (Equation 10) (right).
Figure 7: (a) We show the effect of scaling by comparing model EMA trajectories of the baseline  $(\kappa = 1$ , black dashed) to  $\kappa = 8$  (left) and  $\kappa = 256$  (right), with  $(\rho = \rho_{B}^{\kappa}$ , blue) and without  $(\rho = \rho_{B}$ , red) the EMA Scaling Rule. (b, left) The momentum according for different scaling rules and the empirically optimal  $\rho^{*}$  (Equation 10). (b, right) The approximation error (Equation 10) of trajectories in (b, left) and the target model (orange). Error for  $\rho^{*}$  is computed using a hold-out to mitigate overfitting.
Figure 8: (a) We show the effect of scaling by comparing model EMA trajectories of the baseline  $(\kappa = 1$ , black dashed) to  $\kappa = 8$  (left) and  $\kappa = 256$  (right), with  $(\rho = \rho_{B}^{\kappa}$ , blue) and without  $(\rho = \rho_{B}$ , red) the EMA Scaling Rule. (b, left) The momentum according for different scaling rules and the empirically optimal  $\rho^{*}$  (Equation 10). (b, right) The approximation error (Equation 10) of trajectories in (b, left) and the target model (orange). Error for  $\rho^{*}$  is computed using a hold-out to mitigate overfitting.
(a) Trajectory of the model EMA  $\zeta$  under different scalings  $\kappa$ , with  $\rho_{B} = 0.999$ ,  $\eta_{B} = 10^{-4}$ .
(b) Choices for momentum (left) with corresponding approximation errors (Equation 10) (right).
As described by the scaling error term in Equation 54, the approximation error at a given  $\kappa$  is higher for lower momenta  $\rho$ . For a large range of scalings  $\kappa$ , the EMA Scaling Rule and the optimal momenta  $\rho^{*}$  are consistent. In summary, we see the synthetic experiments validate the results of Section 3.1 for a range of momenta  $\rho$ .
Examples in higher dimensions Our final use of the synthetic noisy parabola will consider an extension to  $D$  dimensions. Consider the optimization of  $\theta \in \mathbb{R}^D$  in a noisy parabola at the origin:
$$
\mathcal {L} (\boldsymbol {\theta}) = \frac {a}{2} \boldsymbol {\theta} ^ {\top} \boldsymbol {\theta}, \quad \boldsymbol {\theta} _ {k + 1} = \boldsymbol {\theta} _ {k} - \eta \mathbf {g} _ {k}, \quad \mathbf {g} _ {k} = a \boldsymbol {\theta} _ {k} + \boldsymbol {\epsilon} _ {k}, \quad \boldsymbol {\epsilon} _ {k} \sim \mathcal {N} \left(\mathbf {0}, \frac {b g _ {k} ^ {2} + c}{\kappa}\right), \tag {71}
$$
for curvature  $a > 0$ , scaled additive  $b > 0$ , and additive  $c > 0$  noise coefficients. The scaling factor  $\kappa$  in the covariance denominator implements the reduction in gradient noise as the scaling (i.e., the batch size) increases (Jastrzebski et al., 2017). Let  $\theta \in \mathbb{R}^D$  be optimized with SGD (Definition 2.1)
and let there be a Polyak-Ruppert average (Definition 3.1)  $\zeta \in \mathbb{R}^D$  with momentum  $\rho = 1 - \beta$  for  $\theta$ . We consider dimensionalities  $D = 2$  (Figure 9),  $D = 16$  (Figure 10), and  $D = 100$  (Figure 11). We observe no significant differences in the EMA scaling behavior as we vary dimensions.
(a) Norm of the model EMA  $\zeta$  under different scalings  $\kappa$ , with  $\rho_{B} = 0.9999$ ,  $\eta_{B} = 10^{-4}$ ,  $D = 2$ .
Figure 9: (a) We show the effect of scaling by comparing model EMA trajectories of the baseline ( $\kappa = 1$ , black dashed) to  $\kappa = 8$  (left) and  $\kappa = 256$  (right), with  $(\rho = \rho_{B}^{\kappa}$ , blue) and without  $(\rho = \rho_{B}$ , red) the EMA Scaling Rule. (b, left) The momentum according for different scaling rules and the empirically optimal  $\rho^{*}$  (Equation 10). (b, right) The approximation error (Equation 10) of trajectories in (b, left) and the target model (orange). Error for  $\rho^{*}$  is computed using a hold-out to mitigate overfitting.
(a) Norm of the model EMA  $\zeta$  under different scalings  $\kappa$ , with  $\rho_{B} = 0.9999$ ,  $\eta_{B} = 10^{-4}$ ,  $D = 16$ .
(b) Choices for momentum (left) with corresponding approximation errors (Equation 10) (right).
Figure 10: (a) We show the effect of scaling by comparing model EMA trajectories of the baseline  $(\kappa = 1$ , black dashed) to  $\kappa = 8$  (left) and  $\kappa = 256$  (right), with  $(\rho = \rho_B^\kappa$ , blue) and without  $(\rho = \rho_B$ , red) the EMA Scaling Rule. (b, left) The momentum according for different scaling rules and the empirically optimal  $\rho^*$  (Equation 10). (b, right) The approximation error (Equation 10) of trajectories in (b, left) and the target model (orange). Error for  $\rho^*$  is computed using a hold-out to mitigate overfitting.
(a) Norm of the model EMA  $\zeta$  under different scalings  $\kappa$ , with  $\rho_{B} = 0.9999$ ,  $\eta_{B} = 10^{-4}$ ,  $D = 100$ .
(b) Choices for momentum (left) with corresponding approximation errors (Equation 10) (right).
Figure 11: (a) We show the effect of scaling by comparing model EMA trajectories of the baseline ( $\kappa = 1$ , black dashed) to  $\kappa = 8$  (left) and  $\kappa = 256$  (right), with  $(\rho = \rho_{B}^{\kappa}$ , blue) and without  $(\rho = \rho_{B}$ , red) the EMA Scaling Rule. (b, left) The momentum according for different scaling rules and the empirically optimal  $\rho^{*}$  (Equation 10). (b, right) The approximation error (Equation 10) of trajectories in (b, left) and the target model (orange). Error for  $\rho^{*}$  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 on a single CPU, and so we do not detail this compute usage as we do in the other experimental sections.
Table 4: Supervised ResNetv2-50 hyperparameters used in Polyak-Ruppert Averaging experiments.
Supervised ResNetv2-50
ImageNet1k Test Top-1	76.27 ± 0.10%
ImageNet1k EMA Test Top-1	76.55 ± 0.07%
Weight initialization	kaiming_normal(relu)
Backbone normalization	BatchNorm
Synchronized BatchNorm over replicas	No
Learning rate schedule	Multi step: ×0.1 at [30, 60, 80] epochs
Learning rate warmup (epochs)	5
Learning rate minimum value	1 × 10-6
Training duration (epochs)	90
Optimizer	SGD + Momentum
SGD momentum	0.9
Optimizer scaling rule	Linear
Base learning rate	0.4
Base batch size	1024
Base Polyak momentum	0.9999
Weight decay	1 × 10-4
Weight decay scaling rule	None
Weight decay skip bias	Yes
Numerical precision	bf16
Augmentation stack	ImageNet
Label smoothing rate	0.1
F.2 Image Classification
Hyperparameters We present the base hyperparameters for our image experiments in Table 4.
Data For large scale vision evaluation, we use the ImageNet1k dataset (Russakovsky et al., 2014), a widely used dataset containing approximately 1.2 million labeled images, distributed almost uniformly across 1000 different object classes, like animals, plants, and vehicles.
The images in ImageNet1k are not consistent in resolution. To handle this, they are resized and cropped to a standard size (in our case,  $224 \times 224$ ), before further processing. This is part of the standard ImageNet augmentation stack for convolutional networks mentioned in Table 4.
Compute usage The compute usage image classification Polyak-Ruppert averaging is summarized in Table 5.
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 to 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	8	35.3	282.4	9	2,541.6
1,024	8	17.1	137.0	3	410.9
2,048	8	13.3	106.7	9	960.6
4,096	8	4.2	33.5	9	301.9
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