Rethinking Learning Rate and Batch Size (III): Muon

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

In the previous two posts, Rethinking Learning Rate and Batch Size (I): Current State and Rethinking Learning Rate and Batch Size (II): Mean Field, we mainly proposed a mean-field method to simplify the calculations relating learning rate and batch size. The optimizers we analyzed at the time were SGD, SignSGD, and SoftSignSGD, and the main goal was simplification — essentially no new conclusions were drawn.

However, in today's feast of optimizers, how could Muon possibly be left off the table? So in this post we'll try to work out the corresponding results for Muon, and see whether its learning rate-batch size relationship reveals any new patterns.

Basic Notation

As is well known, the defining feature of Muon is its non-element-wise update rule, so the element-wise computational methods used previously in How Should the Learning Rate Change as Batch Size Increases? and How Does Adam's Epsilon Affect the Learning Rate Scaling Law? are completely inapplicable here. Fortunately, the mean-field approach introduced in the previous post still works, requiring only a few adjustments to the details. more

Let's first introduce some notation. Let the loss function be $\mathcal{L}(\boldsymbol{W})$, with $\boldsymbol{W}\in\mathbb{R}^{n\times m}$ being a matrix-shaped variable (let $n\geq m$), and $\boldsymbol{G}$ its gradient. The gradient for a single sample is denoted $\tilde{\boldsymbol{G}}$, its mean is $\boldsymbol{G}$, and its variance is $\sigma^2$. When the batch size is $B$, the gradient is denoted $\tilde{\boldsymbol{G}}_B$, whose mean remains $\boldsymbol{G}$, but whose variance becomes $\sigma^2/B$. Note that the variance here is just a scalar $\sigma^2$, rather than the full covariance matrix considered previously.

The core reason for this simplification is that the random variable here is already a matrix, so its corresponding covariance object would actually be a 4th-order tensor, which is cumbersome to discuss. Does simplifying it to a single scalar seriously hurt accuracy? Actually, no — although the previous two posts considered the full covariance matrix $\boldsymbol{\Sigma}$, closer inspection shows that the final result only depends on $\newcommand{tr}{\mathop{\text{tr}}}\tr(\boldsymbol{\Sigma})$, which is equivalent to simplifying it to a scalar from the start.

The Hessian Matrix

Similarly, let the update be $-\eta\tilde{\boldsymbol{\Phi}}_B$, and consider the second-order expansion of the loss function

\begin{equation}\mathcal{L}(\boldsymbol{W} - \eta\tilde{\boldsymbol{\Phi}}_B) \approx \mathcal{L}(\boldsymbol{W}) - \eta \tr(\tilde{\boldsymbol{\Phi}}{}_B^{\top}\boldsymbol{G}) + \frac{1}{2}\eta^2\newcommand{tr}{\mathop{\text{tr}}}\tr(\tilde{\boldsymbol{\Phi}}{}_B^{\top}\boldsymbol{H}\tilde{\boldsymbol{\Phi}}_B)\label{eq:loss-2}\end{equation}

The first two terms should raise no questions; the third term is harder to understand. Like the covariance matrix, the Hessian here, $\boldsymbol{H}$, is a 4th-order tensor, which is troublesome to interpret directly.

The simplest way in is to adopt the linear-operator viewpoint: treat $\boldsymbol{H}$ as a linear operator whose input and output are both matrices. We don't need to know what $\boldsymbol{H}$ looks like, nor how $\boldsymbol{H}$ and $\tilde{\boldsymbol{\Phi}}_B$ are combined algebraically — we only need to know that $\boldsymbol{H}\tilde{\boldsymbol{\Phi}}_B$ is linear in $\tilde{\boldsymbol{\Phi}}_B$. This way, the objects we're dealing with remain matrices, and there's no extra mental burden. Any linear operator satisfying this condition can serve as an approximation to the Hessian, without needing to write out the explicit higher-order tensor form.

