We know that when updating matrix parameters with optimizers like Adam or Muon, the singular values and the left/right singular vectors all change and are usually coupled together. Precisely because of this coupling, we cannot simply regulate the singular values of the matrix parameters. Consequently, when singular values exhibit abnormal growth, we lack a simple and 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 article proposes a Muon variant that updates the left and right singular vectors of a matrix separately—"Rotational Muon (MuonR)". it maintains the singular value distribution of the matrix unchanged, thereby ensuring training stability.

Previous Review #

Since the matrices composed of left and right singular vectors must be orthogonal matrices, let's first briefly review Muon under orthogonal constraints. Let the parameter $\boldsymbol{W}\in\mathbb{R}^{n\times n}$ satisfy $\boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{I}$. Let the update amount be $\Delta\boldsymbol{W}=-\eta \boldsymbol{\Phi}$. We want the parameters to still satisfy orthogonality after the update. The corresponding steepest descent problem in 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}
This can be solved as $\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 anti-symmetrization operator. Considering the retraction operation, 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, it 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 process can be found in "Steepest Descent on Manifolds: 2. Muon + Orthogonality", so we won't expand on it in detail here. Whether it is Equation $\eqref{eq:orth-steepest}$ or $\eqref{eq:orth-steepest-full}$, they are completely analytical and only add a few matrix multiplications on top of Muon. The complexity does not increase significantly, so this result is fully practical.

Now consider a matrix $\boldsymbol{W}\in\mathbb{R}^{n\times m}(n \geq m)$. If it also satisfies $\boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{I}$, then we say $\boldsymbol{W}$ lies on the Stiefel manifold, which is a generalization of the concept of orthogonal matrices. The above results can theoretically be extended to the Stiefel manifold, but for non-square matrices, a system of non-linear equations needs to be solved, which is difficult to implement practically. Details can be found in "Steepest Descent on Manifolds: 3. Muon + Stiefel".

Instantaneous Reparameterization #

Next, we shift our attention to an arbitrary parameter matrix $\boldsymbol{W}\in\mathbb{R}^{n\times m}$. The goal is to keep the singular values unchanged during the update process, thereby eliminating the possibility of abnormal singular value growth.

To achieve this goal, we adopt the idea of "instantaneous reparameterization": before the update starts, first reparameterize the $\boldsymbol{W}$ matrix 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}$, and both are initialized as identity matrices. In this way, $\tilde{\boldsymbol{W}}=\boldsymbol{W}$ at initialization, and by denoting $\boldsymbol{G} = \nabla_{\boldsymbol{W}}\mathcal{L}$, we can write:
\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}
Subsequently, we agree to freeze $\boldsymbol{W}$ and only update $\boldsymbol{L}$ and $\boldsymbol{R}$, while maintaining the orthogonality of $\boldsymbol{L}$ and $\boldsymbol{R}$ during the update process. Thus, the updated $\tilde{\boldsymbol{W}}$ will still have the same singular values as $\boldsymbol{W}$. Now, from the perspective of $\boldsymbol{L}$ and $\boldsymbol{R}$, the problem again becomes steepest descent on the orthogonal manifold, and since $\boldsymbol{L}$ and $\boldsymbol{R}$ are square matrices, the corresponding steepest descent has a complete analytical solution! According to Equation $\eqref{eq:orth-steepest-full}$, the update rules can be written directly as:
\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}})$. By multiplying the new $\boldsymbol{L}$ and $\boldsymbol{R}$ with $\boldsymbol{W}$, we get the complete update rule:
\begin{equation}\boldsymbol{W} \quad \leftarrow\quad \boldsymbol{L}\boldsymbol{W}\boldsymbol{R}\end{equation}
This is the "Rotational Muon (Muon under Rotation, MuonR)" derived from the idea of "instantaneous reparameterization". In practical scenarios, there is usually momentum, which we understand as the smoothed gradient, so we only need to replace the gradient $\boldsymbol{G}$ with momentum $\boldsymbol{M}$.

Some Details #

