Rethinking Learning Rate and Batch Size (I): The Current State of Affairs

重新思考学习率与Batch Size（一）：现状 苏剑林 · 2025-09-01 · Mathematics gradient learning rate optimizer scaling law

In the earlier posts How Should the Learning Rate Change as Batch Size Increases? and How Does Adam's Epsilon Affect the Scaling Law of Learning Rate?, we discussed, from a theoretical standpoint, how the learning rate should scale with batch size. The most classic piece of that discussion is the second-order expansion analysis proposed by OpenAI. However, when we try to apply this analytical framework to non-SGD optimizers, the computations tend to become quite involved, leaving one with a feeling of not knowing where to even start.

In the next few posts, I want to revisit and rethink the relevant details from those earlier posts, try to simplify some of the derivation steps, and offer a more general and lightweight derivation path — while also exploring the possibility of extending it to the Muon optimizer.

Outline of the Method

Let's first recap the previous analytical approach. In How Should the Learning Rate Change as Batch Size Increases?, we introduced several ways of thinking about the relationship between learning rate and batch size, the bulk of which was devoted to the second-order approximation analysis proposed by OpenAI in An Empirical Model of Large-Batch Training. This post follows the same line of thought. more

We first need to introduce some notation. Let the loss function be $\mathcal{L}(\boldsymbol{w})$, let $\boldsymbol{w}\in\mathbb{R}^N$ be the parameter vector, and let $\boldsymbol{g}$ be its gradient. Note that the "ideal" loss function is defined as an expectation over the entire training set, but in practice we can only estimate it from a sampled batch, which means the gradient is itself a random variable. We write the gradient of a single sample as $\tilde{\boldsymbol{g}}$, whose mean is $\boldsymbol{g}$ and whose covariance matrix we denote $\boldsymbol{\Sigma}$. When the batch size is $B$, the gradient is written $\tilde{\boldsymbol{g}}_B$, whose mean is still $\boldsymbol{g}$, but whose covariance matrix becomes $\boldsymbol{\Sigma}/B$.

Further, let the current learning rate be $\eta$ and the update vector be $\tilde{\boldsymbol{\varphi}}_B$, so that the loss after the update becomes

\begin{equation}\begin{aligned} \mathcal{L}(\boldsymbol{w} - \eta\tilde{\boldsymbol{\varphi}}_B) \approx&\, \mathcal{L}(\boldsymbol{w}) - \eta \tilde{\boldsymbol{\varphi}}_B^{\top}\boldsymbol{g} + \frac{1}{2}\eta^2\tilde{\boldsymbol{\varphi}}_B^{\top}\boldsymbol{H}\tilde{\boldsymbol{\varphi}}_B \\ =&\, \mathcal{L}(\boldsymbol{w}) - \eta \tilde{\boldsymbol{\varphi}}_B^{\top}\boldsymbol{g} + \frac{1}{2}\eta^2\newcommand{tr}{\mathop{\text{tr}}}\tr(\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}\boldsymbol{H}) \end{aligned}\end{equation}

On the right-hand side we have expanded to second order via Taylor expansion, where $\boldsymbol{H}$ is the Hessian matrix and $\tr$ denotes the trace of a matrix; the second equality uses the identity $\tr(\boldsymbol{A}\boldsymbol{B})=\tr(\boldsymbol{B}\boldsymbol{A})$. To obtain a deterministic result, we take the expectation of both sides:

\begin{equation}\mathbb{E}[\mathcal{L}(\boldsymbol{w} - \eta\tilde{\boldsymbol{\varphi}}_B)] \approx \mathcal{L}(\boldsymbol{w}) - \eta\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]^{\top}\boldsymbol{g} + \frac{1}{2}\eta^2 \tr(\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]\boldsymbol{H})\end{equation}

We treat the right-hand side as a quadratic function of $\eta$, and assuming the quadratic coefficient is positive (a stronger assumption being that the matrix $\boldsymbol{H}$ is positive definite), we can obtain the location of the minimum:

\begin{equation}\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})}\end{equation}

This is the learning rate that, on average, minimizes the loss most quickly after the update — i.e., the theoretically optimal learning rate. What remains is to work out $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]$ and $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]$ for a specific $\tilde{\boldsymbol{\varphi}}_B$, and from the expression above extract its relationship with the batch size (i.e., $B$).

Warm-up Exercise

As a first example, let's naturally consider the simplest case, SGD, for which we have $\tilde{\boldsymbol{\varphi}}_B=\tilde{\boldsymbol{g}}_B$. It's then straightforward to derive $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]=\boldsymbol{g}$ and $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]=\boldsymbol{g}\boldsymbol{g}^{\top} + \boldsymbol{\Sigma}/B$, giving us

