Rethinking Learning Rate and Batch Size (Part 2): Mean Field Theory

At the end of the previous article "Rethinking Learning Rate and Batch Size (Part 1): Current Landscape", we mentioned that for cases where $\tilde{\boldsymbol{\varphi}}_B$ nonlinearly depends on $\tilde{\boldsymbol{g}}_B$, such as SignSGD and SoftSignSGD, the computational process is mentally burdensome and faces difficulties in generalization. To address this, I invested some effort in attempting to simplify the derivations, and fortunately made some progress. The key idea is the theme of this article—mean field theory.

Mean field theory is a common approximation method in physics that has no fixed form, but the general idea is to move the averaging operation inside the function. In fact, we have already glimpsed the charm of mean field theory in "Why is Adam's Update RMS 0.2?". In this article, we will once again witness its remarkable effectiveness in computing the learning rate scaling laws for SignSGD/SoftSignSGD.

Method Overview#

Following the notation from the previous article, for SignSGD we have $\newcommand{sign}{\mathop{\text{sign}}}\tilde{\boldsymbol{\varphi}}_B=\sign(\tilde{\boldsymbol{g}}_B)$. We first need to compute $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]$ and $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]$, then we can calculate:

where $\boldsymbol{g}$ is the gradient and $\boldsymbol{H}$ is the Hessian matrix. According to our assumptions, the random variable $\tilde{\boldsymbol{g}}_B$ has mean $\boldsymbol{g}$ and covariance matrix $\boldsymbol{\Sigma}/B$. We are primarily concerned with the relationship between $\eta^*$ and batch size $B$. Since $\sign$ is an element-wise operation, we can attempt the analysis starting from a single scalar. The mean field method originated from my sudden discovery one day of a potentially valid approximation:

Readers who have read "How Should Learning Rate Scale with Batch Size?" should be surprised to discover that this result, which can be derived in just one line, differs from the result obtained through numerous assumptions and approximations in the original article only by an insignificant constant $\pi/2$! This fact made me realize that the mean field approximation might be completely sufficient for analyzing the relationship between learning rate and batch size.

Computation Process#

In this section, we will use the mean field approximation to provide the complete computation process for SignSGD. First, for the mean $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]$, the calculation from the previous section is almost complete—we just need to add a few details. Using component notation:

where $\sigma_i^2 = \boldsymbol{\Sigma}_{i,i}$. Since we are ultimately mainly concerned with the relationship between $\eta^*$ and $B$, both of which are scalars, we apply the mean field approximation once more here to separate the denominator part related to $B$ in scalar form:

Here, $\mathcal{B}_{\text{simple}}$ is the same as in the previous article: $\mathcal{B}_{\text{simple}} = \tr(\boldsymbol{\Sigma})/\boldsymbol{g}^{\top}\boldsymbol{g}$, which is also equal to $\mathbb{E}[\sigma_i^2]/\mathbb{E}[g_i^2]$ (where $\mathbb{E}$ denotes averaging over index $i$). That is, it replaces the originally index-dependent $\sigma_i^2/g_i^2$ with an index-independent average $\mathbb{E}[\sigma_i^2]/\mathbb{E}[g_i^2]$. After this approximation, the result is simplified while still preserving the functional form with respect to $B$.

Next, for the second moment $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]$, we reintroduce the component independence assumption to simplify the result. It is possible to compute without this assumption, but the result would be more complex and would require other assumptions to simplify the computation, so it's better to directly introduce the independence assumption. Under the independence assumption, $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]_{i,j}$ is computed in two parts: $i\neq j$ and $i=j$. When $i\neq j$,

When $i=j$, it's even simpler because the square of $\sign$ is necessarily 1, so its expectation is naturally 1. Therefore, the overall result can be concisely written as $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]_{i,j}\approx \mu_i\mu_j + \delta_{i,j}(1 - \mu_i\mu_j)$.

Anomalous Phenomenon#

Substituting the above calculation results into equation (1), we obtain:

