arXiv CSopen access

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

Move on Muon : A Hamiltonian probability gradient flow perspective of Muon optimizer Aratrika Mustafi∗ , Soumya Mukherjee and Bharath K. Sriperumbudur

arXiv:2605.23871v1 [stat.ML] 22 May 2026

Department of Statistics, Pennsylvania State University May 25, 2026 Abstract 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 fromR a single matrix parameter to
finite-particle probability objectives of the form J(ρ) = R F dρ , 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, boundedmomentum, 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.

1

Introduction

Optimization methods for deep neural networks often succeed because they exploit structure that is not visible in the scalar-coordinate view of the parameters. As model sizes grow and parameter blocks become highly structured, the geometry imposed by the optimizer can substantially influence both stability and speed. This is particularly visible for matrix-valued parameters, where the Frobenius geometry used by standard gradient methods is only one possible choice. Standard first-order optimizers such as SGD, Adam, and their variants treat these parameters largely through coordinatewise Euclidean updates. While this has been extremely effective in practice, recent work has shown that matrix geometry can play a more explicit role in optimization, particularly when the update direction is constrained or normalized according to spectral information. The Muon optimizer introduced by Jordan et al. (2024) is a recent and influential example of this matrix-geometric viewpoint. At a high level, Muon maintains a momentum variable for 1

each matrix parameter and updates the parameter in an orthogonalized version of that momentum direction. Empirical studies have shown that such orthogonalized matrix updates can be competitive in language-model training, and recent large-scale implementations have investigated the extent to which Muon can be made practical for LLM training through suitable scaling, weight decay, and implementation choices (Liu et al., 2025). In the idealized form, as considered in this paper, the orthogonalization map keeps the singular vectors of the momentum matrix and replaces its nonzero singular values by one. Thus, rather than moving directly in the raw momentum direction, Muon moves in a polar-factor direction. This operation makes the update insensitive to the scale of the singular values but highly sensitive to their singular subspaces. Practical implementations approximate this orthogonalization numerically, for instance using Newton-Schulz iterations, but the exact polar-factor map captures the central geometric mechanism of the method. These developments motivate a theoretical question : What continuous-time and probabilityspace dynamics are naturally associated with a momentum optimizer whose update direction is produced by spectral orthogonalization? A first answer is already suggested by convex geometry. The polar-factor direction is not an arbitrary normalization: it is a steepest descent direction over a spectral-norm unit ball. Equivalently, the ideal Muon step solves a linear minimization problem under an operator-norm constraint. This perspective has been developed in recent work on non-Euclidean trust-region interpretations of gradient orthogonalization and norm-constrained linear minimization oracles (Kovalev, 2025; Pethick et al., 2025). Related viewpoints also connect Muon-type methods to implicit spectral constraints and broader families of spectral optimizers (Chen et al., 2025a). This paper starts from this existing observation but takes a different route. We take the spectral-norm trust-region interpretation of hard Muon as a starting point and develop from it a probabilistic gradient-flow perspective on Muon with momentum. The key technical obstacle is that the hard orthogonalization map on the momentum variable P , P 7→ Orth(P ) is nonsmooth. Here, Orth(P ) = U V ⊤ with P = U ΣV ⊤ , Σ = diag(σ1 (P ), · · · , σq (P )), is the compact/reduced SVD of P . This nonsmoothness is not merely a technical inconvenience: it changes the nature of the continuous-time limit, replacing an ordinary differential equation by a differential inclusion. To obtain a well-defined smooth flow, we introduce a regularized Muon map ! σi (P ) Orthε (P ) = U diag p V ⊤, ε > 0, 2 2 σi (P ) + ε which keeps the singular directions of P but replaces the hard saturation of singular values by a smooth saturation. As ε ↓ 0, this map converges pointwise to the hard polar-factor map on fixedrank matrices. The role of the present regularization is to enable an exact variational representation of the orthogonalization operation needed for the Hamiltonian probability dynamics developed in the next section. The regularized Muon update then becomes a genuine mirror/prox step in the update or velocity variable, with the momentum serving as the corresponding dual coordinate. This type of regularization has been recently used as a technical tool in Kim et al. (2026). The second goal of the paper is to lift this mirror interpretation from a single matrix update to a probability-space description. We consider functionals of the form Z  J(ρ) = R F (W ) dρ(W ) , where ρ is a probability measure over matrix-valued parameters (or tuples of matrix-valued parameters). 2

This includes finite-particle objectives obtained from empirical measures, and it also provides a convenient mean-field perspective for studying populations of parameters. For an N -particle empirical law, the regularized Muon update induces a coupled particle system in the phase variables (Wi , Pi ). Under the inertial scaling of step-sizes, the discrete dynamics converge to finite-particle ODEs, leading to a McKean-Vlasov continuity equation for µt = Law(Wt , Pt ). The resulting probability flow has a damped Hamiltonian flow structure, instead of being an ordinary Wasserstein gradient flow. Hamiltonian-flow formalisms have been used in literature to study momentum-based acceleration on the space of probability measures (Wang and Li, 2022; Chen et al., 2025b). In our setting, the Hamiltonian structure is specific to Muon and is induced by the mirror map associated with regularized nuclear-norm smoothing. The resulting flow satisfies an exact Hamiltonian dissipation identity, which clarifies why the objective J itself need not decrease monotonically. As in other inertial or accelerated systems, energy can move between the objective and the momentum variable, while the total damped Hamiltonian decreases. Recent work Peyré (2026) considers specral Wasserstein flow whose ODE limit corresponds to a momentum-free Muon update with an extra nuclear norm scaling. Our Hamiltonian probability-flow formulation is closer to the version of Muon largely implemented in practice since it preserves the momentum variable, treats momentum as the dual coordinate of a smooth mirror map, and recovers the hard polar Muon update as ε ↓ 0. We then extend the formulation from a single matrix space to finite product spaces of matrix blocks. This allows the same mirror geometry to act on all blocks of structured models. This is important for neural network layers whose parameters consist of several matrices, and especially for mixture-of-experts models with both expert and router parameters. In the extended product space Θ, the regularized Muon potential is block-separable, and the corresponding orthogonalization (mirror) map OrthΘ ε applies the spectral regularization to each matrix block in a separable manner. This gives a product-space mirror step, a product-space Hamiltonian, and a corresponding dissipation identity. This generalization to product spaces allows us to analyze equally weighted mixture-of-experts where each expert model is a neural network or a transformer, as well as smooth-routing/selection variants of such mixture-of-experts. For transformer mixture-of-experts models, input-dependent routing is encoded inside the Hilbert-valued feature map F . Smooth unnormalized gates and softmax-normalized gates fit directly into the framework, while exact hard top-k routing lies outside the smooth theory and can be treated as a nonsmooth limiting case. This distinction is consistent with prior work on sparse mixture-of-experts routing and noisy router smoothing, emphasizing the discontinuity of sparse routing and the smoothing effect of noisy or softened routers (Shazeer et al., 2017; Fedus et al., 2022; Chen et al., 2022; Vaswani et al., 2017).

2

Muon as a mirror regularized trust-region step

In this section, we isolate the discrete-time Muon update for a single matrix variable and recast it in a form that admits a smooth mirror-map regularization. This reformulation interprets the orthogonalization step in the Muon update as the solution to a variational problem on matrix space. The resulting variational structure then enables the development of a Hamiltonian formulation of the Muon dynamics in later sections.
Let X = Rm×n with Frobenius pairing ⟨A, B⟩F = tr A⊤ B , and let q = min(m, n). For P = U ΣV ⊤ define, Orth(P ) = U V ⊤ , Orth(0) = 0. The idealized Muon update for a smooth objective F : X → R is given by, Pk+1 = βPk + (1 − β)∇F (Wk ) , Wk+1 = Wk − η Orth (Pk+1 ) , 3

(1)

where P denotes the momentum variable and W denotes the spatial variable. The following proposition demonstrates the convex-analytic structure behind this update. Proposition 1 (Spectral trust-region form of Muon). For every P ∈ X , ∥P ∥nuc =

sup ⟨P, G⟩F ,

− Orth(P ) ∈ argmin ⟨P, G⟩F .

∥G∥op ≤1

(2)

∥G∥op ≤1

Equivalently, if Φ0 (G) = ι{∥G∥op ≤1} (G) is the 0/∞-indicator of the spectral-norm unit ball, then Gk+1 ∈ argmin {⟨Pk+1 , G⟩F + Φ0 (G)} , G

Wk+1 = Wk + ηGk+1 . recovers ideal Muon after choosing the canonical minimizer Gk+1 = − Orth (Pk+1 ). The above proposition says that Muon linearizes through the momentum Pk+1 and then takes the steepest direction allowed by a spectral norm unit ball. The non-smoothness enters since Φ∗0 = ∥ · ∥nuc is not differentiable at rank-deficient matrices. To regularize this singular structure and obtain a smooth dynamics we use the scalar  Fenchel (convex) conjugate pair of functions √ √ 2 2 2 ψε (a) = a + ε − ε and ϕε (b) = ε 1 − 1 − b + ι[−1,1] (b) for ε > 0. Lifting through singular values to the matrix space gives the corresponding Fenchel conjugates  ( Pq  p q   2 , ∥G∥ X p ε 1 − 1 − σ (G) i op ≤ 1, i=1 Ψε (P ) = (3) σi (P )2 + ε2 − ε , Φε (G) = +∞, otherwise . i=1 Then Φ∗ε = Ψε and Ψ∗ε = Φε . If P = U diag (σ1 , . . . , σs ) V ⊤ , we define  s σ i  V ⊤. Orthε (P ) = U diag  q 2 2 σi + ε

(4)

i=1

The map is 1/ε-Lipschitz in Frobenius norm, ∥Orthε (P )∥op < 1 and Orthε (P ) → Orth(P ) as ε ↓ 0 for every fixed P . The pair (ϕε , ψε ) should be viewed as a spectral analogue of a smooth saturation. At the scalar level, a 7→ ψε′ (a) = √a2a+ε2 maps the real line into (−1, 1) and tends pointwise to the sign map as ε ↓ 0. Thus Orthε keeps the singular directions of P and replaces each singular value by a softened value in [0, 1). The conjugate ϕε keeps the spectral-ball domain, but it replaces the hard indicator by a smooth barrier-like penalty inside the ball. This is precisely the structure needed to pass between the nonsmooth trust-region interpretation and a  smooth  mirror map. Further, due to the identity ∇Ψε (P ) = U diag

√ σ1 (P2)

σ1 (P ) +ε2

, . . . , √ σs (P2)

σs (P ) +ε2

V⊤ =

Orthε (P ), the soft-orthogonalization operation is exactly a smooth  mirror-map, with its inverse  on e diag the open unit spectral-norm ball being exactly ∇Φε (G) = U

√ εσ1 (G) 2 , . . . , √ εσs (G) 2 1−σ1 (G)

1−σs (G)

Ve ⊤ ,

e diag(σ1 (G), . . . , σs (G))Ve ⊤ where s = rank(G). Hence the regularized with reduced SVD G = U direction is not chosen by an ad hoc smoothing of Orth; it is the primal minimizer associated with the Fenchel-dual kinetic potential Ψε . The same conjugate pair becomes the kinetic term and the mirror geometry in the Hamiltonian probability flow later on.

4

Proposition 2 (Smooth regularized Muon step). For every ε > 0, the problem Gε (P ) = argmin {⟨P, G⟩F + Φε (G)}

(5)

G

has the unique solution Gε (P ) = − Orthε (P ). Hence the regularized Muon update is Pk+1 = βPk + (1 − β)∇F (Wk ) , Wk+1 = Wk − η Orthε (Pk+1 ) .

(6)

Moreover, if Gε,k = − Orthε (Pk ), then Gε,k+1 = argmin {⟨Pk+1 − Pk , G⟩F + DΦε (G, Gε,k )} . G

Remark 1 (Why regularization matters). The hard map P 7→ Orth(P ) is a selected element of ∂∥P ∥nuc and is set-valued at rank-deficient momenta. The smooth family above preserves the spectral trust-region geometry while giving a single-valued Lipschitz vector field. This is the ingredient that makes the ODE, PDE, and propagation-of-chaos arguments standard rather than differential-inclusion arguments. Three equivalent views of the same update. For later use it is helpful to keep three interpretations in parallel. The hard step is a spectral trust-region steepest descent direction. The nonsmooth mirror form says that the same step is generated by the indicator of the spectral unit ball, whose conjugate is the nuclear norm. The regularized step says that Muon is the zero-temperature limit of a smooth mirror family. These are not competing descriptions: the trust-region view explains the geometry, the mirror view provides the variational update, and the smooth Fenchel pair provides the analytic regularity needed for continuous-time and mean-field limits. Why not regularize by adding ε∥G∥2F ? A Euclidean quadratic regularization would also make the direction unique, but it would change the saturation geometry of Muon. The construction above preserves the spectral unit-ball domain and only smooths the singular-value saturation. Consequently, the limit ε ↓ 0 returns the canonical polar-factor direction rather than a Euclidean steepest-descent direction. This is the reason the regularization is tied to a Fenchel pair rather than introduced as a generic numerical smoothing.

3

Probability lift and finite particle Muon

We now minimize a general class of functionals defined over the space of probability measures on X denoted as P(X ) using the dynamics induced by (regularized) Muon. Let  Z J(ρ) = R F (W )dρ(W ) , ρ ∈ P(X ), X

where F P : X → R and R : R → R. For particles W = (W1 , . . . , WN ) define the empirical law 1 N ρW = N N i=1 δWi and the empirical particle lift ! N
1 X N JN (W ) := J ρW = R F (Wi ) . N i=1

1 PN The natural product geometry over particles is the mean-field pairing ⟨U, V ⟩ = avg i=1 ⟨Ui , Vi ⟩F . N
N The mean-field pairing makes the empirical objective JN = J ρW an intensive energy and makes gradavg JN coincide with the particle discretization of the Wasserstein force and keeps the probability objective and particle geometry aligned.

5

R The class J(ρ) = R F dρ is deliberately simple, yet sufficiently expressive to capture population objectives in which a distribution over parameters induces an averaged feature or prediction. For simplicity, we restrict attention in the main text to real-valued functions F . However, ML applications involving neural networks or transformers naturally require extensions to vector-valued outputs and inputs defined on product spaces of matrices. The results and proofs extend to these settings in a standard manner. For deriving a well-posed dynamics, we need to impose some standard regularity conditions. Assumption (A1) (Basic smoothness). F ∈ C 1 (X ), ∇F is globally Lipschitz, and ∥∇F (W )∥F ≤ MF for all W . The derivative R′ is globally Lipschitz with constant LR and is globally bounded with constant MR . Under Assumption (A1), the probability functional J admits a well-defined first variation and Wasserstein gradient that drives the Muon dynamics. R Proposition 3 (First variation and particle gradient). Let mρ = F (W )dρ(W ). A valid first variation of J and its Wasserstein gradient ∇W2 J(ρ)(W ) are given by δJ (ρ)(W ) = R′ (mρ ) F (W ), ∇W2 J(ρ)(W ) = R′ (mρ ) ∇F (W ). δρ P For F N (W ) = N1 N j=1 F (Wj ), the gradient of JN under the natural product geometry is
gradavg JN (W ) = a(W ), ai (W ) = R′ F N (W ) ∇F (Wi ) .

(7)

(8)

The regularized Muon mirror potentials on X N are block-separable: N

ΦN ε (G) =

N

1 X Φε (Gi ) , N

ΨN ε (P ) =

i=1

1 X Ψε (Pi ) . N i=1

With respect to ⟨·, ·⟩avg they remain Fenchel conjugates. The exact finite-particle regularized Muon scheme is therefore Pi,k+1 = βPi,k + (1 − β)R′ (Fk ) ∇F (Wi,k ) , Wi,k+1 = Wi,k − η Orthε (Pi,k+1 ) ,

(9)

P where Fk = N1 N j=1 F (Wj,k ). The scheme is exact for the regularized mirror problem: no approximation has been made beyond the choice of ε > 0. The only coupling among particles is through the empirical scalar Fk , while the orthogonalization is blockwise and particlewise. This separation is central to the mean-field analysis. The interaction enters through the force, while the non-Euclidean geometry enters through the kinetic mirror map P 7→ Orthε (P ). A useful way to understand Equation (9) is as a relaxation system. The momentum Pi,k is a moving average of the current first-variation force, and the position is transported by the mirror-dual velocity generated by that momentum. The limiting ODE derived below preserves this two-time-scale structure.

4

Continuous-time and Hamiltonian probability dynamics

To analyze the continuous-time dynamics, we choose the inertial scaling ηh = h,

βh = 1 − γh + rh , 6

rh /h → 0,

(10)

with γ > 0. The state variable is the matrix and momentum pair Y = (W, P ) belonging to the state space Z := X × X . Let P1 (Z) denote the space of probability measures on Z with finite first moment, equipped with the W1 -Wasserstein topology. Under the stated regularity conditions, we derive the continuous-time evolution both at the particle as well as the distributional law level for the state-space variables. Theorem 1 (Finite-particle ODE limit and phase space PDE). Under Assumption (A1), for fixed N and fixed initial data, the piecewise linear interpolation of Equation (9) converges uniformly on every [0, T ] to the unique global solution of Ẇi (t) = − Orthε (Pi (t))

Ṗi (t) = γ R′ F N (t) ∇F (Wi (t)) − Pi (t) ,

(11)

P 2 where F N (t) = N1 N j=1 F (Wj (t)). If rh = O h , the convergence rate is O(h) on finite horizons. Let Yt = (Wt , Pt ) ∈ Z = X × X , let µt = Law (Yt ), let ρt = (πW )# µt , and set at (W ) =
R ′ R F dρt ∇F (W ). The mean-field law solves ∂t µt + ∇W · (− Orthε (P )µt ) + ∇P · (γ (at (W ) − P ) µt ) = 0.

(12)

The weak solution is unique in C ([0, T ]; P1 (Z)) under the corresponding global Lipschitz hypotheses, and under localized hypotheses as long as the trajectory remains in the bounded region on which the constants are finite. Equation (12) is the probability-flow counterpart of Muon with momentum. The position velocity is the negative regularized orthogonalized momentum and the momentum velocity relaxes toward the Wasserstein force at . Remark 2 (Role of scaling). If β is fixed as h ↓ 0, the momentum variable relaxes on a fast scale and the second-order structure is lost. The scaling chosen in Equation (10) is therefore part of the model and allows a second-order flow dynamics instead of a singular overdamped limit, not merely a technical convenience. We now express the evolution PDE in Theorem 1 as an exact regularized Muon Hamiltonian probability flow. Theorem 2 (Damped Hamiltonian structure). Define Z  Z Hε,γ (µ) = Ψε (P )dµ(W, P ) + γR F (W )dµ(W, P ) Z

(13)

Z

Then Equation (12) is equivalent to       δHε,γ δHε,γ ∂t µt + ∇W · µt −∇P + ∇P · µt ∇W − γP =0 δµt δµt

(14)

Moreover, along solutions, d Hε,γ (µt ) = −γ dt

Z Z

⟨P, Orthε (P )⟩F dµt (W, P ) ≤ 0

The integrand has the singular-value form ⟨P, Orthε (P )⟩F = only if P = 0. 7

Pq

r=1

2

√ σr (P )2

σr (P ) +ε2

(15) and vanishes if and

The theorem explains why the objective J (ρt ) alone need not decrease monotonically, but the Hamiltonian descent drives the “potential energy” component J, which is the target functional, to decrease Ras the dynamics progress through the relation in Equation 13. Momentum stores “kinetic energy” Ψε dµ, and the dissipated quantity is the Hamiltonian defined in Equation 13. Remark 3 (Dissipation as a Fenchel coupling). The integrand in Equation (15) is dε (P ) = ⟨P, Orthε (P )⟩F = Ψε (P ) + Φε (Orthε (P )) , where the second R equality is Fenchel equality. Therefore the dissipated quantity is not simply the kinetic energy Ψε dµ, it is the coupling between the momentum and the mirror-dual velocity. For the chosen spectral potential, 0 ≤ dε (P ) − Ψε (P ) ≤ qε, dε (P ) ≤ ∥P ∥nuc. Hence dε is asymptotically equivalent to the nuclear-norm kinetic energy as ε ↓ 0, but it has the stronger property that it vanishes if and only if the momentum vanishes.

5

Convergence, particle approximation, and hard Muon

Let at (W ) = R′ (mt ) ∇F (W ), J⋆ = inf ρ∈P1 (X ) J(ρ) and define Z Z Z At = ∥at (W )∥2F dµt , Kt = Ψε (P )dµt Dt = ⟨P, Orthε (P )⟩F dµt , Z Ut = J (ρt ) − J⋆ Ht = Kt + γUt , Ct = ⟨at (W ), P ⟩F dµt . The alignment Ct is the term that records whether momentum is consistent with descent. The basic Hamiltonian identity gives Ht′ = −γDt , but this alone does not control the objective gap because Dt controls momentum rather than force. The force norm At is connected to the gap by the PL assumption, while the momentum dissipation is connected to the kinetic energy by coercivity. The bridge between them is the alignment Ct . Positive alignment means that the momentum points in a descent-compatible direction, negative alignment records a transient inertial mismatch. For this reason the convergence proof uses the modified Lyapunov functional Lt = Ht − αCt . The role of α is to reward descent-compatible alignment without letting Ct dominate the Hamiltonian. The upper-gradient condition and kinetic coercivity imply |Ct | ≤ MC Ht , so Lt is equivalent to Ht when αMC < 1. The curvature assumption then guarantees that the derivative of Ct introduces the force norm At up to a controlled remainder. This is the mechanism behind the exponential rate below. Assumption (A2) (Trajectory-level convergence hypotheses). Along the trajectory, Ut ≥ 0 and the following hold. (a) Bounded-momentum kinetic coercivity: for constants κK , κD > 0, χ ≥ 1 and LG < ∞, κK ∥P ∥2F , ⟨P, Orthε (P )⟩F ≥ κD ∥P ∥2F , 2 Ψε (P ) ≤ χ ⟨P, Orthε (P )⟩F , ∥Orthε (P )∥F ≤ LG ∥P ∥F . Ψε (P ) ≥

3/2 For the regularized Muon potential these inequalities hold on ∥P ∥F ≤ BP with κK = ε2 / BP2 + ε2 ,
−1/2 2 2 κD = B P + ε , LG = 1/ε, and χ = 1. (b) Functional PL and upper-gradient bounds: At ≥ 2λUt and At ≤ 2ΛUt for constants λ > 0, Λ < ∞. (c) Curvature control: Ct′ = γAt − γCt + St with |St | ≤ σDt . 8

Theorem 3 (Exponential convergence under explicit assumptions). Let Assumptions (A1) and q Λ (A2) hold true. Let MC = γκK . Choose r ∈ (0, 2) and α > 0 such that αMC < 1, dα,r := n
o dα,r αγ 1 r := 1+αM . Then, γ − ασ − 2rκ > 0. Finally, let us set c min , 2λα 1 − α,r χ 2 D C J (ρt ) − J⋆ ≤

e−cα,r t (H0 − αC0 ) . γ (1 − αMC )

(16)

The proof in Appendix A.7 obtains Theorem 3 as the one-block scalar-output specialization of the Hilbert-domain Theorem 14. Remark 4 (Interpretation of assumptions). The theorem separates geometry from objective landscape. Kinetic coercivity is a property of the regularized Muon potential on bounded momentum sets. PL and upper-gradient conditions are properties of J along the probability trajectory. As observed in practice and demonstrated in our numerical experiments, the exponential convergence behavior holds true. Precise sufficient conditions for the validity of the PL and upper-gradient conditions for J for specific models of interest can be determined. Remark 5. The constant cα,r is the minimum of a momentum-dissipation contribution and a force-gap contribution. Increasing α helps expose the PL force term, but too large an α either destroys equivalence of Lt and Ht or makes the dissipation tradeoff negative. The parameter r is the Young-inequality split between force and momentum. These constraints are not artifacts of notation, they express the fact that accelerated dynamics converge only when momentum alignment, curvature, and dissipation are balanced. For finite N , the particle update Equation (9) gives an interacting approximation of the nonlinear characteristic system Equation (11). The next theorem states the quantitative mean-field consistency result under a second-moment assumption on the initialization law of W, P . Theorem 4 (Propagation of chaos). Let Assumption (A1) hold. Equivalently, after identifying Θ = X and H = R, the scalar data satisfy the global Hilbert smoothness Assumption (A4). Let µ0 ∈ P2 (Z). Couple the N -particle ODE Equation (11)) with i.i.d. nonlinear mean-field copies having law µt and the same initial data. Then, for every T < ∞, there is Cpoc (T, ε) < ∞ such that, for every fixed i, h i C (T, ε) 2 2 poc sup E WiN (t) − W̄i (t) F + PiN (t) − P̄i (t) F ≤ (17) N t≤T Consequently every fixed k particles converge in law to µ⊗k t , uniformly on finite horizons. Remark 6 (Mean-field relevance). Theorem 4 justifies using the nonlinear PDE as the largepopulation limit of the regularized Muon particle system. The constant depends on 1/ε through the Lipschitz constant of Orthε , which is expected since the hard Muon map is not Lipschitz. Finally, the regularized flow has a compactness limit as ε ↓ 0. Let Ψ0 (P ) = ∥P ∥nuc . We discuss this in detail in the Appendix (Theorem 9).

