Steepest Descent on Manifolds: 7. An Analytic Solution for Stiefel

流形上的最速下降：7. Stiefel的解析解 苏剑林 · 2026-08-17 · Mathematics optimizer muon constraint steepest descent

For the orthogonal manifold, our previous conclusion was: in the square-matrix case, we could write down the complete analytic solution for the corresponding steepest descent (see Steepest Descent on Manifolds: 2. Muon + Orthogonal); but for the non-square Stiefel manifold, the steepest descent problem requires solving a nonlinear system of equations, and no analytic solution could be written down directly (see Steepest Descent on Manifolds: 3. Muon + Stiefel).

Recently, however, Muon on the Stiefel Manifold Admits an Exact Closed-Form Update overturned this claim, showing that an explicit solution can also be written for the Stiefel manifold, without needing to solve a system of equations. In this post I'll restate the whole derivation in my own way.

Recap of the Problem

We won't repeat the background in detail; let's go directly to the problem to be solved. Let $\boldsymbol{W}\in\mathbb{R}^{n\times m}(n \geq m)$ be the current parameters, $\boldsymbol{G}$ the gradient of the objective function at $\boldsymbol{W}$, and we want to find the steepest descent direction $\boldsymbol{\Phi}\in\mathbb{R}^{n\times m}$ with step size $\eta$ under the orthogonality constraint, i.e.

\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 \leq 1,\,\, \boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{I}_m,\,\,(\boldsymbol{W} - \eta \boldsymbol{\Phi})^{\top}(\boldsymbol{W} - \eta \boldsymbol{\Phi})=\boldsymbol{I}_m\end{equation}

Following the principle that "a first-order approximation is good enough," we expand the last constraint and drop the $\eta^2$ term, simplifying to

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

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$ would give a completely equivalent result (the optimum of a linear objective is always attained on the boundary), but we keep $\leq$ because $\Vert\boldsymbol{\Phi}\Vert_2 \leq 1$ is a convex set, which gives us more room to maneuver in the argument. Our previous result was

\begin{equation}\boldsymbol{\Phi} = \newcommand{msign}{\mathop{\text{msign}}}\msign(\boldsymbol{G} + \boldsymbol{W}\boldsymbol{X})\label{eq:Phi-ori}\end{equation}

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

\begin{equation}\boldsymbol{W}^{\top}\msign(\boldsymbol{G} + \boldsymbol{W}\boldsymbol{X})+\msign(\boldsymbol{G} + \boldsymbol{W}\boldsymbol{X})^{\top}\boldsymbol{W} = \boldsymbol{0}\label{eq:X-ori}\end{equation}

We already solved the case $n=m$ in Steepest Descent on Manifolds: 2. Muon + Orthogonal, but for the standard Stiefel manifold with $n > m$, solving this equation is nontrivial — even a numerical solution is somewhat troublesome — and that's where the claim "there is no analytic solution in the Stiefel case" came from.

A Weakened Version

The key to cracking this is to eliminate the equality constraint $\boldsymbol{W}^{\top}\boldsymbol{\Phi}+\boldsymbol{\Phi}^{\top}\boldsymbol{W} = \boldsymbol{0}$ via some explicit representation. To this end, let's 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 an arbitrary antisymmetric matrix satisfying $\Vert\boldsymbol{S}\Vert_2\leq 1$, and consider solving

\begin{equation}\max_{\boldsymbol{S}} \tr(\boldsymbol{G}^{\top}\boldsymbol{S}\boldsymbol{W}) \qquad \text{s.t.}\qquad \Vert\boldsymbol{S}\Vert_2 \leq 1\end{equation}

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

