Steepest Descent on Manifolds: 4. Muon + Spectral Sphere

流形上的最速下降：4. Muon + 谱球面 苏剑林 · 2025-08-21 · Mathematics matrix optimizer muon constraint steepest descent

Readers who have made it through the first three posts in this series will by now be familiar with our "routine": first impose a constraint on the update, find the steepest-descent direction under that constraint, then impose a constraint on the parameters themselves and find the new steepest-descent direction. When solving the parameter-constrained problem, we simplify the constraint using the "first-order approximation is good enough" principle, which geometrically corresponds to working in the "tangent space." We then introduce undetermined coefficients to turn the problem into an unconstrained one, write down the analytical form of the solution, and finally solve for the coefficients numerically.

In this post we work through another new example — Muon under a spectral-sphere constraint — which is the analogue and generalization of the first post in the series, Steepest Descent on Manifolds: 1. SGD + Hypersphere. This is worth considering whenever we want the spectral norm of the parameters to stay constant throughout training. Of course, it can also simply be treated as an exercise for practice.

Problem Statement

In Steepest Descent on Manifolds: 2. Muon + Orthogonal and Steepest Descent on Manifolds: 3. Muon + Stiefel, we already discussed in detail how Muon interacts with orthogonality constraints, so we won't repeat that background here and will jump straight to the problem statement:

\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,\,\, \Vert\boldsymbol{W}\Vert_2 = 1,\,\,\Vert\boldsymbol{W} - \eta \boldsymbol{\Phi}\Vert_2=1\end{equation}more

where $\boldsymbol{W},\boldsymbol{\Phi}\in\mathbb{R}^{n\times m}(n \geq m)$, and $\Vert\cdot\Vert_2$ is the spectral norm. Of course, if we like, the latter two spectral norms could be replaced by other norms — for instance, the $F$ norm corresponds to the "Muon + hypersphere" combination.

The "first-order approximation is good enough" principle requires us to compute the gradient of the spectral norm, which we already covered in From the Gradient of the Spectral Norm to a New Kind of Weight Decay and Derivatives of the SVD. The answer is $\nabla_{\boldsymbol{W}}\Vert\boldsymbol{W}\Vert_2=\boldsymbol{u}_1 \boldsymbol{v}_1^{\top}$, where $\boldsymbol{u}_1,\boldsymbol{v}_1$ are the pair of singular vectors corresponding to the largest singular value of $\boldsymbol{W}$, which can be found via power iteration. This result also assumes the largest singular value is unique; we'll discuss the non-unique case later.

If we instead use the $F$ norm, we get $\nabla_{\boldsymbol{W}}\Vert\boldsymbol{W}\Vert_F=\boldsymbol{W}/\Vert\boldsymbol{W}\Vert_F$. In general, regardless of which norm we use, there always exists a matrix $\boldsymbol{\Theta}$ depending only on $\boldsymbol{W}$, such that $\nabla_{\boldsymbol{W}}\Vert\boldsymbol{W}\Vert=\boldsymbol{\Theta}$, so from $\Vert\boldsymbol{W}\Vert = 1$ and $\Vert\boldsymbol{W} - \eta \boldsymbol{\Phi}\Vert=1$ we obtain the first-order approximation $0 = \langle\boldsymbol{\Theta},\boldsymbol{\Phi}\rangle_F = \tr(\boldsymbol{\Theta}^{\top} \boldsymbol{\Phi})$. Hence, under the first-order approximation, the general formulation of this class of problems is:

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

Undetermined Coefficients

The routine remains the same: introduce an undetermined coefficient $\lambda$, and we have

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

The last inequality is exactly the Muon result itself, analogous to the Hölder inequality for two vectors, with equality holding at

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

The remaining task is to solve for a $\lambda$ satisfying the constraint $\tr(\boldsymbol{\Theta}^{\top} \boldsymbol{\Phi})=0$, and we're done.

Because of the presence of $\msign$, $\tr(\boldsymbol{\Theta}^{\top} \boldsymbol{\Phi})=0$ is in fact a nonlinear equation, and the author suspects it has no closed-form solution — so we turn to a numerical approach instead. Fortunately, having gone through Steepest Descent on Manifolds: 3. Muon + Stiefel, we're now well equipped to construct an iterative scheme for equations of this kind.

Iterative Solution

