The Steepest Descent on a Manifold: 6. Muon + Double Rotation

流形上的最速下降：6. Muon + 双旋转 苏剑林 · 2026-06-08 · Mathematics optimizer muon constraint steepest descent

We know that when using optimizers such as Adam or Muon to update matrix parameters, both the singular values and the left/right singular vectors change along with them, and they are usually coupled together. Precisely because of this coupling, we cannot simply control the singular values of a matrix parameter, so when a singular value starts to grow abnormally, we have no simple, effective way to stop it — which can lead to training failure.

Inspired by 《Pion: A Spectrum-Preserving Optimizer via Orthogonal Equivalence Transformation》 (hereafter referred to as Pion), this post proposes a Muon variant that updates the left and right singular vectors of a matrix separately — "Rotation Muon (MuonR)". It keeps the singular value spectrum of the matrix unchanged, thereby guaranteeing training stability.

Recap

Since a matrix composed of left and right singular vectors is necessarily orthogonal, let us first briefly recap Muon under an orthogonality constraint. Let the parameter be $\boldsymbol{W}\in\mathbb{R}^{n\times n}$, satisfying $\boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{I}$, and let the update be $\Delta\boldsymbol{W}=-\eta \boldsymbol{\Phi}$. We want the parameter to remain orthogonal after the update, so the corresponding steepest-descent problem under the spectral norm is

\begin{equation}\newcommand{tr}{\mathop{\text{tr}}}\max_{\boldsymbol{\Phi}} \tr(\boldsymbol{G}^{\top}\boldsymbol{\Phi}) \qquad \text{s.t.}\qquad \Vert\boldsymbol{\Phi}\Vert_2 = 1,\quad(\boldsymbol{W} - \eta \boldsymbol{\Phi})^{\top}(\boldsymbol{W} - \eta \boldsymbol{\Phi})=\boldsymbol{I}\end{equation}

We can solve this to get $\boldsymbol{\Phi} = \boldsymbol{W}\boldsymbol{O}$, where $\newcommand{msign}{\mathop{\text{msign}}}\boldsymbol{O}=\msign([\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{skew}})$, and $[\boldsymbol{X}]_{\text{skew}} = (\boldsymbol{X} - \boldsymbol{X}^{\top})/2$ is the antisymmetrization operator. Taking the retraction operation into account, the complete update rule is:

\begin{equation}\boldsymbol{W} \quad \leftarrow\quad \boldsymbol{W}(\boldsymbol{I} - \eta\boldsymbol{O})\left(\boldsymbol{I} - \boldsymbol{O}^{\top}\boldsymbol{O} + \frac{\boldsymbol{O}^{\top}\boldsymbol{O}}{\sqrt{1+\eta^2}}\right)\label{eq:orth-steepest}\end{equation}

In particular, if $[\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{skew}}$ is full rank, this simplifies to

\begin{equation}\boldsymbol{W} \quad \leftarrow\quad \frac{\boldsymbol{W}(\boldsymbol{I} - \eta\boldsymbol{O})}{\sqrt{1+\eta^2}}\label{eq:orth-steepest-full}\end{equation}

The derivation can be found in 《The Steepest Descent on a Manifold: 2. Muon + Orthogonal》, so we won't go over it in detail here. Whether it's Eq. $\eqref{eq:orth-steepest}$ or $\eqref{eq:orth-steepest-full}$, both are fully analytic, adding only a few matrix multiplications on top of Muon, without any noticeable increase in complexity — so this result is fully practical.

Now consider a matrix $\boldsymbol{W}\in\mathbb{R}^{n\times m}(n \geq m)$: if it simultaneously satisfies $\boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{I}$, we say $\boldsymbol{W}$ lies on the Stiefel manifold, which is a generalization of the concept of an orthogonal matrix. The above result can in principle be extended to the Stiefel manifold, but for non-square matrices this requires solving a nonlinear system of equations, which is rather difficult to make practical. See 《The Steepest Descent on a Manifold: 3. Muon + Stiefel》 for the details.

Instantaneous Reparameterization

Next, let's turn our attention to an arbitrary parameter matrix $\boldsymbol{W}\in\mathbb{R}^{n\times m}$, with the goal of keeping the singular values unchanged during the parameter update, thereby ruling out the possibility of abnormal singular value growth.