\begin{align}1) &\quad \Vert\boldsymbol{\Phi}\Vert_2 = \Vert\boldsymbol{S}\boldsymbol{W}\Vert_2\leq \Vert\boldsymbol{S}\Vert_2 \Vert\boldsymbol{W}\Vert_2\leq 1 \\[5pt] 2) &\quad \boldsymbol{W}^{\top}\boldsymbol{\Phi}+\boldsymbol{\Phi}^{\top}\boldsymbol{W} = \boldsymbol{W}^{\top}\boldsymbol{S}\boldsymbol{W}+\boldsymbol{W}^{\top}\boldsymbol{S}^{\top}\boldsymbol{W} = \boldsymbol{W}^{\top}(\boldsymbol{S} + \boldsymbol{S}^{\top})\boldsymbol{W} = \boldsymbol{0} \end{align}

So the new problem is being solved over a subset of the feasible region of the original problem, and the maximum value obtained will not exceed that of the original problem. Furthermore, using the trace identity and the antisymmetry of $\boldsymbol{S}$, we can transform the objective function into

\begin{equation}\tr(\boldsymbol{G}^{\top}\boldsymbol{S}\boldsymbol{W}) = \tr(\boldsymbol{W}\boldsymbol{G}^{\top}\boldsymbol{S}) = -\tr(\boldsymbol{W}\boldsymbol{G}^{\top}\boldsymbol{S}^{\top}) = -\tr(\boldsymbol{G}\boldsymbol{W}^{\top}\boldsymbol{S}) = \tr([\boldsymbol{W}\boldsymbol{G}^{\top}]_{\text{skew}}\boldsymbol{S})\end{equation}

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

\begin{equation}\max_{\boldsymbol{S}}\tr([\boldsymbol{W}\boldsymbol{G}^{\top}]_{\text{skew}}\boldsymbol{S}) \qquad \text{s.t.}\qquad \Vert\boldsymbol{S}\Vert_2 \leq 1,\,\,\boldsymbol{S} + \boldsymbol{S}^{\top}=\boldsymbol{0}\end{equation}

This already has the shape of the standard Muon problem, so we can directly obtain the result from Muon:

\begin{equation}\boldsymbol{S} = \msign([\boldsymbol{G}\boldsymbol{W}^{\top}]_{\text{skew}})\qquad\Rightarrow\qquad \boldsymbol{\Phi} = \msign([\boldsymbol{G}\boldsymbol{W}^{\top}]_{\text{skew}})\boldsymbol{W}\label{eq:sol}\end{equation}

This is in fact exactly the form of the solution in the square-matrix case, and also the left-side rotation from Steepest Descent on Manifolds: 6. Muon + Double Rotation. It's worth explaining here why we go through a series of identity transformations before applying Muon's $\msign$: it's mainly to exploit the fact that "the $\msign$ of an antisymmetric matrix is still antisymmetric," which guarantees that the final result satisfies the required antisymmetry of $\boldsymbol{S}$.

Full Equivalence

Next, we're going to prove that the optimal solution of the weakened problem solved in the previous section is in fact identical to that of the original problem!

The proof idea is quite direct. The weakened problem is "weak" because the parametrization $\boldsymbol{\Phi} = \boldsymbol{S}\boldsymbol{W}$ we currently use is only a subset of the feasible region of the original problem, which might miss some region of exploration and thus fail to reach the optimum of the original problem. But if we can show that, for any feasible $\boldsymbol{\Phi}$ of the original problem, there exists a suitable $\boldsymbol{S}$ such that $\boldsymbol{\Phi} = \boldsymbol{S}\boldsymbol{W}$, that would mean the two exploration spaces coincide exactly, and hence so do the optima.

The proof proceeds in three steps. (Note: in hindsight, with the right hints, K3 was also able to complete this derivation and proof independently; see here and here.)

Solving the Equation

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

\begin{equation}\boldsymbol{\Phi} = \boldsymbol{S}\boldsymbol{W},\qquad \boldsymbol{S}+\boldsymbol{S}^{\top} = \boldsymbol{0}\end{equation}

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

