Muon on the Stiefel Manifold Admits an Exact Closed-Form Update

Research Question and Approach

Optimization on the Stiefel manifold—the space of matrices with orthonormal columns—is a fundamental task in machine learning and scientific computing. Recent work introduced Muon, a matrix-aware optimizer that outperforms standard methods like Adam by accounting for the underlying matrix structure through a spectral-norm-constrained linear minimization oracle (LMO). However, adapting Muon to the Stiefel manifold has proven difficult, with existing methods relying on heuristic, iterative, or relaxed subproblems that introduce significant computational overhead.

This paper addresses this challenge by reformulating the Stiefel Muon problem (SMP). The authors derive an auxiliary optimization problem over skew-symmetric matrices and demonstrate that it admits an exact closed-form solution. By leveraging this analytical result, they develop Skewon, a practical algorithm that solves the SMP directly in the tangent space of the Stiefel manifold.

The authors establish that the SMP is equivalent to a specific optimization problem over skew-symmetric matrices. They provide an analytical solution for this problem, which allows for an efficient implementation that avoids the need for iterative approximations. The resulting algorithm, Skewon, maintains the same asymptotic computational complexity as the original Muon optimizer. Specifically, for cases where the matrix dimension $n$ is much larger than the rank $p$ ($p \ll n$), the algorithm achieves a per-iteration complexity of $O(np^2)$, making it highly suitable for deep learning applications. Furthermore, the authors provide first-order convergence guarantees for Skewon in the smooth non-convex setting, ensuring its theoretical reliability.

By providing an exact, closed-form solution to the Stiefel Muon problem, this work eliminates the computational bottlenecks associated with previous Riemannian adaptations of Muon. Skewon offers a principled and efficient way to perform orthogonality-constrained optimization, potentially enabling faster and more stable training for neural network architectures that rely on orthogonal weight matrices, such as those used in low-rank adaptation or invariant-subspace computations.