[2604.04891] Muon Dynamics as a Spectral Wasserstein Flow

Abstract:Gradient normalization is central in deep-learning optimization because it stabilizes training and reduces sensitivity to scale. For deep architectures, parameters are naturally grouped into matrices or blocks, so spectral normalizations are often more faithful than coordinatewise Euclidean ones; Muon is the main motivating example of this paper. More broadly, we study a family of spectral normalization rules, ranging from ordinary gradient descent to Muon and intermediate Schatten-type schemes, in a mean-field regime where parameters are modeled by probability measures. We introduce a family of Spectral Wasserstein distances indexed by a norm gamma on positive semidefinite matrices. The trace norm recovers the classical quadratic Wasserstein distance, the operator norm recovers the Muon geometry, and intermediate Schatten norms interpolate between them. We develop the static Kantorovich formulation, prove comparison bounds with W2, derive a max-min representation, and obtain a conditional Brenier theorem. For Gaussian marginals, the problem reduces to a constrained optimization on covariance matrices, extending the Bures formula and yielding a closed form for commuting covariances in the Schatten family. For monotone norms, including all Schatten cases, we prove the equivalence between the static and dynamic Benamou-Brenier formulations, deduce that the resulting transport cost is a genuine metric equivalent to W2 in fixed dimension, and show that the induced Gaussian covariance cost is also a metric. We then interpret the associated normalized continuity equation as a Spectral Wasserstein gradient flow, identify its exact finite-particle counterpart as a normalized matrix flow, obtain first geodesic-convexity results, and show how positively homogeneous mean-field models induce a spectral unbalanced transport on the sphere.

Muon Dynamics as a Spectral Wasserstein Flow Gabriel Peyré CNRS and ENS, PSL Université gabriel.peyre@ens.fr May 11, 2026 arXiv:2604.04891v2 [math.OC] 8 May 2026…

saved by

• Yudhister Joel Kumar

related reading

• Deriving Muonjeremybernste.in
• Modular Manifolds - Thinking Machines Labthinkingmachines.ai
• A Spectral Condition for Feature Learningarxiv.org
• [2604.06366] Stochastic Gradient Descent in the Saddle-to-Saddle Regime of Deep Linear Networksarxiv.org
• Muon: An optimizer for hidden layers in neural networks | Keller Jordan blogkellerjordan.github.io
• Understanding Muonlakernewhouse.com
• An optimization perspective on log-concave sampling and beyond | Sinho Chewichewisinho.github.io
• Understanding the Neural Tangent Kernel – EigenTaleseigentales.com
• Ji-Ha's Blogjiha-kim.github.io
• Notesandyrdt.com
• [2601.08393] Controlled LLM Training on Spectral Spherearxiv.org
• Why Momentum Really Worksdistill.pub