The protagonist of this post is Muon, so we take $\tilde{\boldsymbol{\Phi}}_B=\newcommand{msign}{\mathop{\text{msign}}}\msign(\tilde{\boldsymbol{G}}_B)$ as our approximation for computation. By definition, we can write $\msign(\tilde{\boldsymbol{G}}_B)=\tilde{\boldsymbol{G}}_B(\tilde{\boldsymbol{G}}{}_B^{\top}\tilde{\boldsymbol{G}}_B)^{-1/2}$; from a Newton's-method perspective, this amounts to assuming $\boldsymbol{H}^{-1}\boldsymbol{X} = \eta_{\max}\boldsymbol{X}(\boldsymbol{G}^{\top}\boldsymbol{G})^{-1/2}$, giving us $\boldsymbol{H}\boldsymbol{X} = \eta_{\max}^{-1}\boldsymbol{X}(\boldsymbol{G}^{\top}\boldsymbol{G})^{1/2}$, which we'll use in the calculations below.

Computing the Expectation

Taking the expectation of both sides of $\eqref{eq:loss-2}$, we get

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

First let's compute $\mathbb{E}[\tilde{\boldsymbol{\Phi}}_B]$:

\begin{equation}\mathbb{E}[\tilde{\boldsymbol{\Phi}}_B]=\mathbb{E}[\tilde{\boldsymbol{G}}_B(\tilde{\boldsymbol{G}}{}_B^{\top}\tilde{\boldsymbol{G}}_B)^{-1/2}]\approx\mathbb{E}[\tilde{\boldsymbol{G}}_B](\mathbb{E}[\tilde{\boldsymbol{G}}{}_B^{\top}\tilde{\boldsymbol{G}}_B])^{-1/2} = \boldsymbol{G}(\mathbb{E}[\tilde{\boldsymbol{G}}{}_B^{\top}\tilde{\boldsymbol{G}}_B])^{-1/2}\end{equation}

Writing out $\mathbb{E}[\tilde{\boldsymbol{G}}{}_B^{\top}\tilde{\boldsymbol{G}}_B]$ component-wise, and assuming independence between different components, we get