\begin{equation}\eta^* \approx \frac{\boldsymbol{g}^{\top}\boldsymbol{g}}{\tr((\boldsymbol{g}\boldsymbol{g}^{\top} + \boldsymbol{\Sigma}/B)\boldsymbol{H})} = \frac{\boldsymbol{g}^{\top}\boldsymbol{g}}{\boldsymbol{g}^{\top}\boldsymbol{H}\boldsymbol{g} + \tr(\boldsymbol{\Sigma}\boldsymbol{H})/B} = \frac{\eta_{\max}}{1 + \mathcal{B}_{\text{noise}}/B}\label{eq:eta-sgd}\end{equation}

where

\begin{equation}\eta_{\max} = \frac{\boldsymbol{g}^{\top}\boldsymbol{g}}{\boldsymbol{g}^{\top}\boldsymbol{H}\boldsymbol{g}},\qquad\mathcal{B}_{\text{noise}} = \frac{\tr(\boldsymbol{\Sigma}\boldsymbol{H})}{\boldsymbol{g}^{\top}\boldsymbol{H}\boldsymbol{g}}\end{equation}

There are several ways to interpret the result $\eqref{eq:eta-sgd}$. First, it is a monotonically increasing but bounded function, with the upper bound being $\eta_{\max}$, which tells us that the learning rate cannot grow without bound. This aligns much better with our intuition than a simple linear or square-root scaling law would. When $B \ll \mathcal{B}_{\text{noise}}$, we have

\begin{equation}\eta^* \approx \frac{\eta_{\max}}{1 + \mathcal{B}_{\text{noise}}/B} \approx \frac{\eta_{\max}}{\mathcal{B}_{\text{noise}}/B} = \eta_{\max} B / \mathcal{B}_{\text{noise}}\end{equation}

which shows that, for small batch sizes, the learning rate for SGD does indeed scale linearly with batch size, while also suggesting that $\mathcal{B}_{\text{noise}}$ is a key statistic to track. However, the definition of $\mathcal{B}_{\text{noise}}$ depends on the Hessian matrix $\boldsymbol{H}$, which is essentially impossible to compute exactly for LLMs. In practice, we therefore usually assume it to be (some multiple of) the identity matrix, giving the simplified form

\begin{equation}\mathcal{B}_{\text{simple}} = \frac{\tr(\boldsymbol{\Sigma})}{\boldsymbol{g}^{\top}\boldsymbol{g}}\end{equation}

This result takes the form of "noise strength" ($\tr(\boldsymbol{\Sigma})$) divided by "signal strength" ($\boldsymbol{g}^{\top}\boldsymbol{g}$) — in other words, it is precisely the reciprocal of the signal-to-noise ratio. It tells us that the smaller the signal-to-noise ratio, the larger the batch size needed to make use of the same $\eta_{\max}$, which again matches our intuition. Since $\tr(\boldsymbol{\Sigma})$ depends only on the diagonal elements of $\boldsymbol{\Sigma}$, this means we only need to estimate the mean and variance for each parameter independently, which is practically feasible.

Data Efficiency

Beyond the direct relationship between learning rate and batch size, I think the asymptotic relationship it implies between the amount of training data and the number of training steps is another beautiful part of this story that's well worth studying. In particular, this conclusion seems to be even more general than the learning-rate relation $\eqref{eq:eta-sgd}$ itself, because as we'll see later, SignSGD leads to a conclusion of the same form, even though its learning-rate law is not given by equation $\eqref{eq:eta-sgd}$.

The original paper's treatment of this part is fairly involved; the derivation below has been simplified by me. Specifically, substituting $\eta^*$ back into $\mathcal{L}(\boldsymbol{w} - \eta\tilde{\boldsymbol{g}}_B)$ gives

\begin{equation}\overline{\Delta\mathcal{L}} = \mathcal{L}(\boldsymbol{w}) - \mathbb{E}[\mathcal{L}(\boldsymbol{w} - \eta^*\tilde{\boldsymbol{g}}_B)] \approx \frac{\Delta\mathcal{L}_{\max}}{1 + \mathcal{B}_{\text{noise}}/B}\end{equation}

where $\Delta\mathcal{L}_{\max} = \frac{(\boldsymbol{g}^{\top}\boldsymbol{g})^2}{2\boldsymbol{g}^{\top}\boldsymbol{H}\boldsymbol{g}}$. How should we interpret this result? First, it is a monotonically increasing function of $B$, which equals $\Delta\mathcal{L}_{\max}$ when $B\to\infty$. In other words, if we could use an infinitely large batch size, the loss reduction per step would be $\Delta\mathcal{L}_{\max}$, and the number of training steps required would be minimized, which we denote $S_{\min}$.

If the batch size is finite, then the average loss reduction per step is only $\overline{\Delta\mathcal{L}}$, which means that on average it takes $1 + \mathcal{B}_{\text{noise}}/B$ steps to achieve the same amount of loss reduction that one step would achieve with an infinite batch size. So to reach the same loss, we need to train for $S = (1 + \mathcal{B}_{\text{noise}}/B)S_{\min}$ steps.

