Rethinking Learning Rate and Batch Size (IV): EMA

重新思考学习率与Batch Size（四）：EMA 苏剑林 · 2025-09-22 · Mathematics learning rate optimizer scaling law mean field

In Rethinking Learning Rate and Batch Size (II): Mean Field, we mentioned that one reason for focusing on SignSGD is that we typically use it as a theoretical approximation to Adam — a common simplification strategy for analyzing Adam theoretically. Beyond the learning-rate analysis scenario, we've also used this simplification in Can Configuring Different Learning Rates Push LoRA Even Further? and A First Look at MuP: Cross-Model Scale Transfer Laws for Hyperparameters.

But is SignSGD really a good approximation of Adam? One obvious difference is that SignSGD's Update RMS is always 1, while Adam's is not. The author found that the core reason behind this discrepancy is momentum, which is ubiquitous in optimizers like Adam, Lion, and Muon. So in this post, we examine the effect of momentum — or more broadly, EMA.

Problem Analysis

From Adam's perspective, SignSGD corresponds to the special case $\beta_1=\beta_2=0$, or equivalently to Adam's very first update step (regardless of $\beta_1,\beta_2$). So we'd expect it to share some commonalities with Adam, capturing some general patterns. more

That said, there are also clear differences between them. A typical one is the difference in Update RMS: SignSGD's is always 1, whereas Adam's is often noticeably less than 1. Also, Adam appears closer to SGD — it looks more like an intermediate version between SignSGD and SGD. At first, the author thought this was due to the $\epsilon$ in Adam's denominator, so in How Does Adam's Epsilon Affect the Learning Rate Scaling Law? we specifically computed SoftSignSGD with $\epsilon$.

Later, in Why Is Adam's Update RMS 0.2?, we estimated Adam's Update RMS both via simulation and theory; the mean-field approximation gives $\sqrt{\frac{1-\beta_1}{1+\beta_1}}$, and this was confirmed to match both the simulation results and actual experiments well. This result explicitly depends on $\beta_1$, which clearly points our thinking toward momentum.

That is what led to the analysis below. In summary, we can confirm that the role of $\epsilon$ is indeed secondary; the true protagonist is momentum — the "moving average" of the gradient — which is exactly the subject of this post: EMA (Exponential Moving Average).

Gradient Descent

To analyze the changes brought about by EMA, we start with SGDM, i.e., SGD with momentum. In practice, we almost never use SGD without momentum:

\begin{equation}\begin{aligned} &\boldsymbol{m}_t = \beta_1 \boldsymbol{m}_{t-1} + \left(1 - \beta_1\right) \boldsymbol{g}_t \\[4pt] &\boldsymbol{w}_t = \boldsymbol{w}_{t-1} - \eta_t \boldsymbol{m}_t \end{aligned}\end{equation}

In actual use, $\boldsymbol{g}_t$ is replaced by $\tilde{\boldsymbol{g}}_{B,t}$, a random variable with mean $\boldsymbol{g}_t$ and covariance matrix $\boldsymbol{\Sigma}_t/B$. These basic settings are the same as in Rethinking Learning Rate and Batch Size (I): Current State. The noise here arises from randomly sampling different batches, so we can reasonably assume that $\tilde{\boldsymbol{g}}_{B,t}$ for different values of $t$ are mutually independent.

Our task is to compute

\begin{equation}\newcommand{tr}{\mathop{\text{tr}}}\eta^* \approx \frac{\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]^{\top}\boldsymbol{g}}{\tr(\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]\boldsymbol{H})}\label{eq:eta-opt}\end{equation}

The relevant derivation has already been given in earlier posts, so we won't repeat it here. For SGDM, $\tilde{\boldsymbol{\varphi}}_B = \boldsymbol{m}_t$, and this can be expanded as

\begin{equation}\boldsymbol{m}_t = (1 - \beta_1)\sum\limits_{s=1}^t \beta_1^{t-s}\tilde{\boldsymbol{g}}_{B,s}\end{equation}

Scaling Up the Batch

Now we can compute

