% Source: https://kexue.fm/archives/11772
\documentclass[11pt]{article}
\usepackage[utf8]{inputenc}
\usepackage[T1]{fontenc}
\usepackage{lmodern}
\usepackage{geometry}
\usepackage{hyperref}
\usepackage{graphicx}
\usepackage{amsmath,amssymb}
\usepackage{booktabs,longtable,array}
\usepackage{xcolor}
\usepackage{listings}

\definecolor{skyblue}{rgb}{0.0, 0.47, 0.75}

\title{Why Does the Official Muon Have an Extra $\max(1, \cdot)$ Compared to the MuP Version?}
\author{Su Jianlin}
\date{2026-06-03}

\begin{document}
\maketitle

\section*{Why Does the Official Muon Have an Extra $\max(1, \cdot)$ Compared to the MuP Version?}
\addcontentsline{toc}{section}{Why Does the Official Muon Have an Extra $\max(1, \cdot)$ Compared to the MuP Version?}

In the article \href{https://kexue.fm/archives/11416}{\textit{Muon Optimizer Guide: Quick Start and Key Details}}, we listed several versions of Muon, which differ in scaling factors related to the matrix shape of the learning rate. Among them, the ``official version (KellerJordan version)'' differs from the ``MuP version'' only by an additional $\max(1,\cdot)$ truncation operation. This article will specifically discuss where this truncation comes from.

\section{Several Versions}\label{几个版本}
The update rule for Muon can be uniformly written as
\begin{equation}
\begin{aligned}
\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{aligned}
\end{equation}
The differences among the versions lie 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{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, matrix $\boldsymbol{W}\in\mathbb{R}^{d_{in}\times d_{out}}$ represents the trainable parameters of a linear layer $\boldsymbol{y}=\boldsymbol{x}\boldsymbol{W}$, with input $\boldsymbol{x}\in\mathbb{R}^{d_{in}}$ being a row vector.

We mainly focus on the ``KellerJordan version'' and the ``MuP version''. The former adds a $\max(1, \cdot)$ on top of the latter. According to our analysis in articles such as \href{https://kexue.fm/archives/10795}{\textit{Higher-Order MuP: Simpler but Smarter Spectral Condition Scaling}} and \href{https://kexue.fm/archives/11605}{\textit{Beyond MuP: 2. Linear Layers and Steepest Descent}}, under spectral condition constraints related to MuP, the steepest descent should correspond exactly to the MuP version of Muon. How then should we interpret the extra $\max(1,\cdot)$?

\section{Feature Increment}\label{特征增量}
For simplicity, we omit the subscript $t$ in the following discussion. Without loss of generality, assume the momentum $\boldsymbol{M}$ is full-rank. Then $\boldsymbol{\Phi} = \mathop{\text{msign}}(\boldsymbol{M})$ has all singular values equal to 1. Thus, when $d_{in} \leq d_{out}$, $\boldsymbol{\Phi} \boldsymbol{\Phi}^{\top} = \boldsymbol{I}_{d_{in}}$; when $d_{in} > d_{out}$, $\boldsymbol{\Phi}^{\top} \boldsymbol{\Phi} = \boldsymbol{I}_{d_{out}}$.

Let $\Delta \boldsymbol{W} = \eta\alpha \boldsymbol{\Phi}$. Our task is to find the relationship between $\alpha$ and $d_{in}, d_{out}$. As discussed in \href{https://kexue.fm/archives/11549}{\textit{Why Do We Prefer Isotropy? A Steepest Descent Perspective}}, the parameters are actually by-products of the model; changes at the feature level might be more fundamental. Translating $\Delta \boldsymbol{W}$ to the feature level gives $\Delta \boldsymbol{y} = \boldsymbol{x} \Delta\boldsymbol{W} = \eta\alpha \boldsymbol{x}\boldsymbol{\Phi}$, so $\Vert\Delta \boldsymbol{y}\Vert_{RMS} = \eta\alpha \Vert\boldsymbol{x}\boldsymbol{\Phi}\Vert_{RMS}$.

Now we consider cases separately. First, when $d_{in} \leq d_{out}$, $\boldsymbol{\Phi}$ can be written as $\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}}$, $\boldsymbol{V}\in\mathbb{R}^{d_{out}\times d_{out}}$ are orthogonal matrices. Then
\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 all steps are equalities. Therefore, by setting $\alpha = \sqrt{d_{out}/d_{in}}$, we ensure that the RMS of \textit{every} $\Delta \boldsymbol{y}$ is $\eta\Vert\boldsymbol{x}\Vert_{RMS}$, i.e., the relative update magnitude is consistent across all tokens.

\section{Isotropy}\label{各向同性}
Unfortunately, for the second case $d_{in} > d_{out}$, the above goal of ``perfect consistency'' cannot be achieved. Specifically, in this case, 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}$. Thus,
\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}
Here, $\boldsymbol{x}\boldsymbol{U}$ is a $d_{in}$-dimensional vector with $d_{in} > d_{out}$, so $(\boldsymbol{x}\boldsymbol{U})_{[:d_{out}]}$ only takes the first $d_{out}$ dimensions to compute the RMS. Its RMS is not fixed—it can be as large as $\sqrt{d_{in}/d_{out}}\Vert\boldsymbol{x}\Vert_{RMS}$ (worst case) and as small as 0.

