SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Dan Busbridge* Jason Ramapuram* Pierre Ablin* Tatiana Likhomanenko*
Eeshan Gunesh Dhekane Xavier Suau Russ Webb
{dbusbridge, jramapuram, p_ablin, antares, eeshan, xsuaucuadros, rwebb}@apple.com
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.
With data and models becoming progressively larger (Chen et al., 2020; Kaplan et al., 2020; Bommasani et al., 2021; Srivastava et al., 2022), the ability to reduce training wall-clock time is a requirement 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 scaling 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 a detailed breakdown of author contributions see Appendix J.
Many ML algorithms rely on a model EMA, a functional copy of a target model<sup>2</sup>, whose parameters 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  $\rho$ .
Definition 1.1 (EMA Update). The EMA update for the model EMA parameters  $\zeta_t$  following target model parameters  $\theta_t$  at iteration  $t$  with momentum  $\rho \equiv 1 - \beta_{\rho}$  is
$$
\zeta_ {t + 1} = \rho \zeta_ {t} + (1 - \rho) \theta_ {t} \equiv (1 - \beta_ {\rho}) \zeta_ {t} + \beta_ {\rho} \theta_ {t}. \tag {1}
$$
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 target model, the model EMA moves slowly, making it useful as a stabilizer for networks governing Bellman updates in reinforcement learning, (Lillicrap et al., 2016); and iii) the model EMA is relatively cheap to compute, whilst providing a 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 et al., 2021; Higuchi et al., 2022), to Self-Supervised Learning (SSL) methods like Bootstrap Your Own Latent (BYOL) (Grill et al., 2020), DINO (Caron et al., 2021), and data2vec (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.2) which states how the EMA momentum  $\rho$  hyperparameter should be modified<sup>3</sup>.
Definition 1.2 (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.
In Definition 1.2, the momentum  $\rho$ , which is defined at batch size  $B$ , typically corresponds to a "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 semi-supervised 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 a mixture of a) labeled data and b) data without labels, whose pseudolabels 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 dynamics during optimizer warm-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 BYOL Vision Transformers (ViTs), reducing wall-clock training by  $6 \times$  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.
2The target model usually undergoes gradient-based optimization, but this does not have to be the case. 3We stress that the study of momentum in gradient-based optimizers is not the focus of this work. We refer to Smith & Le (2018); Li et al. (2019) for a discussion on 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 a model with parameters  $\theta_{t}$  at iteration  $t$  updated with SGD (Definition 2.1).
Definition 2.1 (SGD Update). The SGD update for a model with parameters  $\theta_t$  at iteration  $t$  given a minibatch  $\mathbb{B} = \{x^{(b)}\sim P_{\mathbf{x}}:b = 1,2,\dots,B\}$  of  $B = |\mathbb{B}|$  samples with learning rate  $\eta$  is
$$
\boldsymbol {\theta} _ {t + 1} = \boldsymbol {\theta} _ {t} - \eta \times \frac {1}{B} \sum_ {x \in \mathbb {B}} \nabla_ {\boldsymbol {\theta}} \mathcal {L} (x; \boldsymbol {\theta} _ {t}), \tag {2}
$$
where  $\mathcal{L}$  is the loss function,  $\nabla_{\theta}\mathcal{L}(x;\theta_t)$  is the parameter gradient for the sample  $x$  at iteration  $t$ , and the  $x\in \mathbb{B}$  are Independent and Identically Distributed (i.i.d.) from  $P_{\mathbf{x}}$ .
Iterating over a sequence of independent minibatches  $\mathbb{B}_0, \mathbb{B}_1, \ldots, \mathbb{B}_{\kappa - 1}$  produces model parameters
$$
\boldsymbol {\theta} _ {t + \kappa} = \boldsymbol {\theta} _ {t} - \eta \times \frac {1}{B} \sum_ {j = 0} ^ {\kappa - 1} \sum_ {x \in \mathbb {B} _ {j}} \nabla_ {\boldsymbol {\theta}} \mathcal {L} (x; \boldsymbol {\theta} _ {t + j}). \tag {3}
$$
If gradients vary slowly  $\nabla_{\theta}\mathcal{L}(x;\theta_{t + j})\approx \nabla_{\theta}\mathcal{L}(x;\theta_t)$ $j = 0,\dots ,\kappa -1$  , one SGD step with  $\hat{\eta} = \kappa \eta$  on a batch  $\widehat{\mathbb{B}} = \cup_i\mathbb{B}_i$  of size  $\hat{B} = \kappa B$  results in  $\hat{\pmb{\theta}}_{t + 1}\approx \pmb{\theta}_{t + k}$  yielding the SGD Scaling Rule (Definition 2.2).
Definition 2.2 (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).
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  $\kappa B$ .
As discussed in Goyal et al. (2017), although the assumption of slowly changing gradients is strong, if it is true, then  $\theta_{t + k} \approx \hat{\theta}_{t + 1}$  only if  $\hat{\eta} = \kappa \eta$ . 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  $\kappa$  when the SGD Scaling Rule is no longer guaranteed to hold (Li et al., 2021).
To address model parameter EMAs, we first restate the EMA Update (Definition 1.1).
The model EMA parameters  $\zeta$  do not typically receive gradient information, we take the convention that  $\rho$  is close to one, and the  $\beta_{\rho}$  subscript will be omitted where it is clear from the context.
Assuming again that gradients change slowly  $\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 gradient  $\mathbf{g}$ , iterating over  $\kappa$  independent minibatches produces model states (see Appendix E.1 for derivation)
$$
\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 {4}
$$
The first row is the SGD Scaling Rule (Definition 2.2). The third row implements the slowly changing gradients assumption for the first row. The second row is equivalent to a single EMA update (Definition 1.1) with momentum  $\hat{\rho} = \rho^{\kappa}$ ; we can take a single SGD update with batch size  $\hat{B} = \kappa B$  and learning rate  $\hat{\eta} = \kappa \eta$ , and a single EMA update with momentum  $\hat{\rho} = \rho^{\kappa}$ , and we get  $(\hat{\theta}_{t+1}, \hat{\zeta}_{t+1}) \approx (\theta_{t+\kappa}, \zeta_{t+\kappa})$  up to terms  $O(\eta \times \beta_{\rho})$ . This yields the EMA Scaling Rule (Definition 1.2).
The EMA Scaling Rule was derived for SGD, and is extended to other optimizers in the following way. An optimizer scaling rule ensures  $\hat{\theta}_{t + 1} = \theta_{t + \kappa}$ , satisfying identification for the first row. Next, the zeroth order term in  $\eta \times \beta_{\rho}$  in the second row in Equation 4 is optimizer-independent, and therefore unchanged. Finally, the first order terms in  $\eta \times \beta_{\rho}$  in the second row, corresponding to the scaling rule error, are an EMA accumulation of target model  $\theta$  updates under optimization, which is still  $O(\eta \times \beta_{\rho})$ , although its functional form may be different for different optimizers.
The above discussion is intended to give an 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 is 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.1. We consider the following discrete dynamics for EMA:
$$
\theta_ {k + 1} = \theta_ {k} - \eta \mathbf {g} _ {k}, \text {w i t h} \mathbf {g} _ {k} = \nabla f \left(\theta_ {k}, \zeta_ {k}\right) + \sigma \epsilon_ {k}, \text {a n d} \epsilon_ {k} \sim \mathcal {E} _ {\sigma} \left(\theta_ {k}, \zeta_ {k}\right), \tag {5}
$$
$$
\zeta_ {k + 1} = \rho \zeta_ {k} + (1 - \rho) \theta_ {k},
$$
where  $\sigma > 0$  is the noise scale,  $\mathcal{E}_{\sigma}(\theta_k, \zeta_k)$  is the gradient noise distribution, assumed to be zero-mean and variance  $\Sigma(\theta_k, \zeta_k)$  independent of  $\sigma$ , and  $\nabla f(\theta_k, \zeta_k) \equiv \nabla_{\theta} f(\theta_k, \zeta_k)$ . We posit a dependency of the loss  $f$  on the EMA  $\zeta$  in order to cover semi-supervised (Section 3.3) and SSL (Section 3.4). The case of Polyak-Ruppert averaging (Section 3.2), is covered by letting  $f$  be independent of  $\zeta$ .
We aim to obtain an SDE approximation of Equation 5 as  $\eta$  goes to zero. The scaling rule for iterations of  $\theta$  is well known (Li et al., 2021): we let  $\sigma_0 = \sigma \sqrt{\eta}$ . The analysis of Section 2.1 gives the scaling rule  $\hat{\eta} = \eta \kappa$  and  $\hat{\rho} = \rho^{\kappa}$ . Linearizing this rule near  $\eta = 0$  gives  $\hat{\rho} = 1 - \kappa \times (1 - \rho)$ , which is a linear relationship between  $1 - \rho$  and  $\eta$ . We therefore let  $\beta_0 = (1 - \rho) / \eta$  and consider the SDE
$$
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 e r p r o c e s s ,} \tag {6}
$$
$$
d Z _ {t} = \beta_ {0} (\Theta_ {t} - Z _ {t}) d t,
$$
where  $\Theta_t$  and  $Z_t$  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  $\eta$  goes to zero. One way to see this is that an Euler-Maruyama discretization of the SDE with learning rate  $\eta$  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^{\alpha}$  is the set of functions with derivatives up to order  $\alpha$  that have at most polynomial growth (see Definition D.1).
Theorem 2.1 (SDE for  $\mathrm{SGD} + \mathrm{EMA}$ ; informal see Theorem D.1). Assume that  $f$  is continuously differentiable, with  $f \in G^3$ . Let  $\Theta_t, Z_t$  be solutions of Equation 6, and  $\theta_k, \zeta_k$  iterations of Equation 5 with  $\Sigma^{\frac{1}{2}} \in G^2$ . Then, for any time horizon  $T > 0$  and function  $g \in G^2$ , there exists a constant  $c > 0$  independent of  $\eta$  such that
$$
\max  _ {k = 0, \dots , \lfloor T / \eta \rfloor} | \mathbb {E} [ g (\Theta_ {\eta k}, Z _ {\eta k}) ] - \mathbb {E} [ g (\theta_ {k}, \zeta_ {k}) ] | \leq c \times \eta . \tag {7}
$$
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 \in G^3$  and  $\Sigma^{\frac{1}{2}} \in G^2$ . Let  $\theta_k^{(B)}, \zeta_k^{(B)}$  be iterations of 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$ , and  $\hat{\eta}$  determined by the SGD Scaling Rule (Definition 2.2) and  $\hat{\rho}$  determined by the EMA Scaling Rule (Definition 1.2). Then, for any time horizon  $T > 0$  and function  $g \in G^2$ , there exists a constant  $c > 0$  independent of  $\eta$  such that
$$
\max  _ {k = 0, \dots , \lfloor T / \eta \rfloor} | \mathbb {E} [ g (\theta_ {\lfloor k / \kappa \rfloor} ^ {(\kappa B)}, \zeta_ {\lfloor k / \kappa \rfloor} ^ {(\kappa B)}) ] - \mathbb {E} [ g (\theta_ {k} ^ {(B)}, \zeta_ {k} ^ {(B)}) ] | \leq c \times \eta . \tag {8}
$$
Table 1: The role of the model EMA  $\zeta$  in the optimization of  $(\theta, \zeta)$  given a target model  $\theta$  for different techniques, ordered by increasing influence of the EMA model. All statements assume a momentum  $0 \leq \rho < 1$  and that the target model  $\theta$  is subject to stochastic optimization at a batch size  $B$ .
TECHNIQUE	ROLE OF MODEL EMA
POLYAK-RUPPERT AVERAGING, SEC. 3.2	θ undergoes optimization and is tracked by ζ, which does not affect θ. ζ is an estimate of θ with a time horizon and variance determined by B and ρ.
CONTINUOUS PSEUDO-LABELING, SEC. 3.3	Pre-Training is as above in Polyak-Ruppert Averaging. After Pre-Training, ζ (teacher) produces targets for θ (student) from unlabeled data, which is combined with labeled data. The optimization endpoint is dependent on B and ρ.
SELF-SUPERVISED LEARNING, SEC. 3.4	As above in After Pre-Training, except there is no labeled data. The optimization endpoint is dependent on B and ρ.
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.
Now that we have derived and shown the validity of the EMA Scaling Rule, we verify it empirically. The experiments validate the EMA Scaling Rule for a 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 Polyak-Ruppert Averaging (Definition 3.1).
Definition 3.1 (Polyak-Ruppert Average). When optimizing model parameters  $\theta$ , compute their EMA  $\zeta$  (Definition 1.1). Use  $\zeta$  instead of  $\theta$  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  $\theta$  in a noisy parabola whose loss  $\mathcal{L}(\theta)$  is parameterized by coefficients for curvature  $a > 0$ , scaled additive noise  $b \geq 0$ , and additive noise  $c \geq 0$ :
$$
\mathcal {L} (\theta) = \frac {a}{2} \theta^ {2}, \quad \theta_ {k + 1} = \theta_ {k} - \eta \mathrm {g} _ {k}, \quad \mathrm {g} _ {k} = a \theta_ {k} + \epsilon_ {k}, \quad \epsilon_ {k} \sim \mathcal {N} \left(0, \frac {b \mathrm {g} _ {k} ^ {2} + c}{\kappa}\right). \tag {9}
$$
The scaling factor  $\kappa$  in the covariance denominator implements gradient noise reduction as scaling (i.e. batch size) increases (Jastrzebski et al., 2017). Let  $\theta \in \mathbb{R}$  be optimized with SGD (Definition 2.1) and  $\zeta \in \mathbb{R}$  be a Polyak-Ruppert average (Definition 3.1) for  $\theta$  with momentum  $\rho = 1 - \beta$ . At scaling  $\kappa = 1$ , we use  $\beta_{B} = \eta_{B} = 10^{-4}$  and  $I_{B} = 10^{4}$  iterations, to yield a total time  $T = I_{B} \times \eta_{B} = 1$ . To keep gradients  $O(1)$  and gradient noise non-negligible, we take  $a = 1$ ,  $b = 0.5$ , and  $c = 0$ .
First, we observe the effect of scaling on a single run (Figure 1a) by tracking the position of the model EMA. We see that at scaling  $\kappa = 8$  or  $\kappa = 256$ , the runs using the EMA Scaling Rule match the baseline trajectory, whereas the runs using the baseline momentum do not, with a greater deviation induced by greater scaling  $\kappa$ . Even at  $\kappa = 8$ , there is a significant difference between scaled and unscaled trajectories, despite the seemingly small numerical difference of their momenta<sup>4</sup>.
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,  $\operatorname{Err}(\rho, \kappa, g)$ , of a test function  $g$  for a given
<sup>4</sup>Momentum enters optimization exponentially; small changes can lead to very different updates.
scaling  $\kappa$  using momentum  $\rho$ , and the value of the momentum  $\rho^{*}(\kappa, g)$  that minimizes this error:
$$
\rho^ {*} (\kappa , g) = \underset {\rho} {\arg \min } \operatorname {E r r} (\rho , \kappa , g), \quad \operatorname {E r r} (\rho , \kappa , g) \equiv \underset {k = 0, \dots , T / \eta} {\max } \left| \mathbb {E} g \left(\zeta_ {k}\right) - \mathbb {E} g \left(\zeta_ {k / \kappa} ^ {(\kappa , \rho)}\right) \right|. \tag {10}
$$
For scalings  $\kappa \in \{1,2,4,\dots ,1024\}$ , we determine the optimal momentum  $\rho^{*}$  and compare it to the EMA Scaling Rule (Figure 1b, left). The scaling rule tracks the  $\rho^{*}$  until  $\kappa = 256$ , when the  $\rho^{*}$  become systematically higher. We see target model error increase at  $\kappa = 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  $\kappa = 64$ , there is an inflection point in the EMA Scaling Rule approximation error, before the model error grows. This difference indicates the  $O(\eta \times \beta_{\rho})$  terms of Equation 4 are beginning to influence the EMA update. Finally, these observations are true in  $D = 100$  dimensions, (Appendix F.1), and we stress that not changing the momentum at every scaling  $\kappa$  induces large approximation error, indicating there is merit to using the EMA Scaling Rule.
3.2 Supervised learning on real data with Polyak-Ruppert averaging
We now turn to real-world classification where the target model  $\theta$  optimizes a parametric log-likelihood  $\max_{\theta} \log p(\mathbf{y}|\mathbf{x};\theta)$  with inputs and labels  $(x,y)$  drawn from a joint distribution  $p(\mathbf{y},\mathbf{x})$ .