The Steepest Descent on Manifolds: 2. Muon + Orthogonality

流形上的最速下降：2. Muon + 正交 苏剑林 · 2025-08-06 · Mathematics matrix optimizer muon constraint steepest descent

This post continues our series on constrained optimization. In the previous post, The Steepest Descent on Manifolds: 1. SGD + Hypersphere, we revisited the "principle of least action" for optimizers, and proposed that the core difference between different optimizers lies in the different constraints imposed on the update. If this constraint is the Euclidean norm, then the corresponding steepest descent is SGD. Going further, we also discussed the result of additionally imposing a norm constraint on the parameters themselves, which gives us steepest descent on the hypersphere manifold.

However, the previous post was really just a "warm-up," since it dealt with the relatively simple case of optimizing vector parameters. This post moves on to a more challenging part: the optimization parameter changes from a vector to a matrix, and the constraint on the increment is changed to the spectral norm, from which the Muon optimizer emerges. We then further impose an orthogonality constraint on the parameter, which gives us the Muon optimizer on the orthogonal manifold.

Statement of the Problem

Let the parameter to be optimized be a matrix $\boldsymbol{W}\in\mathbb{R}^{n\times m}$, and without loss of generality let $n\geq m$. By the "principle of least action" from the previous post, the steepest-descent increment $\Delta\boldsymbol{W}$ should satisfy

\begin{equation}\min_{\Delta \boldsymbol{W}} \mathcal{L}(\boldsymbol{W} +\Delta\boldsymbol{W}) \qquad \text{s.t.}\qquad \rho(\Delta\boldsymbol{W})\leq \eta\end{equation}

If $\rho$ is taken to be the $F$ norm (Frobenius Norm), then we get the same result as in the previous post, since the $F$ norm is just the vector L2 norm computed by treating the matrix as a vector — so the result is again SGD, treating the matrix as a vector. To get a result that more deeply reveals and better fits the essence of matrices, here we choose the spectral norm, also known as the "$2$ norm," which we also denote $\Vert\cdot\Vert_2$.

As for why we should choose the spectral norm, see Appreciating the Muon Optimizer: The Essential Leap from Vectors to Matrices, Muon Sequel: Why Did We Choose to Try Muon?, and Higher-order MuP: A Simpler Yet More Elegant Spectral-Condition Scaling — we won't repeat that discussion here. In short, the spectral norm is the most compact norm that captures the change induced by a linear layer, which makes it a more suitable measure of "stability" for matrices.

Following the same steps as before, taking a first-order approximation of $\mathcal{L}(\boldsymbol{W} +\Delta\boldsymbol{W})$ gives $\mathcal{L}(\boldsymbol{W}) + \langle \boldsymbol{G}, \Delta\boldsymbol{W}\rangle_F$, where $\boldsymbol{G}=\nabla_{\boldsymbol{W}}\mathcal{L}(\boldsymbol{W})$ — here $\langle\cdot,\cdot\rangle_F$ is just the inner product computed by flattening both matrices into vectors, and it equals $\newcommand{tr}{\mathop{\text{tr}}}\tr(\boldsymbol{G}^{\top}\Delta\boldsymbol{W})$. Then, setting $\Delta\boldsymbol{W} = -\eta \boldsymbol{\Phi}$, the original problem can be simplified to

\begin{equation}\max_{\boldsymbol{\Phi}} \tr(\boldsymbol{G}^{\top}\boldsymbol{\Phi}) \qquad \text{s.t.}\qquad \Vert\boldsymbol{\Phi}\Vert_2 = 1\label{eq:muon-obj}\end{equation}

All the transformation steps up to this point are generic; if you've forgotten the details, please go back and check the previous post.

Basic Result

The solution process for the objective $\eqref{eq:muon-obj}$ was already given in the "Matrix Norm" section of Appreciating the Muon Optimizer: The Essential Leap from Vectors to Matrices, but for completeness we repeat it here. Let the SVD of $\boldsymbol{G}$ be $\boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^{\top} = \sum\limits_{i=1}^r \sigma_i \boldsymbol{u}_i \boldsymbol{v}_i^{\top}$, and let $r$ be the rank of $\boldsymbol{G}$. We have

\begin{equation}\tr(\boldsymbol{G}^{\top}\boldsymbol{\Phi})=\tr\left(\sum_{i=1}^r \sigma_i \boldsymbol{v}_i \boldsymbol{u}_i^{\top}\boldsymbol{\Phi}\right) = \sum_{i=1}^r \sigma_i \boldsymbol{u}_i^{\top}\boldsymbol{\Phi}\boldsymbol{v}_i\end{equation}