Since orthogonal matrices preserve RMS, $\Vert\boldsymbol{x}\boldsymbol{U}\Vert_{RMS}=\Vert\boldsymbol{x}\Vert_{RMS}$. When the distribution of $\boldsymbol{x}$ is sufficiently isotropic, we can assume that each component of $\boldsymbol{x}\boldsymbol{U}$ has an average magnitude of $\Vert\boldsymbol{x}\Vert_{RMS}$. Thus, taking the RMS of the first $d_{out}$ components would also be approximately $\Vert\boldsymbol{x}\Vert_{RMS}$, meaning $\Vert\Delta \boldsymbol{y}\Vert_{RMS}\approx \eta\alpha\Vert\boldsymbol{x}\Vert_{RMS}$. Therefore, setting $\alpha = 1$ achieves an effect similar to the previous section.

\section{Anisotropy}\label{各向异性}
Combining the results from the previous two sections, we obtain
\begin{equation}
\alpha = \sqrt{\max\left(1, \frac{d_{out}}{d_{in}}\right)}
\end{equation}
This is precisely the $\max(1,\cdot)$ that appears in the KellerJordan version of Muon.

However, the conclusion in the previous section relies on the assumption that the input $\boldsymbol{x}$ is sufficiently isotropic. This assumption may approximately hold in the early stages of training, but as training progresses, the feature distribution becomes increasingly anisotropic, tending toward the ``worst case'' that maximizes $\Vert\Delta \boldsymbol{y}\Vert_{RMS}$. In such cases, the average approximation $\Vert\Delta \boldsymbol{y}\Vert_{RMS}\approx \eta\alpha\Vert\boldsymbol{x}\Vert_{RMS}$ becomes less accurate, while the maximum value $\eta\alpha\sqrt{d_{in}/d_{out}}\Vert\boldsymbol{x}\Vert_{RMS}$ becomes more accurate.

Under these conditions, the $\alpha$ that maintains $\Vert\Delta \boldsymbol{y}\Vert_{RMS}\approx \eta\Vert\boldsymbol{x}\Vert_{RMS}$ is $\sqrt{d_{out}/d_{in}}$, which is consistent with the result when $d_{in} \leq d_{out}$, yielding the MuP version again. That is, during the middle and later stages of training, the MuP version of Muon is more appropriate.

To resolve this inconsistency, we have two strategies: 
\begin{enumerate}
    \item Use the MuP version of Muon throughout training. This may slightly reduce convergence speed in the early phase, but the middle and later stages are the ``most critical'' part of training.
    \item Modify 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 smooth transition from the KellerJordan version to the MuP version, at the cost of introducing an additional schedule to tune.
\end{enumerate}

\section{Summary}\label{文章小结}
This article mainly explains the origin of the $\max(1,\cdot)$ in the KellerJordan version from the perspective of uniformity in ``feature increments''.

\textcolor{orange}{\textit{\textbf{Please include this link when reposting: }}\href{https://kexue.fm/archives/11772}{https://kexue.fm/archives/11772}}\par

\textcolor{orange}{\textit{\textbf{For more details about reposting, please see: }}\href{https://kexue.fm/archives/6508}{\textit{Scientific Space FAQ}}}

\end{document}
