[2604.04891] Muon Dynamics as a Spectral Wasserstein Flow

Mathematics > Optimization and Control arXiv:2604.04891 (math) [Submitted on 6 Apr 2026] Title:Muon Dynamics as a Spectral Wasserstein Flow Authors:Gabriel Peyré View a PDF of the paper titled Muon Dynamics as a Spectral Wasserstein Flow, by Gabriel Peyr\'e View PDF HTML (experimental) 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 ...