By 苏剑林 | June 03, 2026
In the article "Muon Optimizer Guide: Quick Start and Key Details", we listed several versions of Muon. Their difference lies in the scaling factor related to the matrix shape of the learning rate. Among them, the "official version (Keller Jordan version)" differs from the "MuP version" only by an additional $\max(1,\cdot)$ truncation operation. This article will specifically discuss: where does this truncation operation come from?
The update rule of Muon can be unified as: \begin{align} \boldsymbol{M}_t =&\, \beta \boldsymbol{M}_{t-1} + \boldsymbol{G}_t \\[5pt] \boldsymbol{W}_t =&\, \boldsymbol{W}_{t-1} - \eta_t (\alpha \mathop{\text{msign}}(\boldsymbol{M}_t) + \lambda \boldsymbol{W}_{t-1}) \end{align} The difference between several versions lies in $\alpha$, which are: $$\alpha = \left\{ \begin{aligned} &1 & \color{skyblue}{(\text{Naive Version})} \\[5pt] & \sqrt{\max(1, d_{out}/d_{in})} & \color{skyblue}{(\text{Keller Jordan 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.$$
Where the matrix $\boldsymbol{W}\in\mathbb{R}^{d_{in}\times d_{out}}$ represents the training parameters of the linear layer $\boldsymbol{y}=\boldsymbol{x}\boldsymbol{W}$, and the input $\boldsymbol{x}\in\mathbb{R}^{d_{in}}$ is a row vector.
This article mainly focuses on the "Keller Jordan version" and the "MuP version." The former adds a $\max(1,)$ on top of the latter. According to our analysis in articles like "Higher-Order MuP: Simpler but Smarter Spectral Condition Scaling" and "Beyond MuP: 2. Linear Layers and Steepest Descent", under the spectral condition constraints related to MuP, the steepest descent should be exactly the MuP version of Muon. How can we explain the extra $\max(1,\cdot)$?
For simplicity, the following discussion omits the subscript $t$. Without loss of generality, we assume the momentum $\boldsymbol{M}$ is full rank, so the singular values of $\boldsymbol{\Phi} = \mathop{\text{msign}}(\boldsymbol{M})$ are all 1. Thus, when $d_{in} \leq d_{out}$, $\boldsymbol{\Phi} \boldsymbol{\Phi}^{\top} = \boldsymbol{I}_{d_{in}}$, and when $d_{in} > d_{out}$, $\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? Understanding Based on Steepest Descent", we know that parameters are actually just a byproduct of the model; 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}$. Then $\Vert\Delta \boldsymbol{y}\Vert_{RMS} = \eta\alpha \Vert\boldsymbol{x}\boldsymbol{\Phi}\Vert_{RMS}$.
Next, we discuss by 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 $\boldsymbol{U}\in\mathbb{R}^{d_{in}\times d_{in}}$ and $\boldsymbol{V}\in\mathbb{R}^{d_{out}\times d_{out}}$ are both orthogonal matrices. Then: \begin{align} \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{align} Note that these are equalities throughout, so we only need to set $\alpha = \sqrt{d_{out}/d_{in}}$ to make the RMS of "every" $\Delta \boldsymbol{y}$ equal to $\eta\Vert\boldsymbol{x}\Vert_{RMS}$, meaning the relative update magnitude for all tokens is consistent.
Unfortunately, for the second case $d_{in} > d_{out}$, the objective of "complete consistency" cannot be achieved. Specifically, the SVD of $\boldsymbol{\Phi}$ is written as $\boldsymbol{U}\begin{bmatrix}\boldsymbol{I}_{d_{out}} \\ \boldsymbol{0}_{(d_{in}-d_{out})\times d_{out}}\end{bmatrix}\boldsymbol{V}^{\top}$. Therefore: \begin{align} \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{align} $\boldsymbol{x}\boldsymbol{U}$ is a $d_{in}$-dimensional vector, and $d_{in} > d_{out}$, so $(\boldsymbol{x}\boldsymbol{U})_{[:d_{out}]}$ just truncates the first $d_{out}$ dimensions of $\boldsymbol{x}\boldsymbol{U}$ to calculate the RMS. At this point, its RMS is uncertain; it can reach a maximum of $\sqrt{d_{in}/d_{out}}\Vert\boldsymbol{x}\Vert_{RMS}$ (Worst Case) and a minimum of 0.
We know that orthogonal matrices do not change the RMS, so $\Vert\boldsymbol{x}\boldsymbol{U}\Vert_{RMS}=\Vert\boldsymbol{x}\Vert_{RMS}$. When the distribution of $\boldsymbol{x}$ is sufficiently isotropic, we can consider that the average scale of each component of $\boldsymbol{x}\boldsymbol{U}$ is $\Vert\boldsymbol{x}\Vert_{RMS}$. Thus, taking the RMS of the first $d_{out}$ components is also approximately equal to $\Vert\boldsymbol{x}\Vert_{RMS}$ on average, 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 the previous section.
Combining the results of the previous two sections, we get: \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 the Keller Jordan version of Muon.
However, the conclusion of the previous section depends on the assumption that the input $\boldsymbol{x}$ is sufficiently isotropic. This might hold approximately 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 this situation, the average approximation $\Vert\Delta \boldsymbol{y}\Vert_{RMS}\approx \eta\alpha\Vert\boldsymbol{x}\Vert_{RMS}$ becomes less accurate, and the maximum value $\eta\alpha\sqrt{d_{in}/d_{out}}\Vert\boldsymbol{x}\Vert_{RMS}$ becomes more accurate.
In this case, the $\alpha$ that makes $\Vert\Delta \boldsymbol{y}\Vert_{RMS}\approx \eta\Vert\boldsymbol{x}\Vert_{RMS}$ is $\sqrt{d_{out}/d_{in}}$, which is consistent with the conclusion when $d_{in} \leq d_{out}$, reclaiming the MuP version result. In other words, for the middle and late stages of training, the MuP version of Muon is more scientific. Regarding this inconsistency, we have two strategies: one is to always use the MuP version of Muon, which will slightly slow down the initial convergence speed, but the middle and late stages are the "top priority" of training after all; the second 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. This achieves a gradual transition from the Keller Jordan version to the MuP version, at the cost of having one more schedule to tune.
This article primarily explains the origin of $\max(1,\cdot)$ in the Keller Jordan version from the perspective of "feature increment" uniformity.