Steepest Descent on Manifolds: 7. Analytical Solution for Stiefel

By 苏剑林 | August 17, 2026

Regarding orthogonal manifolds, our previous conclusion was: in the square matrix case, we can completely write out the analytical solution for the corresponding steepest descent (refer to "Steepest Descent on Manifolds: 2. Muon + Orthogonal"); but for the non-square Stiefel manifold, the steepest descent problem requires solving a system of nonlinear equations, and its analytical solution cannot be directly written out (refer to "Steepest Descent on Manifolds: 3. Muon + Stiefel").

However, recently, "Muon on the Stiefel Manifold Admits an Exact Closed-Form Update" overturned this assertion, pointing out that an explicit solution can also be written for the Stiefel manifold without needing to solve a system of equations. This article will use my own reasoning to recount the entire derivation process.

Problem Review

We won't repeat the problem background too much; let's directly state the problem to be solved: Let $\boldsymbol{W}\in\mathbb{R}^{n\times m}(n \geq m)$ be the current parameters, and $\boldsymbol{G}$ be the gradient of the objective function at $\boldsymbol{W}$. We are looking for the steepest descent direction $\boldsymbol{\Phi}\in\mathbb{R}^{n\times m}$ with step size $\eta$ under orthogonal constraints, namely

According to the "first-order approximation is sufficient" principle, expanding the last constraint and omitting the $\eta^2$ term simplifies it to

where $\Vert\cdot\Vert_2$ is the spectral norm and $\boldsymbol{I}_m$ is the $m\times m$ identity matrix. Here, changing $\Vert\boldsymbol{\Phi}\Vert_2 \leq 1$ to $\Vert\boldsymbol{\Phi}\Vert_2 = 1$ yields completely equivalent results (the optimum of a linear objective function is always reached on the boundary), but we retain $\leq$ here because $\Vert\boldsymbol{\Phi}\Vert_2 \leq 1$ is a convex set, allowing for more maneuverability in argumentation. Our previous solution result was

where $\boldsymbol{X}\in\mathbb{R}^{m\times m}$ is an anti-symmetric matrix satisfying the following matrix equation:

The solution for the $n=m$ case was found in "Steepest Descent on Manifolds: 2. Muon + Orthogonal", but for the standard Stiefel manifold where $n > m$, solving this equation is not trivial. Even numerical solutions are somewhat troublesome, leading to the assertion that "there is no analytical solution in the Stiefel case."

Weakened Version

The key to breaking the deadlock is to remove the equality constraint $\boldsymbol{W}^{\top}\boldsymbol{\Phi}+\boldsymbol{\Phi}^{\top}\boldsymbol{W} = \boldsymbol{0}$ through some explicit representation. To this end, we first introduce a weakened version of the original problem: Let $\boldsymbol{\Phi} = \boldsymbol{S}\boldsymbol{W}$, where $\boldsymbol{S}\in\mathbb{R}^{n\times n}$ is any anti-symmetric matrix satisfying $\Vert\boldsymbol{S}\Vert_2\leq 1$. Then we consider solving

The reason this is called a weakening of the original problem is that $\boldsymbol{\Phi}$ defined this way clearly satisfies the two constraints of the original proposition:

So the new problem is solved on a subset of the feasible region of the original problem, and the maximum value obtained will not exceed the maximum value of the original problem. Furthermore, using trace identities and the anti-symmetry of $\boldsymbol{S}$, the objective function can be transformed into

where $[\boldsymbol{X}]_{\text{skew}} = (\boldsymbol{X} - \boldsymbol{X}^{\top})/2$. Thus the new problem is equivalent to

This is already in the standard Muon problem shape. From Muon's results, we can directly conclude

In fact, this is exactly the form of the solution in the square matrix case and is the left-side rotation in "Steepest Descent on Manifolds: 6. Muon + Double Rotation". It is worth explaining here why we apply Muon's $\msign$ only after a series of identical transformations. This primarily leverages the conclusion that "the $\msign$ of an anti-symmetric matrix is still an anti-symmetric matrix" to ensure the final result satisfies the anti-symmetry requirement of $\boldsymbol{S}$.

Complete Equivalence

Next, we will prove that the solution to the weakened problem we solved in the previous section is actually consistent with the optimal solution of the original problem!

The reasoning is straightforward: the weakened problem is "weak" because the current parameterization $\boldsymbol{\Phi} = \boldsymbol{S}\boldsymbol{W}$ is only a subset of the original problem's feasible region; it might miss some exploration areas, thereby failing to reach the original problem's global optimum. However, if we can prove that for any feasible $\boldsymbol{\Phi}$ of the original problem, we can find an $\boldsymbol{S}$ satisfying the conditions such that $\boldsymbol{\Phi} = \boldsymbol{S}\boldsymbol{W}$, it indicates that their exploration spaces are identical, and thus their optimal solutions are identical as well.

