Muon Dynamics as a Spectral Wasserstein Flow

Gabriel Peyré

math.OCApr 6, 2026

Muon Dynamics as a Spectral Wasserstein Flow

Authors: Gabriel Peyré

Organizations: CNRS and ENS, PSL Université

Abstract

Gradient normalization can stabilize 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 recovers the classical $W_2$ distance, the operator norm gives the Muon geometry, and Schatten norms interpolate between them. We develop a static Kantorovich formulation, a max--min robust-cost representation, and Gaussian reductions extending the Bures formula. For monotone norms, we prove equivalence with a Benamou--Brenier formulation and obtain a gradient-flow interpretation of normalized mean-field training dynamics. Numerical experiments illustrate the resulting flows for MMD objectives, Gaussian reductions, two-layer ReLU models, and shallow attention.

Figures & tables

Figure 1: Static coupling illustrations for two source-target pairs: exact trace assignments (

p=1

) and heuristic operator couplings (

p=\infty

). Only a subsample of matched segments is displayed.

Figure 2: Displacement-interpolation density snapshots induced by the computed static couplings (exact for

p=1

, heuristic for

p=\infty

), for

p=1

(first five panels) and for

p=\infty

(last five panels). We use

n=1500

particles. As in Figure 1 , the first and last snapshots correspond to the same source and target distributions (

t=0

and

t=1

Panel of Figure 3 — Figure 3: Two-layer ReLU MLP trajectories in the reduced parameter plane $(|u|v_{1},|u|v_{2})$ . Dashed rays indicate the teacher neurons. The archived runs use separately normalized nominal horizons; neither intrinsic flow times nor wall-clock runtimes are matched.

Figure 4: MMD particle trajectories on separately normalized nominal horizons. The corresponding loss curves appear in Figure 11 , Appendix G.2 .

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Gaussian one-mode geometry. (a) Trajectories in

(r,s)

for

p\in\{1,2,4,\infty\}

. (b) Interior conserved foliation for

p=1

. After rank loss, trajectories use the continuation of Proposition F.5 ; interior invariants no longer apply.

Figure 6: Gaussian operator leaves and coupled modes. (a) Interior conserved foliation for

p=\infty

. (b) Rank-two trajectories in

(r_{1},r_{2})

, initially

r_{1}=r_{2}=0

and

(s_{1},s_{2})=(0.1,10)

, toward the indicated targets

(\beta_{1},\beta_{2})

. The same rank-boundary continuation is used as in Figure 5 .

Example	$n$	steps $(1,\infty)$	$\Delta t_{1}$	$\Delta t_{\infty}$	step cap
ReLU MLP	1920	$(2000,2000)$	75	50	0.20
Attention	200	$(960,480)$	100	180	0.45
MMD	200	$(6000,6000)$	5	3.2	none

Appendix

Table 1: Nominal notebook clocks. Caps apply to the Frobenius norm of the entire particle-matrix increment, not separately to particles.

Figure 7: ReLU MLP energy on separately normalized nominal horizons: solid

p=1

, dashed

p=\infty

. Colored vertical lines identify the snapshots in Figure 8 ; equal horizontal coordinates do not represent equal intrinsic flow times. Particle trajectories appear in Figure 3 .

Figure 8: Circular histogram snapshots for the ReLU MLP, in chronological order from left to right, with the same angular axis in all panels. Solid curves correspond to

p=1

and dashed curves to

p=\infty

(Muon); colors match the snapshot markers in Figure 7 . Dashed vertical lines mark teacher directions. Only the legacy legend text has been corrected.

Figure 9: Shallow-attention loss curves for five paired trials per regime, with

p=1

in blue and

p=\infty

in red. Each trial redraws both teacher and student. Effective support is approximately 16 tokens in the dense regime and 2.1 in the sparse regime. The horizontal axes use separately normalized nominal horizons, not matched intrinsic flow times.

Figure 10: Teacher-attention histograms for the dense and sparse regimes. Weights are truncated to

[0,0.1]

and renormalized on that interval; sparse weights larger than

0.1

are not displayed. Effective-token statistics in the text instead use all attention weights. Both panels use the same horizontal scale.

Figure 11: Archived MMD objective values for the trajectories in Figure 4 . Nominal horizons are normalized separately for

p=1

and

p=\infty

. These curves do not establish an intrinsic convergence-rate or wall-clock ordering.

Figure 12: Additional static Spectral Wasserstein couplings. Only a subsample of matching segments is displayed.

Figure 13: Additional displacement-interpolation density snapshots for the same shape pairs as Figure 12 .

Notation	Name	First use
$\mathbb{R}^{d}$	Euclidean space	Sec. 2
$\mathbb{S}^{d}$	Symmetric matrices	Sec. 2
$\mathbb{S}_{+}^{d}$	PSD matrices	Sec. 2
$\mathbb{S}_{++}^{d}$	Positive definite matrices	Prop. A.7
$S\preceq T$	Loewner order	Sec. 2
$I$ , $\mathrm{Id}$	Identity matrix/map	Sec. 2

Appendix

Table 2: Classical notation.

Notation	Name	Introduction
$\gamma$	Monotone matrix gauge on $\mathbb{S}_{+}^{d}$	Sec. 2
$\gamma_{p}$	Schatten- $p$ gauge	Sec. 2
$p^{\prime}$	Conjugate exponent of $p$	Rem. A.4
$q=(2p)^{\prime}$	Tangent-space dual exponent	Prop. A.5
$\mathcal{K}_{\gamma}$	Representing set	Def. 3.3
$\mathcal{P}_{\gamma}$	Polar set	App. A

Appendix

Table 3: Main notation introduced in the paper.

Neural Network Training Dynamics

Neural network training dynamics studies how parameters and representations change during optimization and how those changes shape learning, stability, and generalization.

See all 511 papers

Move on Muon : A Hamiltonian probability gradient flow perspective of Muon optimizer
May 22, 2026Aratrika Mustafi, Soumya Mukherjee, Bharath K. SriperumbudurMomentum Methods Muon Optimizer
Department of Statistics, Pennsylvania State University
The Spectral Dynamics and Noise Geometry of Muon
Jun 7, 2026Pierfrancesco Beneventano, Mahmoud Abdelmoneum, Tomaso PoggioMuon Optimizer Matrix Optimization
Massachusetts Institute of Technology
Intrinsic Muon: Spectral Optimization on Riemannian Matrix Manifolds
May 10, 2026Yibang Li, Bihari Lal Pandey, Ravi Sah +4Stiefel Manifold Optimization Riemannian Optimization
Nanyang Technological University, Singapore. · Indian Institute of Technology Bombay, India. · University of Sydney, Australia. +1

May 22, 2026J/K move · Enter open · S save

Move on Muon : A Hamiltonian probability gradient flow perspective of Muon optimizer

Aratrika Mustafi, Soumya Mukherjee, Bharath K. Sriperumbudur

Department of Statistics, Pennsylvania State University

Momentum Methods Muon Optimizer Mean-Field Theory Neural Network Training Dynamics

Open paper AskPDF arXiv

We develop a gradient flow on the space of probability measures defined on matrix-valued parameters induced by regularized Muon, an analytically smoothed version of the idealized Muon optimizer. The key observation is that the regularized orthogonalization map is the gradient of a smooth Fenchel-dual smoothing of the nuclear norm. This identifies the (regularized) Muon update as a mirror/prox step in the update variable, with momentum acting as the dual coordinate. We use this structure to lift Muon from a single matrix parameter to finite-particle probability objectives of the form

J(ρ)=R\left(\int F d ρ\right)

, a setting motivated by mean-field descriptions of neural-network training, and derive the inertial continuous-time limit. Using this structure, we derive the finite-particle continuous-time limit under the inertial scaling of step size and momentum, and then pass to a phase-space mean-field equation over probability laws on parameter-momentum pairs. The resulting flow can be shown to be a damped Hamiltonian probability dynamics whose kinetic energy is induced by the regularized Muon mirror potential. We prove an exact Hamiltonian dissipation identity, showing that the Hamiltonian energy decreases monotonically. While the target objective itself need not be monotone along the inertial Muon dynamics, under additional gradient-dominance, bounded-momentum, and curvature/alignment assumptions, we obtain continuous and discrete-time exponential convergence rates for the objective gap. We also study the well-posedness of the mean-field limit equation and establish propagation of chaos guarantees for the interacting particle system. Finally, we extend the formulation to Hilbert-valued feature maps on product matrix spaces, yielding a blockwise Muon probability flow applicable to smooth transformer mixture-of-experts models.

Abstract

Neural Network Training Dynamics

Explore similar work

Move on Muon : A Hamiltonian probability gradient flow perspective of Muon optimizer

The Spectral Dynamics and Noise Geometry of Muon

Intrinsic Muon: Spectral Optimization on Riemannian Matrix Manifolds