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

2026-08-06

AI Summary

This work addresses the lack of an efficient and exact update mechanism for Muon-type optimization methods on the Stiefel manifold by deriving, for the first time, a closed-form solution to their update rule and proposing a novel algorithm named Skewon. By integrating differential geometry with matrix optimization theory, Skewon fully exploits the underlying manifold structure, thereby circumventing the inefficiencies inherent in existing approaches that rely on heuristic or iterative approximations. Under non-convex smooth settings, Skewon not only enjoys first-order convergence guarantees but also significantly enhances the practicality and performance of Muon methods for orthogonally constrained optimization while maintaining computational efficiency.

Abstract

We study Muon, a recently proposed matrix-aware optimization method, in the context of the Stiefel manifold. This manifold consists of matrices with orthonormal columns and is ubiquitous in machine learning and scientific computing. Existing extensions of Muon to this manifold rely on heuristic, approximate, or iterative updates with varying computational efficiency. We show that the corresponding Stiefel Muon update admits an exact closed-form solution and use this result to develop Skewon, a practical algorithm for orthogonality-constrained optimization with an efficient implementation. We further establish first-order convergence guarantees for Skewon in the smooth non-convex setting.

Problem

• Stiefel manifold
• Muon
• orthogonality-constrained optimization
• closed-form update
• matrix optimization

Innovation

• Stiefel manifold
• closed-form update
• orthogonality-constrained optimization
• Muon
• Skewon