\begin{equation}\boldsymbol{S} = \boldsymbol{U}\begin{bmatrix}\boldsymbol{A} & -\boldsymbol{B}^{\top} \\ \boldsymbol{B} & \boldsymbol{C}\end{bmatrix}\boldsymbol{U}^{\top} = [\boldsymbol{W}, \boldsymbol{W}_{\perp}]\begin{bmatrix}\boldsymbol{A} & -\boldsymbol{B}^{\top} \\ \boldsymbol{B} & \boldsymbol{C}\end{bmatrix}\begin{bmatrix}\boldsymbol{W}^{\top} \\ \boldsymbol{W}_{\perp}^{\top}\end{bmatrix}\label{eq:S-block}\end{equation}

where $\boldsymbol{A}\in\mathbb{R}^{m\times m},\boldsymbol{C}\in\mathbb{R}^{(n-m)\times (n-m)}$ are all antisymmetric matrices, $\boldsymbol{B}\in\mathbb{R}^{(n-m)\times m}$. Then multiplying both sides on the right by $\boldsymbol{W}$ gives

\begin{equation}\boldsymbol{\Phi} = \boldsymbol{S}\boldsymbol{W} = [\boldsymbol{W}, \boldsymbol{W}_{\perp}]\begin{bmatrix}\boldsymbol{A} & -\boldsymbol{B}^{\top} \\ \boldsymbol{B} & \boldsymbol{C}\end{bmatrix}\begin{bmatrix}\boldsymbol{I}_m \\ \boldsymbol{0}\end{bmatrix} = [\boldsymbol{W}, \boldsymbol{W}_{\perp}]\begin{bmatrix}\boldsymbol{A} \\ \boldsymbol{B}\end{bmatrix}\end{equation}

Multiplying both sides on the left by $\boldsymbol{U}^{\top}$, we can directly read off $\boldsymbol{A} = \boldsymbol{W}^{\top}\boldsymbol{\Phi},\boldsymbol{B} = \boldsymbol{W}_{\perp}^{\top}\boldsymbol{\Phi}$, i.e., $\boldsymbol{A},\boldsymbol{B}$ can both be uniquely solved for, leaving $\boldsymbol{C}$ as the remaining degree of freedom. Note that $\boldsymbol{A}$ is antisymmetric, which imposes the compatibility condition

\begin{equation}\boldsymbol{0} = \boldsymbol{A} + \boldsymbol{A}^{\top} = \boldsymbol{W}^{\top}\boldsymbol{\Phi} + \boldsymbol{\Phi}^{\top}\boldsymbol{W}\end{equation}

and this is precisely one of the conditions of the original problem. So up to this point, everything is self-consistent.

Spectral Norm