By definition, when $\Vert\boldsymbol{\Phi}\Vert_2=1$, $\Vert\boldsymbol{\Phi}\boldsymbol{v}_i\Vert_2\leq \Vert\boldsymbol{v}_i\Vert_2=1$, so $\boldsymbol{u}_i^{\top}\boldsymbol{\Phi}\boldsymbol{v}_i\leq 1$, and therefore

\begin{equation}\tr(\boldsymbol{G}^{\top}\boldsymbol{\Phi})\leq \sum_{i=1}^r \sigma_i = \Vert \boldsymbol{G}\Vert_*\end{equation}

Here $\Vert\cdot\Vert_*$ is called the Nuclear Norm of the matrix, and equality holds when all the $\boldsymbol{u}_i^{\top}\boldsymbol{\Phi}\boldsymbol{v}_i$ equal 1, in which case

\begin{equation}\newcommand{msign}{\mathop{\text{msign}}}\boldsymbol{\Phi} = \sum_{i=1}^r \boldsymbol{u}_i \boldsymbol{v}_i^{\top} = \boldsymbol{U}_{[:,:r]}\boldsymbol{V}_{[:,:r]}^{\top} = \msign(\boldsymbol{G})\end{equation}

Note that if $r < m$, then adding $\boldsymbol{u}_{r+1} \boldsymbol{v}_{r+1}^{\top},\boldsymbol{u}_{r+2} \boldsymbol{v}_{r+2}^{\top},\cdots$ on top would also make equality hold, meaning the solution is not unique — but in this case the terms beyond $r$ cannot be uniquely determined, so the expression above can be regarded as a deterministic, minimal solution. Readers who are interested can also try using the "big gun" — the von Neumann trace identity — to derive the general solution under the Schatten-$p$ norm, of which the spectral norm corresponds to the special case $p\to\infty$.

The Orthogonal Manifold

We have thus shown that, for matrix parameters, the fastest-descending direction under a spectral-norm constraint is also not simply the negative gradient direction $-\boldsymbol{G}$, but requires an additional $\msign$ operator, i.e. $-\msign(\boldsymbol{G})$. This is exactly the Muon optimizer used to train Kimi K2, currently one of the most competitive optimizers around — which in turn suggests that the spectral norm is indeed a very fitting stability constraint for matrices.

Of course, everything up to this point is still old news; now let's start tinkering with something new — adding an orthogonality constraint $\boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{I}$ to the parameter $\boldsymbol{W}$ (source: Orthogonal manifold). This splits into two cases: first, $n=m$, in which case $\boldsymbol{W}$ is a genuine orthogonal matrix satisfying $\boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{W}\boldsymbol{W}^{\top}=\boldsymbol{I}$; second, $n > m$, in which case $\boldsymbol{W}\boldsymbol{W}^{\top}=\boldsymbol{I}$ cannot be satisfied — this is usually called a semi-orthogonal matrix, and the corresponding space is called the Stiefel manifold.

Specifically, the problem we now want to solve is:

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

Continuing with the principle that "a first-order approximation is good enough," the last condition can be simplified to $\boldsymbol{W}^{\top}\boldsymbol{\Phi}+\boldsymbol{\Phi}^{\top}\boldsymbol{W} = \boldsymbol{0}$, i.e., $\boldsymbol{W}^{\top}\boldsymbol{\Phi}$ is an antisymmetric matrix:

\begin{equation}\max_{\boldsymbol{\Phi}} \tr(\boldsymbol{G}^{\top}\boldsymbol{\Phi}) \qquad \text{s.t.}\qquad \Vert\boldsymbol{\Phi}\Vert_2 = 1,\,\, \boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{I},\,\,\boldsymbol{W}^{\top}\boldsymbol{\Phi}+\boldsymbol{\Phi}^{\top}\boldsymbol{W} = \boldsymbol{0}\label{eq:muon-obj-orth}\end{equation}

When would we need an orthogonality constraint? In fact there's no shortage of scenarios. For example, in classification problems, if we know in advance that the classes are largely uncorrelated, we might want to impose an orthogonality constraint on the class matrix — though in practice this is often approximated via a regularization term $\Vert\boldsymbol{W}^{\top}\boldsymbol{W}-\boldsymbol{I}\Vert_F^2$ added to the model. Another example is in the LoRA setting, where the parameterization $\boldsymbol{A}\boldsymbol{B}$ actually has redundancy, which can be reduced via an orthogonality constraint (see reference), and so on.

The Solution Process

To solve the objective $\eqref{eq:muon-obj-orth}$, similar to the previous post, we introduce an undetermined coefficient matrix $\boldsymbol{\Lambda}\in\mathbb{R}^{m\times m}$, giving