\begin{equation}\mathbb{E}[\boldsymbol{m}_t] = (1 - \beta_1)\sum_{s=1}^t \beta_1^{t-s}\mathbb{E}[\tilde{\boldsymbol{g}}_{B,s}] = (1 - \beta_1)\sum_{s=1}^t \beta_1^{t-s}\boldsymbol{g}_s\end{equation}

We further assume that once training has settled into a "steady groove," the gradient changes slowly, so we can approximate $\boldsymbol{g}_s$ with the current gradient $\boldsymbol{g}_t$, giving

\begin{equation}\mathbb{E}[\boldsymbol{m}_t] = (1 - \beta_1)\sum_{s=1}^t \beta_1^{t-s}\boldsymbol{g}_t = (1 - \beta_1^t) \boldsymbol{g}_t \approx \boldsymbol{g}_t \qquad (t\to\infty)\end{equation}

As for $\mathbb{E}[\boldsymbol{m}_t \boldsymbol{m}_t^{\top}]$, we use the identity $\mathbb{E}[\boldsymbol{m}_t \boldsymbol{m}_t^{\top}] = \mathbb{E}[\boldsymbol{m}_t] \mathbb{E}[\boldsymbol{m}_t]^{\top} + \mathbb{C}\text{ov}[\boldsymbol{m}_t,\boldsymbol{m}_t]$, then apply additivity of variance to get:

\begin{equation}\mathbb{C}\text{ov}[\boldsymbol{m}_t,\boldsymbol{m}_t] = (1 - \beta_1)^2\sum_{s=1}^t \beta_1^{2(t-s)}\boldsymbol{\Sigma}_s/B\end{equation}

Similarly, assuming the covariance matrix also changes slowly, we get

\begin{equation}\mathbb{C}\text{ov}[\boldsymbol{m}_t] \approx (1 - \beta_1)^2\sum_{s=1}^t \beta_1^{2(t-s)}\boldsymbol{\Sigma}_t/B = (1 - \beta_1)^2\frac{1-\beta_1^{2t}}{1-\beta_1^2}\boldsymbol{\Sigma}_t/B = \frac{1 - \beta_1}{1 + \beta_1}\boldsymbol{\Sigma}_t/B \qquad (t\to\infty)\end{equation}

Substituting into equation $\eqref{eq:eta-opt}$ gives

\begin{equation}\eta^* \approx \frac{\eta_{\max}}{1 + \frac{1 - \beta_1}{1 + \beta_1}\mathcal{B}_{\text{noise}}/B},\qquad \eta_{\max} = \frac{\boldsymbol{g}^{\top}\boldsymbol{g}}{\boldsymbol{g}^{\top}\boldsymbol{H}\boldsymbol{g}},\quad\mathcal{B}_{\text{noise}} = \frac{\tr(\boldsymbol{\Sigma}\boldsymbol{H})}{\boldsymbol{g}^{\top}\boldsymbol{H}\boldsymbol{g}}\end{equation}

This result shows that introducing the momentum mechanism is equivalent to scaling SGD's batch size up by a factor of $\frac{1 + \beta_1}{1 - \beta_1}$. As the author understands it, momentum works by taking an EMA of the gradients along the optimization trajectory, cheaply suppressing gradient noise — so this result is consistent with the intuitive meaning of momentum.

Signed Momentum

Next we consider SignSGDM, which can be viewed as a special case of Lion — essentially SGDM with an extra $\newcommand{sign}{\mathop{\text{sign}}}\sign$ applied:

\begin{equation}\begin{aligned} &\boldsymbol{m}_t = \beta_1 \boldsymbol{m}_{t-1} + \left(1 - \beta_1\right) \boldsymbol{g}_t \\[4pt] &\boldsymbol{w}_t = \boldsymbol{w}_{t-1} - \eta_t \sign(\boldsymbol{m}_t) \end{aligned}\end{equation}

In actual training, $\boldsymbol{g}_t$ is likewise replaced by $\tilde{\boldsymbol{g}}_{B,t}$. For SignSGDM, $\tilde{\boldsymbol{\varphi}}_B = \sign(\boldsymbol{m}_t)$, so by the mean-field approximation we get

\begin{equation}\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B] = \mathbb{E}\bigg[\frac{\boldsymbol{m}_t}{\sqrt{\boldsymbol{m}_t^2}}\bigg]\approx \frac{\mathbb{E}[\boldsymbol{m}_t]}{\sqrt{\mathbb{E}[\boldsymbol{m}_t^2]}}\end{equation}

