When Batch Size Increases, How Should the Learning Rate Change?
With the rapid progress of compute, more and more scenarios hope to "trade compute for time", i.e., shortening model training time by piling on more compute. Ideally, we'd like to invest $n$ times as much compute so that the time to reach the same result shrinks to $1/n$, keeping the total compute cost the same. This "hope" seems reasonable and natural, but in practice it's far from trivial. Even setting aside bottlenecks like communication, once compute exceeds a certain scale, or the model falls below a certain size, adding more compute often can only increase the batch size. But does increasing batch size always shorten training time while keeping the results unchanged?
This is exactly the topic we'll discuss next: as batch size increases, how should the various hyperparameters — especially the learning rate — be adjusted, in order to preserve the original training performance while maximizing training efficiency? We can also call this the scaling law between batch size and learning rate.
The Variance Perspective
Intuitively, as batch size increases, the gradient of each batch becomes more accurate, so we can afford to take bigger steps — that is, increase the learning rate — in order to reach the destination faster and shorten training time. This much is fairly easy to guess. The question is: how much of an increase is actually appropriate?
Square Root Scaling
The earliest answer to this question is probably square-root scaling: if the batch size is scaled up by a factor of $n$, the learning rate should be scaled up by a factor of $\sqrt{n}$. This comes from the 2014 paper One weird trick for parallelizing convolutional neural networks, and the derivation principle is to keep the variance of the SGD increment unchanged. Specifically, denote the gradient of a randomly sampled example as $\tilde{\boldsymbol{g}}$, with mean and covariance $\boldsymbol{g}$ and $\boldsymbol{\Sigma}$ respectively; here $\boldsymbol{g}$ is the gradient over the whole dataset. When we increase the number of samples to $B$, we have
\begin{equation}\tilde{\boldsymbol{g}}_B \triangleq \frac{1}{B}\sum_{i=1}^B \tilde{\boldsymbol{g}}^{(i)},\quad \mathbb{E}[\tilde{\boldsymbol{g}}_B] = \boldsymbol{g},\quad \mathbb{E}[(\tilde{\boldsymbol{g}}_B-\boldsymbol{g})(\tilde{\boldsymbol{g}}_B-\boldsymbol{g})^{\top}]=\frac{\boldsymbol{\Sigma}}{B}\end{equation}
That is, increasing the sample count doesn't change the mean, but shrinks the covariance to $1/B$. For the SGD optimizer, the update increment is $-\eta \tilde{\boldsymbol{g}}_B$, whose covariance is proportional to $\eta^2/B$. We believe that a moderate (neither too much nor too little) amount of noise is necessary during optimization, so as the batch size $B$ changes, we adjust the learning rate $\eta$ to keep the noise intensity of the increment — i.e. its covariance matrix — unchanged, which gives
\begin{equation}\frac{\eta^2}{B} = \text{const}\quad\Rightarrow\quad \eta\propto \sqrt{B}\end{equation}
This is the square-root scaling law between learning rate and batch size. The later paper Train longer, generalize better: closing the generalization gap in large batch training of neural networks also endorses this choice.
Linear Scaling
Interestingly, linear scaling, i.e. $\eta\propto B$, tends to perform better in practice — even the authors of One weird trick for parallelizing convolutional neural networks, who first proposed square-root scaling, pointed this out in their paper, admitting they couldn't offer a convincing explanation.
In a sense, linear scaling matches our intuition better, especially under the assumption made in Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour, that the gradient direction doesn't change much across $n$ consecutive batches — under that assumption linear scaling is almost self-evident. However, this assumption is clearly too strong. Relaxing it requires connecting SGD to SDEs (stochastic differential equations), which was carried out in Stochastic Modified Equations and Dynamics of Stochastic Gradient Algorithms I: Mathematical Foundations, though the paper that first used this connection to point out the scaling relationship between learning rate and batch size should be On the Generalization Benefit of Noise in Stochastic Gradient Descent.
In hindsight, this connection isn't actually hard to understand. Let the model parameters be $\boldsymbol{w}$; then the SGD update rule can be rewritten as
\begin{equation}\boldsymbol{w}_{t+1} =\boldsymbol{w}_t - \eta \tilde{\boldsymbol{g}}_{B,t} =\boldsymbol{w}_t - \eta \boldsymbol{g}_t - \eta (\tilde{\boldsymbol{g}}_{B,t} - \boldsymbol{g}_t)\end{equation}
where $\tilde{\boldsymbol{g}}_{B,t} - \boldsymbol{g}_t$ is the gradient noise. So far we haven't assumed anything about the distribution of this noise, only that its mean is $\boldsymbol{0}$ and its covariance is $\boldsymbol{\Sigma}_t/B$. Next, we assume this noise follows a normal distribution $\mathcal{N}(\boldsymbol{0},\boldsymbol{\Sigma}_t/B)$, and the above iteration can be further rewritten as
\begin{equation}\begin{aligned} \boldsymbol{w}_{t+1} =&\, \boldsymbol{w}_t - \eta \boldsymbol{g}_t - \eta (\tilde{\boldsymbol{g}}_{B,t} - \boldsymbol{g}_t) \\[5pt] =&\, \boldsymbol{w}_t - \eta \boldsymbol{g}_t - \eta \sqrt{\frac{\boldsymbol{\Sigma}_t}{B}}\boldsymbol{z},\quad \boldsymbol{z}\sim \mathcal{N}(\boldsymbol{0},\boldsymbol{I}) \\[5pt] =&\, \boldsymbol{w}_t - \eta \boldsymbol{g}_t - \sqrt{\eta} \sqrt{\frac{\eta\boldsymbol{\Sigma}_t}{B}}\boldsymbol{z},\quad \boldsymbol{z}\sim \mathcal{N}(\boldsymbol{0},\boldsymbol{I}) \end{aligned}\end{equation}
This means that the SGD update scheme $\boldsymbol{w}_{t+1} =\boldsymbol{w}_t - \eta \tilde{\boldsymbol{g}}_{B,t}$ is actually approximately solving the SDE:
\begin{equation}d\boldsymbol{w} = - \boldsymbol{g}_t dt - \sqrt{\frac{\eta\boldsymbol{\Sigma}_t}{B}}d\boldsymbol{z},\quad d\boldsymbol{z}\sim \mathcal{N}(\boldsymbol{0},dt\boldsymbol{I}) \end{equation}
Therefore, for the results to remain essentially unchanged when $B$ varies, the form of the above SDE should stay the same, which gives us linear scaling $\eta\propto B$. The key step in this derivation is that the step size of the noise term in an SDE is the square root of that of the non-noise term, which lets us factor out a term $\sqrt{\eta}$. We've discussed this point before in Rambling on Diffusion Models (V): The SDE Framework — in short, zero-mean Gaussian noise tends to cancel itself out over time, so the step size must be enlarged in order for the noise effect to actually show up.
All the conclusions above are derived for the SGD optimizer. The paper On the SDEs and Scaling Rules for Adaptive Gradient Algorithms extends this to optimizers like RMSProp and Adam, arriving at square-root scaling. Coincidentally, the slightly earlier Large Batch Optimization for Deep Learning: Training BERT in 76 minutes also applied square-root scaling when testing Adam and its variant LAMB. For more, see also the blog post How to Scale Hyperparameters as Batch Size Increases.
Facing the Loss Function Head-On
One thing is certain: whether it's square-root scaling or linear scaling, both can only be locally valid approximations, since both imply the conclusion that "as long as the batch size is large enough, the learning rate can be made arbitrarily large" — which is clearly impossible. Moreover, the work discussed in the previous two sections all revolves around variance, but our fundamental task is to reduce the loss function, so it may be more essential to take the loss function itself as the guiding principle.
Monotonic and Bounded
The classic work from this perspective is OpenAI's An Empirical Model of Large-Batch Training, which analyzes the optimal learning rate for SGD via a second-order approximation of the loss function, arriving at the conclusion that "the learning rate increases monotonically with batch size, but is bounded above." A similar line of thinking also appears in the slightly earlier Dissecting Adam: The Sign, Magnitude and Variance of Stochastic Gradients, though that paper wasn't concerned with the role of batch size.
The key idea throughout the derivation is to treat the learning rate itself as an optimization variable: let the loss function be $\mathcal{L}(\boldsymbol{w})$, and the gradient of the current batch be $\tilde{\boldsymbol{g}}_B$; then the loss after an SGD update is $\mathcal{L}(\boldsymbol{w} - \eta\tilde{\boldsymbol{g}}_B)$. We treat finding the optimal learning rate as the optimization problem
\begin{equation}\eta^* = \mathop{\text{argmin}}_{\eta} \mathbb{E}[\mathcal{L}(\boldsymbol{w} - \eta\tilde{\boldsymbol{g}}_B)]\end{equation}
This objective is intuitive enough: choose the learning rate so that, on average, training efficiency is highest (the loss decreases fastest). To solve this, we expand the loss function approximately to second order:
\begin{equation}\mathcal{L}(\boldsymbol{w} - \eta\tilde{\boldsymbol{g}}_B) \approx \mathcal{L}(\boldsymbol{w}) - \eta\tilde{\boldsymbol{g}}_B^{\top}\underbrace{\frac{\partial \mathcal{L}(\boldsymbol{w})}{\partial\boldsymbol{w}}}_{\text{is}\boldsymbol{g}} + \frac{1}{2}\eta^2 \tilde{\boldsymbol{g}}_B^{\top}\underbrace{\frac{\partial^2 \mathcal{L}(\boldsymbol{w})}{\partial\boldsymbol{w}^2}}_{\text{denote}\boldsymbol{H}}\tilde{\boldsymbol{g}}_B = \mathcal{L}(\boldsymbol{w}) - \eta\tilde{\boldsymbol{g}}_B^{\top}\boldsymbol{g} + \frac{1}{2}\eta^2 \tilde{\boldsymbol{g}}_B^{\top}\boldsymbol{H}\tilde{\boldsymbol{g}}_B\end{equation}
Here $\boldsymbol{H}$ is the Hessian matrix, and $\frac{\partial \mathcal{L}(\boldsymbol{w})}{\partial\boldsymbol{w}}$ is the gradient of the loss function. The ideal objective is computed over the full dataset, which is why its gradient is exactly the mean $\boldsymbol{g}$ of $\tilde{\boldsymbol{g}}_B$. Taking the expectation next, we get
\begin{equation}\mathbb{E}[\mathcal{L}(\boldsymbol{w} - \eta\tilde{\boldsymbol{g}}_B)] \approx \mathbb{E}[\mathcal{L}(\boldsymbol{w}) - \eta\tilde{\boldsymbol{g}}_B^{\top}\boldsymbol{g} + \frac{1}{2}\eta^2 \tilde{\boldsymbol{g}}_B^{\top}\boldsymbol{H}\tilde{\boldsymbol{g}}_B] = \mathcal{L}(\boldsymbol{w}) - \eta\boldsymbol{g}^{\top}\boldsymbol{g} + \frac{1}{2}\eta^2 \mathbb{E}[\tilde{\boldsymbol{g}}_B^{\top}\boldsymbol{H}\tilde{\boldsymbol{g}}_B]\end{equation}
The last term requires a small trick:
\begin{equation}\newcommand{tr}{\mathop{\text{tr}}}\begin{aligned} \mathbb{E}[\tilde{\boldsymbol{g}}_B^{\top}\boldsymbol{H}\tilde{\boldsymbol{g}}_B] =&\, \mathbb{E}[\tr(\tilde{\boldsymbol{g}}_B^{\top}\boldsymbol{H}\tilde{\boldsymbol{g}}_B)]= \mathbb{E}[\tr(\tilde{\boldsymbol{g}}_B\tilde{\boldsymbol{g}}_B^{\top}\boldsymbol{H})] = \tr(\mathbb{E}[\tilde{\boldsymbol{g}}_B\tilde{\boldsymbol{g}}_B^{\top}]\boldsymbol{H})\\[5pt] =&\, \tr((\boldsymbol{g}\boldsymbol{g}^{\top} + \boldsymbol{\Sigma}/B)\boldsymbol{H}) = \boldsymbol{g}^{\top}\boldsymbol{H}\boldsymbol{g} + \tr(\boldsymbol{\Sigma}\boldsymbol{H})/B \end{aligned}\end{equation}
The main tool used in this transformation is $\tr(\boldsymbol{A}\boldsymbol{B}) = \tr(\boldsymbol{B}\boldsymbol{A})$. Now, assuming $\boldsymbol{H}$ is positive definite, the problem reduces to minimizing a quadratic function, which is easy to solve, giving
\begin{equation}\eta^* \approx \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-opt}\end{equation}
This yields the result that the learning rate "increases monotonically with $B$ but is bounded above," 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}
Practical Analysis
When $B \ll \mathcal{B}_{\text{noise}}$, $1 + \mathcal{B}_{\text{noise}}/B\approx \mathcal{B}_{\text{noise}}/B$, so $\eta^* \approx \eta_{\max}B/\mathcal{B}_{\text{noise}}\propto B$ — that is, linear scaling — which again shows that linear scaling is merely a local approximation valid for small batch sizes. When $B > \mathcal{B}_{\text{noise}}$, $\eta^*$ gradually saturates toward the value $\eta_{\max}$, meaning the increase in training cost far outpaces the gain in training efficiency. Thus $\mathcal{B}_{\text{noise}}$ acts as a kind of watershed: once the batch size exceeds this value, there's no longer any point in pouring in more compute to increase it further.
For practical purposes, the crucial question is undoubtedly how to estimate $\eta_{\max}$ and $\mathcal{B}_{\text{noise}}$. In particular, $\mathcal{B}_{\text{noise}}$ directly determines both the learning-rate scaling law and the saturation of training efficiency. Computing them directly involves the Hessian matrix $\boldsymbol{H}$, whose cost is proportional to the square of the parameter count — and given that today's "small" models already have hundreds of millions of parameters, computing the Hessian is clearly infeasible. So we need a more efficient way to compute these quantities.
Let's first look at $\mathcal{B}_{\text{noise}}$, whose formula is $\frac{\tr(\boldsymbol{\Sigma}\boldsymbol{H})}{\boldsymbol{g}^{\top}\boldsymbol{H}\boldsymbol{g}}$. Both numerator and denominator contain a $\boldsymbol{H}$ term, which naturally tempts us to "cancel it out." The simplification idea is exactly that: assuming $\boldsymbol{H}$ is approximately some multiple of the identity matrix, we get
\begin{equation}\mathcal{B}_{\text{noise}} = \frac{\tr(\boldsymbol{\Sigma}\boldsymbol{H})}{\boldsymbol{g}^{\top}\boldsymbol{H}\boldsymbol{g}}\approx \frac{\tr(\boldsymbol{\Sigma})}{\boldsymbol{g}^{\top}\boldsymbol{g}}\triangleq \mathcal{B}_{\text{simple}}\end{equation}
$\mathcal{B}_{\text{simple}}$ is far more tractable computationally, and it turns out empirically to be a good approximation of $\mathcal{B}_{\text{noise}}$, so we choose to estimate $\mathcal{B}_{\text{simple}}$ instead of $\mathcal{B}_{\text{noise}}$. Note that $\tr(\boldsymbol{\Sigma})$ only requires the diagonal elements, so there's no need to compute the full covariance matrix — one only needs to compute the variance of each gradient component individually and sum them. In a data-parallel setting, this can be estimated directly from the gradients computed on each device.
It should be noted that quantities such as $\eqref{eq:eta-opt}$ are in fact dynamic — in theory, $\eta_{\max}$, $\mathcal{B}_{\text{noise}}$, and $\mathcal{B}_{\text{simple}}$ are all different at every training step. So if we want a static rule of thumb, we need to keep training for a while, until the model has entered a "steady state," before the computed $\mathcal{B}_{\text{simple}}$ becomes reliable; alternatively, one can continuously monitor $\mathcal{B}_{\text{simple}}$ during training in order to gauge how far the current setting is from optimal.
As for $\eta_{\max}$, there's actually no need to estimate it from a formula — just perform a grid search over the learning rate at some small batch size to find an approximate $\eta^*$, and combined with the estimated $\mathcal{B}_{\text{simple}}$, we can back out $\eta_{\max}$.
Data Efficiency
Starting from the above results, we can also derive an asymptotic relationship concerning the amount of training data and the number of training steps. The derivation is simple as well: substituting $\eqref{eq:eta-opt}$ into the loss function, we can compute that, under the optimal learning rate, the loss decrease at each iteration is:
\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}\label{eq:Delta-L-sgd}\end{equation}
where $\Delta\mathcal{L}_{\max} = \frac{(\boldsymbol{g}^{\top}\boldsymbol{g})^2}{2\boldsymbol{g}^{\top}\boldsymbol{H}\boldsymbol{g}}$. The key point now is to interpret this result.
When $B\to\infty$, i.e. full-batch SGD, the loss decrease per step reaches its maximum value $\Delta\mathcal{L}_{\max}$, and in that case the fewest training steps (denoted $S_{\min}$) are needed to reach the target. When $B$ is finite, the average loss decrease per step is only $\overline{\Delta\mathcal{L}}$, meaning we need $1 + \mathcal{B}_{\text{noise}}/B$ steps to match the decrease achieved by a single step of full-batch SGD, so the total number of training steps is roughly $S = (1 + \mathcal{B}_{\text{noise}}/B)S_{\min}$.
Since the batch size is $B$, the total number of samples consumed during training is $E = BS = (B + \mathcal{B}_{\text{noise}})S_{\min}$, which is an increasing function of $B$, and when $B\to 0$, $E_{\min} = \mathcal{B}_{\text{noise}}S_{\min}$ — which shows that as long as we train with a sufficiently small batch size, the total number of training samples required, $E$, will decrease accordingly, at the cost of a much larger number of training steps $S$. Using these notations, we can further write down the relationship between them:
\begin{equation}\left(\frac{S}{S_{\min}} - 1\right)\left(\frac{E}{E_{\min}} - 1\right) = 1\label{eq:E-S}\end{equation}
This is the scaling law between the amount of training data and the number of training steps, showing that the smaller the amount of data, the smaller the batch size should be, and the more training steps are needed, in order to have a better chance of reaching a superior solution. The derivation here has been simplified by the author, assuming that $\mathcal{B}_{\text{noise}}$ and $\Delta\mathcal{L}_{\max}$ remain unchanged throughout training. If necessary, one could also, following the original paper's appendix, handle the dynamically-changing case more precisely using an integral (though this requires introducing the assumption $B = \sqrt{r\mathcal{B}_{\text{noise}}}$) — we won't go into that here.
Furthermore, since $\mathcal{B}_{\text{noise}} = E_{\min}/S_{\min}$, the equation above also provides another way to estimate $\mathcal{B}_{\text{noise}}$: run multiple experiments combined with grid search to obtain several $(S,E)$ pairs, then fit the equation above to estimate $E_{\min},S_{\min}$, and from there compute $\mathcal{B}_{\text{noise}}$.
The Adaptive Version
It has to be said, OpenAI truly is one of the pioneers of various scaling laws — the analysis above is quite brilliant, and the results are rich. What's even more remarkable is that the entire derivation isn't complicated at all, giving one the sense that great truths are indeed simple. However, all the conclusions so far were derived for SGD, and it's unclear to what extent they apply to adaptive learning-rate optimizers like Adam. This gap is filled by Surge Phenomenon in Optimal Learning Rate and Batch Size Scaling.
Sign Approximation
The approach to analyzing Adam is the same as for SGD — both rely on a second-order expansion — except that the direction vector changes from $\tilde{\boldsymbol{g}}_B$ to a general vector $\tilde{\boldsymbol{\varphi}}_B$, giving us
\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 now need to determine $\tilde{\boldsymbol{\varphi}}_B$ and compute the corresponding $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]$ and $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]$. Since we only need an asymptotic relationship, following the same approach as Assigning Different Learning Rates — Can LoRA Improve Further?, we choose SignSGD, i.e. $\newcommand{sign}{\mathop{\text{sign}}}\tilde{\boldsymbol{\varphi}}_B = \sign(\tilde{\boldsymbol{g}}_B)$, as an approximation of Adam. This idea probably first appeared in Dissecting Adam: The Sign, Magnitude and Variance of Stochastic Gradients. The reasonableness of this approximation rests on two points:
1. Regardless of the value of $\beta_1,\beta_2$, Adam's first update vector is always $\sign(\tilde{\boldsymbol{g}}_B)$;
2. When $\beta_1=\beta_2=0$, Adam's update vector is always $\sign(\tilde{\boldsymbol{g}}_B)$.
To compute $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]$ and $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]$, we again need to assume, as we did in the "Linear Scaling" section, that $\tilde{\boldsymbol{g}}_B$ follows the distribution $\mathcal{N}(\boldsymbol{g},\boldsymbol{\Sigma}/B)$. And to simplify the calculation further, we additionally assume that $\boldsymbol{\Sigma}$ is a diagonal matrix $\text{diag}(\sigma_1^2,\sigma_2^2,\sigma_3^2,\cdots)$, i.e., that the components are mutually independent, which lets us handle each component separately. By the reparameterization trick, we know that $\tilde{g}_B\sim \mathcal{N}(g, \sigma^2/B)$ is equivalent to $\tilde{g}_B=g + \sigma z/\sqrt{B},z\sim\mathcal{N}(0,1)$, so
\begin{equation}\begin{aligned} \mathbb{E}[\tilde{\varphi}_B] =&\, \mathbb{E}[\sign(g + \sigma z/\sqrt{B})] = \mathbb{E}[\sign(g\sqrt{B}/\sigma + z)] \\[5pt] =&\,\frac{1}{\sqrt{2\pi}}\int_{-\infty}^{\infty} \sign(g\sqrt{B}/\sigma + z) e^{-z^2/2}dz \\[5pt] =&\,\frac{1}{\sqrt{2\pi}}\int_{-\infty}^{-g\sqrt{B}/\sigma} (-1)\times e^{-z^2/2}dz + \frac{1}{\sqrt{2\pi}}\int_{-g\sqrt{B}/\sigma}^{\infty} 1\times e^{-z^2/2}dz \\[5pt] =&\,\text{erf}\left(\frac{g}{\sigma}\sqrt{\frac{B}{2}}\right) \end{aligned}\end{equation}
Here $\text{erf}$ is the error function, an S-shaped function with range $(-1,1)$ similar to $\tanh$, which can serve as a smooth approximation of $\sign$. But $\text{erf}$ itself has no elementary closed form, so it's better to find an elementary approximation in order to more directly observe the trends. We previously discussed this topic in Where Do the Two Elementary-Function Approximations of GELU Come From?, but the approximations there were still too complex (both involving exponentials). Here we'll use a simpler one:
\begin{equation}\text{erf}(x)\approx \sign(x) = \frac{x}{|x|} = \frac{x}{\sqrt{x^2}}\approx \frac{x}{\sqrt{x^2+c}}\end{equation}
We choose $c=\pi/4$ such that the first-order approximation of this expression at $x=0$ matches that of $\text{erf}$. Of course, after so many layers of approximation, the exact value of $c$ is no longer that important — we just need to know that some such $c > 0$ exists. Based on this approximation, we obtain
\begin{equation}\mathbb{E}[\tilde{\varphi}_B] \approx \frac{g/\sigma}{\sqrt{\pi/2B+(g/\sigma)^2}}\quad\Rightarrow\quad\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]_i \approx \frac{g_i/\sigma_i}{\sqrt{\pi/2B+(g_i/\sigma_i)^2}}\triangleq \mu_i\end{equation}
Notably, one clear difference between Adam and SGD is that already at the step $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]$ things depend on $B$. Fortunately though, the second moment is simpler here, since the square of $\sign(x)$ is necessarily 1, so
\begin{equation}\mathbb{E}[\tilde{\varphi}_B^2] = 1\quad\Rightarrow\quad\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]_{i,j} \to\left\{\begin{aligned}&=1, & i = j \\ &\approx\mu_i \mu_j,&\,i\neq j\end{aligned}\right.\end{equation}
Using these results, we can solve for
\begin{gather}\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})} \approx \frac{\sum_i \mu_i g_i}{\sum_i H_{i,i} + \sum_{i\neq j} \mu_i \mu_j H_{i,j}}\label{eq:eta-opt-sign} \\[5pt] \overline{\Delta\mathcal{L}} = \mathcal{L}(\boldsymbol{w}) - \mathbb{E}[\mathcal{L}(\boldsymbol{w} - \eta^*\tilde{\boldsymbol{\varphi}}_B)] \approx \frac{1}{2}\frac{(\sum_i \mu_i g_i)^2}{\sum_i H_{i,i} + \sum_{i\neq j} \mu_i \mu_j H_{i,j}}\label{eq:Delta-L-sign}\end{gather}
Two Special Cases
Compared to SGD's equation $\eqref{eq:eta-opt}$, Adam's equation $\eqref{eq:eta-opt-sign}$ is more complex, to the point that its dependence on $B$ isn't obvious at a glance. So let's start from a few special cases.
First consider $B\to\infty$, in which case $\mu_i = \sign(g_i)$, so
\begin{equation}\eta^* \approx \frac{\sum_i |g_i|}{\sum_i H_{i,i} + \sum_{i\neq j} \sign(g_i g_j) H_{i,j}}\end{equation}
The difference from SGD's $\eta_{\max}$ is that this is not homogeneous in the gradient, but rather proportional to the scale of the gradient.
Next let's consider the case where $\boldsymbol{H}$ is a diagonal matrix, i.e. when $i\neq j$, $H_{i,j}=0$; then
\begin{equation}\eta^* \approx \frac{\sum_i \mu_i g_i}{\sum_i H_{i,i}}=\frac{1}{\sum_i H_{i,i}}\sum_i \frac{g_i^2/\sigma_i}{\sqrt{\pi/2B+(g_i/\sigma_i)^2}}\end{equation}
Each term in this sum is monotonically increasing and bounded above with respect to $B$, so the overall result behaves the same way. To capture the most essential behavior, we can further simplify $\mu_i$ (this is where we start to diverge from the original paper):
\begin{equation}\mu_i = \frac{g_i/\sigma_i}{\sqrt{\pi/2B+(g_i/\sigma_i)^2}} = \frac{\sign(g_i)}{\sqrt{1 + \pi(\sigma_i/g_i)^2/2B}} \approx \frac{\sign(g_i)}{\sqrt{1 + \pi\kappa^2/2B}}\label{eq:mu-approx}\end{equation}
The assumption here is that there exists some constant $\kappa^2$ independent of $i$ (for example, one could take some kind of average over all $(\sigma_i/g_i)^2$ — note that $\kappa^2$ here plays a role similar to the earlier $\mathcal{B}_{\text{simple}}$, and could also be estimated following the definition of $\mathcal{B}_{\text{simple}}$), such that replacing $(\sigma_i/g_i)^2$ with $\kappa^2$ is a good approximation for any $i$, giving
\begin{equation}\eta^* \approx \frac{\sum_i \mu_i g_i}{\sum_i H_{i,i}}\approx \frac{\sum_i |g_i|}{\sum_i H_{i,i}}\frac{1}{\sqrt{1 + \pi\kappa^2/2B}}\label{eq:eta-opt-sign-diag}\end{equation}
When $\pi\kappa^2\gg 2B$, i.e. $B \ll \pi\kappa^2/2$, we can further write the approximation
\begin{equation}\eta^* \approx \frac{\sum_i |g_i|}{\kappa\sum_i H_{i,i}}\sqrt{\frac{2B}{\pi}} \propto \sqrt{B}\end{equation}
This shows that when the batch size itself is small, Adam does indeed follow the square-root scaling law.
Emergent Behavior
If we apply the approximation $\eqref{eq:mu-approx}$ to the original equation $\eqref{eq:eta-opt-sign}$, we find it exhibits some entirely new properties. Specifically, we have
\begin{equation}\eta^* \approx \frac{\sum_i \mu_i g_i}{\sum_i H_{i,i} + \sum_{i\neq j} \mu_i \mu_j H_{i,j}} \approx \frac{\eta_{\max}}{\frac{1}{2}\left(\frac{\beta_{\text{noise}}}{\beta} + \frac{\beta}{\beta_{\text{noise}}}\right)}\label{eq:eta-opt-beta}\end{equation}
where $\beta = (1 + \pi\kappa^2/2B)^{-1/2}$, and
\begin{equation}\beta_{\text{noise}} = \sqrt{\frac{\sum_i H_{i,i}}{\sum_{i\neq j}\sign(g_i g_j) H_{i,j}}},\quad \eta_{\max} = \frac{\sum_i |g_i|}{2\sqrt{\left(\sum_i H_{i,i}\right)\left(\sum_{i\neq j} \sign(g_i g_j) H_{i,j}\right)}}\end{equation}
Note that $\beta$ is a monotonically increasing function of $B$, but the final approximation in equation $\eqref{eq:eta-opt-beta}$ is not a monotonically increasing function of $\beta$ — it first increases and then decreases, reaching its maximum at $\beta=\beta_{\text{noise}}$. This means there exists a corresponding $\mathcal{B}_{\text{noise}}$ such that once the batch size exceeds this $\mathcal{B}_{\text{noise}}$, the optimal learning rate should no longer increase, but should actually decrease! This is precisely the "Surge phenomenon" referred to in the original paper's title. (There's of course one caveat here: $\beta$ is always less than $1$; if $\beta_{\text{noise}} \geq 1$, then the relationship between the optimal learning rate and batch size remains monotonically increasing.)
Regarding Adam's $\eta^*$, OpenAI, in an appendix to their paper, once "conjectured" without proof that Adam's optimal learning rate should be
\begin{equation}\eta^* \approx \frac{\eta_{\max}}{(1 + \mathcal{B}_{\text{noise}}/B)^{\alpha}}\label{eq:openai-adam}\end{equation}
where $0.5 < \alpha < 1$. In hindsight, this form is just the approximate result when the diagonal elements of the Hessian dominate; when the off-diagonal elements can't be neglected, the Surge phenomenon may emerge — where "the learning rate should actually decrease once the batch size is large enough."
How should we intuitively understand the Surge phenomenon? In the author's view, this is essentially a manifestation of the suboptimality of adaptive learning-rate strategies. Take the approximation $\tilde{\boldsymbol{\varphi}}_B = \sign(\tilde{\boldsymbol{g}}_B)$ as an example: the larger $B$ is, the more accurate $\tilde{\boldsymbol{g}}_B$ becomes, and $B\to \infty$ is then $\sign(\boldsymbol{g})$ — but is $\sign(\boldsymbol{g})$ really the most scientifically sound update direction? Not necessarily, especially in the later stages of training, where this adaptive strategy might actually have negative effects. So when $B$ takes on an appropriate value, the noise in $\sign(\tilde{\boldsymbol{g}}_B)$ might actually help correct this suboptimality, whereas as $B$ continues to grow, the noise decreases, reducing the chance for such correction — thus requiring us to be more cautious and lower the learning rate.
Efficiency Relationship
As with the SGD analysis, we can finally consider $\overline{\Delta\mathcal{L}}$: substituting equation $\eqref{eq:eta-opt-beta}$ into equation $\eqref{eq:Delta-L-sign}$, restoring the notation $B$, and simplifying (this simplification requires no further approximations), we get
\begin{equation}\overline{\Delta\mathcal{L}} \approx \frac{\Delta\mathcal{L}_{\max}}{1 + \mathcal{B}_{\text{noise-2}}/B}\label{eq:Delta-L-sign-2}\end{equation}
where
\begin{equation}\Delta\mathcal{L}_{\max} = \frac{\beta_{\text{noise}}\eta_{\max}\sum_i|g_i|}{1 + \beta_{\text{noise}}^2},\quad \mathcal{B}_{\text{noise-2}} = \frac{\pi\kappa^2\beta_{\text{noise}}^2}{2(1 + \beta_{\text{noise}}^2)}\label{eq:beta-B-noise}\end{equation}
Note that $\mathcal{B}_{\text{noise-2}}$ here is a new symbol — it is not the same as $\mathcal{B}_{\text{noise}}$, the latter being the theoretical optimal batch size obtained by inverting $\beta=\beta_{\text{noise}}$, given by
\begin{equation}\mathcal{B}_{\text{noise}} = \frac{\pi\kappa^2\beta_{\text{noise}}^2}{2(1 - \beta_{\text{noise}}^2)}\end{equation}
The relationship between them is
\begin{equation}\frac{1}{\mathcal{B}_{\text{noise-2}}} - \frac{1}{\mathcal{B}_{\text{noise}}} = \frac{4}{\pi\kappa^2}\quad\Rightarrow\quad \mathcal{B}_{\text{noise}} = \left(\frac{1}{\mathcal{B}_{\text{noise-2}}} - \frac{4}{\pi\kappa^2}\right)^{-1}\label{eq:B-1-2}\end{equation}
Since equation $\eqref{eq:Delta-L-sign-2}$ has the same form as SGD's equation $\eqref{eq:Delta-L-sgd}$, the analysis from that section applies equally here, so we can likewise derive equation $\eqref{eq:E-S}$:
\begin{equation}\left(\frac{S}{S_{\min}} - 1\right)\left(\frac{E}{E_{\min}} - 1\right) = 1\end{equation}
except now $E_{\min}/S_{\min} = \mathcal{B}_{\text{noise-2}}$. This gives us a way to estimate $\beta_{\text{noise}}$ and $\mathcal{B}_{\text{noise}}$: run multiple experiments to obtain several $(S,E)$ pairs (and along the way we can also estimate $\kappa^2$), then fit the equation above to get $E_{\min},S_{\min}$, from which we can estimate $\mathcal{B}_{\text{noise-2}}$, and finally solve for $\beta_{\text{noise}}$ from equation $\eqref{eq:beta-B-noise}$.
If $\beta_{\text{noise}} \geq 1$, then there is no optimal $\mathcal{B}_{\text{noise}}$; if $\beta_{\text{noise}} \gg 1$, it indicates that the diagonal elements of the Hessian dominate, in which case the scaling law $\eqref{eq:eta-opt-sign-diag}$ applies — increasing the batch size always allows for a moderate increase in the learning rate. When $\beta_{\text{noise}} < 1$, the optimal $\mathcal{B}_{\text{noise}}$ can be solved from $\eqref{eq:B-1-2}$, and once the batch size exceeds this value, the learning rate should actually decrease instead.
A Note
It should be pointed out that although the starting point and final conclusions of the sections above are broadly similar to those of the original paper Surge Phenomenon in Optimal Learning Rate and Batch Size Scaling, the intermediate approximations differ.
Most of the conclusions in the original paper are approximate results under the assumption $B \ll \pi(\sigma_i/g_i)^2/2$, and so they nearly always end up concluding that the Surge phenomenon will occur — which isn't entirely well-founded. Most notably, the very form of the assumption $B \ll \pi(\sigma_i/g_i)^2/2$ is somewhat problematic: its right-hand side depends on $i$, and we obviously can't assign a separate batch size to every single component. So to arrive at a global result, one is forced into $B \ll \min_i \pi(\sigma_i/g_i)^2/2$, which is rather too strong a requirement.
The approach in this article instead introduces the approximation $\eqref{eq:mu-approx}$, which can be seen as a mean-field approximation. Intuitively this is more reasonable than the pointwise assumption $B \ll \pi(\sigma_i/g_i)^2/2$, so in principle the conclusions should be more precise — for instance, we can obtain the conclusion that "even when the off-diagonal elements of the Hessian cannot be neglected, the Surge phenomenon does not necessarily occur" (depending on $\beta_{\text{noise}}$). Notably, this added precision doesn't come at the cost of simplicity: equation $\eqref{eq:eta-opt-beta}$ remains just as clean and clear, and equation $\eqref{eq:Delta-L-sign-2}$ has the same form as in the original paper, without requiring any extra approximating assumptions, and so on.
Finally, a bit of a reflection: OpenAI's analysis of SGD dates back to 2018, while the paper on the Surge phenomenon was only published in the middle of this year. It's rather surprising that it took six whole years to go from SGD to Adam — largely, I think, because of OpenAI's "prestige" and the conjecture $\eqref{eq:openai-adam}$, which made people feel there was nothing more to explore about Adam. Who would have guessed that Adam might harbor some new characteristics after all. Of course, questions like exactly how reasonable it is to use $\tilde{\boldsymbol{\varphi}}_B = \sign(\tilde{\boldsymbol{g}}_B)$ as an approximation of Adam, and to what extent it reflects reality, are still, in the author's view, worth further thought.
Summary
This article has discussed the classic model-training question of "the scaling law between batch size and learning rate" from multiple perspectives, with particular focus on OpenAI's derivation and conclusions based on a second-order approximation of the loss function, as well as subsequent work that applies the same ideas to analyze the Adam optimizer.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.