\begin{equation}\begin{aligned} \tr(\boldsymbol{G}^{\top}\boldsymbol{\Phi}) =&\, \tr(\boldsymbol{G}^{\top}\boldsymbol{\Phi}) + \tr(\boldsymbol{\Lambda}^{\top}(\boldsymbol{W}^{\top}\boldsymbol{\Phi}+\boldsymbol{\Phi}^{\top}\boldsymbol{W})) \\ =&\, \tr((\boldsymbol{G} + \boldsymbol{W}(\boldsymbol{\Lambda} + \boldsymbol{\Lambda}^{\top}))^{\top}\boldsymbol{\Phi}) \\ \leq &\,\Vert\boldsymbol{G} + \boldsymbol{W}(\boldsymbol{\Lambda} + \boldsymbol{\Lambda}^{\top})\Vert_* \end{aligned}\end{equation}

The second equality uses the trace identity \begin{equation}\tr(\boldsymbol{A}\boldsymbol{B}) = \tr(\boldsymbol{B}\boldsymbol{A}) = \tr(\boldsymbol{A}^{\top}\boldsymbol{B}^{\top}) = \tr(\boldsymbol{B}^{\top}\boldsymbol{A}^{\top})\end{equation}. By the result for Muon in the previous section, the condition for equality is

\begin{equation}\boldsymbol{\Phi} = \msign(\boldsymbol{G} + \boldsymbol{W}(\boldsymbol{\Lambda} + \boldsymbol{\Lambda}^{\top}))\end{equation}

What remains is to find a real symmetric matrix $\boldsymbol{X} = \boldsymbol{\Lambda} + \boldsymbol{\Lambda}^{\top}$ such that $\boldsymbol{W}^{\top}\boldsymbol{\Phi}$ is antisymmetric. This is quite easy to solve when $n=m$, since in that case $\boldsymbol{W}^{\top}$ can be absorbed into $\msign$:

\begin{equation}\boldsymbol{W}^{\top}\boldsymbol{\Phi} = \boldsymbol{W}^{\top}\msign(\boldsymbol{G} + \boldsymbol{W}\boldsymbol{X}) = \msign(\boldsymbol{W}^{\top}(\boldsymbol{G} + \boldsymbol{W}\boldsymbol{X})) = \msign(\boldsymbol{W}^{\top}\boldsymbol{G} +\boldsymbol{X})\end{equation}

Note that $\msign$ has another property: it preserves antisymmetry, i.e., if the square matrix $\boldsymbol{M}$ is antisymmetric, then so is $\msign(\boldsymbol{M})$ (try proving this yourself). So, to make $\boldsymbol{W}^{\top}\boldsymbol{\Phi}$ antisymmetric, it suffices to make $\boldsymbol{W}^{\top}\boldsymbol{G} +\boldsymbol{X}$ antisymmetric. Note that $\boldsymbol{X}$ is symmetric, so this amounts to decomposing $\boldsymbol{W}^{\top}\boldsymbol{G}$ into the sum of a symmetric matrix and an antisymmetric matrix, for which there's a ready-made answer:

\begin{equation}\boldsymbol{W}^{\top}\boldsymbol{G} = \underbrace{\frac{1}{2}(\boldsymbol{W}^{\top}\boldsymbol{G} + \boldsymbol{G}^{\top}\boldsymbol{W})}_{[\boldsymbol{W}^{\top}\boldsymbol{G}] _{\text{sym}}} + \underbrace{\frac{1}{2}(\boldsymbol{W}^{\top}\boldsymbol{G} - \boldsymbol{G}^{\top}\boldsymbol{W})}_{[\boldsymbol{W}^{\top}\boldsymbol{G}] _{\text{skew}}} \end{equation}

where $[\boldsymbol{M}]_{\text{sym}} = (\boldsymbol{M}+\boldsymbol{M}^{\top})/2, [\boldsymbol{M}]_{\text{skew}} = (\boldsymbol{M}-\boldsymbol{M}^{\top})/2$. Based on this identity we can directly obtain $\boldsymbol{X} = -[\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{sym}}$. As for the solution when $n > m$, that's considerably more complicated, and we'll leave it for detailed discussion in the next post — this post aims to fully resolve the case $n=m$.

The Retraction Operation

To summarize, when $n=m$, the final result we obtain is

\begin{equation}\begin{aligned} \boldsymbol{\Phi} =&\, \msign(\boldsymbol{G} - \boldsymbol{W}[\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{sym}}) \\ =&\, \boldsymbol{W}\boldsymbol{W}^{\top}\msign(\boldsymbol{G} - \boldsymbol{W}[\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{sym}}) \\ =&\, \boldsymbol{W}\msign(\boldsymbol{W}^{\top}\boldsymbol{G} - [\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{sym}}) \\ =&\, \boldsymbol{W}\msign([\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{skew}}) \\ \end{aligned}\end{equation}