where vector multiplication defaults to the Hadamard product. We already computed the numerator $\mathbb{E}[\boldsymbol{m}_t]$ in the previous section; the denominator $\mathbb{E}[\boldsymbol{m}_t^2]$ actually equals $\newcommand{diag}{\mathop{\text{diag}}}\diag(\mathbb{E}[\boldsymbol{m}_t \boldsymbol{m}_t^{\top}])$, so we can also substitute the result from the previous section, giving

\begin{equation}\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B] \approx \frac{\boldsymbol{g}_t}{\sqrt{\boldsymbol{g}_t^2 + \frac{1 - \beta_1}{1 + \beta_1}\boldsymbol{\sigma}_t^2/B}} = \frac{\sign(\boldsymbol{g}_t)}{\sqrt{1 + \frac{1 - \beta_1}{1 + \beta_1}(\boldsymbol{\sigma}_t^2/\boldsymbol{g}_t^2)/B}} \approx \frac{\sign(\boldsymbol{g}_t)}{\sqrt{1 + \frac{1 - \beta_1}{1 + \beta_1} \mathcal{B}_{\text{simple}}/B}}\end{equation}

where $\boldsymbol{\sigma}_t^2 = \diag(\boldsymbol{\Sigma}_t), \mathcal{B}_{\text{simple}} = \tr(\boldsymbol{\Sigma}_t)/\boldsymbol{g}_t^{\top}\boldsymbol{g}_t$. The equation above is equivalent to SignSGD with $B$ replaced by $\frac{1 + \beta_1}{1 - \beta_1}B$; if we further compute $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]$, we'll find the same conclusion holds. So, just as with SGDM, momentum is equivalent to scaling up SignSGD's batch size by a factor of $\frac{1 + \beta_1}{1 - \beta_1}$.

In Rethinking Learning Rate and Batch Size (III): Muon, we computed the learning rate scaling law for Muon and found it matches SignSGD's. So we can assert that the role of momentum in Muon is the same as in SignSGDM: it roughly amounts to scaling up the batch size by a factor of $\frac{1 + \beta_1}{1 - \beta_1}$.

Double Moving Average

Finally, let's look at Adam:

\begin{equation}\begin{aligned} &\boldsymbol{m}_t = \beta_1 \boldsymbol{m}_{t-1} + \left(1 - \beta_1\right) \boldsymbol{g}_t\\ &\boldsymbol{v}_t = \beta_2 \boldsymbol{v}_{t-1} + \left(1 - \beta_2\right) \boldsymbol{g}_t^2\\ &\hat{\boldsymbol{m}}_t = \boldsymbol{m}_t\left/\left(1 - \beta_1^t\right)\right.\\ &\hat{\boldsymbol{v}}_t = \boldsymbol{v}_t\left/\left(1 - \beta_2^t\right)\right.\\ &\boldsymbol{\theta}_t = \boldsymbol{\theta}_{t-1} - \eta_t \hat{\boldsymbol{m}}_t\left/\left(\sqrt{\hat{\boldsymbol{v}}_t} + \epsilon\right)\right. \end{aligned}\end{equation}

In actual training, $\boldsymbol{g}_t$ is replaced by $\tilde{\boldsymbol{g}}_{B,t}$. We're considering the state where training has already settled into a "steady groove," i.e., $t\to\infty$, so we don't distinguish between $\boldsymbol{m}_t$ and $\hat{\boldsymbol{m}}_t$, or $\boldsymbol{v}_t$ and $\hat{\boldsymbol{v}}_t$. At the same time, since we're focusing on the effect of EMA, we set $\epsilon = 0$. Then for Adam we have $\tilde{\boldsymbol{\varphi}}_B=\boldsymbol{m}_t/\sqrt{\boldsymbol{v}_t}$, which differs from SignSGDM in that the denominator's $\boldsymbol{m}_t^2$ is replaced by another EMA statistic, $\boldsymbol{v}_t$.

By the mean-field approximation we get

