Intrinsic Muon: Spectral Optimization on Riemannian Matrix Manifolds

TL;DR

This paper introduces intrinsic Muon (iMuon), a Riemannian manifold-based spectral optimization framework that provides closed-form updates and convergence guarantees for various matrix manifolds and norms, improving efficiency and symmetry preservation.

Contribution

It develops a unified intrinsic optimization framework on Riemannian manifolds that yields closed-form solutions and convergence guarantees for multiple matrix constraints and norms.

Findings

iMuon achieves closed-form updates on several manifolds.

Convergence rates depend only on manifold dimension, often only on rank.

Experiments show effective fine-tuning of LLMs, image classification, and subspace learning.

Abstract

Muon and related norm-constrained matrix optimizers have become central to large-scale learning problems. They are formulated as a linear maximization oracle (LMO) over an ambient matrix-norm ball in unconstrained Euclidean space. However, these do not generalize cleanly to manifold-valued parameters such as low-rank factorizations, orthogonality constraints, or symmetric positive definite (SPD) matrices. Naively restricting the Muon LMO to the tangent space (i) breaks quotient symmetries and (ii) couples the tangent-space constraint with an ambient norm bound, thereby obstructing closed-form solutions on various manifolds of interest. We resolve both issues with a single observation: every Riemannian metric canonically lifts a unitarily invariant Euclidean norm to an intrinsic norm on each tangent space, and the resulting intrinsic norm constrained LMO is symmetry preserving. Building on this, we introduce intrinsic Muon (iMuon), a unified framework that yields closed-form updates on the fixed-rank, SPD, Stiefel, and Grassmann manifolds for any unitarily invariant norm, including the spectral, Frobenius, and nuclear norms. We establish convergence guarantees for both deterministic and stochastic iMuon with rate constants that depend only on the manifold dimension. Notably, on the fixed-rank manifold this constant depends only on the rank, making the rate independent of factor conditioning and removing the runtime factor-rescaling required by prior work. Experiments on LoRA finetuning of LLMs, image classification, and subspace learning illustrate the efficacy of the proposed approach.