Muon Dynamics as a Spectral Wasserstein Flow
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 on positive semidefinite matrices: the trace norm gives classical , 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}
}