\begin{equation}\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B] = \mathbb{E}\bigg[\frac{\boldsymbol{m}_t}{\sqrt{\boldsymbol{v}_t}}\bigg]\approx \frac{\mathbb{E}[\boldsymbol{m}_t]}{\sqrt{\mathbb{E}[\boldsymbol{v}_t]}}\end{equation}

We've already computed $\mathbb{E}[\boldsymbol{m}_t]$, so we only need to compute $\mathbb{E}[\boldsymbol{v}_t]$:

\begin{equation}\mathbb{E}[\boldsymbol{v}_t] = (1 - \beta_2)\sum_{s=1}^t \beta_2^{t-s}\mathbb{E}[\tilde{\boldsymbol{g}}_{B,s}^2] = (1 - \beta_2)\sum_{s=1}^t \beta_2^{t-s}(\boldsymbol{g}_s^2 + \boldsymbol{\sigma}_s^2/B)\approx \boldsymbol{g}_t^2 + \boldsymbol{\sigma}_t^2/B\end{equation}

As before, the final approximate equality assumes slowly changing gradients and variance, as well as $t\to\infty$. Thus we have

\begin{equation}\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B] \approx \frac{\boldsymbol{g}_t}{\sqrt{\boldsymbol{g}_t^2 + \boldsymbol{\sigma}_t^2/B}} \approx \frac{\sign(\boldsymbol{g}_t)}{\sqrt{1 + \mathcal{B}_{\text{simple}}/B}}\end{equation}

This result actually coincides with SignSGD, so purely from the perspective of the first moment, SignSGD is a reasonable approximation of Adam. But we still have the second moment $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B \tilde{\boldsymbol{\varphi}}_B^{\top}]$; under the assumption of component-wise independence, we only need to compute $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B^2]$:

\begin{equation}\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B^2] = \mathbb{E}\bigg[\frac{\boldsymbol{m}_t^2}{\boldsymbol{v}_t}\bigg]\approx \frac{\mathbb{E}[\boldsymbol{m}_t^2]}{\mathbb{E}[\boldsymbol{v}_t]} \approx \frac{\boldsymbol{g}_t^2 + \frac{1 - \beta_1}{1 + \beta_1}\boldsymbol{\sigma}_t^2/B}{\boldsymbol{g}_t^2 + \boldsymbol{\sigma}_t^2/B}\label{eq:u2-adam}\end{equation}

Two Special Cases

Let's look at two special cases. First, $\beta_1=0$, in which case the numerator and denominator are identical, and $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B^2]$ is the all-ones vector, consistent with SignSGD. So SignSGD is a good approximation of Adam with $\beta_1=0$ — that is, RMSProp — and as $\beta_1$ increases, the quality of the approximation deteriorates.

When $\beta_1=1$, we have

\begin{equation}\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B^2] \approx \frac{\boldsymbol{g}_t^2}{\boldsymbol{g}_t^2 + \boldsymbol{\sigma}_t^2/B}\approx \mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]^2\end{equation}

From this we obtain $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B \tilde{\boldsymbol{\varphi}}_B^{\top}] \approx \mathbb{E}[\tilde{\boldsymbol{\varphi}}_B] \mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]^{\top}$, and substituting into equation $\eqref{eq:eta-opt}$ gives

\begin{equation}\eta^* \approx \frac{\Vert \boldsymbol{g}\Vert_1 \sqrt{1 + \mathcal{B}_{\text{simple}}/B}}{\sign(\boldsymbol{g})^{\top} \boldsymbol{H} \sign(\boldsymbol{g})}\end{equation}

Note that this is a monotonically decreasing function of $B$, meaning that as the batch size increases, the learning rate should decrease. From this we can conjecture that increasing Adam's $\beta_1$ will hasten the onset of the "Surge phenomenon."

This conclusion may seem a bit puzzling at first, but it's easy to understand from another angle. The "Surge phenomenon" refers to the fact that once the batch size exceeds a certain threshold, the optimal learning rate decreases as the batch size increases further. The earlier results for SGDM and SignSGDM both show that introducing momentum is roughly equivalent to scaling the batch size up by a factor of $\frac{1 + \beta_1}{1 - \beta_1} > 1$, which naturally increases the likelihood of exceeding that threshold.