To achieve this, we adopt the idea of "instantaneous reparameterization": before the update begins, we reparameterize the matrix $\boldsymbol{W}$ as $\tilde{\boldsymbol{W}} = \boldsymbol{L}\boldsymbol{W}\boldsymbol{R}$, where $\boldsymbol{L}\in\mathbb{R}^{n\times n},\boldsymbol{R}\in\mathbb{R}^{m\times m}$, both initialized to the identity matrix. In this way, at initialization we have $\tilde{\boldsymbol{W}}=\boldsymbol{W}$, and by writing $\boldsymbol{G} = \nabla_{\boldsymbol{W}}\mathcal{L}$, we can express

\begin{equation}\nabla_{\boldsymbol{L}}\mathcal{L} = \boldsymbol{G}\boldsymbol{W}^{\top},\qquad \nabla_{\boldsymbol{R}}\mathcal{L} = \boldsymbol{W}^{\top}\boldsymbol{G}\end{equation}

We then stipulate that $\boldsymbol{W}$ is frozen, and only $\boldsymbol{L}$ and $\boldsymbol{R}$ are updated, while maintaining the orthogonality of $\boldsymbol{L},\boldsymbol{R}$ during the update. This way, the updated $\tilde{\boldsymbol{W}}$ still has the same singular values as $\boldsymbol{W}$. Now, viewed from the perspective of $\boldsymbol{L},\boldsymbol{R}$, the problem again becomes steepest descent on the orthogonal manifold, and this time $\boldsymbol{L},\boldsymbol{R}$ are both square matrices, so the corresponding steepest descent problem is fully analytically solvable! Based on Eq. $\eqref{eq:orth-steepest-full}$, we can directly write down the update rule

\begin{gather}\boldsymbol{L}\quad\leftarrow\quad (\boldsymbol{I} - \eta\boldsymbol{O}_L)\left(\boldsymbol{I} - \boldsymbol{O}_L^{\top}\boldsymbol{O}_L + \frac{\boldsymbol{O}_L^{\top}\boldsymbol{O}_L}{\sqrt{1+\eta^2}}\right)\\ \boldsymbol{R}\quad\leftarrow\quad (\boldsymbol{I} - \eta\boldsymbol{O}_R)\left(\boldsymbol{I} - \boldsymbol{O}_R^{\top}\boldsymbol{O}_R + \frac{\boldsymbol{O}_R^{\top}\boldsymbol{O}_R}{\sqrt{1+\eta^2}}\right)\end{gather}

where $\boldsymbol{O}_L = \msign([\boldsymbol{G}\boldsymbol{W}^{\top}]_{\text{skew}}),\boldsymbol{O}_R = \msign([\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{skew}})$. Multiplying the new $\boldsymbol{L},\boldsymbol{R}$ with $\boldsymbol{W}$ together gives us the complete update rule

\begin{equation}\boldsymbol{W} \quad \leftarrow\quad \boldsymbol{L}\boldsymbol{W}\boldsymbol{R}\end{equation}

This is exactly "Rotation Muon (Muon under Rotation, MuonR)", derived from the idea of "instantaneous reparameterization". In practical settings there is usually momentum too, which we can think of as a smoothed version of the gradient, so we simply need to replace the gradient $\boldsymbol{G}$ with the momentum $\boldsymbol{M}$.

Some Details

Since updating $\boldsymbol{L},\boldsymbol{R}$ requires computing $\msign$ once more, even in the most ideal case ($n=m$), MuonR's computational cost is twice that of Muon. However, for sufficiently large models, this doubling of compute has only a mild effect on end-to-end training time and is generally acceptable. If one wants to reduce this overhead, one option is to alternate updates of $\boldsymbol{L},\boldsymbol{R}$, spreading out the computational cost.

In fact, MuonR's biggest issue is that it keeps the entire set of singular values of the matrix unchanged throughout training, which means we must fix the entire singular value spectrum of the parameter at initialization. This is not easy, because matrices at different positions in the network may require different scales, and forcing them all to share the same set of values is likely to be suboptimal.

One feasible approach is: on top of an appropriate random initialization, add an element-wise multiplicative vector before or after each matrix to compensate for the missing scale degree of freedom. For matrices that immediately follow an RMSNorm, the gamma parameter built into RMSNorm already plays this role, so this operation can be omitted for such matrices.

As for how to choose the initial singular values of a matrix, we could consider conventional random initialization, or construct them following Zipf's law. Furthermore, we could try tuning the singular value entropy to match the optimal entropy computed in 《Is a Higher Singular Value Entropy Always Better for Matrix Parameters?》, in the hope of achieving better results.

Of course, if we can indeed determine the singular values of a matrix in advance — for instance, if we expect some parameter to always remain orthogonal — then none of this needs to be considered; we can simply apply MuonR directly.

Switching Midway

Another possible approach is "switching midway", using MuonR only as a "stabilization" measure.

