Why does the official Muon have an extra max(1, ⋅) compared to the MuP version?

为什么官方版Muon比MuP版多出一个max(1, ⋅)？ 苏剑林 · 2026-06-03 · Machine Learning learning rate optimizer muon MuP

In the post Muon Optimizer Guide: Quick Start and Key Details, we listed several versions of Muon, which differ in the matrix-shape-dependent scaling factor applied to the learning rate. Among them, the "official version (Keller Jordan's version)" differs from the "MuP version" only by an extra truncation operation, $\max(1,\cdot)$. This post is devoted to discussing where exactly this truncation operation comes from.

Several versions

Muon's update rule can be uniformly written as

\begin{equation}\begin{aligned}\newcommand{msign}{\mathop{\text{msign}}} \boldsymbol{M}_t =&\, \beta \boldsymbol{M}_{t-1} + \boldsymbol{G}_t \\[5pt] \boldsymbol{W}_t =&\, \boldsymbol{W}_{t-1} - \eta_t (\alpha \msign(\boldsymbol{M}_t) + \lambda \boldsymbol{W}_{t-1}) \end{aligned}\end{equation}more

The different versions differ in $\alpha$, which are respectively:

$$\alpha = \left\{ \begin{aligned} &1 & \color{skyblue}{(\text{naive version})} \\[5pt] & \sqrt{\max(1, d_{out}/d_{in})} & \color{skyblue}{(\text{KellerJordan version})} \\[5pt] & \sqrt{d_{out}/d_{in}} & \color{skyblue}{(\text{MuP version})} \\[5pt] & 0.2\times\sqrt{\max(d_{out},d_{in})} & \color{skyblue}{(\text{Moonlight version})} \end{aligned}\right.$$

Here the matrix $\boldsymbol{W}\in\mathbb{R}^{d_{in}\times d_{out}}$ represents the trainable parameters of the linear layer $\boldsymbol{y}=\boldsymbol{x}\boldsymbol{W}$, where the input $\boldsymbol{x}\in\mathbb{R}^{d_{in}}$ is a row vector.

This post is mainly concerned with the "Keller Jordan version" and the "MuP version"; the former adds an extra $\max(1,)$ on top of the latter. Based on the analysis in posts such as Higher-Order MuP: A Simpler yet Smarter Spectral Condition Scaling and Beyond MuP: 2. Linear Layers and Steepest Descent, under the spectral-condition constraints related to MuP, the steepest descent update should correspond exactly to the MuP version of Muon. So how should we explain this extra $\max(1,\cdot)$?

Feature increment

For simplicity, in the discussion below we will drop the subscript $t$. Without loss of generality, assume that the momentum $\boldsymbol{M}$ is full rank, so that the singular values of $\boldsymbol{\Phi} = \msign(\boldsymbol{M})$ are all 1. Then when $d_{in} \leq d_{out}$, we have $\boldsymbol{\Phi} \boldsymbol{\Phi}^{\top} = \boldsymbol{I}_{d_{in}}$, and when $d_{in} > d_{out}$, we have $\boldsymbol{\Phi}^{\top} \boldsymbol{\Phi} = \boldsymbol{I}_{d_{out}}$.

Let $\Delta \boldsymbol{W} = \eta\alpha \boldsymbol{\Phi}$; what we want to do is find the relationship between $\alpha$ and $d_{in},d_{out}$. From Why Do We Prefer Isotropy? An Understanding Based on Steepest Descent we know that parameters are actually just a byproduct of the model, and that changes at the feature level may be more fundamental. Converting $\Delta \boldsymbol{W}$ to the feature level gives $\Delta \boldsymbol{y} = \boldsymbol{x} \Delta\boldsymbol{W} = \eta\alpha \boldsymbol{x}\boldsymbol{\Phi}$, so that $\Vert\Delta \boldsymbol{y}\Vert_{RMS} = \eta\alpha \Vert\boldsymbol{x}\boldsymbol{\Phi}\Vert_{RMS}$.

Next we need to consider different cases. First, when $d_{in} \leq d_{out}$, $\boldsymbol{\Phi}$ can be written in the form $\boldsymbol{U}[\boldsymbol{I}_{d_{in}}, \boldsymbol{0}_{d_{in}\times (d_{out}-d_{in})}]\boldsymbol{V}^{\top}$, where both $\boldsymbol{U}\in\mathbb{R}^{d_{in}\times d_{in}}, \boldsymbol{V}\in\mathbb{R}^{d_{out}\times d_{out}}$ are orthogonal matrices, so that

\begin{equation}\begin{aligned} \Vert\Delta \boldsymbol{y}\Vert_{RMS} =&\, \eta\alpha\big\Vert\boldsymbol{x}\boldsymbol{U}[\boldsymbol{I}_{d_{in}}, \boldsymbol{0}_{d_{in}\times (d_{out}-d_{in})}]\boldsymbol{V}^{\top}\big\Vert_{RMS} \\[4pt] =&\, \eta\alpha\big\Vert\boldsymbol{x}\boldsymbol{U}[\boldsymbol{I}_{d_{in}}, \boldsymbol{0}_{d_{in}\times (d_{out}-d_{in})}]\big\Vert_{RMS} \\[4pt] =&\, \eta\alpha\big\Vert[\boldsymbol{x}\boldsymbol{U}, \boldsymbol{0}_{d_{out}-d_{in}}]\big\Vert_{RMS} \\[4pt] =&\, \eta\alpha\sqrt{\frac{d_{in}}{d_{out}}}\Vert\boldsymbol{x}\boldsymbol{U}\Vert_{RMS} \\[4pt] =&\, \eta\alpha\sqrt{\frac{d_{in}}{d_{out}}}\Vert\boldsymbol{x}\Vert_{RMS} \\ \end{aligned}\end{equation}

Note that everywhere above is an exact equality, so we only need to set $\alpha = \sqrt{d_{out}/d_{in}}$ in order to make the RMS of "every" $\Delta \boldsymbol{y}$ equal to $\eta\Vert\boldsymbol{x}\Vert_{RMS}$, i.e., to make the relative update magnitude consistent across all tokens.

Isotropy

Unfortunately, for the second case, $d_{in} > d_{out}$, the above goal of "exact consistency" cannot be achieved. Specifically, in this case the SVD of $\boldsymbol{\Phi}$ must be written as $\boldsymbol{U}\begin{bmatrix}\boldsymbol{I}_{d_{out}} \\ \boldsymbol{0}_{(d_{in}-d_{out})\times d_{out}}\end{bmatrix}\boldsymbol{V}^{\top}$, so that

\begin{equation}\begin{aligned} \Vert\Delta \boldsymbol{y}\Vert_{RMS} =&\, \eta\alpha\left\Vert\boldsymbol{x}\boldsymbol{U}\begin{bmatrix}\boldsymbol{I}_{d_{out}} \\ \boldsymbol{0}_{(d_{in}-d_{out})\times d_{out}}\end{bmatrix}\boldsymbol{V}^{\top}\right\Vert_{RMS} \\[5pt] =&\, \eta\alpha\left\Vert\boldsymbol{x}\boldsymbol{U}\begin{bmatrix}\boldsymbol{I}_{d_{out}} \\ \boldsymbol{0}_{(d_{in}-d_{out})\times d_{out}}\end{bmatrix}\right\Vert_{RMS} \\[5pt] =&\, \eta\alpha\big\Vert(\boldsymbol{x}\boldsymbol{U})_{[:d_{out}]}\big\Vert_{RMS} \end{aligned}\end{equation}

$\boldsymbol{x}\boldsymbol{U}$ is a $d_{in}$-dimensional vector, $d_{in} > d_{out}$, so $(\boldsymbol{x}\boldsymbol{U})_{[:d_{out}]}$ is simply computing the RMS over the first $\boldsymbol{x}\boldsymbol{U}$ dimensions of $d_{out}$, and in this case its RMS is not fixed — it can range up to $\sqrt{d_{in}/d_{out}}\Vert\boldsymbol{x}\Vert_{RMS}$ (Worst Case) at most, and down to 0 at least.

We know that orthogonal matrices preserve RMS, so $\Vert\boldsymbol{x}\boldsymbol{U}\Vert_{RMS}=\Vert\boldsymbol{x}\Vert_{RMS}$. When the distribution of $\boldsymbol{x}$ is sufficiently isotropic, we may take this to mean that the average scale of each component of $\boldsymbol{x}\boldsymbol{U}$ is $\Vert\boldsymbol{x}\Vert_{RMS}$, so taking the RMS over the first $d_{out}$ components will, on average, also be approximately $\Vert\boldsymbol{x}\Vert_{RMS}$, i.e., $\Vert\Delta \boldsymbol{y}\Vert_{RMS}\approx \eta\alpha\Vert\boldsymbol{x}\Vert_{RMS}$. Therefore we only need to take $\alpha = 1$ to achieve an effect similar to that in the previous section.

Anisotropy

Combining the results of the previous two sections, we obtain

\begin{equation}\alpha = \sqrt{\max\left(1, \frac{d_{out}}{d_{in}}\right)}\end{equation}

This is exactly the $\max(1,\cdot)$ that appears in Keller Jordan's version of Muon.

However, the conclusion of the previous section relies on the assumption that the input $\boldsymbol{x}$ is sufficiently isotropic, which may approximately hold in the early stages of training. But as training progresses, the distribution of features gradually becomes anisotropic, concentrating on the "Worst Case" that maximizes $\Vert\Delta \boldsymbol{y}\Vert_{RMS}$. In that regime, the average approximation $\Vert\Delta \boldsymbol{y}\Vert_{RMS}\approx \eta\alpha\Vert\boldsymbol{x}\Vert_{RMS}$ is no longer accurate, and it is instead the maximum value $\eta\alpha\sqrt{d_{in}/d_{out}}\Vert\boldsymbol{x}\Vert_{RMS}$ that becomes more accurate.

In this case, the $\alpha$ that makes $\Vert\Delta \boldsymbol{y}\Vert_{RMS}\approx \eta\Vert\boldsymbol{x}\Vert_{RMS}$ hold is $\sqrt{d_{out}/d_{in}}$, which agrees with the conclusion for the case $d_{in} \leq d_{out}$, recovering the MuP-version result. In other words, for the middle and later stages of training, the MuP version of Muon is more principled. To deal with this inconsistency, we have two possible strategies: one is to always use the MuP version of Muon, which slightly slows convergence in the early stage but is arguably worthwhile since the middle-to-late stage is the "main event" of training anyway; the other is to change the scaling factor to

\begin{equation}\alpha = \sqrt{\max\left(\tau_t, \frac{d_{out}}{d_{in}}\right)}\end{equation}

where $\tau_t$ monotonically decays from 1 to 0, thereby implementing a smooth transition from the Keller Jordan version to the MuP version — at the cost of introducing one more schedule to tune.

Summary

This post has mainly explained, from the perspective of the uniformity of "feature increments," where the $\max(1,\cdot)$ in the Keller Jordan version comes from.