Manifold Gradient Descent: 2. Muon with Orthogonal Constraints

This article continues our series on constrained optimization. In the previous article "Manifold Gradient Descent: 1. SGD on Hyperspheres", we revisited the "Least Action Principle" for optimizers, proposing that the core distinction among different optimizers lies in the various constraints imposed on the update vector. If this constraint is the Euclidean norm, the corresponding steepest descent is SGD. Furthermore, we discussed results when simultaneously imposing magnitude constraints on parameters, constituting steepest descent on hyperspherical manifolds.

However, the previous article should be considered merely a "warm-up," as it addressed relatively simple vector parameter optimization. This article formally enters more challenging territory—optimization parameters transition from vectors to matrices, with the incremental constraint changed to the spectral norm, giving rise to the Muon optimizer. Subsequently, we impose orthogonal constraints on parameters, yielding the Muon optimizer on orthogonal manifolds.

Problem Formulation#

As for why the spectral norm is selected, readers may refer to "Muon Optimizer Appreciation: The Essential Leap from Vectors to Matrices", "Muon Sequel: Why Did We Choose to Experiment with Muon?", and "Higher-Order MuP: More Concise yet More Insightful Spectral Conditioning Scaling"; we will not repeat the introduction here. In brief, the spectral norm is the most compact norm revealing linear layer transformations, making it more suitable as a measure of matrix "stability."

Following previous steps, applying a first-order approximation to $\mathcal{L}(\boldsymbol{W} +\Delta\boldsymbol{W})$ yields $\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$ denotes flattening both matrices to vectors and computing their inner product, equivalently $\newcommand{tr}{\mathop{\text{tr}}}\tr(\boldsymbol{G}^{\top}\Delta\boldsymbol{W})$. Setting $\Delta\boldsymbol{W} = -\eta \boldsymbol{\Phi}$, the original problem simplifies to:

These transformation steps are general; if details are forgotten, please refer back to the previous article.

Fundamental Result#

The solution process for objective (2) was already provided in the "Matrix Norms" section of "Muon Optimizer Appreciation: The Essential Leap from Vectors to Matrices". However, for completeness, we restate 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}$, where $r$ is the rank of $\boldsymbol{G}$. We have:

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$, hence $\boldsymbol{u}_i^{\top}\boldsymbol{\Phi}\boldsymbol{v}_i\leq 1$. Therefore:

Here $\Vert\cdot\Vert_*$ is called the Nuclear Norm of the matrix. Equality holds when all $\boldsymbol{u}_i^{\top}\boldsymbol{\Phi}\boldsymbol{v}_i$ equal 1, yielding:

Orthogonal Manifold#

Thus far, we have demonstrated that for matrix parameters, the direction of steepest descent under spectral norm constraints is not the opposite gradient $-\boldsymbol{G}$, but requires an additional $\msign$ operator: $-\msign(\boldsymbol{G})$. This is precisely the Muon optimizer used in training Kimi K2, making it one of the most competitive optimizers currently. This conversely indicates that the spectral norm as a stability constraint for matrices is highly appropriate.

Of course, the results so far are not new. Now we begin exploring novel territory—imposing orthogonal constraints $\boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{I}$ on parameters $\boldsymbol{W}$ (source: "Orthogonal manifold"). This divides into two cases: first, when $n=m$, $\boldsymbol{W}$ is a proper orthogonal matrix satisfying $\boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{W}\boldsymbol{W}^{\top}=\boldsymbol{I}$; second, when $n > m$, $\boldsymbol{W}\boldsymbol{W}^{\top}=\boldsymbol{I}$ cannot be satisfied, typically called a semi-orthogonal matrix, with the corresponding space termed the Stiefel manifold.

When might orthogonal constraints be applied? In fact, scenarios are not rare. For example, in classification problems, if classes are known to have little correlation, orthogonal constraints on class matrices may be considered, though often implemented through regularization terms $\Vert\boldsymbol{W}^{\top}\boldsymbol{W}-\boldsymbol{I}\Vert_F^2$ for approximate orthogonality. Another example is in LoRA scenarios, where parameterization of form $\boldsymbol{A}\boldsymbol{B}$ inherently contains redundancy, which can be reduced via orthogonal constraints (reference), etc.

Solution Process#

To solve objective, similar to the previous article, we introduce an undetermined coefficient matrix $\boldsymbol{\Lambda}\in\mathbb{R}^{m\times m}$, obtaining:

The second equality uses trace identities $\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}$. According to the previous Muon result, equality holds when:

The remaining problem is finding the real symmetric matrix $\boldsymbol{X} = \boldsymbol{\Lambda} + \boldsymbol{\Lambda}^{\top}$ such that $\boldsymbol{W}^{\top}\boldsymbol{\Phi}$ is skew-symmetric. For $n=m$, this is easily solved because $\boldsymbol{W}^{\top}$ can be absorbed into $\msign$:

Retraction Operation#

In summary, for $n=m$, our final result is:

Therefore, the updated variable is:

This is not an orthogonal matrix but is accurate to $\mathcal{O}(\eta^2)$, consistent with our "first-order approximation suffices" principle. To verify this, simply check:

Summary#

In this article, we revisited the conclusion that imposing spectral norm constraints on matrix parameter updates yields the Muon optimizer. We then explored the form of the Muon optimizer when orthogonal constraints are imposed on parameters. If you wish parameters to remain orthogonal matrices during updates, this article may provide some reference value.