\begin{equation}\mathbb{E}[\tilde{\boldsymbol{G}}{}_B^{\top}\tilde{\boldsymbol{G}}_B]_{i,j} = \mathbb{E}\left[\sum_{k=1}^n (\tilde{G}_B)_{k,i}(\tilde{G}_B)_{k,j}\right] = \left\{\begin{aligned} \mathbb{E}\left[\sum_{k=1}^n (\tilde{G}_B)_{k,i}^2\right] = \left(\sum_{k=1}^n G_{k,i}^2\right) + n\sigma^2/B,\quad (i=j) \\[6pt] \sum_{k=1}^n \mathbb{E}[(\tilde{G}_B)_{k,i}] \mathbb{E}[(\tilde{G}_B)_{k,j}] = \sum_{k=1}^n G_{k,i}G_{k,j},\quad (i\neq j) \end{aligned}\right.\end{equation}

Combining these gives $\mathbb{E}[\tilde{\boldsymbol{G}}{}_B^{\top}\tilde{\boldsymbol{G}}_B]=\boldsymbol{G}^{\top}\boldsymbol{G} + (n\sigma^2/B) \boldsymbol{I}$, so

\begin{equation}\mathbb{E}[\tilde{\boldsymbol{\Phi}}_B]\approx \boldsymbol{G}(\boldsymbol{G}^{\top}\boldsymbol{G} + (n\sigma^2/B) \boldsymbol{I})^{-1/2} = \msign(\boldsymbol{G})(\boldsymbol{I} + (n\sigma^2/B) (\boldsymbol{G}^{\top}\boldsymbol{G})^{-1})^{-1/2}\end{equation}

To further simplify the dependence of $B$, we approximate $\boldsymbol{G}^{\top}\boldsymbol{G}$ with $\tr(\boldsymbol{G}^{\top}\boldsymbol{G})\boldsymbol{I}/m$ — that is, we keep only the diagonal part of $\boldsymbol{G}^{\top}\boldsymbol{G}$, and then replace the diagonal entries with their average. This gives us

\begin{equation}\mathbb{E}[\tilde{\boldsymbol{\Phi}}_B]\approx \msign(\boldsymbol{G})(1 + \mathcal{B}_{\text{simple}}/B)^{-1/2}\end{equation}

where $\mathcal{B}_{\text{simple}} = mn\sigma^2/\tr(\boldsymbol{G}^{\top}\boldsymbol{G})= mn\sigma^2/\Vert\boldsymbol{G}\Vert_F$, which is essentially the same as treating $\boldsymbol{G}$ as a vector and computing $\mathcal{B}_{\text{simple}}$ from the previous two posts. The form of the above expression is exactly the same as SignSGD, from which we might guess that Muon won't show much that's new in terms of the learning-rate-batch-size relationship.

The Same Pattern

As for $\mathbb{E}[\tr(\tilde{\boldsymbol{\Phi}}{}_B^{\top}\boldsymbol{H}\tilde{\boldsymbol{\Phi}}_B)]$, we only compute it for the specific assumption for Muon derived above, namely $\boldsymbol{H}\boldsymbol{X} = \eta_{\max}^{-1}\boldsymbol{X}(\boldsymbol{G}^{\top}\boldsymbol{G})^{1/2}$, giving

\begin{equation}\tr(\tilde{\boldsymbol{\Phi}}{}_B^{\top}\boldsymbol{H}\tilde{\boldsymbol{\Phi}}_B) = \eta_{\max}^{-1}\tr(\tilde{\boldsymbol{\Phi}}{}_B^{\top}\tilde{\boldsymbol{\Phi}}_B(\boldsymbol{G}^{\top}\boldsymbol{G})^{1/2})\end{equation}

Note that $\tilde{\boldsymbol{\Phi}}_B$ is the result of $\msign$, so it must be an orthogonal matrix (full rank), meaning $\tilde{\boldsymbol{\Phi}}{}_B^{\top}\tilde{\boldsymbol{\Phi}}_B=\boldsymbol{I}$ — i.e., in this case $\tr(\tilde{\boldsymbol{\Phi}}{}_B^{\top}\boldsymbol{H}\tilde{\boldsymbol{\Phi}}_B)$ is a fixed constant $\eta_{\max}^{-1}\tr((\boldsymbol{G}^{\top}\boldsymbol{G})^{1/2})=\eta_{\max}^{-1}\msign(\boldsymbol{G})^{\top}\boldsymbol{G}$. So we obtain

\begin{equation}\eta^* \approx \frac{\tr(\mathbb{E}[\tilde{\boldsymbol{\Phi}}_B]^{\top}\boldsymbol{G})}{\mathbb{E}[\tr(\tilde{\boldsymbol{\Phi}}{}_B^{\top}\boldsymbol{H}\tilde{\boldsymbol{\Phi}}_B)]}\approx \frac{\eta_{\max}}{\sqrt{1 + \mathcal{B}_{\text{simple}}/B}}\end{equation}

Sure enough, this has exactly the same form as the SignSGD result — nothing new here.

Actually, thinking about it more carefully, this is entirely to be expected: SignSGD directly adds $\newcommand{sign}{\mathop{\text{sign}}}\sign$ to the gradient, whereas Muon's $\msign$ adds $\sign$ to the singular values. Intuitively, this is just adding $\sign$ in a different coordinate system — it introduces a new matrix update rule, but the learning rate $\eta^*$ and batch size $B$ remain scalars, and given that the core mechanism behind both is $\sign$, the asymptotic relationship between these scalars is very unlikely to change noticeably.

Of course, here we've only worked out one special case of $\boldsymbol{H}$; if we consider a more general $\boldsymbol{H}$, it's possible that, as with SignSGD, we could see the surge phenomenon where "as batch size increases, the learning rate should actually decrease." But as we noted in the "Reflections on the Cause" section of the previous post, if the surge phenomenon really is observed, perhaps what's needed is a change of optimizer rather than a correction to the relationship between $\eta^*$ and $B$.

Summary

In this post we attempted a simple analysis of Muon using the mean-field approximation, and concluded that its learning-rate-batch-size relationship matches that of SignSGD, with no new patterns emerging.