The following proof proceeds in three steps. (Note: In hindsight, given appropriate hints, K3 can also independently complete the derivation and proof, refer here and here.)

Solving the Equation

The first thing to do is, given $\boldsymbol{\Phi}, \boldsymbol{W}$, solve the equation

Without loss of generality, consider only $n > m$. We can then complete $\boldsymbol{W}$ into an $n\times n$ orthogonal matrix $\boldsymbol{U} = [\boldsymbol{W}, \boldsymbol{W}_{\perp}]$. Since $\boldsymbol{S}$ is anti-symmetric, $\boldsymbol{U}^{\top}\boldsymbol{S}\boldsymbol{U}$ must also be, so we can write $\boldsymbol{S}$ as

where $\boldsymbol{A}\in\mathbb{R}^{m\times m}, \boldsymbol{C}\in\mathbb{R}^{(n-m)\times (n-m)}$ are both anti-symmetric matrices, and $\boldsymbol{B}\in\mathbb{R}^{(n-m)\times m}$. Multiplying both sides by $\boldsymbol{W}$ on the right gives

Multiplying both sides by $\boldsymbol{U}^{\top}$ on the left, we can directly read $\boldsymbol{A} = \boldsymbol{W}^{\top}\boldsymbol{\Phi}, \boldsymbol{B} = \boldsymbol{W}_{\perp}^{\top}\boldsymbol{\Phi}$. That is, $\boldsymbol{A}$ and $\boldsymbol{B}$ can be uniquely solved, leaving $\boldsymbol{C}$ as a degree of freedom. Note that $\boldsymbol{A}$ being anti-symmetric imposes the compatibility condition

Which is exactly one of the conditions of the original problem. So so far, all results are self-consistent.

Spectral Norm

