Research map / Optimization / Implicit bias and generalization

Muon Dynamics as a Spectral Wasserstein Flow

NeurIPS 2026 · Optimization › Implicit bias and generalization · approach: Gradient flow and mean-field dynamics · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Builds a Spectral Wasserstein distance indexed by monotone matrix gauges on PSD matrices (operator norm gives Muon geometry), proves static-dynamic equivalence and gradient-flow interpretation of normalized mean-field dynamics, illustrated with numerical experiments.

Problem

Muon-type spectrally normalized optimizers act on matrix-shaped parameters, but their mean-field (wide-model) limit lacks a transport geometry; the paper asks whether normalized dynamics are gradient flows of a distance on probability measures.

Approach

Define a Spectral Wasserstein distance W_gamma via a matrix gauge gamma on displacement covariances (trace gives W_2, operator norm gives Muon geometry, Schatten interpolates); prove a max-min cost-robust representation, Gaussian reduction, and for monotone gauges equivalence with a Benamou-Brenier formulation, so normalized mean-field flows are metric gradient flows.

Key results

• Proves W_gamma equals a max over cost matrices of quadratic W_Q (max-min representation) and, for monotone gauges, static equals Benamou-Brenier dynamic distance (Theorem on static-dynamic equivalence).
• For Schatten norms, d^{1/p-1} W_2^2 <= W_{gamma_p}^2 <= W_2^2, with sharpness attained by isotropic Gaussian pairs.
• In two-layer ReLU and shallow attention experiments, p=infinity reaches lower energy than p=1 during the initial phase after time rescaling.

Surprising findings

• For the operator-norm (Muon) geometry the optimal cost matrix is generically rank one, so no Monge map is expected even for measures with density.
• The MMD loss curves for p=1 and p=infinity remain close even though particle trajectories differ markedly.

Limitations

• Global-in-time well-posedness of the spectral Wasserstein gradient-flow PDE is left open.
• Muon's practical benefits for attention layers are explicitly not explained; the work is a formal framework.
• Momentum and Newton-Schulz approximations are outside the gradient-flow theory (momentum yields a phase-space equation).

Experimental setup

Datasets: synthetic Gaussian measures synthetic teacher-student ReLU data synthetic shallow attention teacher-student data synthetic MMD point clouds

Metrics: energy decay MMD objective convergence profile effective number of attended tokens

Baselines: W_2 gradient flow (trace norm, p=1)

Builds on: Muon Wasserstein gradient flow Benamou-Brenier formulation subspace robust Wasserstein distance normalized gradient methods via LMO

Closest papers by meaning

• Phases of Muon: When Muon Eclipses SignSGD (NeurIPS 2026 · similarity 0.88)
• Intrinsic Muon: Spectral Optimization on Riemannian Matrix Manifolds (NeurIPS 2026 · similarity 0.86)
• Denoise First, Orthogonalize Later: Understanding Momentum in Muon via Spectral Filtering (NeurIPS 2026 · similarity 0.86)
• On the Convergence Analysis of Muon (NeurIPS 2026 · similarity 0.85)
• Towards Understanding the Power and Limits of the Muon Optimizer: A River-Valley Perspective (NeurIPS 2026 · similarity 0.85)
• AMUSE: Anytime Muon with Stable Gradient Evaluation (NeurIPS 2026 · similarity 0.84)