First, from the definition $\msign(\boldsymbol{M}) = \boldsymbol{M}(\boldsymbol{M}^{\top}\boldsymbol{M})^{-1/2}$ we can write $\boldsymbol{\Phi}=(\boldsymbol{G} + \lambda\boldsymbol{\Theta})\boldsymbol{Q}^{-1}$, where $\boldsymbol{Q}=((\boldsymbol{G} + \lambda\boldsymbol{\Theta})^{\top}(\boldsymbol{G} + \lambda\boldsymbol{\Theta}))^{1/2}$, so that

\begin{equation}\tr(\boldsymbol{\Theta}^{\top} \boldsymbol{\Phi})=0\qquad\Rightarrow\qquad \lambda = -\frac{\tr(\boldsymbol{\Theta}^{\top}\boldsymbol{G}\boldsymbol{Q}^{-1})}{\tr(\boldsymbol{\Theta}^{\top}\boldsymbol{\Theta}\boldsymbol{Q}^{-1})}\end{equation}

Note that this is not an analytical solution, since $\boldsymbol{Q}$ itself also depends on $\lambda$. But based on the expression above we can build an iterative scheme: plug in an initial $\lambda$, compute $\boldsymbol{Q}= (\boldsymbol{G} + \lambda\boldsymbol{\Theta})^{\top}\boldsymbol{\Phi}$, then substitute back into the equation above to update $\lambda$, and repeat until convergence.

However, although this iterative scheme is workable in principle, it requires computing $\boldsymbol{Q}^{-1}$. Even though we already described an efficient algorithm for this in Efficient Computation of Matrix r-th Roots and Inverse r-th Roots, from an "Occam's razor" standpoint we'd still prefer to avoid introducing any iteration beyond $\msign$ itself. So the author tried to find another iterative scheme that avoids computing $\boldsymbol{Q}^{-1}$. To that end, let's first write

\begin{equation}\boldsymbol{\Theta}^{\top} \boldsymbol{\Phi} = \boldsymbol{\Theta}^{\top}(\boldsymbol{G} + \lambda\boldsymbol{\Theta})\boldsymbol{Q}^{-1}\end{equation}

For our target, the trace of the expression above equals zero. We can explicitly subtract $\tr(\boldsymbol{\Theta}^{\top} \boldsymbol{\Phi})\boldsymbol{I}/m$ from the left-hand side to guarantee this condition holds:

\begin{equation}\boldsymbol{\Theta}^{\top} \boldsymbol{\Phi} - \tr(\boldsymbol{\Theta}^{\top} \boldsymbol{\Phi})\boldsymbol{I}/m = \boldsymbol{\Theta}^{\top}(\boldsymbol{G} + \lambda\boldsymbol{\Theta})\boldsymbol{Q}^{-1}\end{equation}

Now, multiplying both sides by $\boldsymbol{Q}$ and taking the trace lets us solve for $\lambda$:

\begin{equation}\lambda = \frac{\tr(\boldsymbol{\Theta}^{\top} \boldsymbol{\Phi}\boldsymbol{Q}) - \tr(\boldsymbol{\Theta}^{\top}\boldsymbol{\Phi}) \tr(\boldsymbol{Q})/m - \tr(\boldsymbol{\Theta}^{\top}\boldsymbol{G})}{\tr(\boldsymbol{\Theta}^{\top}\boldsymbol{\Theta})}\end{equation}

This way, the iteration no longer requires computing $\boldsymbol{Q}^{-1}$.

Reference Code

Taking $\lambda=-\tr(\boldsymbol{\Theta}^{\top}\boldsymbol{G})/\tr(\boldsymbol{\Theta}^{\top}\boldsymbol{\Theta})$ as the initial value, the test code is as follows:

import numpy as np

def msign(g):
    """奇异值分解精确计算msign
    """
    u, s, vh = np.linalg.svd(g, full_matrices=False)
    return u @ np.diag(np.sign(s)) @ vh

def dot(a, b):
    """恒等于 np.trace(a.T @ b)
    """
    return (a * b).sum()

n, m = 100, 50
w = np.random.randn(n, m) / m**0.5
g = np.random.randn(n, m) / m**0.5
u, s, vh = np.linalg.svd(w, full_matrices=False)
theta = u[:, :1] @ vh[:1]