The remaining question now is: can we find an antisymmetric matrix $\boldsymbol{C}$ such that $\Vert \boldsymbol{S}\Vert_2 \leq 1$? The answer is yes, but the proof requires a "Parrott's lemma" (the original paper cites Davis–Kahan–Weinberger, but on further investigation Parrott's lemma turns out to be 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}

Back to our problem. To have $\Vert \boldsymbol{S}\Vert_2 \leq 1$, it suffices that $\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

\begin{equation}\left\Vert\begin{bmatrix}\boldsymbol{A} \\ \boldsymbol{B}\end{bmatrix}\right\Vert_2 = \left\Vert\begin{bmatrix}\boldsymbol{W}^{\top}\boldsymbol{\Phi} \\ \boldsymbol{W}_{\perp}^{\top}\boldsymbol{\Phi}\end{bmatrix}\right\Vert_2 = \Vert\boldsymbol{U}^{\top}\boldsymbol{\Phi}\Vert_2 = \Vert\boldsymbol{\Phi}\Vert_2 \leq 1\end{equation}

As for $[\boldsymbol{A},-\boldsymbol{B}^{\top}] = [-\boldsymbol{A}^{\top},-\boldsymbol{B}^{\top}]$, it is $$$\begin{bmatrix}\boldsymbol{A} \\ \boldsymbol{B}\end{bmatrix}$$$的负转置，谱范数不变（同样不超过1）。因此，根据Parrott引理，我们能找到矩阵$\boldsymbol{C}$，使得$\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$.

Antisymmetry

But we're not quite done. We also need $\boldsymbol{C}$ to be antisymmetric, and Parrott's lemma doesn't guarantee this.

Fortunately, this fix isn't hard. Let the matrix given by Parrott's lemma be $\boldsymbol{C}_0$; we simply take $\boldsymbol{C} = [\boldsymbol{C}_0]_{\text{skew}}$, and since the remaining part already satisfies the antisymmetry requirement, we have

\begin{equation}\begin{bmatrix}\boldsymbol{A} & -\boldsymbol{B}^{\top} \\ \boldsymbol{B} & [\boldsymbol{C}_0]_{\text{skew}}\end{bmatrix} = \frac{1}{2}\left(\begin{bmatrix}\boldsymbol{A} & -\boldsymbol{B}^{\top} \\ \boldsymbol{B} & \boldsymbol{C}_0\end{bmatrix} - \begin{bmatrix}\boldsymbol{A} & -\boldsymbol{B}^{\top} \\ \boldsymbol{B} & \boldsymbol{C}_0\end{bmatrix}^{\top}\right) \end{equation}

The key point here is that antisymmetrization does not increase the spectral norm: by the triangle inequality we get

\begin{equation}\left\Vert\begin{bmatrix}\boldsymbol{A} & -\boldsymbol{B}^{\top} \\ \boldsymbol{B} & [\boldsymbol{C}_0]_{\text{skew}}\end{bmatrix}\right\Vert_2 \leq \frac{1}{2}\left\Vert\begin{bmatrix}\boldsymbol{A} & -\boldsymbol{B}^{\top} \\ \boldsymbol{B} & \boldsymbol{C}_0\end{bmatrix}\right\Vert_2 + \frac{1}{2}\left\Vert \begin{bmatrix}\boldsymbol{A} & -\boldsymbol{B}^{\top} \\ \boldsymbol{B} & \boldsymbol{C}_0\end{bmatrix}^{\top}\right\Vert_2 \leq \frac{1}{2} + \frac{1}{2} = 1 \end{equation}

So $[\boldsymbol{C}_0]_{\text{skew}}$ is exactly the antisymmetric matrix we were looking for.

With this we have proved the result we originally set out to prove: for every $\boldsymbol{\Phi}$ of the original problem, we can find a $n\times n$, antisymmetric matrix $\boldsymbol{S}$ with spectral norm no greater than 1, such that $\boldsymbol{\Phi}=\boldsymbol{S}\boldsymbol{W}$; hence the exploration space of the weakened version coincides with that of the original problem. In other words, equation $\eqref{eq:sol}$ is the exact solution of the original problem!

Efficient Computation

Theoretically, the problem is now completely resolved; but from a practical standpoint, there's still an efficiency question worth discussing. Muon needs to perform $\msign$ on a $n\times m$ matrix, while the exact solution $\eqref{eq:sol}$ on the Stiefel manifold requires performing $\msign$ on a $n\times n$ matrix. If computed directly, this is clearly much more expensive than Muon when $n \gg m$, so it's worth discussing ways to save computation.

It's not hard to think of one: first, rewrite the core computation in low-rank form

\begin{equation}\msign([\boldsymbol{G}\boldsymbol{W}^{\top}]_{\text{skew}}) = \msign(\boldsymbol{G}\boldsymbol{W}^{\top} - \boldsymbol{W}\boldsymbol{G}^{\top}) = \msign\bigg(\begin{bmatrix}\boldsymbol{G} & \boldsymbol{W}\end{bmatrix} \underbrace{\begin{bmatrix}\boldsymbol{0} & \boldsymbol{I}_m \\ -\boldsymbol{I}_m & \boldsymbol{0}\end{bmatrix}}_{\boldsymbol{J}} \begin{bmatrix}\boldsymbol{G}^{\top} \\ \boldsymbol{W}^{\top}\end{bmatrix}\bigg)\end{equation}

We can see that the rank of $[\boldsymbol{G}\boldsymbol{W}^{\top}]_{\text{skew}}$ is at most $2m$. Suppose $n \geq 2m$; first perform a QR decomposition of $[\boldsymbol{G},\boldsymbol{W}]$ (a polar decomposition of $\msign$ also works), giving $\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 quantity to be computed 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, i.e.

\begin{equation}\msign(\boldsymbol{Q}\boldsymbol{R}\boldsymbol{J}\boldsymbol{R}^{\top} \boldsymbol{Q}^{\top}) = \boldsymbol{Q}\msign(\boldsymbol{R}\boldsymbol{J}\boldsymbol{R}^{\top} )\boldsymbol{Q}^{\top}\end{equation}

So in fact we only need to perform $\msign$ once on a matrix $\boldsymbol{R}\boldsymbol{J}\boldsymbol{R}^{\top}$ of size $2m\times 2m$, which gives a clear speedup for the scenario $n\gg 2m$. But if $n$ and $m$ don't differ by more than an order of magnitude, it's more convenient to just compute using the original formula, since the QR decomposition also has its own cost.

It's also worth discussing how to switch to the momentum setting. There are two possible approaches we could consider: the first is $\boldsymbol{M} = \mathop{\text{EMA}}(\boldsymbol{G})$, in which case we just need to replace $\boldsymbol{G}$ with $\boldsymbol{M}$, leaving 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 with $[\boldsymbol{M}]_{\text{skew}}$. This form retains more information, at the cost of using more memory, and after a long moving average, even if $n\gg 2m$, $\boldsymbol{M}$ generally won't be low-rank anymore, so it can't be accelerated this way.

Of course, in mainstream model architectures, scenarios with $n\gg 2m$ aren't very common (only with Per-Head Muon does each head's matrix barely qualify), so this acceleration trick has limited practical significance — it's enough just to be aware of it.