6

Transformer mixture-of-experts (MoE) optimization using Hamiltonian probability flow

In this section, we discuss how the Hamiltonian probability flow formulation can be extended to optimize ML models with tuples of matrix-valued parameters, such as Transformer mixture-ofexperts (MoE). A shallow transformer MoE uses a product parameter space Θ = Θexp × Θgate = 9

QB

b=1 R

mb ×nb where the expert blocks may include Q, K, V, O, W , W corresponding to the Query, 1 2

Key, Value and Output projection matrices in the attention module, together with the weight matrices in the FFN layer and the gate blocks contain router parameters for the expert routing scheme. The Muon mirror map is applied blockwise to the full tuple of matrices, including the router. For a training input L-token sequence X ∈ RL×d , a single-head expert can be written as   (XQ)(XK)⊤ (18) √ Aω (X) = softmax XV O, ψω (X) = Rout (σ (Aω (X)W1 ) W2 ) . dk A smooth router score sϕ (X) is included in the feature map. For softmax-normalized MoE routing, the empirical output MN (X)t,: =

P

sϕ (X)t i ψωi (X)t,: P sϕj (X)t e j sϕ sϕ

ie

is represented by an augmented Hilbert-space

valued feature map Fsoft (ω, ϕ) = (e ψω , e ) and a smooth normalization map applied after averaging. Another gating choice can be considered as well. For an unnormalized non-negative smooth input-dependent gate gϕ (X), one may take Fun (ω, ϕ) = (gϕ (Xr ) ⊙ ψω (Xr ))nr=1 in logit Hilbert space. Thus
the Transformer MoE model with these two choices of gating fits the form R J(ρ) = R F dρ . The theorem below analyzes the dynamics of parameter optimization for Transformer MoE and uses the following concrete conditions, which are the transformer versions of the localized product-space assumptions in the Appendix: (T1) training inputs are bounded, ∥Xr ∥F ≤ BX ; (T2) the activation and router score maps are C 2 in the parameters on bounded sets; (T3) in the normalized-gate case, denominators remain bounded below by a positive constant on the moment set reached by the trajectory; (T4) the parameter and momentum trajectory remains in a bounded region for the time horizon or asymptotic regime under consideration; (T5) for exponential convergence, the product-space PL and upper-gradient assumptions hold along the trajectory. Conditions (T 1)-(T 4) are analytic well-posedness conditions, while condition (T 5) is an optimization-landscape condition. Separating them prevents the transformer statement from overstating what follows from smoothness alone. Theorem 5 (Transformer MoE consequence). Assume bounded training inputs, C 2 expert and router maps on bounded parameter sets, a denominator lower bound for normalized gates, and bounded parameter/momentum trajectories on the time interval considered. Then, under the inertial scaling ηh = h, βh = 1 − γh + o(h), the regularized Muon scheme on expert-router particles, Pi,k+1 = βPi,k + (1 − β)aN i (θk ) ,

θi,k+1 = θi,k − η OrthΘ ε (Pi,k+1 ) .

  P ∗ −1 where aN = DF (θ ) ∇R N F (θ ) , satisfies the finite-particle ODE limit, phase-space i j i j Hamiltonian PDE, dissipation identity and, under the P L and upper-gradient assumptions, the exponential convergence estimate. The general Hilbert valued product-space PDE and the proof of this theorem are discussed in detail in the Appendix.

10

7

Numerical Experiments

We use two deterministic synthetic experiment classes to test the finite-particle dynamics developed above. The goal is not to benchmark large-scale training, but to isolate the phenomena predicted by the Hamiltonian formulation. The first experiment focuses on a M -particle matrix mean matching problem with the finite objective functional on N particles given by JN (W1 , . . . , WN ) = 2 1 1 PN 1 PM with W⋆ = M i=1 Wi − W̄⋆ j=1 Wj,⋆ . The second experiment is a nonlinear teacher2 N F

student problem on a product matrix space with each particle θi = (Ai , Bi ) ∈ Θ = Rp×r × Rr×d with the (M, N )-particle teacher-student tanh neural network objective JN ((Ai , Bi )N i=1 ) =   2 B√ x 1 PS 1 PN i s − ys trained on S input-output (xs , ys ) pairs. In both cases s=1 N i=1 Ai tanh 2Sp d 2 ∗ J = 0 is attainable. The experimental results are reported in part in Figure 1, with the complete numerical experiments being discussed in the Appendix. The synthetic experiments support the Hamiltonian interpretation developed in the paper. In the matrix mean-matching problem, regularized Muon avoids the finite-step residual floor exhibited by hard polar and Newton-Schulz updates, allowing the objective and Hamiltonian energies to decay smoothly towards zero. In the nonlinear product-space neural network setting, the regularized dynamics preserve the advantages of spectral Muon geometry while exhibiting the stability predicted by the smooth mirror formulation.

Figure 1: Top: Experiment 1 on matrix mean matching with (M, N ) = (4, 32). Bottom: Experiment 2 on product-space teacher-student particles with (d, r, p) = (10, 6, 4) and S = 320 training points and (M, N ) = (3, 12). For both experiments, left panels show JN and right panels show the Hamiltonian K + γJN on a loglinear scale.

11

8

Conclusion

This work identifies Muon as a spectral trust-region mirror method and derives the damped Hamiltonian probability flow induced by its regularized orthogonalization map. It shows that Muon’s matrix-level update admits a coherent variational, mean-field, and Lyapunov structure.

References Kwangjun Ahn and Sinho Chewi. Efficient constrained sampling via the mirror-langevin algorithm. Advances in Neural Information Processing Systems, 34:28405–28418, 2021. Lizhang Chen, Jonathan Li, and Qiang Liu. Muon optimizes under spectral norm constraints. arXiv preprint arXiv:2506.15054, 2025a. Shi Chen, Qin Li, Oliver Tse, and Stephen J Wright. Accelerating optimization over the space of probability measures. Journal of machine learning research, 26(31):1–40, 2025b. Zixiang Chen, Yihe Deng, Yue Wu, Quanquan Gu, and Yuanzhi Li. Towards understanding the mixture-of-experts layer in deep learning. Advances in neural information processing systems, 35: 23049–23062, 2022. Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for overparameterized models using optimal transport. Advances in neural information processing systems, 31, 2018. William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1–39, 2022. Ruichen Jiang, Zakaria Mhammedi, Mehryar Mohri, and Aryan Mokhtari. Adaptive matrix online learning through smoothing with guarantees for nonsmooth nonconvex optimization. arXiv preprint arXiv:2602.08232, 2026. Keller Jordan, Yuchen Jin, Vlado Boza, You Jiacheng, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024. URL https://kellerjordan.github.io/posts/muon/. Juno Kim, Eshaan Nichani, Denny Wu, Alberto Bietti, and Jason D. Lee. Sharp capacity scaling of spectral optimizers in learning associative memory, 2026. URL https://arxiv.org/abs/2603. 26554. Dmitry Kovalev. Understanding gradient orthogonalization for deep learning via non-euclidean trust-region optimization. arXiv preprint arXiv:2503.12645, 2025. Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, Yanru Chen, Huabin Zheng, Yibo Liu, Shaowei Liu, Bohong Yin, Weiran He, Han Zhu, Yuzhi Wang, Jianzhou Wang, Mengnan Dong, Zheng Zhang, Yongsheng Kang, Hao Zhang, Xinran Xu, Yutao Zhang, Yuxin Wu, Xinyu Zhou, and Zhilin Yang. Muon is scalable for llm training, 2025. URL https://arxiv.org/abs/2502.16982.

12

Thomas Pethick, Wanyun Xie, Kimon Antonakopoulos, Zhenyu Zhu, Antonio Silveti-Falls, and Volkan Cevher. Training deep learning models with norm-constrained lmos. arXiv preprint arXiv:2502.07529, 2025. Gabriel Peyré. Muon dynamics as a spectral wasserstein flow, 2026. URL https://arxiv.org/ abs/2604.04891. Louis Sharrock, Lester Mackey, and Christopher Nemeth. Learning rate free sampling in constrained domains. Advances in Neural Information Processing Systems, 36:65380–65415, 2023. Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. 2017. URL https://openreview.net/pdf?id=B1ckMDqlg. Mark Tuddenham, Adam Prügel-Bennett, and Jonathan Hare. Orthogonalising gradients to speed up neural network optimisation. arXiv preprint arXiv:2202.07052, 2022. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017. Yifei Wang and Wuchen Li. Accelerated information gradient flow. Journal of Scientific Computing, 90(1):11, 2022.

13

A

Technical appendices and supplementary material

A.1

Proof of Proposition 1

Let P = U ΣV ⊤ be a reduced SVD with s = rank(P ) and singular values σ1 (P ), . . . , σs (P ). By the duality between the operator norm and the nuclear norm, ∥P ∥nuc =

sup ⟨P, M ⟩F . ∥M ∥op ≤1

For completeness, this follows from von Neumann’s trace inequality: ⟨P, M ⟩F ≤

q X

σi (P )σi (M ) ≤

q X

i=1

σi (P ) = ∥P ∥nuc

i=1

whenever ∥M ∥op ≤ 1. Taking M = U V ⊤ gives equality, since the nonzero singular values of U V ⊤ are all one and the singular directions align with those of P . Therefore Orth(P ) is a maximizer of the support-function problem, and by sign reversal − Orth(P ) is a minimizer of ⟨P, G⟩F over ∥G∥op ≤ 1. Since Φ0 is the indicator of this spectral unit ball, the constrained problem is exactly Equation (2). The canonical choice of minimizer gives the ideal Muon step. When P is rank deficient, the optimizer need not be unique. The subdifferential formula n o ∂∥P ∥nuc = U V ⊤ + Z : U ⊤ Z = 0, ZV = 0, ∥Z∥op ≤ 1 shows the source of nonuniqueness. Muon chooses the canonical element with Z = 0.

A.2

Scalar and spectral Fenchel conjugacy

For ε > 0, define ψε (a) =

  p ϕε (b) = ε 1 − 1 − b2 + ι[−1,1] (b).

p a2 + ε2 − ε,

Both are proper, closed, convex, and even. On their differentiability domains, ψε′ (a) = √

a a2 + ε 2

ϕ′ε (b) = √

,

εb 1 − b2

√ √ These maps are inverse: if b = a/ a2 + ε2 , then a = εb/ 1 − b2 , and conversely. Hence ϕ∗ε = ψε and ψε∗ = ϕε . The spectral lifts Equation (3) are conjugate by the standard conjugacy theorem for unitarily invariant spectral functions. Equivalently, von Neumann’s trace inequality reduces sup {⟨P, G⟩F − Φε (G)} G

P to the scalar singular-value optimization, whose value is i ψε (σi (P )) = Ψε (P ). The reverse conjugacy is identical. The √ gradient formula Equation (4) follows from spectral-function calculus. The derivative of a 7→ a/ a2 + ε2 is ε2 2 3/2

(a2 + ε )

1 ≤ , ε

so Orthε = ∇Ψε is q 1/ε-Lipschitz in Frobenius norm. Finally, each nonzero singular value is transformed as σi /

σi2 + ε2 → 1, proving Orthε (P ) → Orth(P ). 14

A.3

Proof of Proposition 2

The first-order condition for Equation (5) is 0 = P + ∇Φε (Gε (P )) . Since ∇Φε and ∇Ψε are inverse on the interior of the effective domain of Φε , Gε (P ) = ∇Ψε (−P ). The function Ψε is even, so its gradient is odd, and Gε (P ) = −∇Ψε (P ) = − Orthε (P ). For the Bregman form, use ∇Φε (Gε,k ) = −Pk . Then DΦε (G, Gε,k ) = Φε (G) − Φε (Gε,k ) − ⟨∇Φε (Gε,k ) , G − Gε,k ⟩F = Φε (G) + ⟨Pk , G⟩F + Ck where Ck is independent of G. Adding ⟨Pk+1 − Pk , G⟩F gives the same minimizers as ⟨Pk+1 , G⟩F + Φε (G).

A.4

Proof of Proposition 3

Let ν be a finite signed measure with ν(X ) = 0 and set ρs = ρ + sν. Then Z mρs = mρ + s F (W )dν(W ) X

Differentiating J (ρs ) = R (mρs ) at s = 0 gives Z d J (ρs ) = R′ (mρ ) F (W )dν(W ) ds X s=0 This is the defining identity for the first variation, modulo constants independent of W . Taking the spatial gradient gives the Wasserstein gradient in Equation (7). For the particle gradient, let U = (U1 , . . . , UN ). By the chain rule, N
1 X DJN (W )[U ] = R F (W ) ⟨∇F (Wi ) , Ui ⟩F N ′

N

i=1

=

1 N

N X

⟨ai (W ), Ui ⟩F = ⟨a(W ), U ⟩avg .

i=1

Hence gradavg JN (W ) = a(W ).

A.5

Product-space Fenchel duality for particles

Define

N

ΦN ε (G) =

N

1 X Φε (Gi ) , N

ΨN ε (P ) =

i=1

1 X Ψε (Pi ) . N i=1

With respect to ⟨·, ·⟩avg , ) 1 X 1 X ⟨Pi , Gi ⟩F − Φε (Gi ) N N i i 1 X = sup {⟨Pi , Gi ⟩F − Φε (Gi )} = ΨN ε (P ). N Gi

∗ ΦN (P ) = sup ε G

