Muon Dynamics as a Spectral Wasserstein Flow
Not on the map. Why?
Abstract (arXiv)
Gradient normalization stabilizes deep-learning optimization, and spectral normalizations are especially natural for matrix-shaped parameter blocks; Muon is the motivating example. We study an idealized deterministic, continuous-time, vanishing-momentum version of this idea in the mean-field regime, where wide models are represented by probability measures on parameter space. Starting from normalized matrix flows, we introduce Spectral Wasserstein distances indexed by norms γ on positive semidefinite matrices: the trace norm gives classical W2, the operator norm gives the Muon geometry, and Schatten norms interpolate between them. We develop the static Kantorovich formulation, a max-min robust-cost representation, Gaussian reductions extending the Bures formula, and for monotone norms, prove equivalence with a Benamou--Brenier formulation. This yields a gradient-flow interpretation of the mean-field normalized training dynamics. We illustrate these findings by numerical experiments on MMD flows, Gaussian reductions, two-layer ReLU models, and shallow attention.
Key facts. Cited by 3 papers in the map (data through 7 Oct 2026). Builds on 16 of its 29 references that are in the map. First posted on arXiv 6 Apr 2026.
Builds on
16 of 29 references.
References · 16
- 2018Demystifying MMD GANscited by 1,268
- 2018A Mean Field View of the Landscape of Two-Layers Neural Networkscited by 463
- 2018On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transportcited by 422
- 2018Shampoo: Preconditioned Stochastic Tensor Optimizationcited by 413
- 2015Generative Moment Matching Networkscited by 252
- 2025Training Deep Learning Models with Norm-Constrained LMOscited by 177
- 2015Optimal Entropy-Transport problems and a new Hellinger-Kantorovich distance between positive measurescited by 138
- 2020Momentum Improves Normalized SGDcited by 136
- 2019Subspace Robust Wasserstein Distancescited by 89
- 2017Revisiting Normalized Gradient Descent: Fast Evasion of Saddle Pointscited by 28
- 2016Wasserstein Discriminant Analysiscited by 25
- 2018Existence, Duality, and Cyclical monotonicity for weak transport costscited by 21
4 more
- 2023Structured Transforms Across Spaces with Cost-Regularized Optimal Transportcited by 8
- 2013Matrix-valued Monge-Kantorovich Optimal Mass Transportcited by 4
- 2016Matrix Optimal Mass Transport: A Quantum Mechanical Approachcited by 4
- 2023Covariance-modulated optimal transport and gradient flowscited by 3
Built on by
Citing papers per year: 2026: 3.
3 papers.
- 2026From SGD to Muon: Adaptive Optimization via Schatten-p Normscited by 0
- 2026Muon as a Residual Connectioncited by 0
- 2026Physical Muon: Orthogonalization as an Equilibrium Computationcited by 0
- 2026Physical Muon: Orthogonalization as an Equilibrium Computation29 Sep 2026
- 2026Muon as a Residual Connection1 Jul 2026
- 2026From SGD to Muon: Adaptive Optimization via Schatten-p Norms19 May 2026