Paper deep dive
Muon Dynamics as a Spectral Wasserstein Flow
Gabriel Peyré
Intelligence
Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 94%
Last extracted: 4/10/2026, 2:58:28 AM
Summary
The paper introduces 'Spectral Wasserstein' distances, a family of transport metrics indexed by norms on positive semidefinite matrices. This framework generalizes the classical quadratic Wasserstein distance (trace norm) to include Muon-like geometries (operator norm) and intermediate Schatten-type schemes. The authors develop static Kantorovich and dynamic Benamou-Brenier formulations, prove metric properties for monotone norms, and interpret the resulting normalized continuity equations as Spectral Wasserstein gradient flows, providing a theoretical foundation for matrix-aware optimization in deep learning.
Entities (6)
Relation Signals (3)
Spectral Wasserstein distance → generalizes → quadratic Wasserstein distance
confidence 95% · the trace norm recovers the classical quadratic Wasserstein distance
Schatten norms → interpolatesbetween → trace norm and operator norm
confidence 95% · intermediate Schatten norms interpolate between them [trace norm and operator norm]
Muon → isinstanceof → Spectral Wasserstein gradient flow
confidence 90% · the Muon gradient flow is exactly a gradient flow on the space of measures for a Spectral Wasserstein geometry
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Gradient normalization is central in deep-learning optimization because it stabilizes training and reduces sensitivity to scale. For deep architectures, parameters are naturally grouped into matrices or blocks, so spectral normalizations are often more faithful than coordinatewise Euclidean ones; Muon is the main motivating example of this paper. More broadly, we study a family of spectral normalization rules, ranging from ordinary gradient descent to Muon and intermediate Schatten-type schemes, in a mean-field regime where parameters are modeled by probability measures. We introduce a family of Spectral Wasserstein distances indexed by a norm gamma on positive semidefinite matrices. The trace norm recovers the classical quadratic Wasserstein distance, the operator norm recovers the Muon geometry, and intermediate Schatten norms interpolate between them. We develop the static Kantorovich formulation, prove comparison bounds with W2, derive a max-min representation, and obtain a conditional Brenier theorem. For Gaussian marginals, the problem reduces to a constrained optimization on covariance matrices, extending the Bures formula and yielding a closed form for commuting covariances in the Schatten family. For monotone norms, including all Schatten cases, we prove the equivalence between the static and dynamic Benamou-Brenier formulations, deduce that the resulting transport cost is a genuine metric equivalent to W2 in fixed dimension, and show that the induced Gaussian covariance cost is also a metric. We then interpret the associated normalized continuity equation as a Spectral Wasserstein gradient flow, identify its exact finite-particle counterpart as a normalized matrix flow, obtain first geodesic-convexity results, and show how positively homogeneous mean-field models induce a spectral unbalanced transport on the sphere.
Tags
Links
- Source: https://arxiv.org/abs/2604.04891v1
- Canonical: https://arxiv.org/abs/2604.04891v1
Trouble viewing inline? Open PDF directly →
Full Text
87,070 characters extracted from source content.
Expand or collapse full text
Muon Dynamics as a Spectral Wasserstein Flow Gabriel Peyré gabriel.peyre@ens.fr CNRS and ENS, PSL Université (April 2026) Abstract Gradient normalization is a central ingredient of modern deep-learning optimization because it stabilizes training and reduces sensitivity to scale. For deep architectures, parameters are naturally grouped into matrices or blocks, so spectral normalizations are often more faithful than coordinatewise Euclidean ones; Muon is the most iconic recent example and the main motivating application of this paper. The broader purpose of this work is to describe a family of spectral normalization procedures, ranging from ordinary gradient descent to Muon and intermediate Schatten-type rules, in regimes where the number of parameters can be arbitrarily large and is naturally modeled by a probability measure on the space of neurons. For this purpose, we introduce a family of Spectral Wasserstein distances indexed by a norm γ on positive semidefinite matrices, with the trace norm recovering the classical quadratic Wasserstein distance, the operator norm recovering the Muon geometry, and intermediate Schatten norms interpolating between them. We develop the static Kantorovich formulation for arbitrary norms on the positive semidefinite cone, prove comparison estimates with 2 W_2, derive a max-min representation, and obtain a conditional Brenier theorem. For Gaussian marginals, the theory reduces to a positive-semidefinite constrained optimization that defines a covariance cost extending the Bures formula and admits a closed form for commuting covariances in the Schatten family. For monotone norms, which include all Schatten examples, we prove that the static Kantorovich formulation and the dynamic Benamou–Brenier formulation coincide, that the resulting distance is a genuine metric equivalent to 2 W_2 in fixed dimension, and that the induced Gaussian covariance cost is then a genuine metric as well. We then explain how the associated normalized continuity equation should be interpreted as a Spectral Wasserstein gradient flow, identify its exact finite-particle counterpart as a normalized matrix flow, establish first geodesic-convexity results, and show how positively homogeneous mean-field models induce a spectral unbalanced transport on the sphere. Numerical experiments compare static couplings and MMD gradient flows for the trace, Frobenius, and operator norms. 1 Introduction In modern machine learning one often optimizes an objective F(X)F(X) where X∈ℝd×nX ^d× n collects matrix-shaped parameters, for instance the columns of a weight block or a collection of particles. The baseline dynamics is gradient descent, Xk+1=Xk−τ∇F(Xk),X_k+1=X_k-τ∇ F(X_k), but one now often replaces it by spectrally normalized updates that better respect the matrix geometry of the parameter space. A prototypical example is Muon [14]: if Gk=∇F(Xk)G_k=∇ F(X_k) and Gk=UkΣkVk⊤G_k=U_k _kV_k is a singular value decomposition, the orthogonal projection of the gradient is Proj(Gk)=UkVk⊤Proj_O(G_k)=U_kV_k , and the corresponding normalized step reads Xk+1=Xk−τ‖Gk‖S1Proj(Gk).X_k+1=X_k-τ G_k _S_1\,Proj_O(G_k). Its continuous-time counterpart is the scaled gradient flow X˙t=−tr((Gt⊤Gt)1/2)Gt(Gt⊤Gt)†/2,Gt=∇F(Xt). X_t=-tr((G_t G_t)^1/2)\,G_t\,(G_t G_t) /2, G_t=∇ F(X_t). This paper studies this deterministic continuous-time model as an idealized Muon dynamics: it corresponds to the vanishing-momentum limit in which the exponential moving average parameter is suppressed, and it ignores stochasticity. The starting point of this paper is that such matrix dynamics admits a natural measure interpretation. Writing the columns of X as particles xi∈ℝdx_i ^d, one associates the empirical measure μX=1n∑i=1nδxi. _X= 1n _i=1^n _x_i. Our main insight and result is that the Muon gradient flow is exactly a gradient flow on the space of measures for a Spectral Wasserstein geometry. This geometry can be understood as a cost-robust version of Wasserstein distance: instead of transporting Dirac masses independently through a sum of scalar costs, it penalizes a global matrix norm of the displacement covariance and therefore forces the Dirac masses to interact collectively. Although Muon is the main motivating application, the goal of the paper is wider: we develop a mathematical framework for matrix-aware transport geometries that contains both ordinary gradient descent and Muon as special cases, together with intermediate Schatten-type normalizations. 1.1 Previous Works This subsection positions the paper relative to mean-field training, spectral transport costs, and normalized optimization methods for deep models. Mean-field training of neural networks. The idea of describing wide neural networks through probability measures on parameter space is now classical. It underlies the landscape analysis of two-layer networks by Mei et al. [16], the optimal-transport convergence analysis of over-parameterized models by Chizat and Bach [9], and the metric-gradient-flow viewpoint developed in Ambrosio et al. [1]. Our work keeps this mean-field perspective but changes the underlying metric from the Euclidean Wasserstein geometry to matrix-aware Spectral Wasserstein geometries. Generalized transport costs beyond the quadratic Wasserstein distance. On the transport side, our work is closest in spirit to generalized coupling costs and weak transport, as developed for instance by Gozlan et al. [13], Backhoff-Veraguas et al. [2], and Backhoff-Veraguas and Pammer [3]. It is also related to covariance-dependent transport costs [6], to the dynamic viewpoint of Benamou and Brenier [4], and to matrix-valued optimal transport [18, 8]. The Gaussian part of our analysis connects to the Bures–Wasserstein geometry of covariance matrices [5]. What is specific here is that the transported objects remain scalar probability measures, while the cost depends on a global matrix norm applied to the covariance of the displacement field. Robust Wasserstein distances and optimization over costs. Another useful perspective is to robustify Wasserstein distance by letting the cost itself vary in an optimization problem. The closest antecedent to our work is the subspace robust Wasserstein distance of Paty and Cuturi [19], which is obtained by maximizing the transport cost over a family of projected quadratic costs. In this sense it generalizes max-sliced Wasserstein, by optimizing over higher-dimensional subspaces rather than only over one-dimensional directions, and our spectral Wasserstein distance is very close in spirit to that construction, being a generalized PSD-norm version of the same max-over-cost idea. More broadly, maximizing over costs is also connected to metric learning in simple quadratic settings, for instance in Wasserstein discriminant analysis [12] and in ground metric learning [11]. Related max-over-cost formulations also appear in transportation models with congestion, where optimizing transportation costs is used to encode traffic interactions and Wardrop equilibria [7]. There is also an opposite direction, where one minimizes the transport objective with respect to the cost. This is the point of view of Sebbouh et al. [21], who optimize Wasserstein distance with respect to a structured ground metric; because the dependence on the cost is concave, this becomes a concave minimization problem and therefore a globally nonconvex program, which is useful in particular to model Gromov–Wasserstein-type structure. By contrast, our approach performs a robustification through a concave maximization over admissible quadratic costs, and the resulting static problem remains globally convex, which makes it substantially friendlier to analyze and compute. Normalized and spectrally normalized optimization rules. On the optimization side, a growing literature studies normalized first-order methods. The recent framework of Pethick et al. [20] is particularly relevant because it treats norm-constrained linear minimization oracles as a general language for normalized gradient methods and includes spectral normalizations as special cases. Earlier works such as Cutkosky and Mehta [10] and Murray et al. [17] show how normalized gradient methods change the optimization dynamics even in nonconvex settings. For deep architectures, matrix-aware normalizations are especially natural, and Muon has become a leading example of this trend in large-scale training [14, 15]. The model analyzed in this paper should be understood as an idealized Muon limit: we pass to deterministic continuous time, remove stochastic effects, and suppress the auxiliary exponential-moving-average momentum variable. Our contribution is to identify the corresponding mean-field Spectral Wasserstein geometry and to study its static, dynamic, and variational consequences well beyond the Muon case alone. 1.2 Contributions This subsection summarizes the main results of the paper and points to the precise statements proved later on. • Section 2 introduces the static Spectral Wasserstein cost in Kantorovich form and its Monge restriction. Proposition 2.8 proves the comparison with 2 W_2, Theorem 2.10 gives the max-min representation, Corollary 2.11 gives a conditional Brenier statement, and Theorem 2.12 together with Corollary 2.13 characterize the Gaussian case and the induced covariance distance. • Section 3 derives the Benamou–Brenier formulation for the monotone class of norms and proves in Theorem 3.3 that the static and dynamic costs coincide. Corollary 3.4 then yields that the Spectral Wasserstein cost is a genuine distance, and Corollary 3.5 identifies the geodesics explicitly. • Section 4 turns this geometry into a Spectral Wasserstein mean-field gradient flow. Definition 4.1 introduces the duality map, Theorem 4.2 gives its structural form, Proposition 4.3 gives the formal steepest-descent interpretation, Proposition 4.4 gives explicit Schatten selectors, and Corollary 4.5 identifies the corresponding finite-dimensional normalized particle flows. The same section also explains a simple Gaussian-preserving regime through Corollary 4.6. • Section 5 studies geodesic convexity for the new geometry. Definition 5.1 introduces the notion, Theorem 5.2 characterizes linear functionals, and Theorems 5.4–5.5 give relative-entropy convexity results. • Section 6 specializes the discussion to two-layer MLPs and then studies positively two-homogeneous models through a spherical reduction. It naturally leads to a spectral unbalanced transport geometry on the sphere and explains why the classical Wasserstein–Fisher–Rao reduction is recovered only in the trace-norm case. • Section 7 presents two numerical studies: static spectral couplings for the trace, Frobenius, and operator norms, and MMD gradient flows for the same three geometries. The code used to reproduce the numerical experiments is available at https://github.com/gpeyre/spectral-wasserstein. 2 Static Spectral Wasserstein Geometry This section introduces the Spectral Wasserstein transport cost. The key point is that the correct object is coupling-based and depends on a norm γ acting on the global displacement covariance. 2.1 Matrix Norms on the PSD Cone This subsection fixes the matrix-norm framework that will parameterize the whole Spectral Wasserstein geometry. Throughout the paper, γ denotes a norm on the cone +dS_+^d of positive semidefinite matrices. The next example records the benchmark family that interpolates between classical Wasserstein geometry and the Muon geometry. We use the name “Spectral Wasserstein” because the main cases of interest are spectral norms on +dS_+^d, namely norms invariant under orthogonal conjugation and therefore depending only on the eigenvalues of the matrix, as in the Schatten family. Example 2.1 (Schatten norms). The main examples are the Schatten norms restricted to +dS_+^d: γp(S)=‖S‖Sp,1≤p≤∞. _p(S)= S _S_p, 1≤ p≤∞. If λ(S)=(λ1(S),…,λd(S))λ(S)=( _1(S),…, _d(S)) denotes the eigenvalue vector of S⪰0S 0, then γp(S)=‖λ(S)‖ℓp=(∑i=1dλi(S)p)1/p,1≤p<∞,max1≤i≤dλi(S),p=∞. _p(S)= λ(S) _ ^p= cases ( _i=1^d _i(S)^p )^1/p,&1≤ p<∞,\\[3.99994pt] _1≤ i≤ d _i(S),&p=∞. cases In particular, γ1(S)=tr(S) _1(S)=tr(S), γ2(S)=‖S‖F _2(S)= S _F, and γ∞(S)=λmax(S) _∞(S)= _ (S). In this paper, we will show that the choice p=1p=1 leads to the classical quadratic Wasserstein geometry, while the choice p=∞p=∞ leads to the Muon / operator-norm geometry. Intermediate values of p interpolate between these two extremes. The static duality of the paper can be encoded by any convex compact representing set γ⊂dK_γ ^d such that γ(S)=maxQ∈γtr(QS)for every S⪰0.γ(S)= _Q _γtr(QS) every S 0. One canonical choice is the polar set γ≔Q∈d:tr(QS)≤γ(S)for every S⪰0.P_γ \Q ^d:\ tr(QS)≤γ(S)\ for every S 0\. Finite-dimensional convex duality implies that γP_γ is convex and compact and satisfies γ(S)=maxQ∈γtr(QS)for every S⪰0.γ(S)= _Q _γtr(QS) every S 0. Throughout the paper, γK_γ denotes a fixed convex compact representing set for γ; unless specified otherwise, one may simply take γ=γK_γ=P_γ. For Schatten norms γp(S)=‖S‖Sp _p(S)= S _S_p with dual exponent q=p−1q= pp-1, a convenient choice is γp=Q⪰0:‖Q‖Sq≤1for 1<p≤∞.K_ _p=\Q 0:\ Q _S_q≤ 1\ 1<p≤∞. For p=1p=1, one may use the smaller convex compact choice γ1=Id.K_ _1=\Id\. The next proposition characterizes the monotone norms for which representing sets may be chosen inside the positive-semidefinite cone. Proposition 2.2 (PSD representation and monotonicity). The following are equivalent: • γ is monotone on +dS_+^d, namely 0⪯S⪯T⟹γ(S)≤γ(T);0 S T γ(S)≤γ(T); • there exists a convex compact representing set γ⊂+dK_γ _+^d such that γ(S)=maxQ∈γtr(QS)for every S⪰0.γ(S)= _Q _γtr(QS) every S 0. Moreover, when these properties hold, the canonical choice γ=γ∩+dK_γ=P_γ _+^d is admissible. Proof. If γ⊂+dK_γ _+^d, then for 0⪯S⪯T0 S T one has tr(QS)≤tr(QT)tr(QS) (QT) for every Q∈γQ _γ, hence γ(S)≤γ(T)γ(S)≤γ(T) by taking suprema. Conversely, assume γ is monotone and let B≔T⪰0:γ(T)≤1.B \T 0:\ γ(T)≤ 1\. The set B is downward closed. Take Q∈γQ _γ, let P+P_+ be the spectral projector onto the positive eigenspace of Q, and set Q+=P+QP+Q_+=P_+QP_+. For every T∈BT∈ B, define T+≔P+TP+T_+ P_+TP_+. Then 0⪯T+⪯T,0 T_+ T, so T+∈BT_+∈ B, and tr(Q+T)=tr(QT+)≤γ(T+)≤1.tr(Q_+T)=tr(QT_+)≤γ(T_+)≤ 1. Hence Q+∈γ∩+dQ_+ _γ _+^d. Since tr(QS)≤tr(Q+S)tr(QS) (Q_+S) for every S⪰0S 0, taking suprema over γP_γ and then over γ∩+dP_γ _+^d gives the required support formula. Convexity and compactness are inherited from γP_γ. ∎ 2.2 Kantorovich Cost and Monge Restriction This subsection introduces the static transport cost and compares its coupling and map versions. The first definition gives the coupling-based Spectral Wasserstein cost that will serve as the reference static formulation. Definition 2.3 (Generalized static cost). For μ,ν∈2(ℝd)μ,ν _2(R^d), define γ(μ,ν)2≔infπ∈Π(μ,ν)γ(∫ℝd×ℝd(y−x)(y−x)⊤π(x,y)). W_γ(μ,ν)^2 _π∈ (μ,ν)γ\! ( _R^d×R^d(y-x)(y-x) \,dπ(x,y) ). The next definition records the Monge restriction, which is useful for comparison but is not the correct notion in general. Definition 2.4 (Monge restriction). Define γM(μ,ν)2≔infT#μ=νγ(∫ℝd(T(x)−x)(T(x)−x)⊤μ(x)), W_γ^M(μ,ν)^2 _T_\#μ=νγ\! ( _R^d(T(x)-x)(T(x)-x) \,dμ(x) ), with value +∞+∞ if no transport map exists. The next proposition simply records that the map-based problem is a restriction of the coupling-based one. Proposition 2.5 (Monge is a restriction). For every μ,ν∈2(ℝd)μ,ν _2(R^d), γ(μ,ν)≤γM(μ,ν). W_γ(μ,ν)≤ W_γ^M(μ,ν). Proof. Every transport map T induces the coupling (Id,T)#μ(Id,T)_\#μ, and the static cost evaluated on this coupling is exactly the Monge cost associated with T. ∎ The following remark shows on a two-Dirac example that the relaxation can be genuinely strict. Remark 2.6 (Strictness of the Monge restriction). Consider μ=12(δ(−1,0)+δ(1,0)),ν=12(δ(0,−1)+δ(0,1)).μ= 12( _(-1,0)+ _(1,0)), ν= 12( _(0,-1)+ _(0,1)). Any transport map from μ to ν is a bijection between the two atoms. The two possible displacement covariance matrices are (1111)or(1−1−11), pmatrix1&1\\ 1&1 pmatrix pmatrix1&-1\\ -1&1 pmatrix, so for the operator norm and Frobenius norm one gets γM(μ,ν)2=2. W_γ^M(μ,ν)^2=2. By contrast, the split coupling assigning mass 1/41/4 to each source-target pair has displacement covariance equal to the identity matrix, hence γ(μ,ν)2=1for γ=∥⋅∥S∞,γ(μ,ν)2=2for γ=∥⋅∥S2. W_γ(μ,ν)^2=1 γ= · _S_∞, W_γ(μ,ν)^2= 2 γ= · _S_2. Therefore the inequality in Proposition 2.5 is strict already for two-point measures. The next remark explains that, for a completely arbitrary norm on +dS_+^d, the static cost need not yet be a bona fide metric. Remark 2.7 (Triangle inequality may fail for non-monotone norms). For Dirac masses, the unique coupling gives γ(δx,δy)2=γ((y−x)(y−x)⊤). W_γ( _x, _y)^2=γ\! ((y-x)(y-x) ). Hence if γ W_γ were always a distance, then the pointwise cost dγ(x,y)≔γ((y−x)(y−x)⊤)d_γ(x,y) γ\! ((y-x)(y-x) ) would in particular have to define a metric on ℝdR^d. This fails for general norms on +dS_+^d. Fix d≥2d≥ 2 and a parameter M>2M>2, and define γns(S)≔tr(S)+M∑1≤i<j≤d|Sij|,S⪰0. _ns(S) (S)+M _1≤ i<j≤ d S_ij , S 0. This is a norm on +dS_+^d, but it is not spectral since it depends on the matrix entries and not only on the eigenvalues. For the displacement Δ=y−x=(Δ1,…,Δd) =y-x=( _1,…, _d) one has dγns(x,y)2=‖Δ‖22+M∑1≤i<j≤d|ΔiΔj|.d_ _ns(x,y)^2= _2^2+M _1≤ i<j≤ d _i _j . Taking x0=0,x1=e1,x2=e1+e2x_0=0, x_1=e_1, x_2=e_1+e_2 gives dγns(x0,x1)=1,dγns(x1,x2)=1,dγns(x0,x2)=2+M>2,d_ _ns(x_0,x_1)=1, d_ _ns(x_1,x_2)=1, d_ _ns(x_0,x_2)= 2+M>2, so the triangle inequality fails. By contrast, if γ is monotone on +dS_+^d, then Proposition 2.2 allows us to choose γ⊂+dK_γ _+^d, and therefore dγ(x,y)=supQ∈γ‖Q1/2(x−y)‖2,d_γ(x,y)= _Q _γ Q^1/2(x-y) _2, which is a supremum of seminorms and therefore a genuine distance on ℝdR^d. For general measures, however, this pointwise argument does not control the cross terms created by gluing couplings. The full metric property of γ W_γ in the monotone case is therefore proved later, through the dynamic formulation, in Corollary 3.4. 2.3 Comparison with 2 W_2 and Cost-Robust Formulation This subsection compares the Spectral Wasserstein cost with the classical quadratic Wasserstein distance and derives its static dual description. The next proposition quantifies how the Spectral Wasserstein distance is sandwiched between two multiples of 2 W_2. Proposition 2.8 (Norm comparison). Define cγ≔infS⪰0tr(S)=1γ(S),Cγ≔supS⪰0tr(S)=1γ(S).c_γ _ subarraycS 0\\ tr(S)=1 subarrayγ(S), C_γ _ subarraycS 0\\ tr(S)=1 subarrayγ(S). These constants always exist and satisfy 0<cγ≤Cγ<∞0<c_γ≤ C_γ<∞. Equivalently, they are the best comparison constants between γ and the trace norm on +dS_+^d: cγtr(S)≤γ(S)≤Cγtr(S)for every S⪰0.c_γ\,tr(S)≤γ(S)≤ C_γ\,tr(S) every S 0. Therefore cγ2(μ,ν)≤γ(μ,ν)≤Cγ2(μ,ν)for every μ,ν∈2(ℝd). c_γ\, W_2(μ,ν)≤ W_γ(μ,ν)≤ C_γ\, W_2(μ,ν) every μ,ν _2(R^d). For Schatten norms, d1p−12(μ,ν)2≤γp(μ,ν)2≤2(μ,ν)2.d 1p-1\, W_2(μ,ν)^2≤ W_ _p(μ,ν)^2≤ W_2(μ,ν)^2. Equivalently, cγp=d1p−1,Cγp=1.c_ _p=d 1p-1, C_ _p=1. In particular, (cγ2,Cγ2)=(d−1/2,1),(cγ∞,Cγ∞)=(d−1,1).(c_ _2,C_ _2)=(d^-1/2,1), (c_ _∞,C_ _∞)=(d^-1,1). Proof. The slice S⪰0:tr(S)=1\S 0:tr(S)=1\ is compact, so cγc_γ and CγC_γ exist. By homogeneity of γ, the two-sided bound follows for all S⪰0S 0. Applying it to the displacement covariance of any coupling and taking infima yields the comparison with 2 W_2. For Schatten norms, the eigenvalue inequality ‖λ‖ℓp≤‖λ‖ℓ1≤d1−1p‖λ‖ℓp λ _ _p≤ λ _ _1≤ d^1- 1p λ _ _p gives the stated constants. ∎ The following remark records that the lower Schatten bound is sharp. Remark 2.9 (Sharpness of the lower Schatten bound). For γ=γpγ= _p, equality in the lower bound is attained by isotropic Gaussian pairs. Let μ=(0,αId),ν=(0,βId),α,β>0.μ=N(0,α I_d), ν=N(0,β I_d), α,β>0. By the Gaussian computation of Section 2.4, the displacement covariance is (α−β)2Id,( α- β)^2I_d, and hence γp(μ,ν)2=d1p−12(μ,ν)2. W_ _p(μ,ν)^2=d 1p-1 W_2(μ,ν)^2. Hence the lower constant in Proposition 2.8 is optimal for every Schatten norm. The comparison in Proposition 2.8 already implies separation and topological equivalence with 2 W_2. What is not immediate from the static formulation alone is the triangle inequality, because gluing two couplings creates cross terms in the covariance of the composed displacement. We therefore prove the bona fide metric property in Corollary 3.4 after establishing the Benamou–Brenier formulation in Section 3. For every symmetric matrix Q∈dQ ^d, denote by Q W^Q the quadratic transport functional associated with the cost ⟨Q(x−y),x−y⟩ Q(x-y),x-y , namely Q(μ,ν)2≔infπ∈Π(μ,ν)∫(y−x)⊤Q(y−x)π(x,y). W^Q(μ,ν)^2 _π∈ (μ,ν) (y-x) Q(y-x)\,dπ(x,y). In particular, Id=2 W^Id= W_2. The next theorem identifies the static Spectral Wasserstein cost with an anisotropic quadratic transport problem optimized over the dual unit ball. Theorem 2.10 (Max-min and cost-robust representation). For every μ,ν∈2(ℝd)μ,ν _2(R^d), γ(μ,ν)2=maxQ∈γinfπ∈Π(μ,ν)∫(y−x)⊤Q(y−x)π(x,y)=maxQ∈γQ(μ,ν)2. W_γ(μ,ν)^2= _Q _γ _π∈ (μ,ν) (y-x) Q(y-x)\,dπ(x,y)= _Q _γ W^Q(μ,ν)^2. Proof. Using the support-function formula for γ, γ(μ,ν)2=infπ∈Π(μ,ν)maxQ∈γ∫(y−x)⊤Q(y−x)π. W_γ(μ,ν)^2= _π∈ (μ,ν) _Q _γ (y-x) Q(y-x)\,dπ. The coupling set is convex and weakly compact, γK_γ is convex and compact, and the integrand is affine and continuous in each variable, so Sion’s theorem applies. ∎ The following corollary explains when an optimal coupling is induced by a Monge map of Brenier type. Corollary 2.11 (Conditional Brenier theorem). Assume μ is absolutely continuous. If a maximizing matrix Q∗∈γQ_* _γ in Theorem 2.10 is positive definite, then every optimal coupling for γ(μ,ν) W_γ(μ,ν) is induced by a map. More precisely, there exists a convex function u such that T∗(x)=Q∗−1/2∇u(Q∗1/2x)T_*(x)=Q_*^-1/2∇ u(Q_*^1/2x) is optimal. For a general norm γ, this hypothesis is genuinely restrictive because a maximizer Q∗∈γQ_* _γ need not be positive semidefinite and may even be indefinite. If γ is monotone, however, Proposition 2.2 allows us to choose the representing set γK_γ inside +dS_+^d, so the maximizing matrix is automatically positive semidefinite and the remaining assumption is simply that it be invertible. Proof. Fix such a positive definite maximizer Q∗Q_*. For any coupling π between μ and ν, let π~=(Q∗1/2,Q∗1/2)#π, π=(Q_*^1/2,Q_*^1/2)_\#π, and define μ~=(Q∗1/2)#μ,ν~=(Q∗1/2)#ν. μ=(Q_*^1/2)_\#μ, ν=(Q_*^1/2)_\#ν. Then ∫(y−x)⊤Q∗(y−x)π(x,y)=∫|y~−x~|2π~(x~,y~). (y-x) Q_*(y-x)\,dπ(x,y)= y- x ^2\,d π( x, y). Because μ is absolutely continuous and Q∗1/2Q_*^1/2 is invertible, μ~ μ is also absolutely continuous. Brenier’s theorem for the quadratic cost therefore yields a convex potential u such that the unique optimal coupling between μ~ μ and ν~ ν is induced by the map T~(x~)=∇u(x~). T( x)=∇ u( x). Pulling this map back to the original variables gives T∗(x)=Q∗−1/2T~(Q∗1/2x)=Q∗−1/2∇u(Q∗1/2x),T_*(x)=Q_*^-1/2 T(Q_*^1/2x)=Q_*^-1/2∇ u(Q_*^1/2x), and the corresponding coupling is optimal for the inner problem associated with Q∗Q_*. Since Q∗Q_* maximizes Theorem 2.10, this coupling is also optimal for γ(μ,ν) W_γ(μ,ν). ∎ 2.4 Gaussian Marginals and a Generalized Bures Distance This subsection shows that Gaussian marginals compress the transport problem to a finite-dimensional optimization over admissible covariance blocks. The next theorem shows that, for Gaussian marginals, the infinite-dimensional transport problem collapses to an optimization over the cross-covariance matrix. Theorem 2.12 (Gaussian reduction). Let μ=(m0,Σ0)μ=N(m_0, _0) and ν=(m1,Σ1)ν=N(m_1, _1). Then γ(μ,ν)2=infKγ((m1−m0)(m1−m0)⊤+Σ0+Σ1−K−K⊤), W_γ(μ,ν)^2= _Kγ\! ((m_1-m_0)(m_1-m_0) + _0+ _1-K-K ), where the infimum runs over matrices K such that (Σ0K⊤Σ1)⪰0. pmatrix _0&K\\ K & _1 pmatrix 0. In particular, for centered Gaussians the covariance cost γ(Σ0,Σ1)2≔infK:[Σ0K⊤Σ1]⪰0γ(Σ0+Σ1−K−K⊤) B_γ( _0, _1)^2 _K:\, [ smallmatrix _0&K\\ K & _1 smallmatrix ] 0γ( _0+ _1-K-K ) is well defined on the cone of covariance matrices. In the monotone regime of Section 3, it is the restriction of the metric γ W_γ to centered Gaussian laws and therefore defines a metric on covariance matrices. Proof. Let (X,Y)(X,Y) be any coupling of the two Gaussian marginals. Since X and Y are jointly Gaussian if and only if their joint law is determined by first and second moments, every Gaussian coupling is characterized by the cross-covariance matrix K=[(X−m0)(Y−m1)⊤],K=E[(X-m_0)(Y-m_1) ], and the covariance matrix of (X,Y)(X,Y) is (Σ0K⊤Σ1). pmatrix _0&K\\ K & _1 pmatrix. This block matrix must be positive semidefinite. Conversely, every such block positive semidefinite matrix defines a jointly Gaussian vector with marginals μ and ν. For any admissible K, the displacement covariance equals (m1−m0)(m1−m0)⊤+Σ0+Σ1−K−K⊤.(m_1-m_0)(m_1-m_0) + _0+ _1-K-K . Therefore the generalized cost of a Gaussian coupling depends only on K. Optimizing over Gaussian couplings is exactly the same as optimizing over the block positive semidefinite constraint, which proves the formula. ∎ The next corollary shows that commuting covariances lead to a closed form for the Schatten family, exactly as in the classical Bures case. Corollary 2.13 (Commuting covariances). If Σ0 _0 and Σ1 _1 commute, then for Schatten norms γp _p one has γp(Σ0,Σ1)2=γp((Σ01/2−Σ11/2)2). B_ _p( _0, _1)^2= _p\! (( _0^1/2- _1^1/2)^2 ). Equivalently, γp(Σ0,Σ1)2=(∑i=1d|λi(Σ0)−λi(Σ1)|2p)1/p. B_ _p( _0, _1)^2= ( _i=1^d _i( _0)- _i( _1) ^2p )^1/p. When p=1p=1, this is the usual Bures–Wasserstein formula. Proof. If Σ0 _0 and Σ1 _1 commute, they are simultaneously diagonalizable. In that basis the block PSD constraint decouples coordinatewise, and the optimal choice is K=diag(aibi)K=diag( a_ib_i). The resulting displacement covariance is diag((ai−bi)2)diag(( a_i- b_i)^2). ∎ The following remark clarifies the scope of Corollary 2.13. Remark 2.14. The commuting formula is stated for Schatten norms because their value depends only on the eigenvalues of the displacement covariance. For more general norms on +dS_+^d, even when Σ0 _0 and Σ1 _1 commute, one should not expect a closed form depending only on the eigenvalues (ai−bi)2( a_i- b_i)^2, since the norm itself may retain basis-dependent information. 3 Dynamic Formulation and Geodesics We now turn to the Benamou–Brenier side [4]. This section requires the additional monotonicity property satisfied by all Schatten norms and shows that, under this assumption, the same object is obtained dynamically. 3.1 Dynamic and Momentum Formulations This subsection introduces the Benamou–Brenier action and its convex reformulation in momentum variables. From now on in this section, we assume in addition that γ is monotone on +dS_+^d, namely 0⪯S⪯T⟹γ(S)≤γ(T).0 S T γ(S)≤γ(T). By Proposition 2.2, we may and do choose the representing set γK_γ so that γ⊂+d.K_γ _+^d. This positive-semidefinite property is the crucial additional ingredient used throughout the Benamou–Brenier analysis below. The next definition gives the dynamic transport problem associated with the static Spectral Wasserstein cost. Definition 3.1 (Dynamic generalized cost). For μ0,μ1∈2(ℝd) _0, _1 _2(R^d), define γB(μ0,μ1)2≔inf(μt,vt)∫01γ(∫ℝdvt(x)vt(x)⊤μt(x))t, W_γ^B( _0, _1)^2 _( _t,v_t) _0^1γ\! ( _R^dv_t(x)v_t(x) \,d _t(x) )\,dt, where the infimum runs over narrowly continuous curves t↦μt _t and measurable velocity fields satisfying ∂tμt+div(μtvt)=0,μt=0=μ0,μt=1=μ1. _t _t+div( _tv_t)=0, _t=0= _0, _t=1= _1. Passing to momenta linearizes the constraint. If μ=ρλμ=ρλ and m=wλm=wλ for some reference measure λ, define γ(μ,m)≔supQ∈γ∫ℝdw(x)⊤Qw(x)ρ(x)λ(x),A_γ(μ,m) _Q _γ _R^d w(x) Q\,w(x)ρ(x)\,dλ(x), with the usual perspective convention on ρ=0\ρ=0\. The following proposition shows that the momentum formulation is convex and exactly matches the velocity action. Proposition 3.2 (Convex momentum action). The action γ(μ,m)A_γ(μ,m) is intrinsic, convex, and weak-* lower semicontinuous. If m=μvm=μ v, then γ(μ,m)=γ(∫v(x)v(x)⊤μ(x)).A_γ(μ,m)=γ\! ( v(x)v(x) \,dμ(x) ). Hence γB(μ0,μ1)2=inf(μt,mt)∫01γ(μt,mt)t W_γ^B( _0, _1)^2= _( _t,m_t) _0^1A_γ( _t,m_t)\,dt under the linear constraint ∂tμt+div(mt)=0. _t _t+div(m_t)=0. Proof. For m=μvm=μ v, one has w=ρvw=ρ v, hence ∫w⊤Qwρλ=∫v⊤Qvμ=tr(Q∫vv⊤μ). w Qwρ\,dλ= v Qv\,dμ=tr\! (Q v \,dμ ). Taking the supremum over Q∈γQ _γ gives the support-function formula for γ. Convexity and lower semicontinuity follow because the integrand is a supremum of perspective quadratic forms. ∎ 3.2 Static Equals Dynamic This subsection identifies the coupling formulation with the Benamou–Brenier formulation and extracts the metric consequences. The next theorem is the structural core of the paper: it identifies the static Spectral Wasserstein cost with its dynamic Benamou–Brenier counterpart. Theorem 3.3 (Static and dynamic formulations coincide). For every μ0,μ1∈2(ℝd) _0, _1 _2(R^d), γB(μ0,μ1)=γ(μ0,μ1). W_γ^B( _0, _1)= W_γ( _0, _1). Proof. Let π∈Π(μ0,μ1)π∈ ( _0, _1) and consider the displacement interpolation μt=((1−t)x+ty)#π. _t=((1-t)x+ty)_\#π. The velocity along each segment is y−xy-x, so ∫vtvt⊤μt=∫(y−x)(y−x)⊤π v_tv_t \,d _t= (y-x)(y-x) \,dπ for all t. Therefore γB(μ0,μ1)2≤γ(μ0,μ1)2. W_γ^B( _0, _1)^2≤ W_γ( _0, _1)^2. Conversely, take any admissible dynamic plan (μt,vt)( _t,v_t). By the superposition principle, there exists a probability measure η on absolutely continuous paths γ such that μt=(et)#η,γ˙t=vt(γt) _t=(e_t)_\#η, γ_t=v_t( _t) for η-a.e. path and a.e. t. Let π=(e0,e1)#ηπ=(e_0,e_1)_\#η. For every Q∈γQ _γ, ∫01∫vt⊤Qvtμtt=∫01γ˙t⊤Qγ˙ttη(γ). _0^1 v_t Qv_t\,d _t\,dt= _0^1 γ_t Q γ_t\,dt\,dη(γ). This is the key point where monotonicity is used: by Proposition 2.2 we have chosen γ⊂+dK_γ _+^d, so every test matrix satisfies Q⪰0Q 0. Without this reduction one would have to work with possibly indefinite matrices, and the quadratic Jensen inequality below would no longer apply. Applying Jensen to the scalar function t↦Q1/2γ˙t Q^1/2 γ_t gives ∫01γ˙t⊤Qγ˙tt≥(γ1−γ0)⊤Q(γ1−γ0). _0^1 γ_t Q γ_t\,dt≥( _1- _0) Q( _1- _0). Hence ∫01γ(∫vtvt⊤μt)t≥∫(y−x)⊤Q(y−x)π(x,y) _0^1γ\! ( v_tv_t \,d _t )\,dt≥ (y-x) Q(y-x)\,dπ(x,y) for every Q∈γQ _γ. Taking the supremum over Q yields ∫01γ(∫vtvt⊤μt)t≥γ(∫(y−x)(y−x)⊤π). _0^1γ\! ( v_tv_t \,d _t )\,dt≥γ\! ( (y-x)(y-x) \,dπ ). Finally, take the infimum over admissible dynamic plans. ∎ The following corollary records the metric consequences of the static-dynamic equivalence. Corollary 3.4 (Metric properties). The quantity γ W_γ is a bona fide metric on 2(ℝd)P_2(R^d) and induces the same topology as 2 W_2. Proof. Symmetry is obvious. Separation follows from Proposition 2.8. The triangle inequality follows from the dynamic formulation by concatenation and time rescaling of admissible curves. ∎ The next corollary gives the explicit constant-speed geodesics once an optimal coupling is known. Corollary 3.5 (Geodesics). If π∗ _* is any optimal coupling for γ(μ0,μ1) W_γ( _0, _1), then μt=((1−t)x+ty)#π∗ _t=((1-t)x+ty)_\# _* is a constant-speed geodesic and γ(μs,μt)=|t−s|γ(μ0,μ1) W_γ( _s, _t)= t-s \, W_γ( _0, _1) for every 0≤s≤t≤10≤ s≤ t≤ 1. Proof. The upper bound follows by restricting the same displacement interpolation to the time interval [s,t][s,t] and reparameterizing it to [0,1][0,1]. Indeed, the rescaled velocity is (t−s)(y−x)(t-s)(y-x), so Theorem 3.3 gives γ(μs,μt)2≤(t−s)2γ(∫(y−x)(y−x)⊤π∗(x,y))=(t−s)2γ(μ0,μ1)2. W_γ( _s, _t)^2≤(t-s)^2γ\! ( (y-x)(y-x) \,d _*(x,y) )=(t-s)^2 W_γ( _0, _1)^2. Applying the same argument to the segments [0,s][0,s] and [t,1][t,1] yields γ(μ0,μs)≤sγ(μ0,μ1),γ(μt,μ1)≤(1−t)γ(μ0,μ1). W_γ( _0, _s)≤ s W_γ( _0, _1), W_γ( _t, _1)≤(1-t) W_γ( _0, _1). By the triangle inequality from Corollary 3.4, γ(μ0,μ1)≤γ(μ0,μs)+γ(μs,μt)+γ(μt,μ1)≤γ(μs,μt)+(1−(t−s))γ(μ0,μ1). W_γ( _0, _1)≤ W_γ( _0, _s)+ W_γ( _s, _t)+ W_γ( _t, _1)≤ W_γ( _s, _t)+(1-(t-s)) W_γ( _0, _1). Rearranging gives γ(μs,μt)≥(t−s)γ(μ0,μ1), W_γ( _s, _t)≥(t-s) W_γ( _0, _1), hence equality holds. ∎ 4 Spectral Wasserstein Gradient Flows We now return to optimization and explain how the Spectral Wasserstein geometry generates normalized mean-field dynamics. There are in fact two layers in this section. First, for an arbitrary norm γ on +dS_+^d, one can define a local tangent norm, a duality map, and the associated normalized continuity PDE. Second, when γ is monotone and Section 3 applies, the same PDE may be interpreted as a genuine metric gradient flow for the distance γ W_γ. Throughout this section we keep the monotone setting in the background, but we indicate explicitly which statements only use the local norm structure and which ones use the metric theory. 4.1 An Informal Gradient-Flow Picture This subsection introduces the Spectral Wasserstein gradient flow of a functional on measures before specializing to neural-network objectives. Let Ω=ℝd =R^d, let f:2(Ω)→ℝf:P_2( ) be a sufficiently regular functional, and write dεf(μ+εσ)|ε=0=∫Ωδfδμ(x)σ(x) dd f(μ+ σ) |_ =0= _ δ fδμ(x)\,dσ(x) for every signed perturbation σ with zero total mass. The scalar field δf/δμδ f/δμ is the first variation of f, defined up to an additive constant. The associated Wasserstein gradient is the vector field gμ(x)=∇xδfδμ(x).g_μ(x)= _x δ fδμ(x). The tangent cost induced by γ is μ(v)2≔γ(∫Ωv(x)v(x)⊤μ(x)).N_μ(v)^2 γ\! ( _ v(x)v(x) \,dμ(x) ). This construction only uses that γ is a norm on +dS_+^d; no monotonicity is needed at this stage. Informally, the gradient-flow velocity at μ is obtained by minimizing the linearized decrease of f penalized by the squared tangent norm. This leads to the duality map below; to lighten notation, we write JμJ_μ although the map depends on the chosen norm γ. The next definition packages the matrix-normalized steepest-descent direction into a single operator. Definition 4.1 (Duality map). For each μ and force field g, let Jμ(g)J_μ(g) denote any minimizer of v⟼∫Ωg(x)⋅v(x)μ(x)+12μ(v)2.v _ g(x)· v(x)\,dμ(x)+ 12N_μ(v)^2. The next theorem gives the basic structural form of the selector once one chooses an active matrix in the support representation of γ. Theorem 4.2 (Structure theorem for the selector). Let μ∈2(ℝd)μ _2(R^d) and let g∈L2(μ;ℝd)g∈ L^2(μ;R^d) be nonzero. Define the force covariance Sμ(g)≔∫g(x)g(x)⊤μ(x)∈+d.S_μ(g) g(x)g(x) \,dμ(x) _+^d. Choose any active matrix Qμ∗∈argmaxQ∈γtr(QSμ(g)).Q_μ^*∈ *argmax_Q _γtr(QS_μ(g)). If Qμ∗Q_μ^* is invertible, then Jμ(g)(x)=−(Qμ∗)−1g(x)J_μ(g)(x)=-(Q_μ^*)^-1g(x) is a valid selector for Definition 4.1. Proof. Set v(x)≔−(Qμ∗)−1g(x)v(x) -(Q_μ^*)^-1g(x). Then ∫g⋅vμ=−∫g(x)⊤(Qμ∗)−1g(x)μ(x)=−tr((Qμ∗)−1Sμ(g)). g· v\,dμ=- g(x) (Q_μ^*)^-1g(x)\,dμ(x)=-tr\! ((Q_μ^*)^-1S_μ(g) ). Moreover, ∫v(x)v(x)⊤μ(x)=(Qμ∗)−1Sμ(g)(Qμ∗)−1, v(x)v(x) \,dμ(x)=(Q_μ^*)^-1S_μ(g)(Q_μ^*)^-1, so tr(Qμ∗∫vv⊤μ)=tr((Qμ∗)−1Sμ(g)).tr\! (Q_μ^* v \,dμ )=tr\! ((Q_μ^*)^-1S_μ(g) ). Since Qμ∗Q_μ^* is active for Sμ(g)S_μ(g), the support-function formula for γ is attained at Qμ∗Q_μ^* along this direction. Therefore the first-order optimality condition of the convex minimization problem in Definition 4.1 is satisfied by v, which proves the claim. ∎ The corresponding normalized transport flow of f is the continuity equation ∂tμt+div(μtJμt(gμt))=0. _t _t+div\! ( _tJ_ _t(g_ _t) )=0. When γ(S)=tr(S)γ(S)=tr(S), one has Jμ(g)=−gJ_μ(g)=-g, so the equation reduces to the classical 2 W_2 gradient flow. For general γ, the duality map plays the role of a matrix-aware normalization of the force field. When γ is monotone, so that the Benamou–Brenier theory of Section 3 applies, we interpret this same PDE as the metric gradient flow of f for the distance γ W_γ. Without monotonicity, the PDE still makes sense as a normalized steepest-descent transport equation, but the paper does not claim a metric gradient-flow interpretation for γ W_γ. The next proposition records the formal bridge between the local norm μN_μ and the dynamic metric structure from Section 3. Here |μt′|γ _t _ W_γ denotes the metric derivative of the curve t↦μt _t with respect to γ W_γ, namely |μt′|γ≔limh→0γ(μt+h,μt)|h|, _t _ W_γ _h→ 0 W_γ( _t+h, _t) h , whenever the limit exists. The statement is a formal version of the metric-gradient-flow minimal-slope principle. Proposition 4.3 (Formal steepest-descent principle). Assume γ is monotone and let (μt,vt)( _t,v_t) be a sufficiently regular curve satisfying ∂tμt+div(μtvt)=0. _t _t+div( _tv_t)=0. Then the Benamou–Brenier characterization of Theorem 3.3 formally gives |μt′|γ≤μt(vt), _t _ W_γ _ _t(v_t), with equality along smooth characteristic curves. Moreover, if f is differentiable along the curve, then dtf(μt)=∫gμt(x)⋅vt(x)μt(x). ddtf( _t)= g_ _t(x)· v_t(x)\,d _t(x). Consequently, the instantaneous steepest-descent problem for the metric γ W_γ is formally minv∫gμ⋅vμ+12μ(v)2, _v \ g_μ· v\,dμ+ 12N_μ(v)^2 \, which is exactly the definition of Jμ(gμ)J_μ(g_μ). Proof. The dynamic action in Definition 3.1 is the time integral of μt(vt)2N_ _t(v_t)^2, so Theorem 3.3 identifies μN_μ as the infinitesimal metric norm associated with γ W_γ. The chain rule for the first variation gives the displayed identity for dtf(μt) ddtf( _t). Minimizing the linearized decrease penalized by the squared metric norm yields the stated variational problem. ∎ 4.2 Finite-Width Normalized Particle Flows This subsection explains how the abstract duality map becomes an explicit normalized matrix flow for empirical measures. Let μX=1n∑i=1nδxi,X=[x1⊤⋮xn⊤]∈ℝn×d, _X= 1n _i=1^n _x_i, X= bmatrixx_1 \\ \\ x_n bmatrix ^n× d, and define the finite-dimensional objective Fn(X)≔f(μX).F_n(X) f( _X). If the measure flow is transported by characteristics, empirical measures remain empirical and the velocity field becomes an update on the stacked particle matrix X. For Schatten norms, the duality map can be written explicitly at matrix level. If V∈ℝn×dV ^n× d is a matrix of particle velocities, then γp(1nV⊤V)1/2=1n‖V‖S2p. _p\! ( 1nV V )^1/2= 1 n V _S_2p. The next proposition gives the explicit selector associated with each Schatten geometry. Proposition 4.4 (Explicit Schatten selectors). Let γ=γp=∥⋅∥Spγ= _p= · _S_p, let r=2pr=2p, and let q=r−1=2p2p−1q= rr-1= 2p2p-1 be the dual exponent. If G=Udiag(σi)W⊤G=Udiag( _i)W is the singular value decomposition of a gradient matrix, then the empirical duality map is represented by Ξp(G)=−‖G‖Sq2−qUdiag(σiq−1)W⊤. _p(G)=- G _S_q^2-q\,Udiag( _i^q-1)W . Equivalently, if g is evaluated on the support of μX _X and stacked row-wise into the matrix G, then JμX(g)J_ _X(g) is obtained by reading the rows of Ξp(G) _p(G). In particular, • p=1p=1: Ξ1(G)=−G _1(G)=-G, which is the identity normalization and recovers ordinary gradient descent; • p=2p=2: Ξ2(G)=−‖G‖S4/32/3Udiag(σi1/3)W⊤ _2(G)=- G _S_4/3^2/3\,Udiag( _i^1/3)W , which we call the Frobenius-intermediate normalization; • p=∞p=∞: Ξ∞(G)=−‖G‖S1UW⊤=−tr((G⊤G)1/2)G(G⊤G)†/2, _∞(G)=- G _S_1\,UW =-tr((G G)^1/2)\,G\,(G G) /2, which is the Muon normalization. Proof. For empirical measures, the tangent norm is the Schatten r norm on the particle velocity matrix. The displayed formula is the negative gradient of 12‖G‖Sq2 12 G _S_q^2 when 1<q≤21<q≤ 2, while the endpoint q=1q=1 is interpreted through a subgradient selector. ∎ The following corollary turns the measure-valued flow into a finite-dimensional normalized particle dynamics. Corollary 4.5 (Finite-width interpretation). Assume the Spectral Wasserstein flow starting from μX0 _X_0 preserves empirical measures and can be written as μt=1n∑i=1nδxi(t). _t= 1n _i=1^n _x_i(t). Then the stacked particle matrix XtX_t solves X˙t=Ξp(∇Fn(Xt)) X_t= _p\! (∇ F_n(X_t) ) for Schatten-p geometries. Thus p=1p=1 gives the usual particle gradient flow and p=∞p=∞ gives the continuous-time single-block Muon flow. Proof. For empirical measures, the continuity equation reduces to the characteristic system for the atoms. Stacking the particle velocities yields exactly the matrix update of Proposition 4.4. ∎ 4.3 Gaussian-Preserving Gradient Flows This subsection isolates a simple class of spectral gradient flows that can still be analyzed rather explicitly. In general, understanding the long-time behavior of the flow is hard, but when Gaussian measures are preserved the infinite-dimensional evolution reduces to a finite-dimensional ODE on means and covariances. The next corollary explains why Gaussian preservation transfers from the classical 2 W_2 flow to the spectral flow whenever the active matrix remains invertible on the Gaussian class. Corollary 4.6 (From 2 W_2 Gaussian preservation to spectral Gaussian preservation). Assume that for every Gaussian state μ the classical Wasserstein gradient gμ(x)=∇xδfδμ(x)g_μ(x)= _x δ fδμ(x) is affine in x, and that there exists an invertible active matrix Qμ∗∈argmaxQ∈γtr(QSμ(gμ)).Q_μ^*∈ *argmax_Q _γtr(QS_μ(g_μ)). Then the spectral flow ∂tμt+div(μtJμt(gμt))=0 _t _t+div\! ( _tJ_ _t(g_ _t) )=0 preserves Gaussian measures. Proof. Fix a Gaussian state μ and write gμ(x)=bμ+Bμxg_μ(x)=b_μ+B_μx for the affine 2 W_2 gradient provided by the assumption. By Theorem 4.2, any invertible active matrix Qμ∗∈argmaxQ∈γtr(QSμ(gμ))Q_μ^*∈ *argmax_Q _γtr(QS_μ(g_μ)) defines a valid spectral selector through Jμ(gμ)(x)=−(Qμ∗)−1(bμ+Bμx).J_μ(g_μ)(x)=-(Q_μ^*)^-1(b_μ+B_μx). This is again an affine vector field in x. Therefore the spectral continuity equation is of the form vt(x)=at+Atxv_t(x)=a_t+A_tx with coefficients depending on the current Gaussian state. A continuity equation with affine drift preserves the Gaussian class, and the associated mean mtm_t and covariance Σt _t solve the closed ODE system m˙t=at+Atmt,Σ˙t=AtΣt+ΣtAt⊤. m_t=a_t+A_tm_t, _t=A_t _t+ _tA_t . Therefore Gaussian initial data remain Gaussian along the spectral flow. ∎ Two important examples fit into this mechanism. First, for the relative entropy Entν(μ)Ent_ν(μ) with Gaussian target ν, the classical 2 W_2 gradient flow is the Fokker–Planck equation, and on the Gaussian class its Wasserstein gradient is affine in x. Second, whenever f(μ)f(μ) depends only on the mean and covariance of μ, its first variation is a quadratic polynomial and its 2 W_2 gradient is affine, so the same transfer principle applies. This is in particular the case for training a linear two-layer network, obtained by taking σ=Idσ=Id in the MLP model of Section 6 together with a quadratic loss R(H)R(H): the resulting objective depends quadratically on the first two moments of μ. We do not pursue these Gaussian ODE reductions further here, but they provide one of the simplest regimes in which the spectral flow can be studied beyond the formal PDE level. 5 Geodesic Convexity This section studies one of the basic structural notions behind metric gradient flows, namely geodesic convexity, for the Spectral Wasserstein geometries introduced above. 5.1 Definition and General Setup This subsection recalls the relevant notion of convexity along Spectral Wasserstein geodesics and explains why it matters for the variational analysis of the flow. When one studies gradient flows in metric spaces, geodesic convexity plays the same role as ordinary convexity in Euclidean optimization: it governs uniqueness, stability, and the variational structure of the dynamics. Since Corollary 3.5 shows that γ W_γ-geodesics are displacement interpolations associated with optimal couplings, the corresponding convexity inequalities can be tested directly along Euclidean segments. The next definition fixes the notion used in the remainder of this section. Definition 5.1 (Geodesic convexity). A functional F:2(ℝd)→(−∞,+∞]F:P_2(R^d)→(-∞,+∞] is called κ-geodesically convex for γ W_γ if along every constant-speed γ W_γ-geodesic (μt)t∈[0,1]( _t)_t∈[0,1] one has F(μt)≤(1−t)F(μ0)+tF(μ1)−κ2t(1−t)γ(μ0,μ1)2.F( _t)≤(1-t)F( _0)+tF( _1)- κ2t(1-t) W_γ( _0, _1)^2. 5.2 Linear Functionals This subsection shows that for linear functionals the qualitative and quantitative convexity criteria reduce to explicit pointwise conditions on the integrand. For a Borel function h:ℝd→ℝh:R^d , define the linear functional Fh(μ)≔∫ℝdh(x)μ(x).F_h(μ) _R^dh(x)\,dμ(x). Because γ W_γ-geodesics are still given by affine interpolation in space, the convexity question for FhF_h is especially transparent. The next theorem gives the exact criterion. Theorem 5.2 (Linear functionals). Assume γ is monotone and let γK_γ be a convex compact PSD representing set for γ as in Proposition 2.2. 1. The functional FhF_h is geodesically convex for γ W_γ if and only if h is convex on ℝdR^d. 2. If in addition h∈C2(ℝd)h∈ C^2(R^d), then FhF_h is κ-geodesically convex for γ W_γ if and only if ∇2h(z)⪰κQfor every z∈ℝd,Q∈γ.∇^2h(z) κ Q every z ^d,\ Q _γ. Equivalently, ξ⊤∇2h(z)ξ≥κmaxQ∈γξ⊤Qξfor every z,ξ∈ℝd.ξ ∇^2h(z)\,ξ≥κ _Q _γξ Q\,ξ every z,ξ ^d. Proof. Assume first that h is convex and let μt=((1−t)x+ty)#π∗ _t=((1-t)x+ty)_\# _* be any constant-speed γ W_γ-geodesic. Then Fh(μt)=∫h((1−t)x+ty)π∗(x,y).F_h( _t)= h((1-t)x+ty)\,d _*(x,y). By convexity of h, h((1−t)x+ty)≤(1−t)h(x)+th(y),h((1-t)x+ty)≤(1-t)h(x)+th(y), and integrating yields geodesic convexity of FhF_h. Conversely, if FhF_h is geodesically convex, apply the definition to Dirac masses μ0=δx _0= _x and μ1=δy _1= _y. The unique geodesic is μt=δ(1−t)x+ty, _t= _(1-t)x+ty, so h((1−t)x+ty)=Fh(μt)≤(1−t)h(x)+th(y),h((1-t)x+ty)=F_h( _t)≤(1-t)h(x)+th(y), which proves convexity of h. Assume now that h∈C2h∈ C^2 and that ∇2h(z)⪰κQfor every z∈ℝd,Q∈γ.∇^2h(z) κ Q every z ^d,\ Q _γ. Let μt=((1−t)x+ty)#π∗ _t=((1-t)x+ty)_\# _* be any constant-speed γ W_γ-geodesic, write zt=(1−t)x+tyz_t=(1-t)x+ty and Δ=y−x =y-x. Differentiating under the integral sign gives d2dt2Fh(μt)=∫Δ⊤∇2h(zt)Δπ∗(x,y). d^2dt^2F_h( _t)= ∇^2h(z_t) \,d _*(x,y). For every Q∈γQ _γ this implies d2dt2Fh(μt)≥κ∫Δ⊤QΔπ∗(x,y). d^2dt^2F_h( _t)≥κ Q \,d _*(x,y). Taking the maximum over Q∈γQ _γ yields d2dt2Fh(μt)≥κγ(∫ΔΔ⊤π∗(x,y))=κγ(μ0,μ1)2. d^2dt^2F_h( _t)≥κγ\! ( \,d _*(x,y) )=κ W_γ( _0, _1)^2. Integrating twice in t gives the κ-geodesic convexity estimate. Conversely, assume FhF_h is κ-geodesically convex. Testing the definition on the Dirac geodesic between δz _z and δz+ξ _z+ξ gives h(z+tξ)≤(1−t)h(z)+th(z+ξ)−κ2t(1−t)γ(ξξ⊤).h(z+tξ)≤(1-t)h(z)+th(z+ξ)- κ2t(1-t)γ(ξ ). Differentiating at second order in t yields ξ⊤∇2h(z)ξ≥κγ(ξξ⊤)=κmaxQ∈γξ⊤Qξ.ξ ∇^2h(z)\,ξ≥κγ(ξ )=κ _Q _γξ Q\,ξ. This is equivalent to the matrix inequality ∇2h(z)⪰κQfor every Q∈γ.∇^2h(z) κ Q every Q _γ. ∎ The next remark interprets the condition of Theorem 5.2 for the Schatten family. Remark 5.3 (Schatten norms). For the Schatten geometries used throughout the paper, the condition ∇2h(z)⪰κQfor every Q∈γ∇^2h(z) κ Q every Q _γ is equivalent to ∇2h(z)⪰κId.∇^2h(z) . Indeed, for p=1p=1 we may choose γ=IdK_γ=\Id\. For 1<p≤∞1<p≤∞, a convenient choice is γ=Q⪰0:‖Q‖Sq≤1,q=p−1.K_γ=\Q 0:\ Q _S_q≤ 1\, q= pp-1. Every such Q satisfies λmax(Q)≤1 _ (Q)≤ 1, hence Q⪯IdQ , so ∇2h(z)⪰κId∇^2h(z) implies ∇2h(z)⪰κQ∇^2h(z) κ Q for all Q∈γQ _γ. Conversely, every rank-one projector uu⊤u belongs to γK_γ, so the condition for all Q∈γQ _γ implies u⊤∇2h(z)u≥κfor every unit u,u ∇^2h(z)u≥κ every unit u, that is, ∇2h(z)⪰κId∇^2h(z) . Therefore the criterion is exactly the same as for the classical 2 W_2 geometry. 5.3 Relative Entropy This subsection studies the geodesic convexity of relative entropy, which is the basic nonlinear example behind diffusion-type gradient flows. Let dν(x)=Z−1e−V(x)dx,V∈C2(ℝd),dν(x)=Z^-1e^-V(x)\,dx, V∈ C^2(R^d), and define the relative entropy Entν(μ)≔∫log(dμdν)μ,μ≪ν,+∞,otherwise.Ent_ν(μ) cases \! ( dμdν )\,dμ,&μ ν,\\[6.00006pt] +∞,&otherwise. cases The trace case is the classical displacement-convexity theory of entropy. For general spectral geometries the same argument works only when the active quadratic costs remain uniformly elliptic, and otherwise one only gets a weaker statement by approximation. The next theorem records the full all-geodesics result in the uniformly elliptic regime. Theorem 5.4 (Full entropy convexity under uniform ellipticity). Assume γ is monotone and that γK_γ can be chosen so that γ⊂++d.K_γ _++^d. Assume moreover that ∇2V(x)⪰κQfor every x∈ℝd,Q∈γ.∇^2V(x) κ Q every x ^d,\ Q _γ. Then EntνEnt_ν is κ-geodesically convex on (2(ℝd),γ)(P_2(R^d), W_γ). Proof. Let μt=((1−t)x+ty)#π∗ _t=((1-t)x+ty)_\# _* be any constant-speed γ W_γ-geodesic, and set S∗=∫(y−x)(y−x)⊤π∗(x,y).S_*= (y-x)(y-x) \,d _*(x,y). Choose Q∗∈γQ_* _γ such that tr(Q∗S∗)=γ(S∗)=γ(μ0,μ1)2.tr(Q_*S_*)=γ(S_*)= W_γ( _0, _1)^2. By Theorem 2.10, the same coupling π∗ _* is optimal for the quadratic cost cQ∗(x,y)=(y−x)⊤Q∗(y−x),c_Q_*(x,y)=(y-x) Q_*(y-x), and Q∗(μ0,μ1)2=γ(μ0,μ1)2. W^Q_*( _0, _1)^2= W_γ( _0, _1)^2. Since Q∗Q_* is positive definite, let L=Q∗1/2L=Q_*^1/2 and define μ~i=L#μi,ν~=L#ν,π~∗=(L,L)#π∗. μ_i=L_\# _i, ν=L_\#ν, π_*=(L,L)_\# _*. Then π~∗ π_* is optimal for the Euclidean quadratic cost between μ~0 μ_0 and μ~1 μ_1, hence μ~t=((1−t)z0+tz1)#π~∗ μ_t=((1-t)z_0+tz_1)_\# π_* is a 2 W_2-geodesic and 2(μ~0,μ~1)2=γ(μ0,μ1)2. W_2( μ_0, μ_1)^2= W_γ( _0, _1)^2. The reference measure ν~ ν has density dν~(z)=Z~−1e−V~(z)dz,V~(z)=V(L−1z)+constant,d ν(z)= Z^-1e^- V(z)\,dz, V(z)=V(L^-1z)+constant, so ∇2V~(z)=L−1∇2V(L−1z)L−1.∇^2 V(z)=L^-1∇^2V(L^-1z)L^-1. Because Q∗∈γQ_* _γ, the curvature assumption gives ∇2V(x)⪰κQ∗=κL2,∇^2V(x) κ Q_*=κ L^2, hence ∇2V~(z)⪰κId.∇^2 V(z) . The classical entropy convexity theorem for 2 W_2 therefore yields Entν~(μ~t)≤(1−t)Entν~(μ~0)+tEntν~(μ~1)−κ2t(1−t)2(μ~0,μ~1)2.Ent_ ν( μ_t)≤(1-t)Ent_ ν( μ_0)+tEnt_ ν( μ_1)- κ2t(1-t) W_2( μ_0, μ_1)^2. Finally, relative entropy is invariant under the invertible change of variables L, so Entν~(μ~t)=Entν(μt),Ent_ ν( μ_t)=Ent_ν( _t), and similarly at the endpoints. This gives the result. ∎ The next theorem explains what remains true in the more relevant but singular regime where the condition γ⊂++dK_γ _++^d fails. This lack of ellipticity is precisely what happens for Schatten-p geometries with p>1p>1, because the corresponding representing sets contain singular matrices, but one still retains a weaker geodesic-convexity statement by approximation. Theorem 5.5 (Weak entropy convexity). Assume γ is monotone, that γ⊂+dK_γ _+^d is convex compact, and that γ∩++d≠∅.K_γ _++^d≠ . Assume moreover that ∇2V(x)⪰κQfor every x∈ℝd,Q∈γ.∇^2V(x) κ Q every x ^d,\ Q _γ. Then for every μ0,μ1∈2(ℝd) _0, _1 _2(R^d) there exists at least one constant-speed γ W_γ-geodesic (μt)t∈[0,1]( _t)_t∈[0,1] from μ0 _0 to μ1 _1 such that Entν(μt)≤(1−t)Entν(μ0)+tEntν(μ1)−κ2t(1−t)γ(μ0,μ1)2.Ent_ν( _t)≤(1-t)Ent_ν( _0)+tEnt_ν( _1)- κ2t(1-t) W_γ( _0, _1)^2. Proof. Choose Q0∈γ∩++dQ_0 _γ _++^d and define γ,ε≔(1−ε)γ+εQ0,ε∈(0,1).K_γ, (1- )K_γ+ Q_0, ∈(0,1). Then every element of γ,εK_γ, is positive definite. Let γε _ be the support function of γ,εK_γ, , namely γε(S)=maxQ∈γ,εtr(QS)=(1−ε)γ(S)+εtr(Q0S). _ (S)= _Q _γ, tr(QS)=(1- )γ(S)+ (Q_0S). Hence (1−ε)γ(S)≤γε(S)≤γ(S)(1- )γ(S)≤ _ (S)≤γ(S) for all S⪰0S 0, and therefore (1−ε)γ(μ0,μ1)2≤γε(μ0,μ1)2≤γ(μ0,μ1)2.(1- ) W_γ( _0, _1)^2≤ W_ _ ( _0, _1)^2≤ W_γ( _0, _1)^2. By Theorem 5.4, each regularized geometry has full entropy convexity. Choose an optimal coupling πε _ for γε W_ _ and set μtε=((1−t)x+ty)#πε. _t =((1-t)x+ty)_\# _ . Then Entν(μtε)≤(1−t)Entν(μ0)+tEntν(μ1)−κ2t(1−t)γε(μ0,μ1)2.Ent_ν( _t )≤(1-t)Ent_ν( _0)+tEnt_ν( _1)- κ2t(1-t) W_ _ ( _0, _1)^2. By compactness of couplings, one may extract a subsequence πεn⇀π∗ _ _n _*. Setting μt=((1−t)x+ty)#π∗, _t=((1-t)x+ty)_\# _*, the lower semicontinuity of relative entropy and the convergence γεn(μ0,μ1)2→γ(μ0,μ1)2 W_ _ _n( _0, _1)^2→ W_γ( _0, _1)^2 yield the claimed inequality along the limit geodesic. The limit coupling π∗ _* is γ W_γ-optimal by continuity of the displacement covariance and of the support functions on trace-bounded sets. ∎ The following remark explains the scope of the entropy results for the Schatten family. Remark 5.6 (Schatten norms). The full theorem applies to the trace case p=1p=1, where one may choose γ=IdK_γ=\Id\ and recover the classical 2 W_2 displacement-convexity theory. For every Schatten geometry with 1<p≤∞1<p≤∞, the natural representing sets contain singular matrices, so the uniformly elliptic argument above does not apply directly. Theorem 5.5 nevertheless yields a weak geodesic-convexity statement as soon as γK_γ contains one positive definite matrix, which is the case for the standard Schatten choices. 6 Spectral Flow for Training MLPs and Unbalanced Formulation This section explains how the abstract Spectral Wasserstein flow specializes to two-layer MLPs and then how positively two-homogeneous models reduce the ambient transport problem to an unbalanced transport problem on the sphere. The functionals arising in this setting are typically pairwise interaction energies, and geodesic convexity usually fails for such energies even in the classical 2 W_2 geometry except in rather special situations, so proving global convergence to equilibrium is in general difficult. 6.1 Mean-Field Models for 2 Layers MLPs This subsection connects the abstract normalized flow to the standard mean-field parameterization of two-layer MLPs, which is the main modeling bridge to Muon. Consider the mean-field predictor Hμ(z)=∫Ωϕ(z,x)μ(x),f(μ)=R(Hμ).H_μ(z)= _ φ(z,x)\,dμ(x), f(μ)=R(H_μ). For a two-layer network with scalar output and activation σ, one may take parameter x=(u,v)∈ℝd,ϕ(z,x)=uσ(v⋅z).parameter x=(u,v) ^d, φ(z,x)=u\,σ(v· z). This is the standard mean-field representation of a two-layer multilayer perceptron. The notation is a deliberate departure from usual machine-learning conventions: the network input is denoted by z, while the trainable variable is denoted by x so that it matches the transport notation used throughout the paper. Here d is the sum of the block dimensions of u and v. In practical implementations of Muon one often normalizes separate parameter blocks, whereas in our continuum model the whole particle x is treated as a single vector; this is the natural mean-field analogue of a single normalized block. A basic example is the quadratic risk R(H)=12∫|H(z)−y⋆(z)|2ρ(z),R(H)= 12 H(z)-y (z) ^2\,dρ(z), or its empirical version on a dataset. Because HμH_μ depends linearly on μ, this makes f(μ)f(μ) a quadratic interaction functional. In particular, kernelized choices of ϕφ recover interaction energies of MMD type, which is why the MMD experiment of Section 7 is a natural test case for the Spectral Wasserstein flow. Under the regularity assumptions stated below, the characteristic construction shows that empirical initial data remain empirical. Consequently, for discrete measures the Spectral Wasserstein flow is exactly equivalent to the corresponding normalized finite-dimensional flow on X, just as classical 2 W_2 flow corresponds to particle gradient descent and the operator-norm flow corresponds to Muon. Remark 6.1 (A caveat on global existence). The previous discussion identifies the PDE ∂tμt+div(μtJμt(gμt))=0 _t _t+div\! ( _tJ_ _t(g_ _t) )=0 as the natural normalized transport flow associated with the spectral geometry. Proving global existence and uniqueness of solutions by characteristics, however, requires quantitative control of the selected duality map (μ,g)↦Jμ(g)(μ,g) J_μ(g) along the trajectory of the PDE: one needs at least a measurable selection of minimizers in Definition 4.1, together with Lipschitz dependence on (μ,g)(μ,g) and linear growth in the space variable. These requirements are classical in the trace case p=1p=1, since then Jμ(g)=−gJ_μ(g)=-g and the PDE reduces to the usual mean-field Wasserstein gradient-flow equation studied for neural networks, for instance by Chizat and Bach [9], Mei et al. [16]. In that regime, well-posedness boils down to the familiar smoothness, Lipschitz, and growth assumptions on the feature map ϕφ ensuring that the induced Wasserstein gradient field has at most linear growth and sufficient regularity in the particle variable. For general Schatten geometries they become substantially less obvious. Even though Proposition 4.4 gives explicit formulas for the matrix selectors, obtaining uniform Lipschitz bounds for the induced field Jμ(g)J_μ(g) along the nonlinear PDE is delicate for p>1p>1, especially near singular-value multiplicities, rank changes, or degeneracies of the empirical covariance. For this reason, the present paper treats the PDE as a formal metric gradient flow and does not claim a general global existence theorem beyond regimes where such regularity estimates can be verified separately. 6.2 The Positively 22-Homogeneous Case and a Generalized Unbalanced Geometry This subsection explains why positively two-homogeneous models admit a spherical reduction and why the induced spherical dynamics is naturally unbalanced. Let f(μ)=R(∫ℝdΦ(x)μ(x)),Φ(λx)=λ2Φ(x).f(μ)=R\! ( _R^d (x)\,dμ(x) ), (λ x)=λ^2 (x). Writing x=rωx=rω with ω∈d−1ω ^d-1, define the weighted spherical projection ∫d−1ψ(ω)Π2(μ)(ω)≔∫ℝd|x|2ψ(x|x|)μ(x). _S^d-1ψ(ω)\,d _2(μ)(ω) _R^d x ^2ψ\! ( x x )\,dμ(x). The first point is that two-homogeneous models only depend on this weighted spherical projection. Proposition 6.2 (Exact quotient). There exists a functional f¯ f on ℳ+(d−1)M_+(S^d-1) such that f(μ)=f¯(Π2(μ)).f(μ)= f( _2(μ)). Proof. The homogeneity identity Φ(rω)=r2Φ(ω) (rω)=r^2 (ω) immediately gives ∫ℝdΦ(x)μ(x)=∫d−1Φ(ω)Π2(μ)(ω), _R^d (x)\,dμ(x)= _S^d-1 (ω)\,d _2(μ)(ω), which defines f¯ f. ∎ The second point is that the ambient normalized velocity splits into radial and tangential parts on the sphere. If a vector field v is 11-homogeneous, namely v(λx)=λv(x)v(λ x)=λ v(x) for every λ>0λ>0, then for every x=rω≠0x=rω≠ 0 it can be written uniquely as v(rω)=r(b(ω)ω+τ(ω)),τ(ω)∈Tωd−1.v(rω)=r (b(ω)ω+τ(ω) ), τ(ω)∈ T_ωS^d-1. Indeed, one defines b(ω)≔v(ω)⋅ω,τ(ω)≔v(ω)−(v(ω)⋅ω)ω,b(ω) v(ω)·ω, τ(ω) v(ω)- (v(ω)·ω )ω, so that τ(ω)⋅ω=0τ(ω)·ω=0, and then 11-homogeneity gives v(rω)=rv(ω)=r(b(ω)ω+τ(ω)).v(rω)=r\,v(ω)=r (b(ω)ω+τ(ω) ). The decomposition is unique because it is simply the orthogonal splitting of v(ω)v(ω) into its normal and tangential parts on the sphere. In the positively two-homogeneous setting considered here, the first variation is 22-homogeneous and its spatial gradient is therefore 11-homogeneous; correspondingly, the natural steepest-descent velocity fields for the flow belong to this 11-homogeneous class. The next proposition computes the projected spherical PDE. Proposition 6.3 (Projected continuity-reaction equation). If μt _t solves the Spectral Wasserstein flow and νt=Π2(μt) _t= _2( _t), then ∂tνt+divd−1(νtτt)=2btνt. _t _t+div_S^d-1( _t _t)=2b_t _t. Proof. Test the ambient continuity equation against the lifted observable ψ~(rω)=r2ψ(ω) ψ(rω)=r^2ψ(ω) and identify the radial and tangential contributions. ∎ Motivated by Proposition 6.3, define for nonnegative measures ν0,ν1 _0, _1 on d−1S^d-1 γ(ν0,ν1)2≔inf(νt,bt,τt)∫01γ(∫d−1(bt(ω)ω+τt(ω))(bt(ω)ω+τt(ω))⊤νt(ω))t, UW_γ( _0, _1)^2 _( _t,b_t, _t) _0^1γ\! ( _S^d-1(b_t(ω)ω+ _t(ω))(b_t(ω)ω+ _t(ω)) \,d _t(ω) )\,dt, under the continuity-reaction constraint ∂tνt+divd−1(νtτt)=2btνt. _t _t+div_S^d-1( _t _t)=2b_t _t. The next proposition identifies the ambient homogeneous action with this spherical unbalanced action. Proposition 6.4 (Ambient action equals spherical action). For positively two-homogeneous models and 11-homogeneous velocities, the ambient Spectral Wasserstein action equals the spherical action defining γ UW_γ. Proof. Write x=rωx=rω and v(rω)=r(b(ω)ω+τ(ω))v(rω)=r(b(ω)ω+τ(ω)). By definition of Π2(μ) _2(μ), ∫ℝdv(x)v(x)⊤μ(x)=∫d−1(b(ω)ω+τ(ω))(b(ω)ω+τ(ω))⊤Π2(μ)(ω). _R^dv(x)v(x) \,dμ(x)= _S^d-1(b(ω)ω+τ(ω))(b(ω)ω+τ(ω)) \,d _2(μ)(ω). Inserting this identity into the Benamou–Brenier action yields the claim. ∎ The following remark explains how the trace-norm case collapses to the classical Wasserstein–Fisher–Rao geometry. Remark 6.5. When γ(S)=tr(S)γ(S)=tr(S), the integrand becomes ∫d−1(|τt(ω)|2+bt(ω)2)νt(ω). _S^d-1 ( _t(ω) ^2+b_t(ω)^2 )\,d _t(ω). Since the reaction rate is αt=2bt _t=2b_t, this is exactly the classical Wasserstein–Fisher–Rao action. For general γ, one obtains a genuinely new unbalanced transport geometry on the sphere. A static counterpart of this spherical geometry is an interesting open problem. 7 Numerical Experiments This section presents two complementary numerical illustrations: one for the static coupling problem and one for the associated gradient flows. 7.1 Spectral Optimal-Transport Couplings This subsection compares the discrete Monge-type couplings selected by the trace, Frobenius, and operator norms. We consider two discrete measures μ=∑i=1naiδxi,ν=∑j=1mbjδyj,μ= _i=1^na_i _x_i, ν= _j=1^mb_j _y_j, with weights ai,bj≥0a_i,b_j≥ 0 summing to one. As in classical Kantorovich transport, a coupling is represented by a transport matrix P∈ℝ+n×mP _+^n× m satisfying Pm=a,P⊤n=b.P1_m=a, P 1_n=b. The discrete static Spectral Wasserstein problem can then be written as the convex optimization problem minP≥0,Pm=a,P⊤n=bγp(∑i=1n∑j=1mPij(yj−xi)(yj−xi)⊤). _P≥ 0,\;P1_m=a,\;P 1_n=b _p\! ( _i=1^n _j=1^mP_ij(y_j-x_i)(y_j-x_i) ). For p=1p=1 this reduces to the usual linear optimal-transport problem with quadratic cost, hence to a classical assignment problem in the equal-weight case. For p=2p=2 the objective is a convex second-order-cone expression, and for p=∞p=∞ it is a convex semidefinite-representable objective through the maximal eigenvalue. In the numerical code we solve the p=2p=2 and p=∞p=∞ cases with CVXPY. To make the static and dynamic experiments directly comparable, we use the same anisotropic Gaussian source and farther-away Gaussian-mixture target as in the MMD flow experiment below, with n=mn=m and uniform weights ai=bj=1/na_i=b_j=1/n. Figure 1: Static spectral couplings for Schatten p=1,2,∞p=1,2,∞. Red points are the source cloud, blue points are the target cloud, and black segments show a permutation extracted from the optimal coupling for visualization. For consistency with the flow experiment below, we use exactly the same empirical source and target clouds as in the MMD experiment, only viewed now through the static coupling problem. Figure 1 shows how the optimal coupling changes with p. 7.2 MMD Gradient Flows This subsection compares the Spectral Wasserstein gradient flows of the same MMD objective for the three benchmark Schatten norms. We minimize f(μ)=MMD(μ,ν)2f(μ)=MMD(μ,ν)^2 with the energy-distance kernel k(x,y)=−‖x−y‖2.k(x,y)=- x-y _2. For empirical measures μ=1n∑i=1nδxi,ν=1m∑j=1mδyj,μ= 1n _i=1^n _x_i, ν= 1m _j=1^m _y_j, this becomes MMD(μ,ν)2=1n2∑i,i′k(xi,xi′)+1m2∑j,j′k(yj,yj′)−2nm∑i,jk(xi,yj).MMD(μ,ν)^2= 1n^2 _i,i k(x_i,x_i )+ 1m^2 _j,j k(y_j,y_j )- 2nm _i,jk(x_i,y_j). We use n=m=200n=m=200 points in dimension 22, with a farther-away Gaussian-mixture target and an anisotropic Gaussian source. The kernel is smoothed by ‖x−y‖ε=‖x−y‖22+ε2,ε=10−2. x-y _ = x-y _2^2+ ^2, =10^-2. If k∈ℝn×2X_k ^n× 2 is the moving cloud and Gk=∇f(k)G_k= _Xf(X_k), then the three explicit Euler flows are k+1=k+ηpΞp(Gk),p∈1,2,∞.X_k+1=X_k+ _p _p(G_k), p∈\1,2,∞\. The case p=1p=1 is the Euclidean / 2 W_2 flow, p=2p=2 is the Frobenius-intermediate flow, and p=∞p=∞ is the Muon / operator-norm flow. Figure 2: All particle trajectories for the three MMD flows associated with Schatten p=1,2,∞p=1,2,∞. The operator-norm flow is the most globally coordinated, the trace-norm flow is the most local, and the Frobenius flow interpolates between the two. Figure 2 makes the geometry visible. The trace-norm flow reacts locally to the force field, the operator-norm flow concentrates on coherent dominant directions, and the Frobenius flow interpolates between them. The three dynamics optimize the same functional; only the normalized transport geometry changes. 8 Conclusion The main message of this paper is that matrix-normalized optimizers for mean-field neural models are naturally encoded by a family of Spectral Wasserstein distances. A norm γ on positive semidefinite matrices determines a static covariance cost, and for the monotone class that includes all Schatten norms it also determines a dynamic Benamou–Brenier action and a metric gradient flow on measures. The trace norm recovers the classical quadratic geometry, the operator norm recovers Muon, and intermediate Schatten norms interpolate between them. This perspective clarifies both transport and optimization. On the transport side, the geometry is genuinely coupling-based, admits explicit geodesics in the monotone regime, and yields a Gaussian covariance metric extending the Bures formula. On the optimization side, it turns normalized matrix updates into exact steepest descents for a measure-valued metric and provides a continuum interpretation of finite-dimensional normalized training rules. Several directions remain open. A sharper characterization of optimal couplings beyond the conditional Brenier regime would be valuable. On the optimization side, the spherical reduction of positively homogeneous models suggests a spectral unbalanced transport geometry for every γ, but extending the Chizat–Bach global convergence theory beyond the trace norm remains future work. Another natural direction is to incorporate the genuinely blockwise geometries used in practical optimizers. Acknowledgement This work was supported by the European Research Council (ERC project WOLF) and the French government under the management of Agence Nationale de la Recherche as part of the “France 2030” program, reference ANR-23-IACL-0008 (PRAIRIE-PSAI). References [1] L. Ambrosio, N. Gigli, and G. Savaré (2008) Gradient flows in metric spaces and in the space of probability measures. 2 edition, Lectures in Mathematics ETH Zürich, Birkhäuser Basel. Cited by: §1.1. [2] J. D. Backhoff-Veraguas, M. Beiglböck, and G. Pammer (2019) Existence, duality, and cyclical monotonicity for weak transport costs. Calculus of Variations and Partial Differential Equations 58 (6), p. 203. Cited by: §1.1. [3] J. D. Backhoff-Veraguas and G. Pammer (2022) Applications of weak transport theory. Bernoulli 28 (1), p. 370–394. Cited by: §1.1. [4] J. Benamou and Y. Brenier (2000) A computational fluid mechanics solution to the Monge–Kantorovich mass transfer problem. Numerische Mathematik 84 (3), p. 375–393. Cited by: §1.1, §3. [5] R. Bhatia, T. Jain, and Y. Lim (2019) On the Bures–Wasserstein distance between positive definite matrices. Expositiones Mathematicae 37 (2), p. 165–191. Cited by: §1.1. [6] M. Burger, M. Erbar, F. Hoffmann, D. Matthes, and A. Schlichting (2025) Covariance-modulated optimal transport and gradient flows. Archive for Rational Mechanics and Analysis 249 (1), p. 7. Cited by: §1.1. [7] G. Carlier, C. Jimenez, and F. Santambrogio (2008) Optimal transportation with traffic congestion and Wardrop equilibria. SIAM Journal on Control and Optimization 47 (3), p. 1330–1350. Cited by: §1.1. [8] Y. Chen, T. T. Georgiou, and A. Tannenbaum (2018) Matrix optimal mass transport: a quantum mechanical approach. IEEE Transactions on Automatic Control 63 (8), p. 2612–2619. Cited by: §1.1. [9] L. Chizat and F. Bach (2018) On the global convergence of gradient descent for over-parameterized models using optimal transport. In Advances in Neural Information Processing Systems, Vol. 31, p. 3040–3050. Cited by: §1.1, Remark 6.1. [10] A. Cutkosky and H. Mehta (2020) Momentum improves normalized SGD. In Proceedings of the 37th International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 119, p. 2260–2268. Cited by: §1.1. [11] M. Cuturi and D. Avis (2014) Ground metric learning. Journal of Machine Learning Research 15 (1), p. 533–564. Cited by: §1.1. [12] R. Flamary, M. Cuturi, N. Courty, and A. Rakotomamonjy (2018) Wasserstein discriminant analysis. Machine Learning 107 (12), p. 1923–1945. Cited by: §1.1. [13] N. Gozlan, C. Roberto, P. Samson, and P. Tetali (2017) Kantorovich duality for general transport costs and applications. Journal of Functional Analysis 273 (11), p. 3327–3405. Cited by: §1.1. [14] K. Jordan, Y. Jin, V. Boza, J. You, F. Cesista, L. Newhouse, and J. Bernstein (2024) Muon: an optimizer for hidden layers in neural networks. Note: Blog post Cited by: §1.1, §1. [15] J. Liu, J. Su, X. Yao, Z. Jiang, G. Lai, Y. Du, Y. Qin, W. Xu, E. Lu, J. Yan, Y. Chen, H. Zheng, Y. Liu, S. Liu, B. Yin, W. He, H. Zhu, Y. Wang, J. Wang, M. Dong, Z. Zhang, Y. Kang, H. Zhang, X. Xu, Y. Zhang, Y. Wu, X. Zhou, and Z. Yang (2025) Muon is scalable for LLM training. External Links: 2502.16982 Cited by: §1.1. [16] S. Mei, A. Montanari, and P. Nguyen (2018) A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences 115 (33), p. E7665–E7671. Cited by: §1.1, Remark 6.1. [17] R. Murray, B. Swenson, and S. Kar (2019) Revisiting normalized gradient descent: fast evasion of saddle points. IEEE Transactions on Automatic Control 64 (11), p. 4818–4824. Cited by: §1.1. [18] L. Ning, T. T. Georgiou, and A. Tannenbaum (2015) On matrix-valued Monge–Kantorovich optimal mass transport. IEEE Transactions on Automatic Control 60 (2), p. 373–382. Cited by: §1.1. [19] F. Paty and M. Cuturi (2019) Subspace robust wasserstein distances. In Proceedings of the 36th International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 97, p. 5072–5081. Cited by: §1.1. [20] T. Pethick, W. Xie, K. Antonakopoulos, Z. Zhu, A. Silveti-Falls, and V. Cevher (2025) Training deep learning models with norm-constrained LMOs. In Proceedings of the 42nd International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 267, p. 49069–49104. Cited by: §1.1. [21] O. Sebbouh, M. Cuturi, and G. Peyré (2024) Structured transforms across spaces with cost-regularized optimal transport. In Proceedings of The 27th International Conference on Artificial Intelligence and Statistics, Proceedings of Machine Learning Research, Vol. 238, p. 586–594. Cited by: §1.1.