English

Muon Dynamics as a Spectral Wasserstein Flow

Optimization and Control 2026-05-11 v2 Artificial Intelligence Machine Learning

Abstract

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 γ\gamma on positive semidefinite matrices: the trace norm gives classical W2W_2, 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.

Keywords

Cite

@article{arxiv.2604.04891,
  title  = {Muon Dynamics as a Spectral Wasserstein Flow},
  author = {Gabriel Peyré},
  journal= {arXiv preprint arXiv:2604.04891},
  year   = {2026}
}
R2 v1 2026-07-01T11:55:37.777Z