In other words, the conclusion that "as $\beta_1$ increases, the Surge phenomenon becomes more likely to occur" holds even for SignSGDM. Adam does have some new characteristics compared to SignSGDM, but the fact that "the momentum mechanism is roughly equivalent to enlarging the batch size" always holds, so it's not hard to see why the same conclusion emerges.

General Analysis

Let's rewrite equation $\eqref{eq:u2-adam}$:

\begin{equation}\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B^2] \approx \frac{\boldsymbol{g}_t^2 + \frac{1 - \beta_1}{1 + \beta_1}\boldsymbol{\sigma}_t^2/B}{\boldsymbol{g}_t^2 + \boldsymbol{\sigma}_t^2/B} = \frac{2\beta_1}{1+\beta_1}\frac{\boldsymbol{g}_t^2}{\boldsymbol{g}_t^2 + \boldsymbol{\sigma}_t^2/B} + \frac{1 - \beta_1}{1 + \beta_1} \approx \frac{2\beta_1}{1+\beta_1}\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]^2 + \frac{1 - \beta_1}{1 + \beta_1}\end{equation}

From this we can write

\begin{equation}\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B \tilde{\boldsymbol{\varphi}}_B^{\top}] \approx \mathbb{E}[\tilde{\boldsymbol{\varphi}}_B] \mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]^{\top} + \frac{1 - \beta_1}{1 + \beta_1}\diag\left(1 - \mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]^2\right)\end{equation}

and hence

\begin{equation}\eta^* \approx \frac{\sum_i |g_i|}{\frac{1}{\beta}\frac{1 - \beta_1}{1 + \beta_1}\sum_i H_{i,i} + \beta\left(\sum_{i,j} H_{i,j}\sign(g_i g_j) - \frac{1 - \beta_1}{1 + \beta_1}\sum_i H_{i,i}\right)}\end{equation}

Here, $\beta$ without a subscript equals $(1 + \mathcal{B}_{\text{simple}}/B)^{-1/2}$; without careful reading, it might get confused with $\beta_1,\beta_2$. The author apologizes for this, since this is the notation used in the previous two posts, and we've had to stick with it here. Unlike SignSGD, if we assume the Hessian matrix is diagonal, then SignSGD won't exhibit the Surge phenomenon at all — but the equation above still exhibits the Surge phenomenon even under the diagonal-Hessian assumption. In that case:

\begin{equation}\eta^* \approx \frac{\sum_i |g_i|}{\left(\frac{1}{\beta}\frac{1 - \beta_1}{1 + \beta_1} + \beta\frac{2\beta_1}{1 + \beta_1}\right)\sum_i H_{i,i}}\end{equation}

By the mean-value inequality, the expression above attains its maximum at $\beta^*=\sqrt{\frac{1-\beta_1}{2\beta_1}}$. But note that, by the definition of $\beta$, it is $\in(0,1)$, so we also need to check whether $\beta^*\in(0,1)$ holds, i.e., $\beta_1 > 1/3$. When this condition is not satisfied, the maximum is still attained at $\beta=1$, and there is no Surge phenomenon. Conversely, when $\beta_1 > 1/3$ and $\beta > \beta^*$ (i.e., $B > \frac{1-\beta_1}{3\beta_1-1}\mathcal{B}_{\text{simple}}$), the learning rate should decrease as the batch size increases.

This conclusion provides a preliminary explanation for why Muon can support larger batch sizes. From Rethinking Learning Rate and Batch Size (III): Muon, we know that Muon behaves similarly to SignSGDM; under a specific assumption about the Hessian structure, it does not exhibit the Surge phenomenon, which means that increasing the batch size always improves learning efficiency, even though the marginal benefit diminishes.

In contrast, under commonly used settings for Adam (such as $\beta_1=0.9$), the Surge phenomenon occurs even under the diagonal-Hessian assumption, which means that once the batch size exceeds a certain value, learning efficiency starts to decline.

Summary

This post gives a preliminary analysis of how the EMA mechanism in optimizers affects the scaling law between learning rate and batch size. We've confirmed that EMA — and momentum in particular — slightly modifies the scaling law, while Adam, an optimizer with a double EMA operation, exhibits some new characteristics distinct from those of SignSGD.