Open Problem

Here I'd like to pose an open problem: if we drop the orthogonality constraint and consider only steepest descent on the tangent manifold, what would the optimal solution look like? That is, consider

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

This problem can be viewed as a generalization of "Stiefel-style Muon" to general matrices, since it doesn't depend on the orthogonality of $\boldsymbol{W}$ itself, and so it may hold promise for more general settings. The only result I've managed to obtain so far is that equations $\eqref{eq:Phi-ori}$ and $\eqref{eq:X-ori}$ still apply, i.e., the optimal solution has the form

\begin{equation}\boldsymbol{\Phi} = \newcommand{msign}{\mathop{\text{msign}}}\msign(\boldsymbol{G} + \boldsymbol{W}\boldsymbol{X})\end{equation}

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

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

As for how to solve this when $\boldsymbol{W}^{\top}\boldsymbol{W}\neq\boldsymbol{I}_m$, I don't yet know (assuming $n=m$ doesn't seem to simplify things either). Additionally, the construction of the weakened problem also holds for it, requiring only a small modification:

\begin{equation}\boldsymbol{\Phi}_{\text{weak}} = \msign([\boldsymbol{G}\boldsymbol{W}^{\top}]_{\text{skew}})\boldsymbol{W}/\Vert\boldsymbol{W}\Vert_2\end{equation}

But since the equivalence proof above relies heavily on $\boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{I}_m$, it appears that this weakened solution really is just a weakened solution, not the optimum. Any further progress from readers is welcome~

Summary

In this post, we successfully found the analytic solution for the manifold steepest descent problem "Muon + Stiefel," bringing this problem to a satisfying conclusion. From now on, regardless of whether the matrix is square, the Muon optimizer under orthogonality constraints can be computed analytically, and solving it will no longer pose any difficulty.