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

Mikhail Solonko, Molozhavenko Alexander, Maxim Rakhuba

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 ubiq…