Rethinking Learning Rate and Batch Size (Part 2): Mean-Field Approximation
At the end of the previous post, Rethinking Learning Rate and Batch Size (Part 1): The Current State], we noted that for cases like SignSGD and SoftSignSGD, where $\tilde{\boldsymbol{\varphi}}_B$ depends non-linearly on $\tilde{\boldsymbol{g}}_B$, the computational burden becomes quite heavy, and the derivations struggle to generalize. Motivated by this, I invested some effort into trying to simplify the derivation, and fortunately made some progress. The key idea turned out to be the subject of this post — the mean-field approximation.
The mean-field approximation is a common approximation technique in physics. It doesn't have a fixed form, but the general idea is to move an averaging operation inside a function. In fact, we already caught a glimpse of the power of mean-field reasoning in Why Is Adam's Update RMS 0.2?], and in this post we'll see it work its magic again in deriving the learning-rate scaling law for SignSGD/SoftSignSGD.
Outline of the Method
Following the notation of the previous post, 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}]$, from which we can obtain
\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}more
where $\boldsymbol{g}$ is the gradient and $\boldsymbol{H}$ is the Hessian matrix. By assumption, the random variable $\tilde{\boldsymbol{g}}_B$ has mean $\boldsymbol{g}$ and covariance matrix $\boldsymbol{\Sigma}/B$, and what we mainly care about is the relationship between $\eta^*$ and the batch size $B$. Since $\sign$ is an element-wise operation, we can start our exploration from a single scalar. The mean-field method originated from an approximate relation I happened to notice one day, which seemed likely to hold:
\begin{equation}\mathbb{E}[\sign(\tilde{g}_B)] = \mathbb{E}\bigg[\frac{\tilde{g}_B}{\sqrt{\tilde{g}_B^2}}\bigg]\approx \frac{\mathbb{E}[\tilde{g}_B]}{\sqrt{\mathbb{E}[\tilde{g}_B^2}]} = \frac{g}{\sqrt{g^2 + \sigma^2/B}}\end{equation}
Readers who have gone through How Should the Learning Rate Change as Batch Size Increases?] will be surprised to find that this result — derivable in a single line — differs from the result obtained in that post, via a whole battery of assumptions and approximations, by nothing more than an inconsequential constant $\pi/2$! This fact convinced me that the mean-field approximation might well be entirely sufficient for capturing the relationship between learning rate and batch size.
There are several benefits to a mean-field-based derivation. First, it requires fewer assumptions: the original derivation relied on at least three assumptions — component independence, a normal distribution, and approximating $\text{erf}(x)$ by $x/\sqrt{x^2+c}$ — whereas the mean-field approximation dispenses with the assumption about the distribution's form, requiring only that the approximation itself be applicable. Second, the computation is much simpler: above, we completed the calculation in a single line, whereas the original derivation, even under all those assumptions, is far more involved.
The Computation
In this section we use the mean-field approximation to give the full computation for SignSGD. First, the mean $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]$ — actually, the computation in the previous section already got us most of the way there, so we just need to fill in a few details. Writing it component-wise:
\begin{equation}\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]_i = \mathbb{E}[\sign((\tilde{g}_B)_i)] = \mathbb{E}\bigg[\frac{(\tilde{g}_B)_i}{\sqrt{(\tilde{g}_B)_i^2}}\bigg]\approx \frac{\mathbb{E}[(\tilde{g}_B)_i]}{\sqrt{\mathbb{E}[(\tilde{g}_B)_i^2]}} = \frac{g_i}{\sqrt{g_i^2 + \sigma_i^2/B}} = \frac{\sign(g_i)}{\sqrt{1 + (\sigma_i^2/g_i^2)/B}}\end{equation}
where $\sigma_i^2 = \boldsymbol{\Sigma}_{i,i}$. Since what we ultimately care about is the relationship between $\eta^*$ and $B$, both of which are scalars, we apply the mean-field approximation once more here, separating out the denominator term related to $B$ in scalar form:
\begin{equation}\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]_i \approx \frac{\sign(g_i)}{\sqrt{1 + (\sigma_i^2/g_i^2)/B}} \approx \frac{\sign(g_i)}{\sqrt{1 + \mathcal{B}_{\text{simple}}/B}} \triangleq \mu_i\end{equation}
Here, $\mathcal{B}_{\text{simple}}$ is exactly the $\mathcal{B}_{\text{simple}} = \tr(\boldsymbol{\Sigma})/\boldsymbol{g}^{\top}\boldsymbol{g}$ from the previous post, which in turn equals $\mathbb{E}[\sigma_i^2]/\mathbb{E}[g_i^2]$ (this $\mathbb{E}$ denotes averaging over the index $i$). In other words, it replaces the quantity $\sigma_i^2/g_i^2$, which originally depends on the index $i$, with some index-independent average value $\mathbb{E}[\sigma_i^2]/\mathbb{E}[g_i^2]$. This approximation simplifies the result while still preserving the functional dependence on $B$.
Next is the second moment $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]$, for which we reintroduce the component-independence assumption to simplify the result. It's possible to compute this without that assumption, but the result would be more complex and would need additional assumptions to simplify anyway — so we might as well introduce the independence assumption directly. Under independence, $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]_{i,j}$ splits into two parts, $i\neq j$ and $i=j$. When $i\neq j$,
\begin{equation}\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]_{i,j} = \mathbb{E}[(\tilde{\varphi}_B)_i(\tilde{\varphi}_B)_j] = \mathbb{E}[(\tilde{\varphi}_B)_i]\mathbb{E}[(\tilde{\varphi}_B)_j] \approx \mu_i \mu_j\end{equation}
When $i=j$, it's even simpler, since the square of $\sign$ is necessarily 1, and hence so is its expectation. Therefore, the overall result can be written compactly 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)$.
An Anomaly
Substituting the above results into equation $\eqref{eq:eta-opt}$, we get
\begin{equation}\eta^* \approx \frac{\sum_i |g_i|}{\frac{1}{\beta}\sum_i H_{i,i} + \beta\sum_{i\neq j} H_{i,j}\sign(g_i g_j)}\label{eq:eta-opt-sign}\end{equation}
where $\beta = (1 + \mathcal{B}_{\text{simple}}/B)^{-1/2}$. Note that $\beta$ is monotonically increasing in $B$, and $\beta\in(0,1)$, so $\beta$ can be regarded as a normalized batch size. However, with respect to $\beta$ the expression is not always monotonic, so we can end up with the counterintuitive behavior that "as batch size increases, the learning rate should actually decrease." The original paper] calls this the "Surge phenomenon."
Let's understand this step by step. When $B\ll \mathcal{B}_{\text{simple}}$, we have $\beta\approx \sqrt{B/\mathcal{B}_{\text{simple}}}$, so $\beta \ll 1$, meaning the $1/\beta$ term dominates the denominator of equation $\eqref{eq:eta-opt-sign}$, giving
\begin{equation}\eta^* \approx \frac{\sum_i |g_i|}{\sum_i H_{i,i}}\beta \approx \frac{\sum_i |g_i|}{\sum_i H_{i,i}}\sqrt{B/\mathcal{B}_{\text{simple}}}\propto \sqrt{B}\end{equation}
This shows that SignSGD's learning rate follows square-root scaling at small batch sizes. Since we assume positive definiteness of the Hessian in our analysis, we necessarily have $\sum_i H_{i,i} > 0$, so when $\sum_{i\neq j} H_{i,j}\sign(g_i g_j) \leq 0$, equation $\eqref{eq:eta-opt-sign}$ is always monotonically increasing in $\beta$, and hence $\eta^*$ is also monotonically increasing in $B$ — no anomaly appears in this case.
When $\sum_{i\neq j} H_{i,j}\sign(g_i g_j) > 0$, the basic inequality tells us that the denominator of equation $\eqref{eq:eta-opt-sign}$ has a minimum at
\begin{equation}\beta^* = \sqrt{\frac{\sum_i H_{i,i}}{\sum_{i\neq j} H_{i,j}\sign(g_i g_j)}}\end{equation}
Note that $\beta\in(0, 1)$, so there's an additional condition $\beta^*\in(0, 1)$ needed. In this case, $\eta^*$ is no longer monotonically increasing in $B$, but instead first increases then decreases — there exists a critical batch size beyond which the learning rate should actually decrease. This is the "Surge phenomenon."
Reflecting on the Cause
Why does this Surge phenomenon — this anomalous behavior — occur? In fact, it reflects an incompatibility between the optimizer's own assumptions and our method of analysis. Specifically, to estimate the optimal learning rate, we expanded the loss increment to second order and assumed positive definiteness of the Hessian matrix. Under these settings, the optimal update should be that of Newton's method, i.e., $\boldsymbol{H}^{-1}\boldsymbol{g}$.
From the perspective of Newton's method, different optimizers 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}|)$. Of course, in actual training we can only substitute $\boldsymbol{g}$ with $\tilde{\boldsymbol{g}}_B$. The Surge phenomenon in fact reflects the growing deviation, as $B\to\infty$, between the Hessian matrix assumed by SignSGD and the true Hessian matrix.
We know that today's LLMs have parameter counts starting in the hundreds of millions, so computing either the full Hessian matrix or the full covariance matrix is essentially impossible. This is one reason we introduce the independence assumption when computing the second moment $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]$ — under that assumption the covariance matrix becomes purely diagonal, making estimation tractable. The same goes for the Hessian: we can generally only compute Hessians with specific structural constraints.
For example, substituting $\boldsymbol{H}=\eta_{\max}^{-1} \diag(|\boldsymbol{g}|)$ into equation $\eqref{eq:eta-opt-sign}$ gives $\eta^*\approx \eta_{\max} \beta = \eta_{\max} / \sqrt{1 + \mathcal{B}_{\text{simple}}/B}$, which is a quite clean form and exhibits no anomalous behavior. Does this mean the Surge phenomenon won't appear? Not at all — the Surge phenomenon is an objectively real occurrence. What I really want to convey here is: when we observe the Surge phenomenon experimentally, the first order of business probably shouldn't be to fix up the scaling law for $\eta^*$, but rather to consider switching optimizers.
Change in Loss
Given $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]$ and $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]$, we can, just as in the previous post, compute $\overline{\Delta\mathcal{L}}$. What's especially interesting is that it takes the same form as the result for SGD:
\begin{equation}\overline{\Delta\mathcal{L}} = \mathcal{L}(\boldsymbol{w}) - \mathbb{E}[\mathcal{L}(\boldsymbol{w} - \eta^*\tilde{\boldsymbol{g}}_B)] \approx \frac{(\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]^{\top}\boldsymbol{g})^2}{2\tr(\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]\boldsymbol{H})}\approx \frac{\Delta\mathcal{L}_{\max}}{1 + \mathcal{B}_{\text{noise}}/B}\end{equation}
where
\begin{equation}\Delta\mathcal{L}_{\max} = \frac{\frac{1}{2}(\sum_i |g_i|)^2}{\sum_i H_{i,i} + \sum_{i\neq j} H_{i,j}\sign(g_i g_j)},\quad \mathcal{B}_{\text{noise}} = \frac{\mathcal{B}_{\text{simple}}\sum_i H_{i,i}}{\sum_i H_{i,i} + \sum_{i\neq j} H_{i,j}\sign(g_i g_j)}\end{equation}
Note that here we've retained the full Hessian matrix, so the result is actually rather striking — even though the learning rate $\eta^*$ can exhibit the Surge phenomenon, the average loss increment does not exhibit this behavior at all: it remains monotonically increasing in $B$, and it retains the same form as for SGD. This means we can derive the same "training data volume vs. training steps" relationship:
\begin{equation}\left(\frac{S}{S_{\min}} - 1\right)\left(\frac{E}{E_{\min}} - 1\right) = 1\end{equation}
A question worth pondering further: why is it that, even though the updates of SGD and SignSGD are entirely different — including the visibly different behaviors of their learning rates $\eta^*$ — the relationship between $\overline{\Delta\mathcal{L}}$ and $B$ nonetheless takes the same form? Is this purely coincidental, or is there some deeper principle at work?
A General Pattern
Starting again from the mean-field approximation, I found an answer that leans toward the latter. Whether for $\eta^*$ or $\overline{\Delta\mathcal{L}}$, the core difficulty always 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 find a unified computational rule covering both.
Let us generally set $\tilde{\boldsymbol{\varphi}}_B=\tilde{\boldsymbol{H}}{}_B^{-1}\tilde{\boldsymbol{g}}_B$, where $\tilde{\boldsymbol{H}}_B$ is some positive semi-definite matrix. Then we can write
\begin{equation}\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B] = \mathbb{E}[\tilde{\boldsymbol{H}}{}_B^{-1}\tilde{\boldsymbol{g}}_B]\approx \underbrace{\mathbb{E}[\tilde{\boldsymbol{H}}_B]^{-1}}_{\text{denote}\hat{\boldsymbol{H}}{}^{-1}}\mathbb{E}[\tilde{\boldsymbol{g}}_B] = \hat{\boldsymbol{H}}{}^{-1}\boldsymbol{g}\end{equation}
as well as
\begin{equation}\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}] = \mathbb{E}[\tilde{\boldsymbol{H}}{}_B^{-1}\tilde{\boldsymbol{g}}_B\tilde{\boldsymbol{g}}_B^{\top}\tilde{\boldsymbol{H}}{}_B^{-1}]\approx \mathbb{E}[\tilde{\boldsymbol{H}}_B]^{-1}\mathbb{E}[\tilde{\boldsymbol{g}}_B\tilde{\boldsymbol{g}}_B^{\top}]\mathbb{E}[\tilde{\boldsymbol{H}}_B]^{-1} = \hat{\boldsymbol{H}}{}^{-1}(\boldsymbol{g}\boldsymbol{g}^{\top} + \boldsymbol{\Sigma}/B)\hat{\boldsymbol{H}}{}^{-1} \end{equation}
Substituting into the expression for $\overline{\Delta\mathcal{L}}$, we obtain
\begin{equation}\overline{\Delta\mathcal{L}} \approx \frac{1}{2}\frac{(\boldsymbol{g}^{\top}\hat{\boldsymbol{H}}{}^{-1}\boldsymbol{g})^2}{\boldsymbol{g}^{\top}\hat{\boldsymbol{H}}{}^{-1}\boldsymbol{H}\hat{\boldsymbol{H}}{}^{-1}\boldsymbol{g} + \tr(\boldsymbol{\Sigma}\hat{\boldsymbol{H}}{}^{-1}\boldsymbol{H}\hat{\boldsymbol{H}}{}^{-1})/B}\end{equation}
Note that the above expression is homogeneous in $\hat{\boldsymbol{H}}$. If we assume the relationship between $\hat{\boldsymbol{H}}$ and $B$ can be factored out into a separate scalar form such as $\hat{\boldsymbol{H}}\approx f(B) \boldsymbol{G}$, where $f(B)$ is a scalar function of $B$, and $\boldsymbol{G}$ has no strong relationship with $B$, then $f(B)$ can be cancelled simultaneously from numerator and denominator, and the final relationship in terms of $B$ can be organized into the following form:
\begin{equation}\overline{\Delta\mathcal{L}} \approx \frac{\Delta\mathcal{L}_{\max}}{1 + \mathcal{B}_{\text{noise}}/B}\end{equation}
This proves that $\overline{\Delta\mathcal{L}}$ has the same asymptotic behavior with respect to $B$, and the core reason is the homogeneity in $\hat{\boldsymbol{H}}$. By contrast, $\eta^*$ has no such unified result, because it is not homogeneous in $\hat{\boldsymbol{H}}$.
The Scope of the Analysis
By this point, I imagine readers have a good sense of the mean-field method. Its main feature is computational simplicity — or, put more fundamentally, mean-field reasoning means always heading toward whatever direction is simplest and most tractable to compute, which grants it tremendous flexibility. Flexibility, though, is often also a drawback: it means it's very hard to anticipate what comes next.
As for explaining why this approach works, that's even harder — one really has to analyze case by case, and in some cases the analysis may not even be tractable. My personal feeling is that the mean-field method is three parts calculation, three parts luck, three parts intuition, plus one part black magic. That said, there's no harm in trying, so let's take the SignSGD computation above as an example and attempt some analysis.
Clearly, the core computation in SignSGD is $\mathbb{E}[\sign(x)]$. Let's write $\mathbb{E}[x]=\mu,\mathbb{E}[x^2]=\mu^2 + \sigma^2$, and then
\begin{equation}\sign(x) = \frac{x}{\sqrt{x^2}} = \frac{x}{\sqrt{\mu^2 + \sigma^2 + (x^2 - \mu^2 - \sigma^2)}}\end{equation}
Assuming $x^2 - \mu^2 - \sigma^2$ is small, we do a Taylor expansion:
\begin{equation}\sign(x) = \frac{x}{\sqrt{\mu^2 + \sigma^2}} - \frac{1}{2}\frac{x(x^2 - \mu^2 - \sigma^2)}{（\mu^2 + \sigma^2)^{3/2}} + \frac{3}{8}\frac{x(x^2 - \mu^2 - \sigma^2)^2}{（\mu^2 + \sigma^2)^{5/2}}-\cdots \end{equation}
Now the denominator no longer depends on $x$, and the numerator is a polynomial in $x$, so taking expectations on both sides, the first term is exactly the mean-field approximation result $\mu/\sqrt{\mu^2 + \sigma^2}$. To gauge the reasonableness of the mean-field approximation, let's compute the second term:
\begin{equation}\frac{1}{2}\frac{\mathbb{E}[x(x^2 - \mu^2 - \sigma^2)]}{（\mu^2 + \sigma^2)^{3/2}} = \frac{1}{2}\frac{\mathbb{E}[x^3] - (\mu^3 + \mu\sigma^2)}{（\mu^2 + \sigma^2)^{3/2}} \end{equation}
This brings in $\mathbb{E}[x^3]$, a new statistic that turns out to be the key factor governing the mean-field error. We can get a feel for it using the normal distribution $\mathcal{N}(x;\mu,\sigma^2)$, in which case $\mathbb{E}[x^3]=\mu^3 + 3\mu\sigma^2$, and substituting into the above gives
\begin{equation}\frac{\mu\sigma^2}{（\mu^2 + \sigma^2)^{3/2}} = \frac{\sigma^2/\mu^2}{（1 + \sigma^2/\mu^2)^{3/2}}\end{equation}
The right-hand side is a bounded expression, attaining its maximum at $\sigma^2/\mu^2=2$, giving the value $2/3^{3/2}=0.3849\cdots$. This indicates that the error of the mean-field approximation is very likely bounded, and the error term tends to 0 as both $\sigma\to 0$ and $\sigma\to\infty$ go to 0 — which to some extent demonstrates the validity of the mean-field approximation.
A More General Approximation
One reason for choosing to analyze SignSGD is that we usually use it as a theoretical proxy for Adam. In How Does Adam's Epsilon Affect the Scaling Law of the Learning Rate?], we computed a theoretically better approximation, SoftSignSGD, which takes into account the effect of $\epsilon$:
\begin{equation}\sign(x)=\frac{x}{\sqrt{x^2}}\quad\to\quad\newcommand{softsign}{\mathop{\text{softsign}}}\softsign(x)=\frac{x}{\sqrt{x^2+\epsilon^2}}\end{equation}
Here $\tilde{\boldsymbol{\varphi}}_B = \softsign(\tilde{\boldsymbol{g}}_B)$. Let's dive straight in:
\begin{equation}\begin{aligned} &\,\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]_i = \mathbb{E}[\softsign((\tilde{g}_B)_i)] = \mathbb{E}\bigg[\frac{(\tilde{g}_B)_i}{\sqrt{(\tilde{g}_B)_i^2 + \epsilon^2}}\bigg]\approx \frac{\mathbb{E}[(\tilde{g}_B)_i]}{\sqrt{\mathbb{E}[(\tilde{g}_B)_i^2]+ \epsilon^2}} \\[8pt] =&\, \frac{g_i}{\sqrt{g_i^2 + \sigma_i^2/B + \epsilon^2}} = \frac{\softsign(g_i)}{\sqrt{1 + \sigma_i^2/(g_i^2 + \epsilon^2)/B}}\approx \frac{\softsign(g_i)}{\sqrt{1 + \mathcal{B}_{\text{simple}}/B}}\triangleq \nu_i\beta \end{aligned}\end{equation}
Here $\mathcal{B}_{\text{simple}}$ is a little different — it equals $\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$; and as for the final $\nu_i=\softsign(g_i), \beta = (1 + \mathcal{B}_{\text{simple}}/B)^{-1/2}$. Next we compute $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]$. Under the independence assumption, when $i\neq j$ we can still take means separately, so we get $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]_{i,j}=\nu_i \nu_j \beta^2$, meaning we only need to handle the case $i=j$:
\begin{equation}\begin{aligned} &\,\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]_{i,i} = \mathbb{E}[\softsign((\tilde{g}_B)_i)^2] = \mathbb{E}\bigg[\frac{(\tilde{g}_B)_i^2}{(\tilde{g}_B)_i^2 + \epsilon^2}\bigg]\approx \frac{\mathbb{E}[(\tilde{g}_B)_i^2]}{\mathbb{E}[(\tilde{g}_B)_i^2]+ \epsilon^2} \\[8pt] =&\, \frac{g_i^2 + \sigma_i^2/B}{g_i^2 + \sigma_i^2/B + \epsilon^2} = 1 - \frac{1 - \softsign(g)^2}{1 + \sigma_i^2/(g_i^2 + \epsilon^2)/B}\approx 1 - \frac{1 - \softsign(g)^2}{1 + \mathcal{B}_{\text{simple}}/B} \end{aligned}\end{equation}
which can be written uniformly 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)$, giving us
\begin{equation}\eta^* \approx \frac{\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]^{\top}\boldsymbol{g}}{\text{Tr}(\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]\boldsymbol{H})} \approx \frac{\beta\sum_i \nu_i g_i}{\sum_i H_{i,i} + \beta^2(\sum_{i,j} \nu_i \nu_j H_{i,j} - \sum_i H_{i,i})}\end{equation}
Apart from $\beta$, everything else in the above expression is independent of $B$, so we've now obtained an explicit relationship between $\eta^*$ and $B$, whose form is much the same as that of SignSGD. The remaining analysis can either follow How Does Adam's Epsilon Affect the Scaling Law of the Learning Rate?] or be carried out by mimicking the discussion above.
Summary
In this post, we used the mean-field approximation to recompute the results for SignSGD and SoftSignSGD, greatly simplifying the relevant derivations, and took a first look at the general patterns underlying these computations.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.