Specifically, we start off using regular Muon and monitor the spectral norm/F-norm of the matrix. Once the matrix's norm exceeds the range we expect, we switch to MuonR. Since both variants of Muon depend on the same gradient/momentum and differ only in how they are computed, this kind of switch is permissible. MuonR keeps the singular values of the matrix unchanged, so the spectral norm/F-norm stops growing, which makes it a perfect tool for "stabilization".

However, we need to make sure the magnitude of the update is aligned as closely as possible before and after the switch, to avoid introducing a "jump". To this end, we consider the first-order approximation of MuonR

\begin{equation}\boldsymbol{L}\boldsymbol{W}\boldsymbol{R} \approx (\boldsymbol{I} - \eta\boldsymbol{O}_L) \boldsymbol{W} (\boldsymbol{I} - \eta\boldsymbol{O}_R) \approx \boldsymbol{W} - \eta(\boldsymbol{O}_L \boldsymbol{W} + \boldsymbol{W} \boldsymbol{O}_R)\end{equation}

Since the singular values of $\boldsymbol{O}_L,\boldsymbol{O}_R$ never exceed 1 (note that we cannot guarantee that $[\boldsymbol{G}\boldsymbol{W}^{\top}]_{\text{skew}}$ and $[\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{skew}}$ are both full rank, so we cannot directly use the orthogonality of $\boldsymbol{O}_L,\boldsymbol{O}_R$), we have $\Vert\boldsymbol{O}_L \boldsymbol{W}\Vert_F \leq \Vert \boldsymbol{W}\Vert_F$ and $\Vert \boldsymbol{W}\boldsymbol{O}_R\Vert_F\leq \Vert\boldsymbol{W}\Vert_F$, and thus

\begin{equation}\Vert\boldsymbol{O}_L \boldsymbol{W} + \boldsymbol{W} \boldsymbol{O}_R\Vert_F \leq \Vert\boldsymbol{O}_L \boldsymbol{W}\Vert_F + \Vert\boldsymbol{W} \boldsymbol{O}_R\Vert_F \leq 2\Vert\boldsymbol{W}\Vert_F\end{equation}

Regular Muon is $\boldsymbol{W} - \eta \msign(\boldsymbol{G})$, and the F-norm of $\msign(\boldsymbol{G})$ is generally $\sqrt{\min(n,m)}$. Therefore, in order to align the F-norm of the update, when switching from Muon to MuonR, we should roughly multiply the learning rate by $\frac{\sqrt{\min(n,m)}}{2\Vert\boldsymbol{W}\Vert_F}$.

In practice, the first inequality above may not be tight enough; $\boldsymbol{O}_L \boldsymbol{W}$ and $\boldsymbol{W} \boldsymbol{O}_R$ are more nearly orthogonal to each other, so by the Pythagorean theorem, the result should be approximately equal to $\sqrt{2}\Vert\boldsymbol{W}\Vert_F$, meaning this multiplier should be scaled up further by $\sqrt{2}$. However, considering that $\sqrt{2}$ and $1$ are not particularly different, and in order to ensure usability in extreme cases, it is still recommended to keep the form above.

Comparative Analysis

At the outset of this post, we stated clearly that MuonR was inspired by Pion; let's now examine their connections and differences.

First, the idea of restricting the update rule to this double-rotation form — left- and right-multiplication by orthogonal matrices — mainly comes from Pion. Once this update form is settled on, obtaining the corresponding gradient via "instantaneous reparameterization" is a fairly natural next step. From there, Pion and MuonR "part ways":

1. Pion achieves orthogonality via the matrix exponential $\exp(\text{antisymmetric matrix})$, which in actual computation is approximated by expanding to second order;
2. Pion follows the Adam route, maintaining separate moving averages of the gradients with respect to $\boldsymbol{L},\boldsymbol{R}$, which means it requires 4 sets of cache variables;
3. MuonR follows the Muon route, caching only the momentum just like Muon, which allows us to switch between the two at any time;
4. MuonR is based on the analytic solution of steepest descent on the orthogonal manifold, requiring only a limited number of extra steps to achieve orthogonality exactly.

Overall, Pion's design for orthogonality is somewhat empirical, and its four sets of cache variables are somewhat daunting; MuonR, on the other hand, is a relatively natural outgrowth of a series of works on Muon and steepest descent on the orthogonal manifold. The author believes it is, on the whole, more consistent with first principles.

Summary

This post proposed MuonR, a Muon variant that constrains the update to be a left/right rotation matrix form. It is able to keep the singular value spectrum of the matrix unchanged, making it a simple training scheme for maintaining training stability.