Since the updates for both $\boldsymbol{L}$ and $\boldsymbol{R}$ require calculating $\msign$, the computational cost of MuonR is double that of Muon even in the most ideal case ($n=m$). However, for a sufficiently large model, this doubling of computation has a very slight impact on the end-to-end training time and is usually acceptable. If you want to reduce this overhead, you could consider alternating updates between $\boldsymbol{L}$ and $\boldsymbol{R}$ to amortize the computation.

In fact, the biggest problem with MuonR is that it maintains the entire set of singular values of the matrix from start to finish. This means we must determine the entire set of singular values for the parameters at initialization. This is not easy because matrices at different positions may require different scales; forcing them all to the same set of values is likely sub-optimal.

A feasible approach is to add an element-wise multiplication vector before or after each matrix on the basis of appropriate random initialization to compensate for the scale's degrees of freedom. For matrices immediately following an RMSNorm, the gamma parameter inherent in RMSNorm already plays this role, so this operation can be omitted for those matrices.

As for how to select the initial singular values of the matrix, we can consider conventional random initialization or constructing them in the form of Zipf's law. Furthermore, we can also try to adjust the singular value entropy to the optimal entropy calculated in "Is Higher Singular Value Entropy of Matrix Parameters Always Better?", in hopes of achieving better results.

Of course, if we can indeed determine the singular values of the matrix beforehand—for example, if we expect certain parameters to always remain orthogonal—then there's no need to consider these issues; just use MuonR directly.

Midway Switching #

Another alternative is "midway switching," using MuonR only as a means of "stabilization."

Specifically, we start with conventional Muon and monitor the spectral norm/Frobenius norm of the matrix. Once the norm of the matrix exceeds our desired range, we switch to MuonR. Since both types of Muon rely on the same gradient/momentum and only the calculation differs, this switch is permissible. MuonR does not change the singular values of the matrix, so neither the spectral norm nor the Frobenius norm will grow further, making it useful as a "stabilization" tool.

However, it should be noted that we need to align the update magnitudes before and after the switch as much as possible to avoid introducing "mutations." 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$ and $\boldsymbol{O}_R$ do not 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$ and $\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$, 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}
For conventional Muon, the update is $\boldsymbol{W} - \eta \msign(\boldsymbol{G})$, and the Frobenius norm of $\msign(\boldsymbol{G})$ is generally $\sqrt{\min(n,m)}$. Therefore, to align the update F-norm, when switching from Muon to MuonR, the learning rate should roughly be multiplied by $\frac{\sqrt{\min(n,m)}}{2\Vert\boldsymbol{W}\Vert_F}$.

In practice, the first inequality above might not be tight; $\boldsymbol{O}_L \boldsymbol{W}$ and $\boldsymbol{W} \boldsymbol{O}_R$ are more likely to be nearly orthogonal. According to the Pythagorean theorem, the result should be approximately $\sqrt{2}\Vert\boldsymbol{W}\Vert_F$, so this multiple should be multiplied by an additional $\sqrt{2}$. However, considering that $\sqrt{2}$ is not particularly different from $1$, and to ensure usability in extreme cases, it is suggested to keep the form mentioned above.

Similarity and Difference Analysis #

As stated at the beginning, MuonR is inspired by Pion. Let's look at the connections and differences between them.

First, the idea of restricting the update rule to a double rotation form by multiplying left and right orthogonal matrices is mainly from Pion. Once this update form was determined, obtaining the corresponding gradients through "instantaneous reparameterization" is quite natural. Then, Pion and MuonR start to "part ways":

1. Pion achieves orthogonality through the matrix exponential $\exp(\text{skew-symmetric matrix})$, and in actual calculation, it approximates it by expanding to the second order;
2. Pion follows the Adam route and performs a moving average of the gradients of $\boldsymbol{L}$ and $\boldsymbol{R}$ separately, which results in having 4 sets of buffer variables;
3. MuonR follows the Muon route, and like Muon, it only buffers momentum, which allows us to switch with Muon at any time;
4. MuonR is based on the analytical solution for steepest descent on an orthogonal manifold, requiring only a few extra steps to achieve orthogonality precisely.

Overall, Pion's orthogonality design is more empirical, and the four sets of buffer variables are somewhat "daunting." MuonR is a relatively natural product of a series of works like Muon and steepest descent on orthogonal manifolds. In the author's opinion, it fits the first principles better overall.

Article Summary #