(

i

15

The reverse conjugacy is identical. Thus the finite-particle regularized Muon step is the blockwise minimization stated in Equation (9).

A.6

Proof of Theorem 1 and 2

Theorem 6 (Finite- N ODE limit). . Assume the vector field W 7→ a(W ) is Lipschitz on the region visited by the discrete
and continuous trajectories, and Orthε is 1/ε-Lipschitz. Under ηh = h and βh = 1 − γh + O h2 , the piecewise-linear interpolation of the particle scheme converges uniformly on compact intervals to Ẇi = − Orthε (Pi ) ,

Ṗi = γ (ai (W ) − Pi ) .

If the Lipschitz bound is global, the solution is global and the convergence rate is O(h) when
βh = 1 − γh + O h2 . Proof. Let BN (W, P ) = (− Orthε (Pi ) , γ (ai (W ) − Pi ))N i=1 ,
where ai (W ) = R′ F N (W ) ∇F (Wi ). The map P 7→ Orthε (P ) is globally 1/ε-Lipschitz. Under Assumption (A1), a(W ) is globally Lipschitz in the mean-field norm. Indeed, for two particle f, configurations W, W  
f ) ≤ R′ F N (W ) fi ai (W ) − ai (W ∇F (Wi ) − ∇F W F F    
′ N ′ N f fi + R F (W ) − R F (W ) ∇F W . F

The first term is controlled by the global bound |R′ | ≤ MR and the Lipschitzness of ∇F . The second term is controlled by the Lipschitzness of R′ , the boundedness of ∇F , and the Lipschitzness of F , which follows from ∥∇F ∥ ≤ MF . Thus the ODE has a unique global solution. The global √ existence follows from bounded W -velocity, ∥Orthε (Pi )∥F ≤ q, and at most linear growth in Pi . Then, Pi,k+1 − Pi,k = (1 − βh ) (ai (Wk ) − Pi,k ) = (γh − rh ) (ai (Wk ) − Pi,k ) , which is hγ (ai (Wk ) − Pi,k ) + o(h) locally uniformly on bounded sets. Also, Wi,k+1 − Wi,k = −h Orthε (Pi,k+1 ) = −h Orthε (Pi,k ) − h (Orthε (Pi,k+1 ) − Orthε (Pi,k )) .
The last term is O h2 + o(h)h on bounded sets by the Lipschitzness of Orth ε . The one-step
consistency error therefore tends to zero after division by h, and it is O(h) when rh = O h2 . The standard Euler convergence estimate with discrete Gronwall gives uniform convergence on [0, T ] and rate O(h) in the second-order-consistent case. Remark 7. The statement is intentionally phrased in a local form. For neural-network parameterizations, global Lipschitz constants are often unavailable, while finite-horizon bounded-trajectory constants are natural. Let µt ∈ P1 (Z) with Z = X × X and ρt = (πW )# µt . Define bµt (W, P ) = (− Orthε (P ), γ (aρt (W ) − P )) . 16

Theorem 7 (Well-posed McKean-Vlasov equation). Under global Lipschitz assumptions on bµ in ( W, P ) and in µ with respect to W1 , for every µ0 ∈ P1 (Z) there is a unique solution µ ∈ C ([0, ∞); P1 (Z)) of ∂t µt + ∇W · (− Orthε (P )µt ) + ∇P · (γ (aρt (W ) − P ) µt ) = 0. Equivalently, for every ζ ∈ Cc∞ (Z), Z Z d ζ dµt = [⟨∇W ζ, − Orthε (P )⟩ + ⟨∇P ζ, γ (aρt (W ) − P )⟩] dµt dt Proof. The nonlinear characteristic system is Z t Orthε (Ps ) ds, Wt = W0 −

Z t γ (aρs (Ws ) − Ps ) ds

Pt = P0 + 0

0

where ms = EF (Ws ). The drift is Lipschitz in ( W, P ) and in the law with respect to W1 under the preceding estimates, with at most linear growth in P . Picard iteration for McKean-Vlasov ODEs gives existence and uniqueness. Setting µt = Law (Wt , Pt ) and applying the chain rule to any ζ ∈ Cc∞ (Z) gives D E D E d ζ (Wt , Pt ) = ∇W ζ (Wt , Pt ) , Ẇt + ∇P ζ (Wt , Pt ) , Ṗt dt = ⟨∇W ζ, − Orthε (Pt )⟩ + ⟨∇P ζ, γ (ai (W ) − Pi )⟩ . Conversely, the superposition principle for Lipschitz continuity equations represents any weak solution by characteristics, and uniqueness of the nonlinear characteristic equation gives uniqueness of the PDE solution. Remark 8. For the scalar functional, Lipschitz dependence on the measure follows from Lipschitzness of R′ and F together with boundedness of ∇F and the boundedness/local boundedness of R′ needed for the spatial Lipschitz term. Theorem 8 (Hamiltonian form and dissipation). Let Z  Z Hε,γ (µ) = Ψε (P )dµ(W, P ) + γR F (W )dµ(W, P ) A first variation is δHε,γ (µ)(W, P ) = Ψε (P ) + γR′ (mµ ) F (W ). δµ Therefore Equation (12) is the damped Hamiltonian equation. Along weak solutions for which the cutoff argument is justified, Z d Hε,γ (µt ) = −γ ⟨P, Orthε (P )⟩F dµt ≤ 0. dt Proof. Let ν be a finite signed measure on Z with ν(Z) = 0 and finite first moment. For small s, formally set µs = µ + sν. Then

Z mµs =

Z F (W )d(µ + sν)(W, P ) = mµ + s

Z

F (W )dν(W, P ) Z

17

Note that, Hε,γ (µs ) =

F (W )dν . Z

Z

Z



Z

Ψε (P )dν+γR mµ + s

Ψε (P )dµ+s

Ψε (P )d(µ+sν)+γR (mµs ) = Z



Z

Z

Z

Differentiating at s = 0, we have that d Hε,γ (µs ) = ds s=0

Z



 Ψε (P ) + γR′ (mµ ) F (W ) dν(W, P ).

Z

Hence a valid first variation is δHε,γ (µ)(W, P ) = Ψε (P ) + γR′ (mµ ) F (W ), δµ up to an additive constant depending on µ but independent of (W, P ). Such constants do not affect ∇W or ∇P . At a fixed µ = µt , the scalar mt := mµt is constant with respect to the variables (W, P ). Therefore δHε,γ ∇P (µt )(W, P ) = ∇Ψε (P ) δµ and ∇W

δHε,γ (µt )(W, P ) = γR′ (mt ) ∇F (W ). δµ

Let Gε (P ) := ∇Ψε (P ), and define at (W ) := R′ (mt ) ∇F (W ). Then the PDE in Equation (12) is the continuity equation with velocity field bt (W, P ) = (bW (t, W, P ), bP (t, W, P )) , where bW (t, W, P ) = −Gε (P )

bP (t, W, P ) = γ (at (W ) − P ) .

and

Thus ∂t µt + ∇W · (µt bW ) + ∇P · (µt bP ) = 0. In weak form, for every smooth compactly supported test function ζ ∈ Cc∞ (Z), Z Z d ζ(W, P )dµt (W, P ) = [⟨∇W ζ, bW ⟩F + ⟨∇P ζ, bP ⟩F ] dµt . dt Z Z

(19)

The functions F (W ) and Ψε (P ) are not compactly supported, but they have at most linear growth and bounded gradients: |F (W )| ≤ |F (0)| + MF ∥W ∥F , and 0 ≤ Ψε (P ) ≤ ∥P ∥nuc ≤

√

q∥P ∥F ,

while ∥∇F (W )∥F ≤ MF ,

∥∇Ψε (P )∥F ≤

√

q.

Hence F and Ψε can be used as test functions by a standard cutoff argument. For completeness, here is the cutoff justification. Let χ ∈ Cc∞ ([0, ∞)) satisfy 0 ≤ χ ≤ 1,

χ(r) = 1 for r ≤ 1, 18

χ(r) = 0 for r ≥ 2.

√ Define χR (W, P ) := χ

∥W ∥2F +∥P ∥2F R

 . For a C 1 function ζ with at most linear growth and

bounded gradient, set ζR := χR ζ. Then ζR ∈ Cc1 (Z), ζR → ζ pointwise, and ∇ζR = χR ∇ζ + ζ∇χR The first term converges to ∇ζ. The second term is supported on the annulus q R ≤ ∥W ∥2F + ∥P ∥2F ≤ 2R, and satisfies C (1 + ∥W ∥F + ∥P ∥F ) 1{R≤∥(W,P )∥≤2R} ≤ C1{∥(W,P )∥≥R} . R Since the vector field has at most linear growth in P , and µt has finite first moment, the contribution of this annulus vanishes as R → ∞. Therefore Equation 19 remains valid for ζ = F (W ) and ζ = Ψε (P ). R By definition, mt = Z F (W )dµt (W, P ). Take ζ(W, P ) = F (W ). Then |ζ∇χR | ≤

∇W ζ(W, P ) = ∇F (W ),

∇P ζ(W, P ) = 0.

Using Equation 19, d mt = dt

Z Z

⟨∇F (W ), bW (t, W, P )⟩F dµt (W, P ).

Since bW = −Gε (P ), d mt = − dt

Z

Z ⟨∇F (W ), Gε (P )⟩F dµt (W, P ) = −

Z

Z

⟨∇F (W ), ∇Ψε (P )⟩F dµt (W, P ).

(20)

Since both ∇F and ∇Ψε are bounded, mt is absolutely continuous. Further, R ∈ C 1 and mt is absolutely continuous, the ordinary chain rule and Equation 20 gives Z d ′ ′ ′ [γR (mt )] = γR (mt ) mt = −γR (mt ) ⟨∇F (W ), Gε (P )⟩F dµt . dt Z Since R′ (mt ) is a scalar independent of (W, P ), this is Z d [γR (mt )] = −γ R′ (mt ) ∇F (W ), Gε (P ) F dµt . dt Z Using at (W ) = R′ (mt ) ∇F (W ), d [γR (mt )] = −γ dt Define

Z Z

⟨at (W ), Gε (P )⟩F dµt .

Z Kε (t) :=

Ψε (P )dµt (W, P ) Z

Take ζ(W, P ) = Ψε (P ). Then ∇W ζ(W, P ) = 0,

∇P ζ(W, P ) = Gε (P ). 19

(21)

Using Equation 19, d Kε (t) = dt

Z Z

⟨Gε (P ), bP (t, W, P )⟩F dµt

Since bP (t, W, P ) = γ (at (W ) − P ) , we obtain d Kε (t) = γ dt

Z Z

⟨Gε (P ), at (W ) − P ⟩F dµt .

Expanding the inner product, Z Z d Kε (t) = γ ⟨Gε (P ), at (W )⟩F dµt − γ ⟨Gε (P ), P ⟩F dµt . dt Z Z Therefore

d Kε (t) = γ dt

Z

Z

Z

⟨at (W ), Gε (P )⟩F dµt − γ

⟨P, Gε (P )⟩F dµt .

(22)

Z

The total Hamiltonian energy is Hε,γ (t) = Kε (t) + γR (mt ) . Using Equations 21 and 22,  Z  Z Z d Hε,γ (t) = γ ⟨at (W ), Gε (P )⟩F dµt − γ ⟨P, Gε (P )⟩F dµt − γ ⟨at (W ), Gε (P )⟩F dµt dt Z Z Z The mixed force-transport terms cancel exactly: Z Z γ ⟨at (W ), Gε (P )⟩F dµt − γ ⟨at (W ), Gε (P )⟩F dµt = 0 Z

Z

Hence, we obtain d Hε,γ (t) = −γ dt

Z

Z ⟨P, Gε (P )⟩F dµt = −γ

Z

Z

⟨P, ∇Ψε (P )⟩F dµt (W, P ).

Let P = U diag (σ1 , . . . , σs ) V ⊤ be a reduced SVD of P , with σi > 0. Then, we have that ! σ1 σs ∇Ψε (P ) = U diag p 2 ,..., p V ⊤. 2 2 2 σ + ε σ1 + ε s Therefore

h i ⟨P, ∇Ψε (P )⟩F = tr P ⊤ ∇Ψε (P ) .

Substituting the SVD expressions, we have that P⊤ = V diag (σ1 , . . . , σs ) U ⊤ , so P ⊤ ∇Ψε (P ) =  V diag (σ1 , . . . , σs ) U ⊤ U diag ⊤

√ σ21

σ1 +ε2

, . . . , √ σ2s

P ∇Ψε (P ) = V diag

σs +ε2

σ12

V ⊤ . Since U ⊤ U = I and V ⊤ V = I, σs2

p ,..., p σs2 + ε2 σ12 + ε2 20

! V ⊤.

Taking the trace, ⟨P, ∇Ψε (P )⟩F =

s X

σ2 q i . σi2 + ε2 i=1

If we use the full singular-value list σ1 , . . . , σq , with zero padding beyond the rank, this becomes ⟨P, ∇Ψε (P )⟩F =

q X

σ (P )2 p i . 2 + ε2 σ (P ) i i=1

Each term is nonnegative. Therefore ⟨P, ∇Ψε (P )⟩F ≥ 0. Moreover, σ2 q i = 0 ⇐⇒ σi = 0. σi2 + ε2 Thus ⟨P, ∇Ψε (P )⟩F = 0

⇐⇒

σi (P ) = 0 for all i

⇐⇒

P = 0.

Hence the dissipation functional satisfies Dε (µ) = 0

⇐⇒

P = 0 µ-a.e.

⇐⇒

µP = δ0 .

Consequently, we have that d Hε,γ (t) ≤ 0. dt

A.7

Proof of Theorem 3 as a scalar Hilbert-space corollary

Proof. We apply the Hilbert-domain convergence theorem, Theorem 14, with the one-block scalaroutput choice Θ = X = Rm×n , H = R, where Θ carries the Frobenius inner product and H carries the usual Euclidean inner product on R. The product-space block number is B = 1, so the product norm ∥ · ∥Θ is exactly ∥ · ∥F . We use the same scalar feature map F : X → R and the same outer loss R : R → R. Therefore Z mρ = F (W ) dρ(W ) Θ

is the same moment as in the main text, and the Hilbert-space objective R(mρ ) is exactly J(ρ). We next identify the Hilbert-space force with the scalar matrix-space force. For u ∈ R and V ∈ X , the adjoint of DF (W ) : X → R is characterized by ⟨DF (W )∗ u, V ⟩F = u DF (W )[V ] = u ⟨∇F (W ), V ⟩F = ⟨u∇F (W ), V ⟩F . Hence DF (W )∗ u = u∇F (W ).

21

Taking u = ∇R(mρ ) = R′ (mρ ) gives aρ (W ) = DF (W )∗ ∇R(mρ ) = R′ (mρ )∇F (W ), which is precisely the force at (W ) used in Theorem 3. Thus the quantities At , Ut , Ct in (83) reduce exactly to the main-text quantities in Section 5. Likewise, since B = 1, ΨΘ ε (P ) = Ψε (P ),

OrthΘ ε (P ) = Orthε (P ),

and therefore Kt , Dt , Ht in Theorem 14 reduce exactly to the corresponding Kt , Dt , Ht in Theorem 3. The assumptions also specialize exactly. The kinetic estimates (85)-(88) are precisely Assumption (A2)(a). For completeness, when the momentum support satisfies ∥P ∥F ≤ BP , these constants are obtained as follows. For s ∈ [0, BP ],   Z s p r d r ε2 ε2 √ √ s2 + ε2 − ε = dr, = 2 ≥ . dr (r + ε2 )3/2 (BP2 + ε2 )3/2 r 2 + ε2 r 2 + ε2 0 √ Since r/ r2 + ε2 vanishes at r = 0, integration gives p κK 2 s2 + ε2 − ε ≥ s , 2

κK =

ε2 . (BP2 + ε2 )3/2

Summing over singular values gives Ψε (P ) ≥ (κK /2)∥P ∥2F . Similarly, √

s2 s2 ≥q = κD s2 , 2 s2 + ε2 2 B +ε

κD = (BP2 + ε2 )−1/2 ,

P

which gives ⟨P, Orthε (P )⟩F ≥ κD ∥P ∥2F . The inequality p s2 s2 + ε2 − ε ≤ √ s2 + ε2 √ is equivalent to ε ≤ s2 + ε2 , so Ψε (P ) ≤ ⟨P, Orthε (P )⟩F , i.e. χ = 1. Finally, Orthε (0) = 0 and the 1/ε-Lipschitz property imply ∥ Orthε (P )∥F ≤ ε−1 ∥P ∥F , so LG = 1/ε. Assumption (A2)(b) is exactly the Hilbert-domain PL and upper-gradient Assumption (A6) after the above identification. Assumption (A2)(c) is exactly the curvature assumption (99) in the scalar matrix-space notation. If one wants primitive sufficient conditions for this curvature assumption, they are obtained from Lemma 3: if along the trajectory |R′ (mt )| ≤ MR ,

|R′′ (mt )| ≤ MR,2 ,

∥∇F (W )∥F ≤ MF ,

∥∇2 F (W )[V ]∥F ≤ MF,2 ∥V ∥F ,

then, in the Hilbert notation, MD = MF and MD,2 = MF,2 , and Lemma 3 gives Ct′ = γAt − γCt + St ,

|St | ≤

22

LG MF2 MR,2 + MR MF,2 Dt . κD

Thus the scalar curvature condition is precisely the one-block instance of the Hilbert curvature condition. All hypotheses of Theorem 14 are therefore satisfied with the same λ, Λ, κK , κD , χ, LG , σ. In this one-block specialization the alignment constant in Theorem 14 is s Λ MC = , γκK which is exactly the constant appearing in Theorem 3. Applying Theorem 14 gives J(ρt ) − J⋆ ≤ with

1 min cα,r = 1 + αMC



exp(−cα,r t) (H0 − αC0 ), γ(1 − αMC )

  dα,r r , , 2λα 1 − χ 2

dα,r = γ − ασ −

αγ . 2rκD

This is exactly (16), and the proof is complete.

A.8

Proof of Theorem 4 as a scalar Hilbert-space corollary

Proof. We apply Theorem 16 with Θ = X = Rm×n ,

H = R,

again using the Frobenius inner product on Θ. The phase space Z = Θ × Θ is therefore the same as Z = X × X in the main text, and ∥ · ∥Θ = ∥ · ∥F . We first verify the assumptions. Assumption (A1) gives ∥∇F (W )∥F ≤ MF and global Lipsf ∈ X, chitzness of ∇F . Hence, for all W, W f )| ≤ MF ∥W − W f ∥F , |F (W ) − F (W so the Hilbert-space Lipschitz constant LF in Assumption (A4) may be taken to be MF . Moreover, DF (W )[V ] = ⟨∇F (W ), V ⟩F , so ∥DF (W )∥Θ→R = ∥∇F (W )∥F ≤ MF , and the Hilbert constant MD may be taken to be MF . The global Lipschitz constant of DF is the global Lipschitz constant of ∇F . Finally, the theorem assumes in addition that R′ is globally bounded, while Assumption (A1) assumes that R′ is globally Lipschitz. Since ∇R(z) = R′ (z) in R, these two bounds are exactly (36)-(37). Thus, the scalar data satisfy Assumption (A4). Under this identification, the Hilbert-space force in (31) becomes aµ (W ) = DF (W )∗ ∇R(mµ ) = R′ (mµ )∇F (W ), as shown in the proof of Theorem 3. Therefore the finite-particle ODE (64) becomes     N X 1 ẆiN (t) = − Orthε (PiN (t)), ṖiN (t) = γ R′  F (WjN (t)) ∇F (WiN (t)) − PiN (t) , N j=1

which is exactly Equation (11). The nonlinear characteristic system (73) becomes
˙ (t) = − Orth (P̄ (t)), W̄ P̄˙i (t) = γ R′ (mt )∇F (W̄i (t)) − P̄i (t) , mt = EF (W̄i (t)), i ε i 23

which is the mean-field copy system stated in Theorem 4. Theorem 16 now gives, for each fixed i and every T < ∞, a constant Cpoc (T, ε) < ∞ such that   Cpoc (T, ε) sup E ∥WiN (t) − W̄i (t)∥2F + ∥PiN (t) − P̄i (t)∥2F ≤ . N 0≤t≤T This is precisely (17). The k-particle Wasserstein estimate (156) in Theorem 16 implies convergence in law of every fixed k particles to µ⊗k t , uniformly on finite time horizons. This proves Theorem 4.

A.9

Hard-Muon subsequential limit

Theorem 9. Let µεt solve Equation (14) with common initial law µ0 ∈ P1 (Z). For every T < ∞, every sequence εk ↓ 0 has a subsequence converging in C ([0, T ]; P1 (Z)) to a curve µt . There exists a Borel field Vt (W, P ) such that Vt (W, P ) ∈ ∂∥P ∥nuc

µt dt-a.e.

and

∂t µt + ∇W · (−Vt µt ) + ∇P · γ R′ (mt ) ∇F (W ) − P µt = 0. If P has full rank almost everywhere, then Vt = Orth(P ) almost everywhere. Proof. The limiting potential is Ψ0 (P ) = ∥P ∥nuc . If P = U ΣV ⊤ has rank r, then n o ∂∥P ∥nuc = U V ⊤ + Z : U ⊤ Z = 0, ZV = 0, ∥Z∥op ≤ 1 Thus Orth(P ) = U V ⊤ is a selected subgradient, and the subdifferential is singleton exactly when the rank is full in the smaller dimension. √ For the compactness argument, note that ∥Orthε (P )∥F ≤ q uniformly in ε, and the force R′ (mt ) ∇F (W ) is bounded under Assumption (A1) on finite horizons after the same moment estimates as above. Along characteristics, √

d ∥P ε ∥ ≤ γ (C + ∥Ptε ∥F ) . dt t F Gronwall’s inequality and the common initial law µ0 ∈ P1 (Z) give uniform integrability of the first moments of {µεt : 0 ≤ t ≤ T, ε > 0}. Together with the equicontinuity estimate in W1 , this yields relative compactness in C ([0, T ]; P1 (Z)) by the Arzelà-Ascoli criterion for W1 -continuous probability curves. The fluxes Orthεk (P )µεt k dt are uniformly bounded vector-valued measures, so along a subse√ quence they converge weak-star to Vt (W, P )µt dt for a Borel field V with ∥V ∥F ≤ q. Convexity gives, for every Q ∈ X , ∥Wtε − Wsε ∥F ≤

q|t − s|,

Ψεk (Q) ≥ Ψεk (P ) + ⟨Orthεk (P ), Q − P ⟩F . Further Ψε → Ψ0 locally uniformly. To pass to the limit with the unbounded test function P , first multiply by a compactly supported cutoff in P , pass to the limit using local uniform convergence Ψε → Ψ0 , and then remove the cutoff using the uniform first-moment bound. Therefore, using this standard cutoff-based strategy to handle the unboundedness of P and passing to the limit in the integrated inequality with arbitrary nonnegative test weights yields Ψ0 (Q) ≥ Ψ0 (P ) + ⟨Vt (W, P ), Q − P ⟩F 24

for µt dt-almost every (W, P ). This is precisely Vt (W, P ) ∈ ∂Ψ0 (P ). Passing to the limit in the weak formulation gives the hard-Muon inclusion. If P has full rank almost everywhere, the subdifferential formula forces Vt = Orth(P ). Remark 9 (Why only subsequential uniqueness). At rank deficient P, ∂∥P ∥nuc is not a singleton. The regularized maps select limits of subgradients, but without an additional selection or uniqueness principle the hard limit is naturally a differential inclusion.

A.10

Hilbert-valued probability functionals on product spaces

The scalar F -output theory in the main paper is based on a probability distribution over a single matrix space. In particular, the one-hidden layer neural network teacher-student mean-field and the transformer mixture-of-experts setting requires a probability distribution over parameter tuples and an expert-router parameter tuple, respectively. The parameter space is therefore enlarged before the variational and dynamical constructions are introduced. A.10.1

Extended product parameter space

Let Θ = Θexp × Θgate =

B Y

Rmb ×nb

(23)

b=1

be a finite-dimensional real Hilbert space. The first Bexp blocks may represent expert parameters, while the remaining Bgate = B − Bexp blocks represent router or gating parameters. A generic point is written as   θ = (ω, ϕ) = θ(1) , . . . , θ(B) , (24) where ω ∈ Θexp and ϕ ∈ Θgate . The product Hilbert inner product is ⟨θ, ϑ⟩Θ :=

B D E X θ(b) , ϑ(b) , F

b=1

∥θ∥2Θ :=

B X

θ(b)

b=1

2 F

.

(25)

The Wasserstein spaces below are built from the metric induced by ∥·∥Θ . The routing parameter is part of the state variable. Consequently, any force field, momentum variable, mirror map, and Hamiltonian gradient is defined on the same extended space Θ and not only on the expert subspace. Vector-valued and scalar router parameters can be included in (23) by treating vectors as one-column matrices and scalars as 1 × 1 matrices. Biases may also be absorbed into matrix blocks by appending a homogeneous coordinate to the input representation. Thus the product matrix formulation covers the usual affine router scores without introducing a separate notation. A.10.2

Hilbert-valued feature maps with routing included

Let H be a real Hilbert space and let F : Θ → H,

R : H → R.

(26)

For ρ ∈ P1 (Θ) define the Bochner integral Z mρ :=

F (θ) dρ(θ) Θ

25

(27)

whenever it is finite, and set J(ρ) := R(mρ ).

(28)

The map F is allowed to contain the router score, the gate weight, the expert output, and any Hilbert-valued augmentation needed to express normalized gates. In particular, if θ = (ω, ϕ) and ψω is an expert output while sϕ is an input-dependent router score, then maps of the form F (ω, ϕ) = “input-indexed function involving sϕ and ψω ”

(29)

are ordinary Hilbert-valued feature maps on Θ. Proposition 4 (First variation on the extended gated space). Assume that F is Bochner integrable under ρ and that R is Frechet differentiable at mρ . Then a valid first variation of J at ρ is δJ (ρ)(θ) = ⟨∇R(mρ ), F (θ)⟩H , δρ

(30)

up to an additive constant independent of θ. If F is Frechet differentiable, then the Wasserstein force on the extended expert-router parameter is aρ (θ) := ∇θ

δJ (ρ)(θ) = DF (θ)∗ ∇R(mρ ) ∈ Θ. δρ

(31)

If θ = (ω, ϕ), then
gate aρ (θ) = aexp ρ (ω, ϕ), aρ (ω, ϕ) ,

(32)

where the two components are the projections of DF (ω, ϕ)∗ ∇R(mρ ) onto Θexp and Θgate , respectively. Proof. Let ν be a finite signed measure on Θ with ν(Θ) = 0 and define ρs = ρ + sν for s in an interval on which the perturbation is meaningful. By linearity of the Bochner integral, Z mρs = mρ + s F (θ) dν(θ). Θ

Frechet differentiability of R gives Z  d J(ρs ) = DR(mρ ) F (θ) dν(θ) ds Θ s=0   Z = ∇R(mρ ), F (θ) dν(θ) Θ H Z = ⟨∇R(mρ ), F (θ)⟩H dν(θ). Θ

This is precisely the defining identity for the first variation, modulo constants independent of θ. If F is differentiable, then for every v ∈ Θ,   Dθ ⟨∇R(mρ ), F (θ)⟩H [v] = ⟨∇R(mρ ), DF (θ)[v]⟩H = ⟨DF (θ)∗ ∇R(mρ ), v⟩Θ . The Riesz representation in the Hilbert space Θ gives (31). The decomposition (32) follows from the orthogonal product decomposition Θ = Θexp × Θgate . 26

A.10.3

Regularity assumptions

The smooth ODE and PDE theory can be stated under either global assumptions or localized finite-horizon assumptions. The localized version is the natural one for transformer models, because attention, feed-forward, and router score maps are smooth on bounded parameter sets, while global Lipschitz constants on all of parameter space are generally unavailable. Assumption (A3) (Localized Hilbert smoothness on the extended space). For every RΘ < ∞ and every RH < ∞, the following quantities are finite on the indicated balls: MF (RΘ ) :=

sup ∥θ∥Θ ≤RΘ

MD (RΘ ) :=

sup ∥θ∥Θ ≤RΘ

MR (RH ) :=

sup ∥z∥H ≤RH

∥F (θ)∥H ,

LF (RΘ ) :=

∥F (θ) − F (ϑ)∥H , ∥θ − ϑ∥Θ ∥θ∥Θ ,∥ϑ∥Θ ≤RΘ sup θ̸=ϑ

∥DF (θ)∥Θ→H ,

LD (RΘ ) :=

∥DF (θ) − DF (ϑ)∥Θ→H , ∥θ − ϑ∥Θ ∥θ∥Θ ,∥ϑ∥Θ ≤RΘ sup θ̸=ϑ

∥∇R(z)∥H ,

LR (RH ) :=

∥∇R(z) − ∇R(z ′ )∥H . ∥z − z ′ ∥H ∥z∥H ,∥z ′ ∥H ≤RH sup

z̸=z ′

When second-order convergence estimates are invoked, F is twice Frechet differentiable on bounded Θ-balls and R is twice Frechet differentiable on bounded H-balls, with bounded second derivatives on those balls. Assumption (A4) (Global Hilbert smoothness). There exist constants LF , MD , LD , MR , LR < ∞ such that for all θ, ϑ ∈ Θ and all z, z ′ ∈ H, ∥F (θ) − F (ϑ)∥H ≤ LF ∥θ − ϑ∥Θ ,

(33)

∥DF (θ)∥Θ→H ≤ MD ,

(34)

∥DF (θ) − DF (ϑ)∥Θ→H ≤ LD ∥θ − ϑ∥Θ ,

(35)

∥∇R(z)∥H ≤ MR , ′

∇R(z) − ∇R(z ) H ≤ LR z − z

(36) ′ H

.

(37)

Lemma 1 (Force estimates). Under Assumption (A4), for every ρ, ρ̃ ∈ P1 (Θ) and θ, θ̃ ∈ Θ, aρ (θ) − aρ̃ (θ̃)

Θ

≤ LD MR θ − θ̃

Θ

+ MD LR LF W1 (ρ, ρ̃).

(38)

Furthermore, ∥aρ (θ)∥Θ ≤ MD MR .

(39)

Under Assumption (A3), the same estimates hold with the corresponding localized constants whenever θ, θ̃ and the supports of ρ, ρ̃ lie in a common bounded Θ-ball and mρ , mρ̃ lie in a common bounded H-ball. Proof. Let uρ := ∇R(mρ ). Then   aρ (θ) − aρ̃ (θ̃) = DF (θ)∗ − DF (θ̃)∗ uρ + DF (θ̃)∗ (uρ − uρ̃ ) . The first term has norm at most LD θ − θ̃

Θ

MR . The second has norm at most

MD LR ∥mρ − mρ̃ ∥H . 27

For every coupling π of ρ and ρ̃, Z (F (ξ) − F (η)) dπ(ξ, η),

mρ − mρ̃ = Θ×Θ

so

Z ∥ξ − η∥Θ dπ(ξ, η).

∥mρ − mρ̃ ∥H ≤ LF

Taking the infimum over π yields ∥mρ − mρ̃ ∥H ≤ LF W1 (ρ, ρ̃). This proves (38). The bound (39) follows directly from ∥DF (θ)∥ ≤ MD and ∥uρ ∥ ≤ MR . The localized statement is identical after restricting every estimate to the relevant bounded balls.

A.11

Regularized Muon mirror geometry on the extended product space

The regularized Muon mirror map is defined on the full product space Θ, including both expert and router blocks. The construction is block-separable, and therefore the Fenchel conjugacy and mirror-update interpretation are inherited from the single-matrix spectral construction. A.11.1

Single block Fenchel pair

Fix a matrix block Rmb ×nb and set qb = min(mb , nb ). For P ∈ Rmb ×nb define Ψε,b (P ) :=

qb   X p σr (P )2 + ε2 − ε ,

(40)

r=1

where the singular values are padded by zeros. For G ∈ Rmb ×nb define  qb   X p  ε 1 − 1 − σr (G)2 , ∥G∥op ≤ 1, Φε,b (G) := r=1   +∞, ∥G∥op > 1.

(41)

On the relative interior ∥G∥op < 1, Φε,b is differentiable. If P = U diag(σ1 , . . . , σs )V ⊤ is a reduced SVD, then ! σ1 σs ∇Ψε,b (P ) = U diag p 2 ,..., p V ⊤ =: Orthε,b (P ). (42) σs2 + ε2 σ1 + ε 2 Moreover, ∥Orthε,b (P )∥op < 1,

∥Orthε,b (P )∥F ≤

√

Orthε,b (P ) − Orthε,b (P̃ )

qb ,

F

≤

1 P − P̃ . ε F (43)

The scalar pair a 7→

p a2 + ε2 − ε,

b 7→ ε(1 −

p 1 − b2 ) + ι[−1,1] (b)

is Fenchel conjugate. By von Neumann’s trace inequality and the standard spectral-function conjugacy theorem, this scalar conjugacy lifts to Φ∗ε,b = Ψε,b ,

Ψ∗ε,b = Φε,b .

28

(44)

A.11.2

Product Fenchel pair, including router blocks

For P = (P (1) , . . . , P (B) ) ∈ Θ and G = (G(1) , . . . , G(B) ) ∈ Θ, define ΨΘ ε (P ) :=

B X

Ψε,b (P (b) ),

ΦΘ ε (G) :=

b=1

B X

Φε,b (G(b) ).

(45)

b=1

The product regularized Muon map is   Θ (1) OrthΘ ), . . . , Orthε,B (P (B) ) . ε (P ) := ∇Ψε (P ) = Orthε,1 (P

(46)

The total block rank parameter is B X

qΘ :=

qb .

(47)

b=1

Equations (43) imply Θ OrthΘ ε (P ) − Orthε (P̃ )

Θ

≤

1 P − P̃ , ε Θ

OrthΘ ε (P ) Θ ≤

√

qΘ .

(48)

Proposition 5 (Fenchel duality and mirror update on the extended expert-router space). With respect to the product pairing (25), ∗ Θ (ΦΘ ε ) = Ψε ,

∗ Θ (ΨΘ ε ) = Φε .

(49)

For every P ∈ Θ, the variational problem  Gε (P ) = argmin ⟨P, G⟩Θ + ΦΘ ε (G)

(50)

Θ Gε (P ) = −∇ΨΘ ε (P ) = − Orthε (P ).

(51)

G∈Θ

has the unique solution Equivalently, if Gε,k = − OrthΘ ε (Pk ), then n o Gε,k+1 = argmin ⟨Pk+1 − Pk , G⟩Θ + DΦΘ (G, G ) , ε,k ε G∈Θ

where Θ Θ DΦΘ (G, H) = ΦΘ ε (G) − Φε (H) − ∇Φε (H), G − H Θ ε

whenever H lies in the differentiability domain of ΦΘ ε. Proof. By separability and (44), ∗ (ΦΘ ε ) (P ) = sup G∈Θ

=

=

( B XD

B X

P

(b)

b=1 G B X

,G

b=1

sup

(b)

E F

−

B X

) (b)

Φε,b (G )

b=1

nD E o P (b) , G(b) − Φε,b (G(b) ) F

(b) ∈Rmb ×nb

Ψε,b (P (b) ) = ΨΘ ε (P ).

b=1

29

(52)

∗ Θ The proof of (ΨΘ ε ) = Φε is identical. The first-order optimality condition for (50) is

0 = P + ∇ΦΘ ε (Gε (P )). Θ Θ Since ∇ΦΘ ε and ∇Ψε are inverse maps between the interior of the effective domain of Φε and Θ, and since ΨΘ ε is even, its gradient is odd. Therefore Θ Gε (P ) = ∇ΨΘ ε (−P ) = −∇Ψε (P ).

This gives (51). For (52), use ∇ΦΘ ε (Gε,k ) = −Pk . Then Θ DΦΘ (G, Gε,k ) = ΦΘ ε (G) − Φε (Gε,k ) − ⟨−Pk , G − Gε,k ⟩Θ ε

= ΦΘ ε (G) + ⟨Pk , G⟩Θ + Ck , where Ck is independent of G. Adding ⟨Pk+1 − Pk , G⟩Θ gives ⟨Pk+1 , G⟩Θ + ΦΘ ε (G) + Ck . Thus the minimizer is exactly the minimizer of (50) with P = Pk+1 . Remark 10 (Persistence of the mirror-map interpretation). The mirror-map interpretation survives on the extended gated product space because the mirror potential is defined on every block of Θ, including the router blocks. If P = (P exp , P gate ), then
exp exp OrthΘ ), Orthgate (P gate ) , ε (P ) = Orthε (P ε and the Bregman update (52) is a single mirror step in the full expert-router Hilbert space. There is no separate post-hoc extension of the mirror map after the gate is introduced; the router geometry is part of the product Fenchel pair from the beginning.

A.12

Finite-particle lift and exact regularized Muon scheme

Let N ∈ N and let

θ = (θ1 , . . . , θN ) ∈ ΘN ,

θi = (ωi , ϕi ).

The empirical measure is N

ρN θ :=

1 X δθi . N

(53)

i=1

The lifted finite-particle objective is JN (θ) := J(ρN θ )=R

! N 1 X F (θi ) . N

(54)

i=1

The mean-field inner product on ΘN is N

⟨U, V ⟩avg :=

1 X ⟨Ui , Vi ⟩Θ , N i=1

Define

N

∥U ∥2avg :=

1 X ∥Ui ∥2Θ . N

(55)

i=1

N

mN (θ) :=

1 X F (θj ), N

uN (θ) := ∇R(mN (θ)).

j=1

30

(56)

Proposition 6 (Particle gradient on the extended space). Assume F and R are differentiable. The gradient of JN with respect to (55) is
N gradavg JN (θ) = aN (θ) = aN (57) 1 (θ), . . . , aN (θ) , where ∗ aN i (θ) := DF (θi ) uN (θ) ∈ Θ.

(58)

In particular, the gate component of aN i is the gradient of the lifted objective with respect to the router blocks of particle i. Proof. For U = (U1 , . . . , UN ) ∈ ΘN , the chain rule gives + * N 1 X DF (θi )[Ui ] DJN (θ)[U ] = ∇R(mN (θ)), N i=1

=

1 N

N X

H

⟨DF (θi )∗ uN (θ), Ui ⟩Θ

i=1 N

= a (θ), U avg . This is the defining identity for the mean-field gradient. The particle-level mirror potentials are N

ΨN ε (P ) :=

N

1 X Θ Ψε (Pi ), N

ΦN ε (G) :=

i=1

1 X Θ Φε (Gi ). N

(59)

i=1

∗ N N ∗ N With respect to ⟨·, ·⟩avg , the conjugacy (ΦN ε ) = Ψε and (Ψε ) = Φε follows by the same separability proof as Proposition 5. The exact finite-particle regularized Muon scheme on the extended expert-router space is

Pi,k+1 = βPi,k + (1 − β)aN i (θ k ),  Gi,k+1 = argmin ⟨Pi,k+1 , G⟩Θ + ΦΘ ε (G) , G∈Θ

(60)

θi,k+1 = θi,k + ηGi,k+1 . Using Proposition 5, this is equivalently 

 N X 1 Pi,k+1 = βPi,k + (1 − β)DF (θi,k )∗ ∇R  F (θj,k ) , N j=1

(61)

θi,k+1 = θi,k − η OrthΘ ε (Pi,k+1 ). Writing θi = (ωi , ϕi ) and Pi = (Piexp , Pigate ), the second line updates expert and router blocks simultaneously: exp ωi,k+1 = ωi,k − η Orthexp ε (Pi,k+1 ),

gate ϕi,k+1 = ϕi,k − η Orthgate (Pi,k+1 ). ε

31

(62)

A.13

Finite-particle continuous-time limit

Let the inertial scaling be βh = 1 − γh + rh ,

ηh = h,

rh → 0, h

(63)

where γ > 0 is fixed. The common second-order-consistent choice is rh = O(h2 ). Theorem 10 (Finite-N ODE limit on the extended gated space). Fix N ∈ N, ε > 0, γ > 0, and initial data (θ 0 , P0 ) ∈ ΘN × ΘN . Suppose Assumption (A4) holds. Then the ODE system θ̇i (t) = − OrthΘ ε (Pi (t)), 



Ṗi (t) = γ DF (θi (t))∗ ∇R 

N 1 X

N





(64)

F (θj (t)) − Pi (t) ,

j=1

for i = 1, . . . , N , has a unique global solution. Let (θ hk , Pkh ) be the iterates of (61) with scaling (63), and let (θ h (t), P h (t)) be the continuous piecewise-linear interpolation satisfying (θ h (kh), P h (kh)) = (θ hk , Pkh ). Then, for every T < ∞,   h h sup θ (t) − θ(t) + P (t) − P (t) →0 (65) avg

0≤t≤T

avg

as h ↓ 0. If rh = O(h2 ), the convergence rate is O(h) on every finite time interval. Proof. Define aN by (58). The vector field on ΘN × ΘN is
Θ N N BN (θ, P ) = − OrthΘ ε (P1 ), . . . , − Orthε (PN ), γ(a1 (θ) − P1 ), . . . , γ(aN (θ) − PN ) .

(66)

The map P 7→ OrthΘ ε (P ) is globally 1/ε-Lipschitz by (48). Lemma 1 implies a global Lipschitz estimate for aN in the mean-field norm. Indeed, for two particle configurations θ, θ̃ ∈ ΘN , N

mN (θ) − mN (θ̃)

H

≤

1 X LF θj − θ̃j ≤ LF θ − θ̃ . N Θ avg j=1

Consequently, N aN i (θ) − ai (θ̃)

Θ

≤ LD MR θi − θ̃i

Θ

+ MD LR LF θ − θ̃

avg

.

Squaring, averaging over i, and using (x + y)2 ≤ 2x2 + 2y 2 yields a global Lipschitz bound for aN . Hence BN is globally Lipschitz, and the Picard-Lindelof theorem gives a unique global solution of (64). The discrete momentum update satisfies   h h h h Pi,k+1 − Pi,k = (1 − βh ) aN (θ ) − P i k i,k   h h = (γh − rh ) aN (θ ) − P i k i,k   h h = hγ aN i (θ k ) − Pi,k + o(h)

32

locally uniformly on bounded sets. For the position update, h h h θi,k+1 − θi,k = −h OrthΘ ε (Pi,k+1 )

  h Θ h Θ h = −h OrthΘ ε (Pi,k ) − h Orthε (Pi,k+1 ) − Orthε (Pi,k ) . h h 2 Since OrthΘ ε is 1/ε-Lipschitz and Pi,k+1 − Pi,k = O(h) on bounded sets, the second term is O(h ). Therefore the one-step local truncation error of the scheme relative to the Euler discretization of (64) is O(h2 ) + o(h)h. A standard discrete Gronwall estimate for globally Lipschitz ODEs gives uniform convergence on [0, T ]. If rh = O(h2 ), the local truncation error is O(h2 ) and the global error is O(h).

Remark 11 (Localized finite-horizon form). The transformer specialization generally satisfies Assumption (A3), not Assumption (A4). The conclusion of Theorem 10 remains valid on each finite interval [0, T ] under Assumption (A3). Indeed, the position velocity obeys θ̇i (t) so

Θ

= OrthΘ ε (Pi (t)) Θ ≤

√

qΘ ,

√ ∥θi (t)∥Θ ≤ ∥θi (0)∥Θ + T qΘ .

All force-field constants needed on [0, T ] are therefore evaluated on a bounded parameter ball. Once θi (t) is confined to that ball, the comparison inequality ! d sup aN ∥Pi (t)∥Θ ≤ γ i (θ(s)) Θ − ∥Pi (t)∥Θ dt 0≤s≤T precludes finite-time momentum blow-up. Local Lipschitzness on the resulting bounded phase-space region gives uniqueness and the same Euler convergence argument.

A.14

Phase-space mean-field equation

Set Z := Θ × Θ,

z = (θ, p).

(67)

For µ ∈ P1 (Z) let ρ = (πθ )# µ and Z mµ :=

Z F (θ) dµ(θ, p) =

Z

F (θ) dρ(θ).

(68)

Θ

Define aµ (θ) := DF (θ)∗ ∇R(mµ )

(69)

bµ (θ, p) := − OrthΘ ε (p), γ(aµ (θ) − p) .

(70)

and the phase-space drift

Theorem 11 (Well-posed phase-space mean-field equation). Assume Assumption (A4). For every µ0 ∈ P1 (Z) there exists a unique curve µ ∈ C([0, ∞); P1 (Z))

33

such that, for every ζ ∈ Cc∞ (Z), Z Z   d ζ(θ, p) dµt (θ, p) = ∇θ ζ, − OrthΘ ε (p) Θ + ⟨∇p ζ, γ(aµt (θ) − p)⟩Θ dµt . dt Z Z

(71)

Equivalently,
∂t µt + ∇θ · − OrthΘ ε (p)µt + ∇p · (γ(aµt (θ) − p)µt ) = 0.

(72)

The solution is the law of the nonlinear characteristic system Θ̇t = − OrthΘ ε (Pt ), Ṗt = γ (DF (Θt )∗ ∇R (EF (Θt )) − Pt ) ,

(73)

with (Θ0 , P0 ) ∼ µ0 . Proof. The drift is globally Lipschitz in (θ, p) and Lipschitz in the measure argument with respect to W1 . The Lipschitz estimate in (θ, p) follows from (48) and Lemma 1. For the measure argument, ∥aµ (θ) − aν (θ)∥Θ ≤ MD LR LF W1 ((πθ )# µ, (πθ )# ν) ≤ MD LR LF W1 (µ, ν). The drift has at most linear growth in p and bounded growth in the θ component. The standard Picard iteration for McKean-Vlasov ODEs with Lipschitz drift gives a unique nonlinear process solving (73). Setting µt = Law(Θt , Pt ) and applying the chain rule to ζ(Θt , Pt ) gives (71). Conversely, the superposition principle for Lipschitz continuity equations implies that any weak solution is transported by the same characteristic flow. The fixed-point uniqueness for the nonlinear characteristic equation therefore gives uniqueness of the PDE solution. Remark 12 (Localized compact-support form). On a finite horizon [0, T ], Assumption (A4) can √ be replaced by Assumption (A3) when µ0 has compact support. The bound Θ̇t ≤ qΘ confines Θ the position support to a bounded Θ-ball, and the momentum comparison estimate confines the momentum support to a bounded ball depending on T , the initial support, and the localized force bound. On this bounded region the drift is Lipschitz, and the proof of Theorem 11 applies without modification.

A.15

Hamiltonian formulation and dissipation identity

Define the regularized Muon Hamiltonian on phase-space probability measures by Z  Z Θ Hε,γ (µ) := Ψε (p) dµ(θ, p) + γR F (θ) dµ(θ, p) . Z

(74)

Z

Proposition 7 (First variation of the Hamiltonian). A valid first variation of Hε,γ is δHε,γ (µ)(θ, p) = ΨΘ ε (p) + γ ⟨∇R(mµ ), F (θ)⟩H , δµ

(75)

up to an additive constant independent of (θ, p). Consequently, ∇p

δHε,γ = OrthΘ ε (p), δµ

∇θ

34

δHε,γ = γaµ (θ). δµ

(76)

Proof. Let ν be a finite signed measure on Z with ν(Z) = 0 and set µs = µ + sν. Then   Z Z d Hε,γ (µs ) = ΨΘ (p) dν(θ, p) + γ ∇R(m ), F (θ) dν(θ, p) µ ε ds Z s=0 H ZZ  Θ  Ψε (p) + γ ⟨∇R(mµ ), F (θ)⟩H dν(θ, p). = Z

Θ This proves (75). Differentiating with respect to p gives ∇ΨΘ ε (p) = Orthε (p). Differentiating with respect to θ and using Proposition 4 gives the second identity in (76).

Theorem 12 (Damped Hamiltonian form). The phase-space PDE (72) is equivalent to       δHε,γ δHε,γ ∂t µt + ∇θ · µt −∇p + ∇ p · µt ∇ θ − γp = 0. δµt δµt

(77)

Proof. Substitution of (76) into (77) gives θ-velocity − OrthΘ ε (p) and p-velocity γaµt (θ) − γp, which is exactly (72). Theorem 13 (Dissipation identity on the extended gated space). Let µt solve (72). Assume that the first moment is finite on compact time intervals and that the test-function cutoff argument is justified by either Assumption (A4) or the localized compact-support conditions of Remark 12. Then t 7→ Hε,γ (µt ) is absolutely continuous on compact intervals and, for almost every t, Z d Hε,γ (µt ) = −γ p, OrthΘ (78) ε (p) Θ dµt (θ, p) ≤ 0. dt Z Moreover, Θ dΘ ε (p) := p, Orthε (p) Θ =

qb B X X

σ (p(b) )2 p r , σr (p(b) )2 + ε2 b=1 r=1

(79)

Θ so dΘ ε (p) ≥ 0, and dε (p) = 0 if and only if p = 0.

Proof. Let G(p) = OrthΘ ε (p) and at (θ) = aµt (θ). The PDE is the continuity equation with velocity bt (θ, p) = (−G(p), γ(at (θ) − p)). For smooth compactly supported ζ, Z Z   d ζ dµt = ⟨∇θ ζ, −G(p)⟩Θ + ⟨∇p ζ, γ(at (θ) − p)⟩Θ dµt . dt The functions θ 7→ F (θ) and p 7→ ΨΘ ε (p) are admissible by the stated cutoff hypothesis. For the feature moment, Z Z d ṁt = F (θ) dµt (θ, p) = DF (θ)[−G(p)] dµt (θ, p). (80) dt Therefore d [γR(mt )] = γ ⟨∇R(mt ), ṁt ⟩H dt Z

⟨DF (θ)∗ ∇R(mt ), G(p)⟩Θ dµt

= −γ Z = −γ

⟨at (θ), G(p)⟩Θ dµt . 35

(81)

For the kinetic term, Z Z d Θ Ψε (p) dµt = ⟨G(p), γ(at (θ) − p)⟩Θ dµt dt Z Z = γ ⟨at (θ), G(p)⟩Θ dµt − γ ⟨p, G(p)⟩Θ dµt .

(82)

Adding (81) and (82) cancels the mixed term and gives (78). Formula (79) follows by inserting the blockwise SVDs into (46). Each summand is nonnegative and vanishes exactly when the corresponding singular value is zero. Hence the sum vanishes if and only if all blocks of p vanish.

A.16

Continuous-time convergence under functional assumptions

The accelerated Muon flow is not an ordinary Wasserstein gradient flow for J(ρt ), and J(ρt ) need not be monotone. The dissipated Lyapunov quantity is the damped Hamiltonian. Exponential convergence of the objective gap follows from the standard force-norm assumptions, kinetic coercivity on bounded momentum sets, and a curvature estimate for the force-momentum alignment. Throughout this section define Z Z Θ Kt := Ψε (p) dµt (θ, p), Dt := p, OrthΘ ε (p) Θ dµt (θ, p), Z

Z

Jt := J(ρt ) = R(mt ),

Ut := Jt − J⋆ , J⋆ := inf J(ρ), ρ∈P1 (Θ) Z At := ∥at (θ)∥2Θ dµt (θ, p),

Ht := Kt + γUt , Z Ct := ⟨at (θ), p⟩Θ dµt (θ, p),

Z

at (θ) = DF (θ)∗ ∇R(mt ).

(83)

Z

A.16.1

Kinetic coercivity on bounded momentum sets

Lemma 2 (Kinetic constants for extended product Muon). Assume ∥p∥Θ ≤ BP on the support of µt for all times under consideration. Then κK :=

ε2 , (BP2 + ε2 )3/2

1 , κD := q BP2 + ε2

1 LG := , ε

χ := 1

(84)

satisfy, for all such p, κK ∥p∥2Θ , 2 2 p, OrthΘ ε (p) Θ ≥ κD ∥p∥Θ , ΨΘ ε (p) ≥

Θ ΨΘ ε (p) ≤ χ p, Orthε (p) Θ , OrthΘ ε (p) Θ ≤ LG ∥p∥Θ .

Proof. For s ∈ [0, BP ],

Z s p r √ s2 + ε2 − ε = dr. r 2 + ε2 0

(85) (86) (87) (88)

√ The derivative of r 7→ r/ r2 + ε2 is ε2 /(r2 + ε2 )3/2 , which is at least κK on [0, BP ]. Since the derivative vanishes at 0, integration gives p κK 2 s2 + ε2 − ε ≥ s . 2 36

Summing this over all singular values in all blocks gives (85). Similarly, √

s2 s2 ≥q = κD s2 , 2 s2 + ε2 2 B +ε P

which gives (86). The inequality p s2 s2 + ε2 − ε ≤ √ s2 + ε2 is equivalent to ε ≤ (88). A.16.2

√

s2 + ε2 , and summing gives (87). Finally, OrthΘ ε (0) = 0 and (48) gives

Curvature identity for the alignment term

Assumption (A5) (Second-order trajectory regularity). Along the trajectory there are constants MD , MD,2 , MR , MR,2 < ∞ such that ∥DF (θ)∥Θ→H ≤ MD ,

D2 F (θ) Θ×Θ→H ≤ MD,2 ,

(89)

D2 R(mt ) H→H ≤ MR,2 .

(90)

for µt -almost every θ, and ∥∇R(mt )∥H ≤ MR ,

Lemma 3 (Curvature remainder on the extended space). Assume the hypotheses of Lemma 2 and Assumption (A5). Then Ct is absolutely continuous and Ct′ = γAt − γCt + St , where

Z St = Z

(91)

∂t at (θ) − Dθ at (θ)[OrthΘ ε (p)], p Θ dµt (θ, p),

and |St | ≤ σDt ,

σ :=

LG 2 MD MR,2 + MR MD,2 . κD

(92)

(93)

Proof. Use the characteristic representation. Along a characteristic, θ̇t = −G(pt ),

ṗt = γ(at (θt ) − pt ),

G = OrthΘ ε .

The chain rule gives D E d ⟨at (θt ), pt ⟩Θ = ∂t at (θt ) + Dθ at (θt )[θ̇t ], pt + ⟨at (θt ), ṗt ⟩Θ dt Θ = ⟨∂t at (θt ) − Dθ at (θt )[G(pt )], pt ⟩Θ + γ ∥at (θt )∥2Θ − γ ⟨at (θt ), pt ⟩Θ . Integration with respect to the law of the characteristic gives (91) and (92). It remains to bound St . Write ut = ∇R(mt ), so at (θ) = DF (θ)∗ ut . For v ∈ Θ, Dθ at (θ)[v] = D2 F (θ)[v, ·]∗ ut , and hence ∥Dθ at (θ)[G(p)]∥Θ ≤ MD,2 MR ∥G(p)∥Θ . 37

(94)

Also, ∂t at (θ) = DF (θ)∗ D2 R(mt )[ṁt ]. By (80), Z ṁt = −

DF (θ)[G(p)] dµt (θ, p), Z

so

Z ∥G(p)∥Θ dµt .

∥ṁt ∥H ≤ MD

(95)

Consequently, 2 ∥∂t at (θ)∥Θ ≤ MD MR,2 ∥ṁt ∥H ≤ MD MR,2

Z ∥G(p)∥Θ dµt .

(96)

Using (86), (88), and Cauchy’s inequality, Z Z LG 1/2 ∥G(p)∥Θ dµt ≤ LG ∥p∥Θ dµt ≤ √ Dt , κD Z 1 1/2 ∥p∥Θ dµt ≤ √ Dt , κD Z Z LG ∥G(p)∥Θ ∥p∥Θ dµt ≤ LG ∥p∥2Θ dµt ≤ Dt . κD Combining (94) and (96) gives Z  Z  Z 2 |St | ≤ MD MR,2 ∥G(p)∥Θ dµt ∥p∥Θ dµt + MR MD,2 ∥G(p)∥Θ ∥p∥Θ dµt ≤

LG 2 MD MR,2 + MR MD,2 Dt . κD

This is (93). A.16.3

Exponential convergence theorem

Assumption (A6) (Functional PL and upper-gradient conditions). Along the trajectory there exist constants λ > 0 and Λ < ∞ such that At ≥ 2λUt ,

(97)

At ≤ 2ΛUt .

(98)

and Theorem 14 (Exponential convergence under general Hilbert domain). Assume that the flow (72) is defined for all t ≥ 0 and that Ut ≥ 0. Assume that the kinetic estimates (85)-(88), the curvature identity Ct′ = γAt − γCt + St , |St | ≤ σDt , (99) and the PL/upper-gradient Assumption (A6) hold for all t ≥ 0. Lemma 2 and Lemma 3 give sufficient conditions for the kinetic estimates and the curvature identity, respectively. Define s Λ MC := . (100) γκK 38

Choose r ∈ (0, 2) and α > 0 such that

Let

αγ > 0. 2rκD

(101)

  dα,r r . , 2λα 1 − χ 2

(102)

dα,r := γ − ασ −

αMC < 1,



1 cα,r := min 1 + αMC

Then J(ρt ) − J⋆ ≤

exp(−cα,r t) [H0 − αC0 ] . γ(1 − αMC )

(103)

Proof. First, bound the alignment term. By Cauchy’s inequality, (85), and (98), 1/2 1/2 Z Z 2 2 |Ct | ≤ ∥p∥Θ dµt ∥at (θ)∥Θ dµt ≤ (2ΛUt )

1/2



√ √ √ Since 2 Ut Kt ≤ (Kt + γUt )/ γ = Ht / γ,

2Kt κK

1/2

r =2

Λ p Ut Kt . κK

|Ct | ≤ MC Ht .

(104)

Define Lt := Ht − αCt . Equation (104) implies (1 − αMC )Ht ≤ Lt ≤ (1 + αMC )Ht .

(105)

By Theorem 13, Ht′ = Kt′ + γUt′ = −γDt . By the assumed curvature identity (99), L′t = Ht′ − αCt′ = −γDt − α(γAt − γCt + St ) = −γDt − αγAt + αγCt − αSt . Use |St | ≤ σDt and Young’s inequality: Z Z r 1 r 1 Dt . Ct ≤ ∥at (θ)∥Θ ∥p∥Θ dµt ≤ At + ∥p∥2Θ dµt ≤ At + 2 2r 2 2rκD Therefore

   αγ r γ − ασ − Dt − αγ 1 − At = −dα,r Dt − aα,r At , 2rκD 2 where aα,r = αγ(1 − r/2) > 0. By (87), Dt ≥ Kt /χ. By (97), At ≥ 2λUt . Hence L′t ≤ −

(106)

dα,r Kt + 2λaα,r Ut χ ≥ qα,r (Kt + γUt ) = qα,r Ht ,

dα,r Dt + aα,r At ≥

where

 qα,r = min

dα,r 2λaα,r , χ γ



 = min

  dα,r r , 2λα 1 − . χ 2

Combining this with (106) gives L′t ≤ −qα,r Ht . Since Lt ≤ (1 + αMC )Ht , one has Ht ≥ Lt /(1 + αMC ), and therefore qα,r L′t ≤ − Lt = −cα,r Lt . 1 + αMC Gronwall’s inequality gives Lt ≤ exp(−cα,r t)L0 . Finally, γUt ≤ Ht ≤ Lt /(1 − αMC ), and (103) follows. 39

Remark 13 (Meaning of the PL condition). The PL condition (97) is a property of the probability functional J along the probability-flow trajectory. It is not a consequence of cross-entropy smoothness alone and is not automatic for arbitrary transformer mixture-of-experts parameterizations. In the transformer specialization below, smoothness and bounded-gradient properties of the loss and feature map verify the analytic assumptions needed for well-posedness, ODE limits, Hamiltonian dissipation, and curvature control. Exponential convergence additionally requires (97) and (98), as in the scalar theory. Corollary 1 (Criticality of compact omega-limit points). Assume the trajectory {µt : t ≥ 0} is relatively compact in P1 (Z) and Hε,γ is bounded below. Then Z ∞ Dt dt < ∞. 0

Every invariant omega-limit point is supported on {p = 0} and satisfies aµ (θ) = 0 for µ-almost every (θ, p). Proof. The dissipation identity gives Z T γ

Dt dt = Hε,γ (µ0 ) − Hε,γ (µT ) ≤ Hε,γ (µ0 ) − inf Hε,γ .

0

Letting T → ∞ proves integrability. Since Dt = 0 if and only if p = 0 by Theorem 13, any invariant limit must be supported on p = 0. On this support, the p-velocity is γaµ (θ). Invariance forces aµ (θ) = 0 on the support; otherwise the measure immediately leaves {p = 0}.

A.17

Discrete-time convergence of the regularized Hamiltonian-Muon map

The preceding convergence result, Theorem 14, is a continuous-time statement for the damped Hamiltonian probability flow. We now prove a genuine fixed-step discrete-time convergence result for the natural semi-implicit Euler discretization of that flow. The result is stated on the extended product space Θ because this notation covers the gated transformer setting. The scalar matrix-valued theory in the main paper is recovered by taking Θ = X = Rm×n , H = R, and aµ (W ) = R′ (mµ )∇F (W ),

OrthΘ ε = Orthε .

Throughout this section the regularized Muon velocity and its Fenchel dissipation density are denoted explicitly by Θ Gε (p) := OrthΘ ε (p) = ∇Ψε (p),

dε (p) := ⟨p, Gε (p)⟩Θ .

(107)

Thus every appearance of Gε , dε , and ΨΘ ε depends on the fixed regularization parameter ε > 0. No unregularized hard-Muon map is used in the discrete convergence proof. For a phase-space law µ ∈ P(Θ × Θ) let Z ρ = (πθ )# µ, mµ = F (θ) dµ(θ, p), aµ (θ) = DF (θ)∗ ∇R(mµ ). (108) Θ×Θ

40

A.17.1

The law-level discrete map and the correct Hamiltonian scaling

Let ηh > 0 be the position step and let βh ∈ (0, 1) be the momentum retention coefficient. Set δh := 1 − βh . Given µk , define the force ak (θ) := aµk (θ), and update the momentum first: p+ k (θ, p) := βh p + δh ak (θ).

(109)

The position is then transported using the regularized Muon direction evaluated at the updated momentum:
+ Tkh (θ, p) := θ − ηh Gε (p+ µk+1 := (Tkh )# µk . (110) k (θ, p)), pk (θ, p) , Equivalently, if (Θk , Pk ) ∼ µk , then Pk+1 = βh Pk + (1 − βh )aµk (Θk ),

Θk+1 = Θk − ηh Gε (Pk+1 ),

µk+1 = Law(Θk+1 , Pk+1 ). (111) P If µk = N −1 i δ(θi,k ,Pi,k ) is empirical, then (111) is exactly the finite-particle regularized Muon scheme (61). The scaling needed to approximate the damped Hamiltonian probability flow with finite damping parameter γ > 0 is 1 (112) ηh = h, 1 − βh = γh, 0<h< . γ Under this scaling, (111) is the semi-implicit Euler discretization of the characteristic equations θ̇ = −Gε (p),

ṗ = γ(aµt (θ) − p).

More generally, the same fixed-step convergence mechanism is stable under the second-orderconsistent inertial scaling ηh = h + O(h2 ),

1 − βh = γh + O(h2 ).

(113)

The exact relation 1 − βh = γηh is used first because it exposes the Hamiltonian cancellation. The perturbative case (113) is discussed in Remark 14. Keeping βh fixed independently of h is not a finite-damping Hamiltonian discretization: in that case Pk+1 − Pk 1 − βh = (aµk (Θk ) − Pk ) , h h so the relaxation rate diverges as h ↓ 0. For the iterates generated by (110), define Z Kk := ΨΘ ε (p) dµk (θ, p),

Z Dk :=

Jk := J(ρk ), Hk := Kk + γUk , Z := Ck ⟨ak (θ), p⟩Θ dµk (θ, p), 41

dε (p) dµk (θ, p),

Uk := Jk − J⋆ , Z Ak := ∥ak (θ)∥2Θ dµk (θ, p), Lk := Hk − αCk .

(114)

Here J⋆ = inf ρ∈P1 (Θ) J(ρ). Since µk+1 is the pushforward of µk by Tkh , the post-momentum dissipation satisfies Z
Dk+1 = dε p+ (115) k (θ, p) dµk (θ, p). Assumption (A7) (Uniform assumptions for the discrete trajectory). The law sequence (µk )k≥0 generated by (110) satisfies the following properties. 1. The objective gap is nonnegative: Uk ≥ 0 for every k. 2. The PL and upper-gradient conditions hold at every discrete iterate: Ak ≥ 2λUk ,

Ak ≤ 2ΛUk ,

(116)

where λ > 0 and Λ < ∞. 3. The momenta p under µk and the post-momenta p+ k (θ, p) under µk remain in a common bounded momentum ball on which Lemma 2 holds with constants κK , κD , LG , χ. 4. The one-step segment θk,s (θ, p) := θ − sηh Gε (p+ k (θ, p)),

0 ≤ s ≤ 1,

remains in a region on which Assumption (A5) holds with constants MD , MD,2 , MR , MR,2 . Set 2 Bcurv := MR MD,2 + MD MR,2 ,

A.17.2

σ :=

LG Bcurv . κD

(117)

One-step estimates

Lemma 4 (Second-order upper estimate for J along a transport step). Let µ ∈ P(Θ × Θ), let ρ = (πθ )# µ, and let aµ (θ) = DF (θ)∗ ∇R(mµ ). Let v ∈ L2 (µ; Θ) and define Z θs = θ + sv(θ, p), ρs = (θs )# µ, ms = F (θs ) dµ(θ, p). Assume the segment {θs : 0 ≤ s ≤ 1} lies in a region where the constants in Assumption (A5) are valid. Then Z Z Bcurv J(ρ1 ) − J(ρ0 ) ≤ ⟨aµ (θ), v(θ, p)⟩Θ dµ(θ, p) + ∥v(θ, p)∥2Θ dµ(θ, p). (118) 2 Proof. For 0 ≤ s ≤ 1, define us := ∇R(ms ) and as (θs ) := DF (θs )∗ us . By the chain rule, d J(ρs ) = ds

d ∇R(ms ), ms ds

Z =

∗

⟨DF (θs ) us , v⟩Θ dµ =

Z ⟨as (θs ), v⟩Θ dµ.

H

We next compare as (θs ) with a0 (θ). The exact decomposition is as (θs ) − a0 (θ) = (DF (θs )∗ − DF (θ)∗ ) u0 + DF (θs )∗ (us − u0 ). 42

(119)

Since ∥u0 ∥H ≤ MR and D2 F ≤ MD,2 on the segment, ∥(DF (θs )∗ − DF (θ)∗ ) u0 ∥Θ ≤ MR MD,2 s ∥v∥Θ .

(120)

Also, Z ∥ms − m0 ∥H =

(F (θs ) − F (θ)) dµ

1/2

Z

Z ≤

∥v∥2Θ dµ

MD s ∥v∥Θ dµ ≤ MD s

.

H

Hence

1/2

Z

∥v∥2Θ dµ

.

Z

1/2 Z

∥us − u0 ∥H ≤ MR,2 ∥ms − m0 ∥H ≤ MD MR,2 s

(121)

Using ∥DF (θs )∥ ≤ MD and Cauchy’s inequality, Z

Z

2 MR,2 s ∥v∥2Θ dµ + MD

⟨as (θs ) − a0 (θ), v⟩Θ dµ ≤ MR MD,2 s Z ≤ sBcurv ∥v∥2Θ dµ. Integrating (119) from s = 0 to s = 1 and using

∥v∥2Θ dµ

∥v∥Θ dµ

R1

0 s ds = 1/2 yields (118).

Lemma 5 (Discrete Hamiltonian increment). Assume the exact scaling (112), so that δh = γηh and βh = 1 − δh . Under Assumption (A7),   γ 2 ηh2 ηh2 L2G γ 2 Hk+1 − Hk ≤ −γηh Dk+1 + Ak + + γBcurv Dk+1 . (122) 2βh 2κD βh Proof. All integrals below are with respect to µk . For readability write p+ = p+ k (θ, p),

g + = Gε (p+ ),

a = ak (θ),

η = ηh ,

δ = δh .

Since ΨΘ ε is convex and differentiable, + Θ + + + ΨΘ ε (p ) − Ψε (p) ≤ Gε (p ), p − p Θ = δ g , a − p Θ .

Therefore

Z Kk+1 − Kk ≤ δ

g + , a − p Θ dµk .

The position displacement is v = −ηg + . Applying Lemma 4 gives Z Z Bcurv η 2 2 + g + Θ dµk . Uk+1 − Uk ≤ −η a, g Θ dµk + 2

(123)

(124)

(125)

Multiplying (125) by γ and adding (124), and then using δ = γη, gives the cancellation of the mixed force-transport term: Z Z γBcurv η 2 2 + Hk+1 − Hk ≤ −δ p, g Θ dµk + g + Θ dµk . (126) 2 Now p+ = p + δ(a − p), hence p = p+ − δ(a − p),

a − p+ = βh (a − p). 43

Consequently, Z −δ

p, g

Z

+

dµk = −δ

Θ

+

p ,g

+ Θ

δ2 = −δDk+1 + βh

dµk + δ

Z

2

Z

a − p, g + Θ dµk

a − p+ , g + Θ dµk .

By Young’s inequality and the nonnegativity of ⟨p+ , g + ⟩Θ , Z Z Z 1 1 + + + a − p , g Θ dµk ≤ a, g Θ dµk ≤ Ak + 2 2

(127)

2

g + Θ dµk .

(128)

By (86) and (88), Z

2 g + Θ dµk ≤ L2G

Z

2

p+ Θ dµk ≤

L2G Dk+1 . κD

(129)

Combining (126)-(129) and using δ = γη proves (122). Lemma 6 (Discrete alignment increment). Under Assumption (A7), Ck+1 − Ck = δh (Ak − Ck ) + Rk ,

(130)

where the remainder obeys |Rk | ≤ ηh σDk+1 ,

σ=

LG Bcurv . κD

(131)

+ + + + Proof. Again write p+ = p+ k (θ, p), g = Gε (p ), a = ak (θ), and θ = θ − ηh g . Since µk+1 = h (Tk )# µk , Z

ak+1 (θ+ ), p+ Θ dµk .

Ck+1 = Therefore Z Ck+1 − Ck =

Z

+

a, p − p Θ dµk +

ak+1 (θ+ ) − a, p+ Θ dµk

= δh (Ak − Ck ) + Rk , where

Z Rk :=

(132)

ak+1 (θ+ ) − ak (θ), p+ Θ dµk .

It remains to bound Rk . Put uk = ∇R(mk ) and uk+1 = ∇R(mk+1 ). Then
ak+1 (θ+ ) − ak (θ) = DF (θ+ )∗ − DF (θ)∗ uk + DF (θ+ )∗ (uk+1 − uk ). The first term satisfies
DF (θ+ )∗ − DF (θ)∗ uk Θ ≤ MR MD,2 ηh g + Θ .

(133)

For the second term, Z ∥mk+1 − mk ∥H =

+

Z

F (θ ) − F (θ) dµk

≤ MD ηh H

so + ∗

DF (θ ) (uk+1 − uk )

2 ≤ MD MR,2 ηh Θ

44

Z

g + Θ dµk ,

g + Θ dµk .

(134)

(135)

By Cauchy’s inequality, (133), and (135), Z Z + + 2 |Rk | ≤ MR MD,2 ηh g Θ p Θ dµk + MD MR,2 ηh

g

 Z

+ Θ

dµk



+

p

Θ

dµk .

Using (86) and (88), Z g and

+

p

Θ

Z g

Z

+ Θ

dµk ≤ LG  Z

+ Θ

dµk

2

p+ Θ dµk ≤ 

+

p

Θ

dµk

≤

LG Dk+1 , κD

LG Dk+1 . κD

Combining the last three displays gives (131). A.17.3

Discrete modified-Lyapunov contraction

Theorem 15 (Fixed-step discrete convergence of the regularized Hamiltonian-Muon map). Assume the exact scaling (112) and Assumption (A7). Define s Λ MC := . (136) γκK Choose r ∈ (0, 2) and α > 0 such that αMC < 1. For a fixed step size h set ηh = h, βh = 1 − γh, and define   L2G γ 2 αγ −h + γBcurv , dh := γ − ασ − 2rβh κD 2κD βh (137)   r γ2h := ah αγ 1 − − . 2βh 2βh Assume dh > 0, ah > 0, and define   d β2κ h h D , qh := min  LG

 

ah , LG γ 2 h2  γ 2 2λ + βh

ch :=

qh . 1 + αMC

(138)

If hch ≤ 1, then for every k ≥ 0, Lk+1 ≤ (1 − hch )Lk .

(139)

Consequently, J(ρk ) − J⋆ ≤

(1 − hch )k exp(−ch kh) [H0 − αC0 ] ≤ [H0 − αC0 ] . γ(1 − αMC ) γ(1 − αMC )

(140)

Moreover, the above positivity conditions are guaranteed by an explicit small-step bound. Suppose in addition that  αγ r := := d0 γ − ασ − > 0, a0 αγ 1 − > 0. (141) 2rκD 2 Set Bd :=

L2 αγ 2 + G 2γ 2 + γBcurv , rκD 2κD

45

Ba := γ 2 (1 + αr),

(142)

with the convention that d0 /(2Bd ) = +∞ if Bd = 0 and a0 /(2Ba ) = +∞ if Ba = 0. Define ( ) a0 d0 κD q⋆
, , , (143) q⋆ := min c⋆ := γ 2 8LG 2 2λ + 4LG γ 1 + αMC and



1 d0 a0 1 h⋆ := min 1, , , , 2γ 2Bd 2Ba c⋆

 .

(144)

Then every h ∈ (0, h⋆ ] satisfies the hypotheses 0 < h < 1/γ, dh > 0, ah > 0, and hc⋆ ≤ 1, and the estimate (1 − hc⋆ )k exp(−c⋆ kh) J(ρk ) − J⋆ ≤ [H0 − αC0 ] ≤ [H0 − αC0 ] (145) γ(1 − αMC ) γ(1 − αMC ) holds for all k ≥ 0. Proof. The proof has five steps. Step (i): equivalence between Lk and Hk . By Cauchy’s inequality, the upper-gradient condition in (116), and the kinetic lower bound (85), Z |Ck | ≤

∥ak (θ)∥2Θ dµk

1/2 Z

∥p∥2Θ dµk

1/2

≤ (2ΛUk )1/2



2Kk κK

1/2 ≤ MC Hk .

(146)

√ √ √ The final inequality follows from 2 Uk Kk ≤ (Kk + γUk )/ γ = Hk / γ. Thus (1 − αMC )Hk ≤ Lk ≤ (1 + αMC )Hk .

(147)

Step (ii): one-step decay of Lk in terms of Dk+1 and Ak . By Lemmas 5 and 6, and by δh = γh, Lk+1 − Lk = Hk+1 − Hk − α(Ck+1 − Ck ) ≤ Hk+1 − Hk − αγhAk + αγhCk + αhσDk+1 . Since p+ = βh p + γhak (θ), p=

(148)

p+ − γhak (θ) . βh

Therefore Z 1 γh ak (θ), p+ Θ dµk − Ak βh βh Z Z 1 r 1 ≤ ak (θ), p+ Θ dµk ≤ Ak + βh 2βh 2rβh r 1 ≤ Ak + Dk+1 . 2βh 2rβh κD

Ck =

2

p+ Θ dµk (149)

Substituting (149) and (122) into (148) gives Lk+1 − Lk ≤ −hdh Dk+1 − hah Ak , with dh and ah defined in (137). 46

(150)

Step (iii): Dk+1 and Ak dominate Hk . Because Gε (0) = 0 and Gε is LG -Lipschitz, Z 1 Z 1 LG Θ LG s ∥p∥2Θ ds = ∥p∥2Θ . Ψε (p) = ⟨Gε (sp), p⟩Θ ds ≤ 2 0 0

(151)

Moreover, ∥p∥2Θ =

1 2 2γ 2 h2 2 + + 2 ≤ + p − γha (θ) p ∥ak (θ)∥2Θ . k Θ Θ βh2 βh2 βh2

Using (86) and the PL inequality, Hk = Kk + γUk LG ≤ 2 Dk+1 + βh κD



LG γ 2 h2 γ + 2 2λ βh

 Ak .

(152)

By the definition of qh in (138), (152) implies dh Dk+1 + ah Ak ≥ qh Hk .

(153)

Lk+1 − Lk ≤ −hqh Hk .

(154)

Combining (150) with (153) yields

Step (iv): contraction and objective decay for a fixed admissible h. Since Lk ≤ (1 + αMC )Hk , we have Hk ≥ Lk /(1 + αMC ). Hence (154) gives   qh Lk+1 ≤ 1 − h Lk = (1 − hch )Lk . 1 + αMC If hch ≤ 1, iteration proves (139). Finally, γUk ≤ Hk ≤ Lk /(1 − αMC ) by (147), so Uk ≤

(1 − hch )k L0 Lk ≤ . γ(1 − αMC ) γ(1 − αMC )

Since L0 = H0 − αC0 and (1 − x)k ≤ exp(−kx) for x ∈ [0, 1], (140) follows. Step (v): explicit sufficient upper bound on h. Assume (141). If h ≤ 1/(2γ), then βh = 1 − γh ≥ 1/2 and γh 1 −1= ≤ 2γh. βh 1 − γh Using (137),     L2G γ 2 1 αγ dh = d0 − −1 −h + γBcurv 2rκD βh 2κD βh  2 
L2G αγ 2 ≥ d0 − h + 2γ + γBcurv = d0 − hBd . rκD 2κD Hence h ≤ d0 /(2Bd ) implies dh ≥ d0 /2. Similarly,   αγr 1 γ2h ah = a0 − −1 − ≥ a0 − αrγ 2 h − γ 2 h = a0 − hBa . 2 βh 2βh 47

Thus h ≤ a0 /(2Ba ) implies ah ≥ a0 /2. If also h ≤ 1, then LG γ 2 h2 ≤ 4LG γ 2 , βh2

1 βh2 ≥ . 4

Consequently, for every h ≤ h⋆ ,  qh ≥ min

a0 /2 (d0 /2)(1/4)κD , γ 2 LG 2λ + 4LG γ

Using (154) with q⋆ in place of qh gives  Lk+1 ≤ 1 − h

q⋆ 1 + αMC

 = q⋆ .

 Lk = (1 − hc⋆ )Lk .

The final entry h ≤ 1/c⋆ in (144) makes this contraction factor nonnegative. Iterating and using γUk ≤ Lk /(1 − αMC ) proves (145). Remark 14 (Second-order momentum scaling). The exact identity 1 − βh = γηh is not essential. Suppose instead δh := 1 − βh = γηh + ζh ,

|ζh | ≤ Cβ ηh2 ,

ηh = h + O(h2 ).

In the proof of Lemma 5, the cancellation in (126) leaves the additional term Z (δh − γηh ) ak (θ), Gε (p+ k (θ, p)) Θ dµk . By Young’s inequality and (129), its absolute value is bounded by   L2G 2 1 Cβ η h Ak + Dk+1 . 2 2κD All other appearances of δh are also γηh + O(ηh2 ). Consequently the coefficients dh and ah in (137) are changed only by O(h) terms. Thus Theorem 15 remains true, with slightly smaller positive constants and a correspondingly smaller explicit threshold, under the standard inertial scaling (113). This is the valid scaling for a finite-damping Hamiltonian limit; keeping βh fixed as h ↓ 0 instead gives a singular overdamped relaxation of the momentum variable. P Corollary 2 (Finite-particle and scalar matrix-space forms). If µk = N −1 N i=1 δ(θi,k ,Pi,k ) , then the law update (110) is equivalent to Pi,k+1 = βh Pi,k + (1 − βh )aN i (θ k ),

θi,k+1 = θi,k − ηh OrthΘ ε (Pi,k+1 ),

and Theorem 15 gives the same exponential bound for the finite-particle objective JN (θ k ) − J⋆ . In the original scalar matrix-space setting, this reads Pi,k+1 = βh Pi,k + (1 − βh )R′ (F̄k )∇F (Wi,k ),

Wi,k+1 = Wi,k − ηh Orthε (Pi,k+1 ),

with ηh = h and 1 − βh = γh + O(h2 ). Therefore the discrete-time convergence theorem applies directly to the natural valid discretization of the Hamiltonian probability flow (12)-(14) and to its extended-space analogue (72)-(77). 48

A.18

Propagation of chaos

The finite-N ODE (64) is the interacting-particle approximation of the nonlinear characteristic equation (73). Theorem 16 (Propagation of chaos on the Hilbert space domain). Assume Assumption (A4) and µ0 ∈ P2 (Z). Let (θiN (t), PiN (t))N i=1 solve the N -particle ODE (64) with i.i.d. initial data distributed according to µ0 . Let (θ̄i (t), P̄i (t))i≥1 be i.i.d. nonlinear mean-field copies solving (73) with the same initial data, i.e. (θiN (0), PiN (0)) = (θ̄i (0), P̄i (0)). Then for every T < ∞ there exists Cpoc (T, ε) < ∞ such that, for each fixed i, sup E

h

0≤t≤T

i C (T, ε) 2 2 poc θiN (t) − θ̄i (t) Θ + PiN (t) − P̄i (t) Θ ≤ . N

(155)

Consequently, for every fixed k ∈ N,  kC (T, ε)  poc ≤ . sup W22 Law((θ1N (t), P1N (t)), . . . , (θkN (t), PkN (t))), µ⊗k t N 0≤t≤T

(156)

Furthermore, if N

µN t =

N

1 X δ(θN (t),P N (t)) , i i N

µ̄N t =

i=1

then sup EW22 (µN t , µt ) ≤

0≤t≤T

1 X δ(θ̄i (t),P̄i (t)) , N i=1

2Cpoc (T, ε) + 2 sup EW22 (µ̄N t , µt ), N 0≤t≤T

and the second term tends to zero as N → ∞. Proof. Let eθi (t) = θiN (t) − θ̄i (t), and set

 u(t) = E

ePi (t) = PiN (t) − P̄i (t),

eθi (t)

2 Θ

+

2 ePi (t) Θ

 ,

which is independent of i by exchangeability. Since OrthΘ ε is 1/ε-Lipschitz, D E d θ 2 N Θ ei = −2 eθi , OrthΘ (P ) − Orth ( P̄ ) i ε i ε dt Θ Θ ≤ eθi

2 Θ

+ ε−2 ePi

2 . Θ

For the momentum error, d P e dt i

2 = 2γ Θ

N P ePi , aN i (θ ) − aµt (θ̄i ) − ei Θ 2

N P ≤ γ aN i (θ ) − aµt (θ̄i ) Θ − γ ei 2

N ≤ γ aN i (θ ) − aµt (θ̄i ) Θ .

49

2 Θ

(157)

By the force decomposition used in Lemma 1, there exists a constant Ca depending only on the global constants such that 2

N θ aN i (θ ) − aµt (θ̄i ) Θ ≤ Ca ei

where

2

2

Θ

+ C a mN N (t) − mt H ,

(158)

N

mN N (t) =

1 X F (θjN (t)), N

mt = EF (θ̄i (t)).

j=1

Decompose   N N X X
1 1 mN F (θjN (t)) − F (θ̄j (t)) +  F (θ̄j (t)) − mt  . N (t) − mt = N N j=1

j=1

By Jensen’s inequality and the Lipschitz property of F , 2

N

N

1 X (F (θjN ) − F (θ̄j )) E N j=1

≤

2 1 X 2 ≤ L2F u(t). LF E eθj N Θ j=1

H

The random variables F (θ̄j (t)) are i.i.d.; hence 2

N

1 X E F (θ̄j (t)) − mt N j=1

=

1 2 E F (θ̄i (t)) − mt H . N

H

On [0, T ] this variance is finite because F is Lipschitz and µ0 ∈ P2 (Z). Therefore 2

E mN N (t) − mt H ≤ CT u(t) +

CT . N

Combining the previous differential inequalities gives CT , u(0) = 0. u′ (t) ≤ CT u(t) + N Gronwall’s inequality proves (155). The k-particle estimate (156) follows by coupling each interacting particle with its mean-field copy and summing squared errors. Finally, 2 N N 2 N W22 (µN t , µt ) ≤ 2W2 (µt , µ̄t ) + 2W2 (µ̄t , µt )  N  2 2 X 2 ≤ eθi (t) + ePi (t) Θ + 2W22 (µ̄N t , µt ). N Θ i=1

Taking expectations and using (155) proves (157). The convergence of the i.i.d. empirical term, for every fixed t, follows from the law of large numbers in W2 on finite-dimensional spaces with finite second moment. To obtain the supremum over t ∈ [0, T ], use the L2 -continuity of the nonlinear characteristics. The drift has at most linear growth and the second moments remain bounded on [0, T ], hence E ∥Zt − Zs ∥2 ≤ CT |t − s|2 . The same estimate holds for the empirical nonlinear system µ̄N t . For a grid 0 = t0 < · · · < tM = T , this gives

N sup W2 µ̄N t , µt ≤ max W2 µ̄tℓ , µtℓ + CT ∆t t∈[0,T ]

ℓ

in expectation, up to the standard empirical and population continuity terms. First, we let N → ∞ for fixed grid, then we let ∆t ↓ 0. 50

A.19

Hard Muon limit on the extended gated space

The limiting product potential is ΨΘ 0 (p) :=

B X

p(b)

nuc

b=1

.

(159)

For a block p(b) = U ΣV ⊤ of rank r, the nuclear-norm subdifferential is n o = U V ⊤ + Z : U ⊤ Z = 0, ZV = 0, ∥Z∥op ≤ 1 . ∂ p(b) nuc

Thus ∂ΨΘ 0 (p) =

B Y

∂ p(b)

b=1

nuc

.

(160)

(161)

The canonical hard Muon map is   OrthΘ (p) = Orth(p(1) ), . . . , Orth(p(B) ) ,

(162)

which is a selected element of ∂ΨΘ 0 (p). Theorem 17 (Subsequential hard-Muon limit on the extended gated space). Let µ0 ∈ P1 (Z) and assume that the force field satisfies the global bounds of Assumption (A4). Let µεt solve (72) with Θ OrthΘ ε = ∇Ψε . For every T < ∞, every sequence εk ↓ 0 has a subsequence, not relabeled, such that µεt k → µt

in C([0, T ]; P1 (Z)).

There exists a Borel vector field Vt (θ, p) with Vt (θ, p) ∈ ∂ΨΘ 0 (p)

µt dt-a.e.

(163)

such that µt solves ∂t µt + ∇θ · (−Vt µt ) + ∇p · (γ(aµt (θ) − p)µt ) = 0.

(164)

If every block p(b) has full rank for µt dt-almost every (θ, p), then Vt (θ, p) = OrthΘ (p) almost everywhere and the limit is the canonical hard Muon flow. √ qΘ for all ε > 0, and the force is bounded Proof. The velocity OrthΘ ε (p) is uniformly bounded by by Lemma 1. Along characteristics, ∥Θεt − Θεs ∥Θ ≤ and

√

qΘ |t − s|,

d ∥P ε ∥ ≤ γ(MD MR + ∥Ptε ∥Θ ). dt t Θ Gronwall’s inequality gives a first-moment bound on Ptε , uniformly for t ∈ [0, T ] and ε. These estimates imply tightness and equicontinuity of {µε }ε>0 in C([0, T ]; P1 (Z)), hence subsequential compactness. √ ε Let Gε (p) = OrthΘ qΘ implies, after extraction, weakε (p). The uniform bound ∥G (p)∥Θ ≤ star convergence of the fluxes Gεk (p)µεt k dt to a vector-valued measure absolutely continuous with √ respect to µt dt. Its density is denoted by Vt (θ, p) and satisfies ∥Vt ∥ ≤ qΘ . 51

It remains to identify Vt (θ, p) as a subgradient of ΨΘ 0 . For every q ∈ Θ and every p ∈ Θ, Θ convexity of Ψε gives Θ Θ ΨΘ (165) ε (q) ≥ Ψε (p) + Orthε (p), q − p Θ . Θ The functions ΨΘ ε converge locally uniformly to Ψ0 as ε ↓ 0. Passing to the limit in the integrated form of (165), with p replaced by the momentum coordinate and with arbitrary bounded nonnegative test weights, gives Θ ΨΘ µt dt-a.e. 0 (q) ≥ Ψ0 (p) + ⟨Vt (θ, p), q − p⟩Θ

This inequality for every q ∈ Θ is exactly Vt (θ, p) ∈ ∂ΨΘ 0 (p). Passing to the limit in the weak formulation of (72) gives (164); the p-velocity term passes by the force-field continuity established in Lemma 1 and the convergence of µεk to µ. The final statement follows from (160): if p(b) has full rank, the orthogonality constraints U ⊤ Z = 0 and ZV = 0 force Z = 0, so the block subdifferential is the singleton {Orth(p(b) )}.

A.20

Transformer MoE specialization with input-dependent routing

The abstract theory specializes to transformer mixture-of-experts models by choosing Θ to contain both an expert tuple and router parameters. The input dependence of the gate is encoded in the Hilbert-valued feature map F . Softmax normalization across experts is represented by augmenting the Hilbert output with numerator and denominator features. A.20.1

Expert-router parameter space

Let an input sequence be X ∈ RL×d .

(166)

ω = (Q, K, V, O, W1 , W2 ),

(167)

A single expert is parameterized by

where Q, K ∈ Rd×dk ,

V ∈ Rd×dv ,

O ∈ Rdv ×d ,

W1 ∈ Rd×df ,

W2 ∈ Rdf ×d .

(168)

Thus Θexp = Rd×dk × Rd×dk × Rd×dv × Rdv ×d × Rd×df × Rdf ×d .

(169)

Let the router parameter space be a finite product of matrix spaces Bgate

Θgate =

Y

Rar ×br .

(170)

r=1

A linear per-token router, for example, can be represented by a matrix G ∈ Rd×1 and score sG (X)t = ⟨Xt,: , G⟩Rd ; more general smooth router MLPs correspond to several matrix blocks in (170). The full particle parameter is θ = (ω, ϕ) ∈ Θ := Θexp × Θgate .

(171)

The Hilbert norm is the product Frobenius norm Bgate

∥θ∥2Θ = ∥Q∥2F + ∥K∥2F + ∥V ∥2F + ∥O∥2F + ∥W1 ∥2F + ∥W2 ∥2F +

X r=1

52

ϕ(r)

2 F

.

(172)

The regularized Muon mirror map is OrthΘ ε (PQ , PK , PV , PO , PW1 , PW2 , Pϕ(1) , . . .)
= Orthε (PQ ), Orthε (PK ), Orthε (PV ), Orthε (PO ), Orthε (PW1 ), Orthε (PW2 ), Orthε (Pϕ(1) ), . . . . (173) with each Orthε evaluated at the corresponding rectangular block. Proposition 5 applies directly to this space. A.20.2

Smooth single-head attention-plus-FFN expert and router

Let σ : R → R be a smooth activation applied rowwise. Define the single-head attention map   (XQ)(XK)⊤ √ Aω (X) := softmax XV O ∈ RL×d , (174) dk where softmax is applied rowwise. Define the feed-forward map Bω (X) := σ(Aω (X)W1 )W2 ∈ RL×d .

(175)

Let Rout : RL×d → RL×C be a fixed output projection to token logits. The expert logit function is ψω (X) := Rout Bω (X) ∈ RL×C .

(176)

sϕ (X) ∈ RL

(177)

Let be a smooth tokenwise router-score map. Sequence-level routing is obtained as the special case in which all coordinates of sϕ (X) are identical or only one scalar score is stored. Lemma 7 (Local smoothness of transformer expert-router features). Assume the input set is bounded, ∥X∥F ≤ BX , and the activation σ and router score map ϕ 7→ sϕ (X) have bounded derivatives up to order two on bounded parameter sets. Then, for every R < ∞, the maps ω 7→ ψω (X),

ϕ 7→ sϕ (X)

are C 2 on {∥θ∥Θ ≤ R} with first and second derivative bounds uniform over ∥X∥F ≤ BX . The same statement holds as a map into L2 (PX ) when X is almost surely bounded. Proof. On {∥θ∥Θ ≤ R} and ∥X∥F ≤ BX , every matrix product in (174) and (175) is bounded by a constant depending only on R, BX , L, d, dk , dv , df . Matrix multiplication is polynomial in the entries of ω and is therefore smooth with bounded derivatives on bounded sets. The rowwise softmax is C ∞ and all derivatives are bounded on bounded subsets of its input space. The activation σ has bounded derivatives on the bounded interval reached by Aω (X)W1 . The composition and product rules for Frechet derivatives imply that ω 7→ ψω (X) is C 2 with uniform derivative bounds. The assumed smoothness of the router score gives the corresponding conclusion for ϕ 7→ sϕ (X). The L2 (PX ) statement follows from the pointwise derivative bounds and dominated convergence.

53

A.20.3

Token-level cross-entropy loss

For a finite training set {(Xr , Yr )}nr=1 with Yr,t ∈ {1, . . . , C}, set
n Hlogit := RL×C

(178)

with averaged inner product n

⟨f, g⟩Hlogit :=

L

1 XX ⟨fr,t , gr,t ⟩RC . nL

(179)

r=1 t=1

The token-level cross-entropy risk is " !# n L C X 1 XX RCE (f ) := −fr,t,Yr,t + log . efr,t,c nL r=1 t=1

(180)

c=1

This is the standard negative log-likelihood/cross-entropy objective used with transformer token logits; label smoothing replaces the one-hot target below by a fixed target distribution. The transformer architecture and training objective with cross-entropy and label smoothing were introduced in Vaswani et al. (2017); sparse transformer mixture-of-experts layers and switch routing are developed in Shazeer et al. (2017) and Fedus et al. (2022). Theoretical motivation for router-driven specialization and the role of smoothing/noisy routing appears in Chen et al. (2022). Proposition 8 (Cross-entropy satisfies the Hilbert loss assumptions). The map RCE : Hlogit → R is C ∞ , lower bounded by 0, and its Hilbert gradient with respect to (179) is (∇RCE (f ))r,t = softmax(fr,t ) − eYr,t , where ey is the yth coordinate vector. Moreover, √ ∥∇RCE (f )∥Hlogit ≤ 2, ∥∇RCE (f ) − ∇RCE (g)∥Hlogit ≤ ∥f − g∥Hlogit .

(181)

(182)

The same conclusions hold with label smoothing, replacing eYr,t by any target vector in the probability simplex. Proof. For one token define ℓy (z) = −zy + log

C X

ezc .

c=1

Then ∇ℓy (z) = softmax(z) − ey . The Hessian is ∇2 ℓy (z) = diag(s) − ss⊤ ,

s = softmax(z),

which is the covariance matrix of a categorical random variable with probability vector s. Hence it is positive semidefinite √ and its operator norm is at most 1. Thus ∇ℓy is 1-Lipschitz. Also, ∥softmax(z) − ey ∥2 ≤ 2 because both vectors lie in the probability simplex. Averaging P over (r, t) with the inner product (179) proves (181) and (182). Nonnegativity follows from log c ezc ≥ zy . Smoothness follows from smoothness of log-sum-exp. The label-smoothed case is identical because the target vector remains in the probability simplex. 54

A.20.4

Hilbert feature maps for input-dependent gates

Two smooth gate encodings are directly covered by (28). Unnormalized smooth gates. Let gϕ (X) ∈ RL be a nonnegative smooth input-dependent gate. Define the empirical logit-space feature map Fun (ω, ϕ) := (gϕ (Xr ) ⊙ ψω (Xr ))nr=1 ∈ Hlogit ,

(183)

where ⊙ denotes tokenwise multiplication, broadcasting the scalar gate at token t across the C logits. Then Z mρ =

Fun (ω, ϕ) dρ(ω, ϕ)

(184)

Θ

is an input-dependent gated population output. Softmax-normalized gates over the expert distribution. A finite softmax MoE with particles (ωi , ϕi )N i=1 has tokenwise output PN MN (X)t,: =

sϕi (X)t ψωi (X)t,: i=1 e . PN sϕ (X)t j e j=1

(185)

For a probability distribution ρ over Θ, the corresponding population output is R s (X) e ϕ t ψω (X)t,: dρ(ω, ϕ) . Mρ (X)t,: = Θ R s (X) t dρ(ω, ϕ) ϕ Θe

(186)

Although (186) is not a single unnormalized average of logits, it is exactly of the form (28) after Hilbert-output augmentation. Define
n Haug := Hlogit ⊕ RL (187) with product Hilbert inner product, and define  n Fsoft (ω, ϕ) := esϕ (Xr ) ⊙ ψω (Xr )

r=1

Then msoft = ρ

 n  , esϕ (Xr ) ∈ Haug . r=1

(188)

Z Fsoft (ω, ϕ) dρ(ω, ϕ) = (Nρ , Dρ ).

(189)

Θ

On the open set Dr,t > 0, define Γ(N, D)r,t,c :=

Nr,t,c . Dr,t

(190)

The softmax-gated risk is Rsoftgate (N, D) := RCE (Γ(N, D)). Therefore

Z J(ρ) = Rsoftgate

(191)

 Fsoft (θ) dρ(θ)

(192)

Θ

P is exactly of the Hilbert-valued form (28). For the empirical measure ρN = N −1 i δ(ωi ,ϕi ) , the factors 1/N cancel between numerator and denominator, and (192) gives exactly (185). 55

Lemma 8 (Local smoothness of the softmax-gate feature and loss). Assume ∥Xr ∥F ≤ BX for all training inputs, and assume the expert and router maps satisfy Lemma 7. On each bounded parameter ball ∥θ∥Θ ≤ R, the maps Fun and Fsoft are C 2 with bounded first and second derivatives. In the normalized case, if Dr,t ≥ δ > 0 on the moment set reached by the trajectory, then Rsoftgate is C 2 with bounded gradient and bounded Hessian on that set. Proof. The maps ψω (Xr ) and sϕ (Xr ) are C 2 with bounded derivatives on bounded parameter balls by Lemma 7. The exponential map has bounded derivatives on bounded score intervals. Products of esϕ with ψω therefore have bounded derivatives up to order two. This proves the statement for Fun and Fsoft . For the normalized risk, Γ(N, D) = N/D is C ∞ on Dr,t > 0. On a set where Dr,t ≥ δ and N, D are bounded, its first and second derivatives are bounded by constants depending on δ and the bounds on N, D. Proposition 8 gives bounded first and second derivatives of RCE on logit space. The chain rule gives the stated bounds for Rsoftgate = RCE ◦ Γ. A.20.5

Hard top-k routing and noisy routing

Exact hard top-k or switch routing uses the discontinuous map that selects the largest router scores. This map is not covered by the smooth ODE, Hamiltonian, curvature, and propagation-of-chaos theorems above. There are two mathematically consistent ways to connect hard routing to the present theory. First, hard routing can be replaced by a smooth relaxation, such as softmax with temperature, Gumbel-softmax, or a differentiable expected router. In the Gumbel-max case, the expected top-1 selection probabilities are precisely softmax probabilities, so the augmented feature (188) applies directly. More generally, if the smoothed expected router weight is a C 2 function of (X, θ) on bounded sets, then it can be included in F exactly as in (183) or (188). Second, exact hard routing can be treated as a nonsmooth limit. The limiting evolution is then a differential inclusion, analogous to the hard-Muon limit in Section A.19. At score ties, the router subdifferential or selection correspondence is set-valued. A smooth-flow convergence proof must therefore be replaced by compactness plus graph-convergence arguments for the router selection map. The MoE analysis of Chen et al. (2022) emphasizes that sparse top-1 routing is discontinuous and that injected random noise smooths the routing probabilities. Their smoothing lemma gives Lipschitz dependence of the route probabilities on the router scores when the noise density is bounded. Such a smoothed expected router is compatible with the present Hilbert-valued framework whenever the resulting expected gate is used as the gate component of F . A.20.6

Specialization theorem

Theorem 18 (Transformer MoE specialization with router parameters included). Consider the extended parameter space (171) with product Frobenius geometry (172) and blockwise regularized Muon map (173). Consider the single-head attention-plus-FFN expert (174)-(176), a smooth router score map (177), and the token-level cross-entropy loss (180). Let F be either the unnormalized gate feature (183) or the augmented softmax-normalized feature (188). Assume: (T1) the training inputs are bounded, ∥Xr ∥F ≤ BX ; (T2) the activation and router score maps are C 2 in the parameters on bounded parameter sets; (T3) in the normalized-gate case, the denominator satisfies Dρ,r,t ≥ δ > 0 on the trajectory; 56

(T4) the trajectory remains in a bounded parameter and momentum region for the time interval or asymptotic regime under consideration; (T5) for exponential convergence, the PL and upper-gradient assumptions (97)-(98) hold along the trajectory. Then all constructions and results in Sections A.10-A.19 apply to this transformer MoE. In particular: 1. the first variation is δJ (ρ)(θ) = ⟨∇R(mρ ), F (θ)⟩H , δρ with H = Hlogit in the unnormalized case and H = Haug in the normalized case; 2. the Wasserstein force on the expert-router tuple is aρ (θ) = DF (θ)∗ ∇R(mρ ) ∈ Θ; 3. the finite-N regularized Muon scheme updates both expert and router blocks by Pi,k+1 = βPi,k + (1 − β)aN i (θ k ),

θi,k+1 = θi,k − η OrthΘ ε (Pi,k+1 );

4. under the inertial scaling (63), the finite-particle ODE limit, mean-field PDE, Hamiltonian formulation, dissipation identity, and hard-Muon subsequential limit hold as stated above on the bounded region specified by (T1)-(T4), while, propagation of chaos, as stated above, holds under the global Lipschitz Assumption (A4), or under an explicitly stated localized compact-support version with uniform support bounds; 5. under (T5) and the curvature assumptions, the convergence estimate (103) holds for the transformer MoE cross-entropy objective. Proof. The parameter space (171) is a finite product of matrix spaces. Proposition 5 therefore gives the product Fenchel duality and mirror-map interpretation on the full expert-router space. Proposition 8 verifies the bounded-gradient and Lipschitz-gradient properties of token-level crossentropy on logit space. Lemma 7 gives local C 2 regularity of the expert and router score maps on bounded parameter regions. Lemma 8 gives local C 2 regularity of the gated Hilbert feature maps and, in the normalized case, of the normalization-composed risk under the denominator lower bound. These statements verify Assumption (A3) and Assumption (A5) on the region reached by the trajectory. Assumption (T4) provides the bounded trajectory region needed for the localized finite-horizon versions of well-posedness and ODE convergence. Thus the first variation, particle gradient, discrete scheme, finite-N ODE limit, mean-field equation, Hamiltonian formulation, dissipation identity, and hard-limit statements follow from Propositions 4, 6, and Theorems 10, 11, 12, 13, and 17, using their localized forms where appropriate. If global constants are imposed instead of localized constants, Theorem 16 gives propagation of chaos. Under (T5), Theorem 14 gives (103).

57

A.20.7

Consequences for the gated transformer-MoE model

A distribution over transformer experts with input-dependent routing is represented as a distribution over extended particles θi = (Qi , Ki , Vi , Oi , W1,i , W2,i , ϕi ) ∈ Θexp × Θgate . The empirical law is N

ρN θ =

1 X δθi . N i=1

R For smooth unnormalized routing, the model output is encoded by mρ = Fun (θ) dρ(θ). For softmax-normalized routing across experts, the numerator and denominator are encoded by the R augmented Hilbert moment msoft = F (θ) dρ(θ) and the output is recovered by the smooth soft ρ normalization map Γ. Thus input dependence of the gate is not an obstruction: it is part of the input-indexed Hilbert feature map. The only obstruction to the smooth theory is discontinuity of exact hard top-k selection; that case requires smoothing or a nonsmooth differential-inclusion treatment. The mirror map and Fenchel conjugate are defined on the full space Θ rather than only on the expert matrices. Hence the regularized Muon step is a genuine product-space mirror step for the pair consisting of expert and router parameters. The Hamiltonian probability flow evolves on phase space Θ × Θ and dissipates according to Z d Hε,γ (µt ) = −γ p, OrthΘ ε (p) Θ dµt (θ, p) ≤ 0. dt Θ×Θ Under the standard PL, upper-gradient, bounded-momentum, and bounded-curvature assumptions, the objective gap decays exponentially according to (103). Finite particles converge to the meanfield law by propagation of chaos under the global Lipschitz version of the assumptions, and the hard Muon dynamics is recovered as a subsequential nonsmooth limit as ε ↓ 0.

A.21

Synthetic experiments on finite-particle Muon dynamics

We use two deterministic synthetic experiment classes to test the finite-particle dynamics developed above. The goal is not to benchmark large-scale training, but to isolate the phenomena predicted by the Hamiltonian formulation. All reported runs are full-batch deterministic runs with seed 0, double precision, zero initial momentum, h = η = 0.01, γ = 1, and therefore β = 1 − γh = 0.99. The modified Lyapunov diagnostic uses α = 0.01. For each setting we compare four update rules. After the common momentum update Pi,k+1 = βPi,k + (1 − β)ai,k ,

(193)

θi,k+1 = θi,k − ηG(Pi,k+1 ).

(194)

we update The regularized Muon choice is G(P ) = Orthε (P ), applied blockwise on product spaces. We also include the ideal hard polar factor G(P ) = Orth(P ), a Newton-Schulz approximation to the polar factor using five fifth-order iterations, and the Euclidean momentum baseline G(P ) = P . For regularized Muon, the plotted ε values are chosen to show both the smooth regime and the nearly-hard regime.

58

A.21.1

Experiment 1: matrix mean matching

The first experiment is a linear mean-matching problem on a single matrix block. The parameter space is Θ = X = R16×8 with Frobenius inner product. We set 1 2 R(A) = A − W̄⋆ F , 2

F (W ) = W, Thus

1 J(ρ) = 2

M

1 X W̄⋆ = Wj,⋆ . M

(195)

j=1

2

Z W dρ(W ) − W̄⋆

(196) F

and the finite-particle objective is N

2

1 1 X Wi − W̄⋆ JN (W1 , . . . , WN ) = 2 N i=1

.

(197)

i = 1, . . . , N,

(198)

F

The mean-field force has the closed form N

ai,k =

1 X Wℓ,k − W̄⋆ , N ℓ=1

so all particles see the same force. This makes the experiment a clean test of the momentum-toupdate map rather than of modeling complexity. Since J⋆ = 0 is attainable, the plotted objective is also the objective gap. We report two target/approximation-particle choices: a single target particle with an overparameterized 10-particle approximation, (M, N ) = (1, 10), and a four-particle target with a 32-particle approximation, (M, N ) = (4, 32). √ The target and initial matrices are sampled from centered Gaussian ensembles scaled by 1/ 8. Each run is executed for 10,000 iterations. For (M, N ) = (1, 10) the regularized Muon values are ε ∈ {1, 3 · 10−2 , 10−3 , 10−8 }; for (M, N ) = (4, 32) they are ε ∈ {1, 10−1 , 10−2 , 10−4 }. A.21.2

Experiment 2: product-space teacher-student particles

The second experiment is a nonlinear teacher-student problem on a product matrix space. Each particle is θi = (Ai , Bi ) ∈ Θ = Rp×r × Rr×d , (d, r, p) = (10, 6, 4). (199) The product-space inner product is D E D E D E e B) e e + B, B e . (A, B), (A, = A, A Θ

F

F

(200)

We draw frozen inputs xs ∈ Rd , s = 1, . . . , S, with S = 320, and generate a frozen teacher output   M Bj,⋆ x 1 X √ ys = f⋆ (xs ), f⋆ (x) = Aj,⋆ tanh . (201) M d j=1 √ √ The teacher matrices are sampled with Aj,⋆ scaled by 1/ r and Bj,⋆ scaled by 1/ d. The feature map is the empirical prediction tensor   S Bxs F (A, B) = A tanh √ ∈ RS×p , (202) d s=1 59

Table 1: Experimental configurations used in the main figures. M is the number of target/teacher particles and N is the number of approximation particles. All runs use h = η = 0.01, γ = 1, β = 0.99, α = 0.01, full-batch gradients, and seed 0. Class

Parameter space

Exp. 1 Exp. 1 Exp. 2 Exp. 2

16×8

R R16×8 R4×6 × R6×10 R4×6 × R6×10

M

N

Iters.

Regularized ε values shown

1 4 3 10

10 32 12 10

10,000 10,000 2,000 2,000

1, 3 · 10−2 , 10−3 , 10−8 1, 10−1 , 10−2 , 10−4 10−1 , 10−3 , 10−4 , 10−5 10−1 , 10−3 , 10−5 , 10−7

and the outer loss is the squared empirical L2 loss S

1 X ∥Zs − ys ∥22 . R(Z) = 2Sp

(203)

s=1

Consequently, JN ((Ai , Bi )N i=1 ) =

2   S N 1 X 1 X Bi xs Ai tanh √ − ys . 2Sp N d s=1

i=1

(204)

2

The Muon map is applied separately to the A- and B-momentum blocks. Unlike Experiment 1, this objective is nonconvex and the particles do not share a common force. We use two particle choices. The first is an overparameterized approximation to a small teacher, (M, N ) = (3, 12). The second is a matched-particle comparison, (M, N ) = (10, 10). In both cases the student is initialized randomly, rather than near the teacher, with initial scale 0.1. Each run is executed for 2,000 iterations. For (M, N ) = (3, 12) the regularized Muon values are ε ∈ {10−1 , 10−3 , 10−4 , 10−5 }; for (M, N ) = (10, 10) they are ε ∈ {10−1 , 10−3 , 10−5 , 10−7 }. A.21.3

Results and interpretation

Figure 2 shows the objective and Hamiltonian traces for the matrix mean-matching experiment. The main qualitative point is that the smooth regularization changes the finite-step behavior near equilibrium. The hard polar map and the Newton-Schulz polar approximation keep an essentially fixed normalized update direction whenever the momentum is nonzero; as a result, for this finite step size they settle into a small residual floor around 10−3 in objective value. In contrast, the regularized map satisfies Orthε (P ) ≈ P/ε when ∥P ∥ is small relative to ε, so the update size shrinks near the optimum and the objective can decay to numerical precision. The Euclidean momentum baseline also reaches numerical precision here because Experiment 1 is a quadratic mean-matching problem with a shared force; thus this experiment should be read as a diagnostic for the regularization mechanism, not as evidence that Muon is always faster than Euclidean momentum. Figure 3 shows the nonlinear product-space teacher-student experiment. This setting is more representative of the product-space theory: the force depends on each particle, the parameter has two matrix blocks, and the loss is nonconvex. The Euclidean momentum baseline barely decreases the objective over the plotted horizon. Hard Muon and Newton-Schulz Muon reduce the loss by several orders of magnitude, confirming that spectral normalization is useful in this product-space particle model. The regularized runs display the expected ε tradeoff. A large value such as ε = 10−1 is too smooth and behaves closer to a damped Euclidean update, while very small values approach the hard polar map. Intermediate small values, especially ε = 10−5 in the plotted neural settings, 60

Figure 2: Experiment 1: matrix mean matching. Top: (M, N ) = (1, 10). Bottom: (M, N ) = (4, 32). Left panels show JN and right panels show the Hamiltonian K +γJN on a logarithmic scale. Smooth regularized Muon with moderate ε reaches numerical precision, while the hard and Newton-Schulz polar directions plateau at a finite-step residual floor. give the best final losses among the displayed regularized runs. The Hamiltonian panels follow the same qualitative decay pattern and provide the energy diagnostic predicted by the damped Hamiltonian identity. The results support three conclusions. First, the regularized operator is not merely a numerical perturbation of hard Muon: at finite step size it removes the non-vanishing update floor of the hard polar map near equilibrium. Second, in the nonlinear product-space problem, the spectral Muon geometry is substantially more effective than raw Euclidean momentum. Third, ε has an interpretable role. Large ε over-smooths the map; extremely small ε behaves like hard Muon; intermediate small ε preserves the spectral acceleration while still allowing a smooth Hamiltonian interpretation. These observations are consistent with the theoretical picture in which ε > 0 supplies the smooth mirror map and the Hamiltonian dissipation diagnostic, while the hard Muon dynamics are recovered only as a singular limit.

61

Figure 3: Experiment 2: product-space teacher-student particles with (d, r, p) = (10, 6, 4) and S = 320 frozen inputs. Top: overparameterized case (M, N ) = (3, 12). Bottom: matched-particle case (M, N ) = (10, 10). Left panels show JN and right panels show K + γJN . Spectral Muon-type updates substantially outperform Euclidean momentum in this nonlinear product-space setting; the regularized map interpolates between overly smooth Euclidean-like behavior and the hard polar regime.

Table 2: Final objective values at the last plotted iteration. The “best regularized” column selects the best value among the plotted ε values for that row. These are single-seed deterministic diagnostics, so the table should be interpreted as a mechanistic comparison rather than a statistical benchmark. Setting Exp. 1, (M, N ) = (1, 10) Exp. 1, (M, N ) = (4, 32) Exp. 2, (M, N ) = (3, 12) Exp. 2, (M, N ) = (10, 10)

GD+momentum

hard Muon

NS Muon

best regularized Muon

best ε

2.4 · 10−29 8.5 · 10−31 1.1 · 10−2 4.5 · 10−3

1.1 · 10−3 1.1 · 10−3 4.4 · 10−6 5.4 · 10−6

1.9 · 10−3 1.4 · 10−3 2.8 · 10−6 4.6 · 10−6

2.0 · 10−32 1.5 · 10−33 1.3 · 10−7 6.0 · 10−7

3 · 10−2 10−2 10−5 10−5

62

A.22

Consequences of this probabilistic view and future directions

The probability-flow perspective gives several concrete consequences that are difficult to see from the discrete update alone. First, it identifies the correct energy. For Euclidean gradient flow, the objective decreases directly. For Muon with momentum, the objective can increase transiently because momentum can be misaligned with the force. The Hamiltonian Kt + γJ (ρt ) is the quantity that dissipates exactly. This distinction is not a weakness of the method- it is the signature of acceleration. Second, the theory explains the role of ε. A positive ε gives a smooth, Lipschitz vector field and therefore a classical ODE/PDE theory. The hard optimizer is recovered as ε ↓ 0, but the limiting equation is set-valued at rank-deficient momenta. Thus ε has a dual interpretation - algorithmically it is a soft orthogonalization parameter, and analytically it is the regularity parameter that selects a smooth Hamiltonian flow before taking a nonsmooth limit. P Third, the particle formulation makes the mean-field limit explicit. The empirical law N −1 i δ(Wi ,Pi ) is not an auxiliary construction, it is the object whose limit is the phase-space law µt . Propagation of chaos then says that any fixed number of particles behaves asymptotically like independent nonlinear characteristics. This is the mathematical bridge between a finite collection of matrix blocks and a population optimization model. Fourth, the alignment term Ct gives an interpretable measure of whether momentum helps or hurts descent. Positive Ct means the momentum is aligned with the force in the sense used by the Lyapunov proof. Negative Ct means the system has kinetic energy but is initially moving against the current force. The convergence proof does not assume this term is always positive, instead, it subtracts a small multiple of Ct and controls the resulting Lyapunov function. Fifth, the transformer extension shows that input-dependent routing is compatible with probability measures over parameters. The route score does not have to be external to the measure. Once the particle is enlarged to θ = (ω, φ), both expert outputs and router scores are ordinary components of a Hilbert-valued feature map. Softmax normalization across experts is represented by augmenting the feature moment
with numerator and denominator components. This keeps R the objective in the form R F dρ and lets the same first-variation and Hamiltonian machinery apply. Several directions remain open. The most important is to verify or replace the PL and uppergradient assumptions in concrete transformer regimes. A second direction is to prove uniqueness or selection principles for the hard-Muon differential inclusion at rank-deficient momenta. A third is to analyze stochastic mini-batch noise, finite Newton-Schulz approximation error, and adaptive choices of ε within the Hamiltonian framework.

63