Steepest Descent on Manifolds: 3. Muon + Stiefel
Last time, we saw that once we shift the optimization target from vector parameters to matrix parameters, and adopt the spectral norm constraint that is better suited to matrices, the Muon optimizer emerges quite naturally. Going further, we considered the steepest descent direction under an additional orthogonality constraint on the parameters, splitting the discussion into the square-matrix case and the non-square-matrix case. The square-matrix case was already solved in the previous post, but the non-square case was left open.
The goal of this post is to fill in that gap for the non-square case, so that optimization under the orthogonality constraint is fully resolved.
Recap
Let's briefly recall the results from Steepest Descent on Manifolds: 2. Muon + Orthogonality]. The objective we want to solve 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,\,\, \boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{I},\,\,(\boldsymbol{W} - \eta \boldsymbol{\Phi})^{\top}(\boldsymbol{W} - \eta \boldsymbol{\Phi})=\boldsymbol{I}\end{equation}
where $\boldsymbol{W},\boldsymbol{\Phi}\in\mathbb{R}^{n\times m}(n \geq m)$, and $\Vert\cdot\Vert_2$ is the spectral norm. Following the principle that "a first-order approximation is good enough," this 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,\,\, \boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{I},\,\,\boldsymbol{W}^{\top}\boldsymbol{\Phi}+\boldsymbol{\Phi}^{\top}\boldsymbol{W} = \boldsymbol{0}\label{eq:ori-obj}\end{equation}
where the set of all $\boldsymbol{\Phi}$ satisfying $\boldsymbol{W}^{\top}\boldsymbol{\Phi}+\boldsymbol{\Phi}^{\top}\boldsymbol{W} = \boldsymbol{0}$ is also called the "tangent space" of $\boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{I}$. In the previous post we already worked out the general form of the solution
\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 to be determined.
The remaining challenge is to find a way of computing the symmetric matrix $\boldsymbol{X}$ such that $\boldsymbol{W}^{\top}\boldsymbol{\Phi}$ is skew-symmetric. Once this is solved, the corresponding $\boldsymbol{\Phi}$ is automatically the optimal solution. For $n=m$ we already have a closed-form solution $\boldsymbol{X}=-[\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{sym}}$; the truly difficult case is $n > m$, which is also known as the "Stiefel manifold" — this is exactly the open problem left by Orthogonal manifold].
Reformulating the equation
Put plainly, our task now is to solve the system of equations:
\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:start}\end{equation}
When $n=m$, $\boldsymbol{W}^{\top}$ can simply be absorbed into $\msign$, which simplifies the solution. However, when $n > m$, this absorption trick no longer works, which is precisely where the difficulty lies. My own suspicion is that when $n > m$, there is no simple explicit solution, so we turn to numerical algorithms instead.
By definition $\msign(\boldsymbol{M})=\boldsymbol{M}(\boldsymbol{M}^{\top}\boldsymbol{M})^{-1/2}$, we can write
\begin{equation}\boldsymbol{W}^{\top}\msign(\boldsymbol{G} + \boldsymbol{W}\boldsymbol{X}) = \boldsymbol{W}^{\top}(\boldsymbol{G} + \boldsymbol{W}\boldsymbol{X})\boldsymbol{Q}^{-1} = (\boldsymbol{W}^{\top}\boldsymbol{G} + \boldsymbol{X})\boldsymbol{Q}^{-1}\end{equation}
where $\boldsymbol{Q} = ((\boldsymbol{G} + \boldsymbol{W}\boldsymbol{X})^{\top}(\boldsymbol{G} + \boldsymbol{W}\boldsymbol{X}))^{1/2}$. Under this new notation, the system of equations becomes
\begin{equation}(\boldsymbol{W}^{\top}\boldsymbol{G} + \boldsymbol{X})\boldsymbol{Q}^{-1} + \boldsymbol{Q}^{-1}(\boldsymbol{G}^{\top}\boldsymbol{W} + \boldsymbol{X}) = \boldsymbol{0}\end{equation}
Multiplying both sides on the left and right by $\boldsymbol{Q}$ gives
\begin{equation}\boldsymbol{Q}(\boldsymbol{W}^{\top}\boldsymbol{G} + \boldsymbol{X}) + (\boldsymbol{G}^{\top}\boldsymbol{W} + \boldsymbol{X})\boldsymbol{Q} = \boldsymbol{0}\label{eq:r-x}\end{equation}
where $\boldsymbol{Q}$ also satisfies
\begin{equation}\boldsymbol{Q} = (\boldsymbol{G} + \boldsymbol{W}\boldsymbol{X})^{\top}\msign(\boldsymbol{G} + \boldsymbol{W}\boldsymbol{X})\label{eq:r-q}\end{equation}
Iterative solution
My idea now is: starting from some initial value of $\boldsymbol{X}$, substitute it into equation $\eqref{eq:r-q}$ to obtain $\boldsymbol{Q}$, then plug $\boldsymbol{Q}$ into the system $\eqref{eq:r-x}$ to solve for a new $\boldsymbol{X}$, and repeat this iteration until convergence. Given $\msign$, equation $\eqref{eq:r-q}$ can be computed explicitly, so the only real difficulty is solving the system $\eqref{eq:r-x}$.
We can rearrange equation $\eqref{eq:r-x}$ as:
\begin{equation}\boldsymbol{Q}\boldsymbol{X} + \boldsymbol{X}\boldsymbol{Q} = -2[\boldsymbol{Q}\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{sym}}\label{eq:r-xx}\end{equation}
Given $\boldsymbol{Q}$, this is actually a linear system in $\boldsymbol{X}$, known as the "continuous Lyapunov equation]," which can also be viewed as a special case of the "Sylvester equation]." If we're only computing on CPU, Scipy already provides a built-in solver scipy.linalg.solve_continuous_lyapunov for this equation, which we can call directly.
As for the choice of initial value, we can consider the solution $-[\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{sym}}$ from the square-matrix case, which is clearly a natural transition from square to non-square matrices. We can also examine the reasonableness of the initial value $-[\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{sym}}$ by looking at an equivalent form of equation $\eqref{eq:r-xx}$:
\begin{equation}\boldsymbol{Q}(\boldsymbol{X} + [\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{sym}}) + (\boldsymbol{X} + [\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{sym}})\boldsymbol{Q} =[\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{skew}}\boldsymbol{Q} -\boldsymbol{Q}[\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{skew}}\end{equation}
So how accurate $-[\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{sym}}$ is depends on how commutative the multiplication between $[\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{skew}}$ and $\boldsymbol{Q}$ is — the closer they are to commuting matrices, the more accurate $-[\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{sym}}$ becomes. However, the empirical results later show that our iterative algorithm is not particularly sensitive to the initial value; even starting from an all-zero matrix causes no real problem.
Doing it ourselves
We just mentioned that Scipy has a built-in solver for the Lyapunov equation, so we can call it directly without worrying about the solving process. But this is only available for CPU-based Scipy. As far as I can tell, neither Torch nor Jax has an equivalent function, so if we want to compute on GPU, we have to roll our own.
There are two ways to write our own solver for equation $\eqref{eq:r-xx}$. The first follows the approach in What can the matrix sign function mcsgn compute?], using $\newcommand{mcsgn}{\mathop{\text{mcsgn}}}\mcsgn$ (not $\msign$) to solve it:
\begin{equation}\boldsymbol{X} = \mcsgn\left(\begin{bmatrix}-\boldsymbol{Q} & -[\boldsymbol{Q}\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{sym}} \\ \boldsymbol{0} & \boldsymbol{Q}\end{bmatrix}\right)_{[:m,m:]}\end{equation}
The second is based on SVD, a method we already used when computing the gradient of $\msign$ in The derivative of msign]; here we'll re-introduce it in the context of equation $\eqref{eq:r-xx}$. Since $\boldsymbol{Q}$ is positive definite and symmetric by definition, it admits an eigendecomposition $\boldsymbol{V}\boldsymbol{\Sigma}\boldsymbol{V}^{\top}$, where $\boldsymbol{V}$ is orthogonal and $\boldsymbol{\Sigma}=\mathop{\text{diag}}(\sigma_1,\cdots,\sigma_m)$ is diagonal. Substituting into equation $\eqref{eq:r-xx}$ and rearranging gives
\begin{equation}\boldsymbol{\Sigma}(\boldsymbol{V}^{\top}\boldsymbol{X}\boldsymbol{V}) + (\boldsymbol{V}^{\top}\boldsymbol{X}\boldsymbol{V})\boldsymbol{\Sigma} = -2\boldsymbol{V}^{\top}[\boldsymbol{Q}\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{sym}}\boldsymbol{V}\end{equation}
The left-hand side can be written as $(\boldsymbol{V}^{\top}\boldsymbol{X}\boldsymbol{V})\otimes \boldsymbol{S}$, where $\otimes$ is the Hadamard product and $\boldsymbol{S}_{i,j} = \sigma_i + \sigma_j$. From this we can solve for
\begin{equation}\boldsymbol{X} = -2\boldsymbol{V}((\boldsymbol{V}^{\top}[\boldsymbol{Q}\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{sym}}\boldsymbol{V})\oslash \boldsymbol{S})\boldsymbol{V}^{\top}\end{equation}
where $\oslash$ is the Hadamard (element-wise) quotient. What's interesting here is that eigendecomposing $\boldsymbol{Q}$ is essentially equivalent to computing the SVD of $\boldsymbol{G} + \boldsymbol{W}\boldsymbol{X}$, and the SVD of $\boldsymbol{G} + \boldsymbol{W}\boldsymbol{X}$ can also be used to compute $\msign(\boldsymbol{G} + \boldsymbol{W}\boldsymbol{X})$ — so a single SVD suffices to compute both $\msign$ and the solution to equation $\eqref{eq:r-xx}$.
Each approach has its own characteristics. The first approach requires computing $\msign$ for a matrix of shape $m\times m$, then $\mcsgn$ for a matrix of shape $2m\times 2m$. Although both can be computed efficiently via Newton–Schulz iteration, this still comes at a considerable cost. Moreover, we need to choose coefficients that both converge and are sufficiently accurate (I recommend the results from Newton–Schulz iteration for the msign operator (part 2)]), otherwise neither $\mcsgn$ nor $\msign$ will even converge, let alone $\boldsymbol{X}$.
The second approach relies on SVD. Although SVD has higher complexity and often requires forcing FP32 precision, for this particular problem each iteration only needs a single SVD to compute both $\msign$ and $\boldsymbol{X}$ simultaneously, so the overall efficiency isn't too bad. If we only have a small number of matrix parameters that need the orthogonality constraint, SVD may well be the simplest choice.
Related work
Before this post, @leloy] proposed, in his blog post Heuristic Solutions for Steepest Descent on the Stiefel Manifold], two heuristic methods for solving the original objective $\eqref{eq:ori-obj}$. Here "heuristic" means that in most cases it produces a fairly good solution, but there is no guarantee of optimality — let's take a look at these as well.
The first method can be described as purely geometric. First, define the projection operation:
\begin{equation}\newcommand{proj}{\mathop{\mathcal{P}}}\proj\nolimits_{\boldsymbol{W}}(\boldsymbol{M}) = \boldsymbol{M} - \boldsymbol{W}[\boldsymbol{W}^{\top}\boldsymbol{M}]_{\text{sym}}\end{equation}
One can verify that $\boldsymbol{W}^{\top}\proj\nolimits_{\boldsymbol{W}}(\boldsymbol{M})$ is always skew-symmetric, meaning that $\proj\nolimits_{\boldsymbol{W}}(\boldsymbol{M})$ always lies in the tangent space. So we can regard this as the operation that projects an arbitrary matrix $\boldsymbol{M}$ onto the tangent space of $\boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{I}$.
Starting from the gradient $\boldsymbol{G}$, $\proj\nolimits_{\boldsymbol{W}}(\boldsymbol{M})$ is certainly in the tangent space, but we know that Muon's update must be an orthogonal matrix (when full rank), while $\proj\nolimits_{\boldsymbol{W}}(\boldsymbol{M})$ is not necessarily orthogonal. So we can use $\msign$ to find the nearest orthogonal matrix, i.e. $\msign(\proj\nolimits_{\boldsymbol{W}}(\boldsymbol{M}))$. But after $\msign$, the result may no longer be in the tangent space, so we project it back into the tangent space, then again find the nearest orthogonal matrix, and repeat:
\begin{equation}\boldsymbol{\Phi} = (\msign\circ\proj\nolimits_{\boldsymbol{W}}\circ\cdots\circ\msign\circ\proj\nolimits_{\boldsymbol{W}})(\boldsymbol{M})\end{equation}
This is @leloy's first idea: alternately projecting onto the tangent space and the orthogonal space until convergence — quite intuitive, really. And in fairly random test cases, it turns out to be remarkably close to the optimal solution, sometimes matching it to four decimal places, which initially led me to believe it was in fact the exact solution. But after further searching, I found cases where it deviates substantially from the optimum, confirming that the earlier agreement was just a coincidence rather than genuine optimality.
The second method can be called a line search. Specifically, when $n > m$, we can consider completing $\boldsymbol{W}$ into a standard $n\times n$ orthogonal matrix $[\boldsymbol{W},\overline{\boldsymbol{W}}]$, and then decompose the quantity we want to solve for, $\boldsymbol{\Phi}$, into two parts $\boldsymbol{W}^{\top}\boldsymbol{\Phi}$ and $\overline{\boldsymbol{W}}{}^{\top}\boldsymbol{\Phi}$. @leloy then makes a greedy approximation: first solve for the optimal $\boldsymbol{W}^{\top}\boldsymbol{\Phi}$, then for the optimal $\overline{\boldsymbol{W}}{}^{\top}\boldsymbol{\Phi}$, and introduces a line search between the two to improve accuracy.
This whole procedure does yield a reasonably good approximate solution, and it is guaranteed to lie in the tangent space while satisfying orthogonality. The solving process requires computing the spectral norm, $\msign$, and a Cholesky decomposition] — see the author's post for the details. Also, when $m=2$, this method can in principle find the exact optimum, because the skew-symmetric matrices of size $2\times 2$ have only a single free parameter, and the line search provides exactly one degree of freedom.
Testing it out
Below we empirically test the methods above in Numpy. The main purpose is to verify the correctness of the methods themselves, so we directly implement $\msign$ and $\mcsgn$ using singular value decomposition and eigendecomposition.
import numpy as np
import scipy as sp
def mcsgn(x):
"""特征值分解精确计算mcsgn
"""
s, v = np.linalg.eig(x)
return v @ np.diag(np.sign(s)) @ np.linalg.inv(v)
def msign(g):
"""奇异值分解精确计算msign
"""
u, s, vh = np.linalg.svd(g, full_matrices=False)
return u @ np.diag(np.sign(s)) @ vh
def sym(x):
"""对称化
"""
return (x + x.T) * 0.5
def skew(x):
"""反对称化
"""
return (x - x.T) * 0.5
def proj(g, w):
"""投影到正交的切空间
"""
return g - w @ sym(w.T @ g)
def jianlin_by_mcsgn(g, w, steps=20):
"""通过mcsgn来构建本文的迭代
"""
n, m = g.shape
x = -sym(w.T @ g)
for i in range(1, steps + 1):
phi = msign(z := g + w @ x)
print('step:', i, ', inner product:', (phi * g).sum(), ', tangent error:', np.abs(sym(w.T @ phi)).mean())
if i == steps:
return phi
q = z.T @ phi
x = mcsgn(np.block([[-q, -sym(q @ w.T @ g)], [np.zeros_like(q), q]]))[:m, m:]
# x = -2 * sp.linalg.solve_continuous_lyapunov(q, sym(q @ w.T @ g))
def jianlin_by_svd(g, w, steps=20):
"""通过svd来构建本文的迭代
"""
x = -sym(w.T @ g)
for i in range(1, steps + 1):
u, s, vh = np.linalg.svd(z := g + w @ x, full_matrices=False)
phi = (u * np.sign(s)) @ vh
print('step:', i, ', inner product:', (phi * g).sum(), ', tangent error:', np.abs(sym(w.T @ phi)).mean())
if i == steps:
return phi
x = -2 * vh.T @ (vh @ sym(z.T @ phi @ w.T @ g) @ vh.T / (s + s[:, None])) @ vh
def leloy_v1(g, w, steps=20):
"""交替投影到切空间和正交空间
"""
phi = g
for i in range(1, steps + 1):
phi = msign(proj(phi, w))
print('step:', i, ', inner product:', (phi * g).sum(), ', tangent error:', np.abs(sym(w.T @ phi)).mean())
return phi
def leloy_v2(g, w, steps=20):
"""分部贪心求解 + 线搜索（形式经过笔者的简化）
"""
n, m = g.shape
taus = np.linspace(0, 1, steps + 2)[1:-1]
p_max, tau_opt, phi_opt = 0, 0, None
for tau in taus:
b = (b := skew(w.T @ g)) * tau / max(np.linalg.norm(b, ord=2), 1e-8)
r = np.linalg.cholesky(np.eye(m) - b.T @ b)
c = msign((np.eye(n) - w @ w.T) @ g @ r) @ r
phi = w @ b + c
print('tau:', tau, ', inner product:', p := (phi * g).sum())
if p > p_max:
p_max, tau_opt, phi_opt = p, tau, phi
print('best inner product:', p_max, ', tau:', tau_opt)
return phi_opt
w = np.array([[ 0.69453734, -0.26590866, -0.44721806, 0.2753041 ],
[-0.11738148, -0.5588003 , -0.17580748, 0.3218624 ],
[-0.4515288 , -0.23489913, -0.26683152, -0.25739142],
[ 0.02392521, 0.02664689, 0.48423648, 0.6193399 ],
[ 0.45194831, -0.25206333, 0.27654836, -0.60242337],
[ 0.21197332, -0.09174792, 0.24521762, -0.08484317],
[-0.15496767, -0.26446804, -0.34942415, -0.01877318],
[-0.16181251, -0.6474956 , 0.45243263, -0.01776086]])
g = np.array([[-17.85745 , -10.758921 , -2.9583392 , 6.245008 ],
[-28.883093 , 19.772121 , 8.086545 , -21.564013 ],
[ -1.6274693 , -14.96859 , 3.4465332 , 3.1070817 ],
[ -7.8890743 , 1.5304767 , -8.949573 , 9.579629 ],
[ 2.246596 , 14.46572 , 12.8451 , -2.7370298 ],
[ -0.9496974 , 6.9879804 , 2.849277 , 1.1148484 ],
[ -8.115278 , -18.054405 , -0.19287404, 7.0389237 ],
[-15.062008 , -15.02901 , 2.9083247 , 21.706533 ]])
phi1 = jianlin_by_mcsgn(g, w, steps=100)
phi2 = jianlin_by_svd(g, w, steps=100)
phi3 = leloy_v1(g, w, steps=100)
phi4 = leloy_v2(g, w, steps=100)
assert np.allclose(phi1, phi2)
w = np.linalg.qr(np.random.randn(100, 50))[0]
g = np.random.randn(100, 50)
phi1 = jianlin_by_mcsgn(g, w, steps=10)
phi2 = jianlin_by_svd(g, w, steps=10)
phi3 = leloy_v1(g, w, steps=10)
phi4 = leloy_v2(g, w, steps=10)
assert np.allclose(phi1, phi2)
For the first set of $\boldsymbol{W},\boldsymbol{G}$ given in the code, the optimal $\tr(\boldsymbol{G}^{\top} \boldsymbol{\Phi})$ found by my method is approximately $90$, and $\mcsgn$ agrees exactly with the SVD-based result. @leloy's first method gives a result of roughly $70$, and the second method gives roughly $80$ — both of which differ noticeably from the optimal solution.
That said, this particular set of $\boldsymbol{W},\boldsymbol{G}$ was deliberately searched for as an extreme example to highlight the gap between the three methods. If we switch to more random values, the solution from this post and @leloy's first method actually turn out to be very close, and the number of iterations required drops significantly (5–10 steps); in that setting, @leloy's second method deviates more from the optimum. Readers are welcome to construct their own test cases.
Further thoughts
This concludes our discussion of solving the original problem $\eqref{eq:ori-obj}$. Let's wrap up with a few additional details that readers may find puzzling.
First, for ease of exposition, the iterative solution procedure described earlier carries an implicit assumption: that $\boldsymbol{G} + \boldsymbol{W}\boldsymbol{X}$ remains full rank (rank $m$) throughout. Otherwise the matrix $\boldsymbol{S}$ would have zero components, making $\oslash\boldsymbol{S}$ awkward to handle. But this difficulty isn't fundamental, because equation $\eqref{eq:start}$ is guaranteed to have a solution — so whenever a denominator is zero, the corresponding numerator must also be zero. We can therefore simply replace the zero components of $\boldsymbol{S}$ with a small positive number to get the correct result.
From a numerical computation standpoint, we rarely encounter singular values that are exactly zero anyway, so this isn't really something to worry too much about — we can just assume $\boldsymbol{G} + \boldsymbol{W}\boldsymbol{X}$ is full rank by default. Under this assumption, the retraction operation becomes quite simple, because
\begin{equation}(\boldsymbol{W} - \eta\boldsymbol{\Phi})^{\top}(\boldsymbol{W} - \eta\boldsymbol{\Phi}) = \boldsymbol{W}^{\top} \boldsymbol{W} - \eta(\boldsymbol{W}^{\top} \boldsymbol{\Phi} + \boldsymbol{\Phi}^{\top}\boldsymbol{W}) + \eta^2 \boldsymbol{\Phi}^{\top}\boldsymbol{\Phi}\end{equation}
By the definition of the Stiefel manifold, the first term on the right is $\boldsymbol{I}$; by the tangent space condition, the second term is $\boldsymbol{0}$; and finally, since $\msign$ is full rank, the result of the operation is again a matrix on the Stiefel manifold, so the third term is $\eta^2 \boldsymbol{I}$. Altogether we get $(1+\eta^2)\boldsymbol{I}$, and we only need to divide by $\sqrt{1+\eta^2}$ to obtain the retraction:
\begin{equation}\boldsymbol{W}\quad\leftarrow\quad\frac{\boldsymbol{W} - \eta\boldsymbol{\Phi}}{\sqrt{1+\eta^2}}\end{equation}
At this point, readers may notice a deeper issue lurking here: whether for the relatively simple orthogonal manifold or the more complex Stiefel manifold, what precision should we actually use for computation? "Orthogonality" is an exact, quantitative constraint — $\boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{I}$ encodes $m(m+1)/2$ equality constraints — and it's easy to imagine that iterating the above formula at low precision will, over time, drift significantly away from orthogonality, not to mention the errors accumulated while solving $\boldsymbol{\Phi}$.
For this reason, I believe that unless we periodically apply an orthogonalization operation to the parameters (i.e. $\boldsymbol{W}\leftarrow\msign(\boldsymbol{W})$) to pull them back onto the orthogonal manifold, the computation during the solving process needs at least FP32 precision. Given that the number of parameters requiring orthogonality constraints is usually not very large, this generally isn't too costly.
Summary
This post has generalized "Muon + orthogonal manifold" from the previous post to the more general "Muon + Stiefel manifold," with the main contribution being an iterative algorithm for solving the corresponding update.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.