Now the remaining question is whether we can find an anti-symmetric matrix $\boldsymbol{C}$ such that $\Vert \boldsymbol{S}\Vert_2 \leq 1$. The answer is yes, but the proof requires the use of "Parrott's Lemma" (the original paper cites Davis–Kahan–Weinberger, but research shows Parrott's Lemma is earlier):

Given a block matrix $\begin{bmatrix}\boldsymbol{A} & \boldsymbol{C} \\ \boldsymbol{B} & \boldsymbol{D}\end{bmatrix}$, with $\boldsymbol{A}, \boldsymbol{B}, \boldsymbol{C}$ fixed, we can always find a matrix $\boldsymbol{D}$ such that \begin{equation}\left\Vert\begin{bmatrix}\boldsymbol{A} & \boldsymbol{C} \\ \boldsymbol{B} & \boldsymbol{D}\end{bmatrix}\right\Vert_2 = \max\left\{\left\Vert\begin{bmatrix}\boldsymbol{A} \\ \boldsymbol{B}\end{bmatrix}\right\Vert_2, \left\Vert\begin{bmatrix}\boldsymbol{A} & \boldsymbol{C}\end{bmatrix}\right\Vert_2\right\}\end{equation}

Returning to our problem. To have $\Vert \boldsymbol{S}\Vert_2 \leq 1$, we only need $\left\Vert\begin{bmatrix}\boldsymbol{A} & -\boldsymbol{B}^{\top} \\ \boldsymbol{B} & \boldsymbol{C}\end{bmatrix}\right\Vert_2\leq 1$. For this block matrix, we have

As for $[\boldsymbol{A},-\boldsymbol{B}^{\top}] = [-\boldsymbol{A}^{\top},-\boldsymbol{B}^{\top}]$, it is the negative transpose of $\begin{bmatrix}\boldsymbol{A} \\ \boldsymbol{B}\end{bmatrix}$, so its spectral norm remains unchanged (also not exceeding 1). Therefore, according to Parrott's Lemma, we can find a matrix $\boldsymbol{C}$ such that $\Vert \boldsymbol{S}\Vert_2 = \left\Vert\begin{bmatrix}\boldsymbol{A} & -\boldsymbol{B}^{\top} \\ \boldsymbol{B} & \boldsymbol{C}\end{bmatrix}\right\Vert_2 \leq 1$.

Anti-symmetry

But we aren't done yet. We also have an anti-symmetry requirement for $\boldsymbol{C}$, which Parrott's Lemma does not guarantee.

Fortunately, this fix is not difficult. Let the matrix given by Parrott's Lemma be $\boldsymbol{C}_0$. We directly take $\boldsymbol{C} = [\boldsymbol{C}_0]_{\text{skew}}$. Since the rest of the matrix already satisfies the anti-symmetry requirement, we have

The key point here is that anti-symmetrization does not increase the spectral norm: from the triangle inequality, we have

Thus $[\boldsymbol{C}_0]_{\text{skew}}$ is exactly the anti-symmetric matrix we are looking for.

At this point, we have proven the result we initially set out to prove: for every $\boldsymbol{\Phi}$ of the original problem, we can find an $n\times n$ anti-symmetric matrix $\boldsymbol{S}$ with a spectral norm not exceeding 1 such that $\boldsymbol{\Phi}=\boldsymbol{S}\boldsymbol{W}$. Thus, the exploration space of the weakened version is identical to the exploration space of the original problem. In other words, equation \eqref{eq:sol} is the exact solution to the original problem!

Efficient Computation

Theoretically, the problem is perfectly solved; but practically, there is still an efficiency issue worth discussing. Muon requires an $\msign$ on an $n\times m$ matrix, while the exact solution \eqref{eq:sol} on the Stiefel manifold requires an $\msign$ on an $n\times n$ matrix. If computed directly, it is obviously much more expensive than Muon when $n \gg m$. Therefore, it is necessary to explore schemes to save computational load.

Actually, it's not hard to think of. First, rewrite the core operation in low-rank form

It can be seen that the rank of $[\boldsymbol{G}\boldsymbol{W}^{\top}]_{\text{skew}}$ is at most $2m$. Assuming $n \geq 2m$, first perform a QR decomposition on $[\boldsymbol{G}, \boldsymbol{W}]$ to get $\boldsymbol{Q}\boldsymbol{R}$, where $\boldsymbol{Q}\in\mathbb{R}^{n\times 2m}$ satisfies $\boldsymbol{Q}^{\top}\boldsymbol{Q}=\boldsymbol{I}_{2m}$, and $\boldsymbol{R}\in\mathbb{R}^{2m\times 2m}$ is an upper triangular matrix. Then the part to be calculated becomes $\msign(\boldsymbol{Q}\boldsymbol{R}\boldsymbol{J}\boldsymbol{R}^{\top} \boldsymbol{Q}^{\top})$. Note that the $\msign$ operation is covariant with respect to orthogonal matrices, namely

So in fact, we only need to perform one $\msign$ on the $2m\times 2m$ matrix $\boldsymbol{R}\boldsymbol{J}\boldsymbol{R}^{\top}$. For scenarios where $n \gg 2m$, the speedup is significant. However, if $n$ and $m$ are not different by more than an order of magnitude, calculating according to the original formula is more convenient, as QR decomposition also has a cost.

Additionally, how to switch to a momentum scenario is also worth discussing. There are two schemes we can consider: the first is $\boldsymbol{M} = \mathop{\text{EMA}}(\boldsymbol{G})$, in which case you just replace $\boldsymbol{G}$ with $\boldsymbol{M}$ and leave everything else unchanged; the other is $\boldsymbol{M} = \mathop{\text{EMA}}(\boldsymbol{G}\boldsymbol{W}^{\top})$, in which case $[\boldsymbol{G}\boldsymbol{W}^{\top}]_{\text{skew}}$ is replaced by $[\boldsymbol{M}]_{\text{skew}}$. This form preserves more information, but the cost is more space consumption, and after a long-term moving average, $\boldsymbol{M}$ usually won't be low-rank even if $n \gg 2m$, so it cannot be accelerated.

Of course, in mainstream model architectures, there aren't many $n \gg 2m$ scenarios (only when using Per-Head Muon might the matrix of each Head barely qualify). Therefore, the practical significance of this acceleration trick is limited, and we can just understand it briefly.

Open Questions

The author proposes an open question here: If we remove the orthogonality constraint and only consider the steepest descent on the tangent manifold, what would the optimal solution be? That is, consider

This problem can be seen as a generalization of the "Stiefel version of Muon" to general matrices. It doesn't depend on the orthogonality of $\boldsymbol{W}$ itself and is thus expected to be applicable to more general scenarios. The only result the author can currently obtain is that equations \eqref{eq:Phi-ori} and \eqref{eq:X-ori} still apply, i.e., the optimal solution has the following form:

where $\boldsymbol{X}\in\mathbb{R}^{m\times m}$ is an anti-symmetric matrix satisfying the following matrix equation:

As for how to solve it when $\boldsymbol{W}^{\top}\boldsymbol{W}\neq\boldsymbol{I}_m$, it is currently unknown (assuming $n=m$ doesn't seem to simplify it either). Furthermore, the construction of the weakened problem also holds for it, just with a slight modification:

However, since the previous proof of equivalence strongly relies on $\boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{I}_m$, it appears that this weakened solution is truly just a weakened solution, not the optimal one. Everyone is welcome to provide further progress~

Article Summary

In this article, we successfully found an analytical solution for the steepest descent on the "Muon + Stiefel" manifold, bringing a satisfying conclusion to this problem. Since then, regardless of whether it's a square matrix, the Muon optimizer under orthogonal constraints can be calculated analytically, and the solution itself will no longer be a difficulty.