lamb = - dot(theta, g) / dot(theta, theta)
for i in range(10):
    phi = msign(z := g + lamb * theta)
    print('step:', i, ', inner product:', dot(phi, g), ', tangent error:', dot(theta, phi))
    q, x = z.T @ phi, theta.T @ phi
    lamb = (dot(x, q) - np.trace(x) * np.trace(q) / m - dot(theta, g)) / dot(theta, theta)

Other Details

As with the previous three posts, because we rely on the "first-order approximation is good enough" principle, the spectral norm of $\boldsymbol{W} - \eta\boldsymbol{\Phi}$ is only accurate to order $1 + \mathcal{O}(\eta^2)$ and generally cannot reach exactly 1. So we still need to apply a spectral normalization step:

\begin{equation}\boldsymbol{W}\quad\leftarrow\quad \frac{\boldsymbol{W} - \eta\boldsymbol{\Phi}}{\Vert\boldsymbol{W} - \eta\boldsymbol{\Phi}\Vert_2}\end{equation}

Fortunately, the spectral norm can be computed efficiently via power iteration, so this is not a particularly expensive computation (compared with the iteration for $\msign$ itself).

It's also worth analyzing the case where the largest singular value is not unique. In practice, this special case can usually be ignored during numerical computation, but for theoretical completeness it should be included in the analysis. In this case the corresponding singular vectors are also not unique, which is equivalent to saying there are multiple distinct tangent spaces, and the actual feasible space is the intersection of these tangent spaces. Taking the case of two largest singular values as an example, the problem becomes

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

where $\boldsymbol{\Theta}_1=\boldsymbol{u}_1 \boldsymbol{v}_1^{\top}, \boldsymbol{\Theta}_2=\boldsymbol{u}_2 \boldsymbol{v}_2^{\top}$. Introducing two undetermined coefficients $\lambda_1,\lambda_2$, we can solve to get

\begin{equation}\boldsymbol{\Phi} = \msign(\boldsymbol{G} + \lambda_1\boldsymbol{\Theta}_1+ \lambda_2\boldsymbol{\Theta}_2)\end{equation}

Next we need to solve the system of equations $\tr(\boldsymbol{\Theta}_1^{\top} \boldsymbol{\Phi})=0,\tr(\boldsymbol{\Theta}_2^{\top} \boldsymbol{\Phi})=0$. Similarly, introducing

\begin{equation}\boldsymbol{Q}=((\boldsymbol{G} + \lambda_1\boldsymbol{\Theta}_1+ \lambda_2\boldsymbol{\Theta}_2)^{\top}(\boldsymbol{G} + \lambda_1\boldsymbol{\Theta}_1+ \lambda_2\boldsymbol{\Theta}_2))^{1/2} = (\boldsymbol{G} + \lambda_1\boldsymbol{\Theta}_1+ \lambda_2\boldsymbol{\Theta}_2)^{\top}\boldsymbol{\Phi}\end{equation}

we can write the system of equations

\begin{equation}\begin{gathered} \boldsymbol{\Theta}_1^{\top} \boldsymbol{\Phi} - \tr(\boldsymbol{\Theta}_1^{\top} \boldsymbol{\Phi})\boldsymbol{I}/m = \boldsymbol{\Theta}_1^{\top}(\boldsymbol{G} + \lambda_1\boldsymbol{\Theta}_1+ \lambda_2\boldsymbol{\Theta}_2)\boldsymbol{Q}^{-1} \\ \boldsymbol{\Theta}_2^{\top} \boldsymbol{\Phi} - \tr(\boldsymbol{\Theta}_2^{\top} \boldsymbol{\Phi})\boldsymbol{I}/m = \boldsymbol{\Theta}_2^{\top}(\boldsymbol{G} + \lambda_1\boldsymbol{\Theta}_1+ \lambda_2\boldsymbol{\Theta}_2)\boldsymbol{Q}^{-1} \\ \end{gathered}\end{equation}

Multiplying both sides by $\boldsymbol{Q}$ and taking the trace turns this into a linear system of two equations in the two unknowns $\lambda_1,\lambda_2$, which can then be solved to construct an iterative scheme. We won't go into the details here; readers interested in practicing are welcome to fill them in themselves.

Summary

This post mainly considered the form Muon takes when the parameters are constrained to a spectral norm sphere, or more generally to some other norm sphere. Building on the previous three posts, this one presents no particularly new technical difficulties, and readers can simply treat it as a supplementary exercise for practice.