so the new variable is

\begin{equation}\boldsymbol{W} - \eta \boldsymbol{\Phi} = \boldsymbol{W}(\boldsymbol{I} - \eta\,\underbrace{\msign([\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{skew}})}_{\text{denote}\boldsymbol{O}})\label{eq:updated-W}\end{equation}

This is not itself an orthogonal matrix, but it is accurate up to $\mathcal{O}(\eta^2)$, which is consistent with our "a first-order approximation is good enough" principle. To see this, we just need to verify

\begin{equation}\begin{aligned} (\boldsymbol{I} - \eta\boldsymbol{O})^{\top}\boldsymbol{W}^{\top}\boldsymbol{W}(\boldsymbol{I} - \eta\boldsymbol{O}) =&\,(\boldsymbol{I} - \eta\boldsymbol{O})^{\top}(\boldsymbol{I} - \eta\boldsymbol{O}) \\ =&\,\boldsymbol{I} - \eta(\boldsymbol{O}^{\top} + \boldsymbol{O}) + \eta^2\boldsymbol{O}^{\top}\boldsymbol{O} \\ =&\,\boldsymbol{I} + \eta^2\boldsymbol{O}^{\top}\boldsymbol{O} \\ \end{aligned}\label{eq:orth-check}\end{equation}

If $[\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{skew}}$ is full rank, then $\boldsymbol{O}$ is an orthogonal matrix, i.e., $\boldsymbol{O}^{\top}\boldsymbol{O}=\boldsymbol{I}$, and in that case simply dividing by $\sqrt{1+\eta^2}$ suffices to make $\eqref{eq:updated-W}$ satisfy orthogonality. However, when $[\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{skew}}$ is not full rank, there's no simple transformation that will make it satisfy orthogonality — the standard approach in that case is to find the nearest orthogonal matrix, which is exactly what $\msign$ does (see here)! So the complete update rule is

\begin{equation}\boldsymbol{W} \quad \leftarrow\quad \msign(\boldsymbol{W} - \eta \boldsymbol{\Phi}) = \msign(\boldsymbol{W}(\boldsymbol{I} - \eta\boldsymbol{O})) = \boldsymbol{W}\msign(\boldsymbol{I} - \eta\boldsymbol{O})\end{equation}

But this requires computing $\msign$ twice, so let's try to simplify it a bit. From the definition and equation $\eqref{eq:orth-check}$ we get

\begin{equation}\msign(\boldsymbol{I} - \eta\boldsymbol{O}) = (\boldsymbol{I} - \eta\boldsymbol{O})(\boldsymbol{I} + \eta^2\boldsymbol{O}^{\top}\boldsymbol{O})^{-1/2}\end{equation}

Note that regardless of whether it's full rank or not, we always have $(\boldsymbol{O}^{\top}\boldsymbol{O})^2 = \boldsymbol{O}^{\top}\boldsymbol{O}$. Letting $(1+\eta^2 x)^{-1/2}=1 + a_1 x + a_2 x^2 + a_2 x^3 + \cdots $, we get

\begin{equation}\begin{aligned} (\boldsymbol{I} + \eta^2\boldsymbol{O}^{\top}\boldsymbol{O})^{-1/2} =&\, \boldsymbol{I} + a_1 (\boldsymbol{O}^{\top}\boldsymbol{O}) + a_2 (\boldsymbol{O}^{\top}\boldsymbol{O})^2 + a_3 (\boldsymbol{O}^{\top}\boldsymbol{O})^3 + \cdots \\ =&\, \boldsymbol{I} + a_1 (\boldsymbol{O}^{\top}\boldsymbol{O}) + a_2 (\boldsymbol{O}^{\top}\boldsymbol{O}) + a_3 (\boldsymbol{O}^{\top}\boldsymbol{O}) + \cdots \\ =&\, \boldsymbol{I} - \boldsymbol{O}^{\top}\boldsymbol{O} + \underbrace{(1 + a_1 + a_2 + a_3 + \cdots)}_{(1+\eta^2 x)^{-1/2}\text{substitute}x=1}\boldsymbol{O}^{\top}\boldsymbol{O} \\ \end{aligned}\end{equation}

This eliminates one computation of $\msign$; the simplified, complete result 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)\end{equation}

Summary

In this post, we revisited the conclusion that imposing a spectral-norm constraint on the increment of a matrix parameter gives the Muon optimizer, and then explored the form of the Muon optimizer once we additionally impose an orthogonality constraint on the parameter itself. If you want your parameters to remain orthogonal matrices throughout the update process, hopefully this post offers some useful reference.