Since the batch size is $B$, it follows readily that the total amount of training data consumed is $E = BS = (B + \mathcal{B}_{\text{noise}})S_{\min}$. From this result we can see that as we increase the batch size, in order to reach the same effect we also need to appropriately increase the amount of data $E$; when $B\to 0$, the amount of data required is minimized, at $E_{\min} = \mathcal{B}_{\text{noise}}S_{\min}$. Using this notation, we can write

\begin{equation}\left(\frac{S}{S_{\min}} - 1\right)\left(\frac{E}{E_{\min}} - 1\right) = 1\end{equation}

This is the classic relationship between the amount of training data and the number of training steps. It has two parameters, $S_{\min},E_{\min}$, which we can also fit by running experiments over multiple values of $(S,E)$ to estimate $S_{\min},E_{\min}$, and from there estimate $\mathcal{B}_{\text{noise}} = E_{\min} / S_{\min}$. For more details of this analysis, please see the earlier post How Should the Learning Rate Change as Batch Size Increases? or OpenAI's original paper An Empirical Model of Large-Batch Training.

Where the Difficulty Lies

Everything discussed so far is still confined to SGD. From a computational standpoint, SGD is a trivial case; the real complexity arises when $\tilde{\boldsymbol{\varphi}}_B$ depends nonlinearly on $\tilde{\boldsymbol{g}}_B$ — for example, SignSGD, which corresponds to $\newcommand{sign}{\mathop{\text{sign}}}\tilde{\boldsymbol{\varphi}}_B=\sign(\tilde{\boldsymbol{g}}_B)$, and which is often used in theoretical analysis as an approximation to Adam. A more accurate approximation is SoftSignSGD, which takes $\epsilon$ into account, and which we attempted to analyze in How Does Adam's Epsilon Affect the Scaling Law of Learning Rate?.

In these nonlinear settings, computing $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]$ and $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]$ tends to be quite difficult, even if we assume $\tilde{\boldsymbol{g}}_B$ follows a simple normal distribution (note that in the SGD analysis, we didn't need to make any normality assumption at all about its distribution). For instance, in the earlier post, in order to compute $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]$ for SignSGD with $\tilde{\boldsymbol{\varphi}}_B=\sign(\tilde{\boldsymbol{g}}_B)$, we had to go through the following steps:

1. Assume the components of $\tilde{\boldsymbol{g}}_B$ are mutually independent, reducing the problem to the expectation of a single (non-bold) component $\tilde{\varphi}_B=\sign(\tilde{g}_B)$;
2. Assume $\tilde{g}_B$ (now a scalar) follows a normal distribution, which lets us compute $\mathbb{E}[\tilde{\varphi}_B]$, whose answer must be expressed in terms of the $\newcommand{erf}{\mathop{\text{erf}}}\erf$ function;
3. Approximate the $\erf$ function with a function of the form $x/\sqrt{x^2+c}$, in order to simplify the result.

In other words, we had to wind our way through a whole series of convoluted steps just to grind out an approximate result that could even be analyzed further (this procedure first appeared in Tencent's paper Surge Phenomenon in Optimal Learning Rate and Batch Size Scaling) — and that was already the easy case, because SoftSignSGD is even more complicated:

1. Assume the components of $\tilde{\boldsymbol{g}}_B$ are mutually independent, reducing the problem to the expectation of a single component $\tilde{\varphi}_B=\newcommand{softsign}{\mathop{\text{softsign}}}\softsign(\tilde{g}_B, \epsilon)$;
2. Approximate the $\softsign$ function with a piecewise linear function, which is needed in order to compute the following integral;
3. Assume $\tilde{g}_B$ follows a normal distribution and, combined with the approximation from step 2, compute $\mathbb{E}[\tilde{\varphi}_B]$, whose answer is a complicated function involving $\erf$;
4. Approximate the complicated function with a function of the form $x/\sqrt{x^2+c}$, in order to simplify the result.

And that's still not the end of it. After going through all that trouble and piling on all those assumptions, we've only barely managed to compute $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]$ — and we still have $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]$ left to compute, which is often even more complicated (SignSGD is an exception here, since $\sign(x)^2$ is always 1, which actually makes things simpler). But the computational complexity is really only a secondary concern; the bigger issue is that none of these steps seem to follow any pattern that could be generalized — it looks as though every case just has to be handled on its own terms, which is exhausting to contemplate.

To Be Continued

To keep this post from running too long, I'll stop here for now, having mainly given a brief recap of the existing analytical results and the computational difficulties involved. In the next post, I'll introduce some of the approaches I've tried in order to reduce the mental overhead of these derivations.