where $\beta = (1 + \mathcal{B}_{\text{simple}}/B)^{-1/2}$. Note that $\beta$ is monotonically increasing with respect to $B$, and $\beta\in(0,1)$, so $\beta$ can be viewed as a normalized batch size. However, $\eta^*$ is not always monotonic with respect to $\beta$, which could lead to the anomalous behavior where "increasing batch size should actually decrease the learning rate"—termed the "Surge phenomenon" in the original paper.

Let's understand this step by step. When $B\ll \mathcal{B}_{\text{simple}}$, we have $\beta\approx \sqrt{B/\mathcal{B}_{\text{simple}}}$, and thus $\beta \ll 1$. In this case, the $1/\beta$ term in the denominator of equation (6) dominates, giving:

This indicates that SignSGD's learning rate follows square root scaling for small batch sizes. Since we assume positive definiteness of the Hessian matrix during analysis, we necessarily have $\sum_i H_{i,i} > 0$. Then when $\sum_{i\neq j} H_{i,j}\sign(g_i g_j) \leq 0$, equation (6) is always monotonically increasing with respect to $\beta$, so $\eta^*$ is also monotonically increasing with respect to $B$, showing no anomalous behavior.

When $\sum_{i\neq j} H_{i,j}\sign(g_i g_j) > 0$, using basic inequalities we can find that the denominator of equation (6) has a minimum point at:

Note that $\beta\in(0, 1)$, so there is an additional condition $\beta^*\in(0, 1)$. In this case, $\eta^*$ is no longer monotonically increasing with respect to $B$, but rather increases first and then decreases. There exists a critical batch size beyond which the learning rate should actually decrease—this is the "Surge phenomenon".

Causal Reflection#

Why does this anomalous Surge phenomenon occur? In fact, it reflects an incompatibility between the optimizer's underlying assumptions and our analytical method. Specifically, to estimate the optimal learning rate, we expanded the loss increment to second-order approximation and assumed positive definiteness of the Hessian matrix. Under these settings, the optimal update should be Newton's method: $\boldsymbol{H}^{-1}\boldsymbol{g}$.

From the perspective of Newton's method, different optimizers essentially correspond to different assumptions about the Hessian matrix. For example, SGD corresponds to assuming $\boldsymbol{H}=\eta_{\max}^{-1} \boldsymbol{I}$, while SignSGD corresponds to assuming $\newcommand{diag}{\mathop{\text{diag}}}\boldsymbol{H}=\eta_{\max}^{-1} \diag(|\boldsymbol{g}|)$, though in practice we can only replace $\boldsymbol{g}$ with $\tilde{\boldsymbol{g}}_B$. The Surge phenomenon actually reflects that as $B\to\infty$, the deviation between SignSGD's assumed Hessian matrix and the actual Hessian matrix increases.

We know that today's LLM models have parameters numbering in the billions. Computing either the full Hessian matrix or the covariance matrix is nearly impossible, which is one reason why we introduce the independence assumption when computing the second moment $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]$—only then does the covariance matrix become diagonal, making estimation feasible. The Hessian matrix is similar; we can often only compute it for specific structures.

For example, substituting $\boldsymbol{H}=\eta_{\max}^{-1} \diag(|\boldsymbol{g}|)$ into equation (6) yields $\eta^*\approx \eta_{\max} \beta = \eta_{\max} / \sqrt{1 + \mathcal{B}_{\text{simple}}/B}$. This form is very concise and shows no anomalous behavior. Does this mean the Surge phenomenon won't occur? Not exactly. The Surge phenomenon objectively exists; the point here is more that: when we observe the Surge phenomenon in experiments, perhaps the first thing to consider is not correcting the variation pattern of $\eta^*$, but rather changing the optimizer.

Loss Variation#

With $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]$ and $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]$, we can also compute $\overline{\Delta\mathcal{L}}$ as in the previous article. Interestingly, it has the same format as the SGD result:

where:

Note that this retains the full Hessian matrix, so the result is actually quite interesting—although the learning rate $\eta^*$ may exhibit the Surge phenomenon, the average loss increment does not show this phenomenon. It is always monotonically increasing with respect to $B$ and maintains the same form as SGD, meaning we can derive the same "training data volume–training steps" relationship:

A more thought-provoking question is: why do SGD and SignSGD have completely different update magnitudes, including明显 different behaviors in learning rate $\eta^*$, yet $\overline{\Delta\mathcal{L}}$ has the same form with respect to $B$? Is this merely a coincidence, or is there a deeper principle underlying it?

General Patterns#

Starting again from the mean field approximation, I obtained an answer leaning toward the latter. Whether for $\eta^*$ or $\overline{\Delta\mathcal{L}}$, the core difficulty lies in computing $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]$ and $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]$, so our goal is to explore unified computational patterns for both.

We generally set $\tilde{\boldsymbol{\varphi}}_B=\tilde{\boldsymbol{H}}{}_B^{-1}\tilde{\boldsymbol{g}}_B$, where $\tilde{\boldsymbol{H}}_B$ is some positive semidefinite matrix. Then we can write:

and:

Substituting into the expression for $\overline{\Delta\mathcal{L}}$, we obtain:

Note that the above expression is homogeneous with respect to $\hat{\boldsymbol{H}}$. If we assume that the relationship between $\hat{\boldsymbol{H}}$ and $B$ can be separated into a scalar form such as $\hat{\boldsymbol{H}}\approx f(B) \boldsymbol{G}$, where $f(B)$ is a scalar function of $B$ and $\boldsymbol{G}$ is not明显 related to $B$, then $f(B)$ can be canceled out in both numerator and denominator. Ultimately, the relationship with $B$ can be organized into the following form:

This proves that $\overline{\Delta\mathcal{L}}$ has the same asymptotic pattern with respect to $B$, with the core being the homogeneity with respect to $\hat{\boldsymbol{H}}$. In contrast, $\eta^*$ does not have such a unified result because it is not homogeneous with respect to $\hat{\boldsymbol{H}}$.

Validity Analysis#

By this point, I believe everyone has some understanding of the mean field method. Its main characteristic is computational simplicity, or more fundamentally, mean field theory chooses the simple, computable direction for computation, which gives it great flexibility. Flexibility is often also a drawback, meaning it's difficult to grasp the patterns for the next step.

As for explaining why this approach is valid, that's even harder—it can only be analyzed case by case, and sometimes even case-by-case analysis is difficult. My feeling is that the mean field method is three parts computation, three parts luck, three parts intuition, plus one part metaphysics. Of course, trying it out is fine. Let's take the previous SignSGD computation as an example and attempt some analysis.

Generalized Approximation#

One reason for choosing to analyze SignSGD is that we typically use it as a theoretical approximation for Adam. In "How Does Adam's Epsilon Affect Learning Rate Scaling Laws?", we computed a theoretically better approximation: SoftSignSGD, which considers the effect of $\epsilon$.

In this case, $\tilde{\boldsymbol{\varphi}}_B = \softsign(\tilde{\boldsymbol{g}}_B)$. Let's proceed directly to the main topic:

Here, $\mathcal{B}_{\text{simple}}$ is slightly different: it is $\tr(\boldsymbol{\Sigma})/(\boldsymbol{g}^{\top}\boldsymbol{g} + N\epsilon^2)$, where $N$ is the total number of model parameters, i.e., $\boldsymbol{g}\in\mathbb{R}^N$. As for the final terms: $\nu_i=\softsign(g_i)$ and $\beta = (1 + \mathcal{B}_{\text{simple}}/B)^{-1/2}$. Next, compute $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]$. Under the independence assumption, when $i\neq j$ we can still compute the means separately, so $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]_{i,j}=\nu_i \nu_j \beta^2$. Thus we only need to compute the case $i=j$:

This can be uniformly written as $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]_{i,j}\approx \nu_i \nu_j\beta^2 + \delta_{i,j}(1-\beta^2)$. Therefore:

Except for $\beta$, all other parts of the above expression are independent of $B$, so we have already obtained the explicit relationship between $\eta^*$ and $B$, with a form大同小异 to that of SignSGD. The remaining analysis can refer to "How Does Adam's Epsilon Affect Learning Rate Scaling Laws?" or follow the previous content.

Summary#

In this article, we used mean field approximation to recompute the conclusions for SignSGD and SoftSignSGD, greatly simplifying the related computational processes, and preliminarily considered the general patterns of these computations.