← Forthcoming posts

Alpha preview · Work in progress

An early version of a forthcoming blog post. The text, proofs, and examples may still evolve.

Spectral Wasserstein distances

Keep the directions of motion, then choose how to measure them.

Quadratic optimal transport minimizes the average squared distance traveled by mass. But two movements can have the same average energy and very different geometry: one may spread its effort evenly across directions, while the other concentrates it along a single axis. The usual Wasserstein distance adds all these contributions together.

In Muon Dynamics as a Spectral Wasserstein Flow [1], I keep the entire matrix of displacement second moments and apply a matrix gauge to it. This gives the spectral Wasserstein distance \(\mathsf W_\gamma\). It connects a static coupling problem, an adversary choosing a quadratic ground cost, and a Benamou–Brenier action for moving distributions.

The construction extends the subspace robust Wasserstein distances of François-Pierre Paty and Marco Cuturi [2]. We will prove the robust and dynamic formulations, then explore the Schatten family numerically. There is a useful subtlety: the distance and dynamic formula remain meaningful down to \(p=1/2\), even though the convex robust-cost interpretation stops at \(p=1\).

Replace a scalar transport cost by a displacement matrix

Let \(\mu,\nu\in\mathcal P_2(\mathbb R^d)\): probability measures with finite second moments. A coupling \(\pi\in\Pi(\mu,\nu)\) is a joint law of \((X,Y)\) with these marginals. Define

\[ S_\pi=\mathbb E[(Y-X)(Y-X)^\top]. \]

This positive semidefinite matrix records directional mean-square displacement: for a unit vector \(u\), \(u^\top S_\pi u=\mathbb E[\langle u,Y-X\rangle^2]\). Although often called the displacement covariance, it is an uncentered second moment: motion of the mean is included.

Given a continuous, positive one-homogeneous gauge \(\gamma\), positive on nonzero positive semidefinite matrices, set

\[ \mathsf W_\gamma(\mu,\nu)^2 =\min_{\pi\in\Pi(\mu,\nu)}\gamma(S_\pi). \]

For the main theory, \(\gamma\) is convex and monotone in the Loewner order: \(0\preceq A\preceq B\) implies \(\gamma(A)\le\gamma(B)\). The monotone matrix norms in the paper satisfy these assumptions. The minimum exists: couplings with fixed marginals are weakly compact, and their displacement matrices vary continuously under weak convergence, since those marginals have fixed finite second moments.

The Schatten family is especially transparent. If \(\lambda_1,\ldots,\lambda_d\) are the eigenvalues of \(S\succeq0\), define

\[ \gamma_p(S)=\left(\sum_i\lambda_i^p\right)^{1/p}, \qquad 0\lt p\lt \infty, \] \[ \gamma_\infty(S)=\lambda_{\max}(S). \]

At \(p=1\), \(\gamma_1(S)=\operatorname{tr}S\), so \(\mathsf W_{\gamma_1}=\mathsf W_2\). At \(p=2\) we measure the Frobenius norm of the displacement matrix. At \(p=\infty\), the cost is the largest directional mean-square displacement. For now take \(p\ge1\); we return to smaller exponents below. The subscript in \(\gamma_p\) indexes a matrix gauge, not the exponent of the usual scalar \(p\)-Wasserstein cost.

All these gauges agree on a rank-one matrix. Thus two Dirac masses have the same Euclidean distance for every \(p\). Their difference appears when motion uses several directions. For instance, the unique coupling from \(\delta_0\) to a centered law with second-moment matrix \(\sigma^2I_d\) gives

\[ \mathsf W_{\gamma_p}^2=\sigma^2d^{1/p}, \qquad \mathsf W_{\gamma_\infty}^2=\sigma^2. \]

A robust quadratic cost

A convex monotone gauge has a positive-semidefinite support representation [1]:

\[ \gamma(S)=\max_{Q\in K_\gamma}\operatorname{tr}(QS), \qquad S\succeq0, \]

where \(K_\gamma\) is compact and convex, and every \(Q\in K_\gamma\) is positive semidefinite. One may take the set of all such \(Q\) satisfying \(\operatorname{tr}(QS)\le\gamma(S)\) for every \(S\succeq0\). The explicit sets below suffice for all our examples.

For Schatten \(p\ge1\), Hölder duality gives

\[ K_p=\{Q\succeq0:\lVert Q\rVert_{S_q}\le1\}, \qquad \frac1p+\frac1q=1. \]

At \(p=1\), this is \(0\preceq Q\preceq I\), and the single matrix \(I\) already represents the trace. At \(p=\infty\), it is \(Q\succeq0\) with \(\operatorname{tr}Q\le1\); matrices of trace exactly one also suffice.

Each \(Q\) defines a quadratic ground cost and its transport value:

\[ c_Q(x,y)=(y-x)^\top Q(y-x), \qquad \mathsf W_Q^2=\min_{\pi\in\Pi(\mu,\nu)}\int c_Q\,d\pi. \]

It stretches and compresses directions before measuring transport. More precisely, \(\mathsf W_Q=\mathsf W_2((Q^{1/2})_\#\mu,(Q^{1/2})_\#\nu)\), where \(T_\#\mu\) denotes the law of \(T(X)\) for \(X\sim\mu\). Singular \(Q\) may erase some directions.

Theorem 1 (cost-robust reformulation). Under the support-representation assumptions above,

\[ \mathsf W_\gamma(\mu,\nu)^2 =\max_{Q\in K_\gamma}\mathsf W_Q(\mu,\nu)^2. \]

The transport plan minimizes its cost while an adversary selects the most demanding admissible quadratic geometry. The matrix is shared by all transported pairs.

Proof: a support function and a minimax exchange

For \(F(\pi,Q)=\operatorname{tr}(QS_\pi)\), the support formula gives

\[ \mathsf W_\gamma^2=\min_\pi\max_{Q\in K_\gamma}F(\pi,Q). \]

The coupling set and \(K_\gamma\) are convex and compact, and \(F\) is continuous and affine in each variable. Sion's minimax theorem [3] exchanges the two optimizations:

\[ \min_\pi\max_Q F(\pi,Q) =\max_Q\min_\pi F(\pi,Q) =\max_Q\mathsf W_Q^2. \]
Why finite second moments suffice for continuity

The marginal terms \(\mathbb E[XX^\top]\) and \(\mathbb E[YY^\top]\) are fixed, so only cross moments vary with the coupling. Truncating a cross term makes it bounded and continuous. The omitted tail is uniformly small, since

\[ \mathbb E[|X||Y|\mathbf1_{|X|>R}] \le\sqrt{\mathbb E[|X|^2\mathbf1_{|X|>R}] \mathbb E|Y|^2}, \]

and similarly with \(X,Y\) exchanged. These expectations depend only on the fixed marginals. Letting \(R\) grow proves continuity of the displacement matrix on the coupling set.

Why the projected transport value equals the quadratic-cost value

Every coupling of \(\mu,\nu\) pushes forward to a coupling of their images under \(Q^{1/2}\), with the same cost. Conversely, any coupling of those images lifts to a coupling of the original laws by sampling independently from the conditional laws of \(X\) and \(Y\) on their respective fibers. Its quadratic cost is unchanged. This proves equality even when \(Q\) is singular.

The same representation proves the triangle inequality: each \(\mathsf W_Q\) is a pseudometric, and their maximum satisfies the triangle inequality. Separation follows from \(c\operatorname{tr}S\le\gamma(S)\le C\operatorname{tr}S\) for some positive constants. Hence \(\sqrt c\,\mathsf W_2\le\mathsf W_\gamma\le\sqrt C\,\mathsf W_2\): this is a genuine distance with the usual \(\mathsf W_2\) topology.

The connection with Paty and Cuturi

For the Ky Fan gauge \(\gamma_k(S)=\sum_{i=1}^k\lambda_i(S)\), with eigenvalues in decreasing order, the representing matrices satisfy

\[ 0\preceq Q\preceq I, \qquad \operatorname{tr}Q=k. \]

This is exactly the subspace robust Wasserstein construction of Paty and Cuturi [2]. Their \(k=1\) case gives the operator-norm endpoint; \(k=d\) gives classical quadratic transport. Their later work develops a broader connection between regularized transport and adversarial ground costs [4].

The convex matrix set matters. Restricting the adversary to actual rank-\(k\) orthogonal projectors gives the distinct projection robust problem, studied for instance by Lin, Fan, Ho, Cuturi, and Jordan [5]. In particular, spectral operator-norm transport need not equal the maximum of the one-dimensional projected Wasserstein distances.

A Benamou–Brenier formulation

A moving law \(\rho_t\) with a Borel velocity field \(v_t\) satisfies the continuity equation

\[ \partial_t\rho_t+\nabla\cdot(\rho_tv_t)=0. \]

This means that, for smooth test functions \(\psi\) compactly supported in space and smooth up to the endpoint times,

\[ \int_0^1\!\int(\partial_t\psi+\nabla\psi\cdot v_t)\,d\rho_t\,dt =\int\psi(1)\,d\nu-\int\psi(0)\,d\mu. \]

At each time form the global velocity second moment and its spectral action:

\[ S_t=\int v_t(x)v_t(x)^\top\,d\rho_t(x), \qquad \mathcal E_\gamma(\rho,v)=\int_0^1\gamma(S_t)\,dt. \]

The gauge is applied after averaging over space. Applying a Schatten gauge separately to each rank-one matrix \(v_tv_t^\top\) would simply give \(|v_t|^2\), recovering ordinary quadratic transport for every \(p\). The spectral geometry comes from measuring collective motion.

Theorem 2 (spectral Benamou–Brenier formula). For a convex monotone gauge as above,

\[ \mathsf W_\gamma(\mu,\nu)^2 =\min_{\substack{\partial_t\rho+\nabla\cdot(\rho v)=0\\ \rho_0=\mu,\ \rho_1=\nu}} \mathcal E_\gamma(\rho,v). \]

Use narrowly continuous curves of probabilities and velocities with \(\int_0^1\int|v_t|^2\,d\rho_t\,dt\lt \infty\). This is Theorem 4.2 of the paper [1], extending the classical formulation of Benamou and Brenier [6]. Here the energy has no factor \(1/2\), so its value is the squared distance.

Proof: conditional Jensen in space, Jensen in time

From a coupling to a flow. Draw \((X,Y)\sim\pi\), put \(D=Y-X\), and move each pair along \(Z_t=X+tD\). Let \(\rho_t\) be the law of \(Z_t\). Its Eulerian velocity is

\[ v_t(z)=\mathbb E[D\mid Z_t=z]. \]

Differentiating test functions along the segments proves the continuity equation. Conditional Jensen gives the matrix inequality

\[ S_t=\mathbb E[\mathbb E[D\mid Z_t]\, \mathbb E[D\mid Z_t]^\top] \preceq\mathbb E[DD^\top]=S_\pi. \]

Thus \(\mathcal E_\gamma(\rho,v)\le\gamma(S_\pi)\) by monotonicity. The conditional mean is essential: several particle trajectories may pass through the same location.

From a flow to a coupling. The superposition principle [7] represents a finite-energy solution as a probability law on absolutely continuous paths \(Z_t\), with \(\dot Z_t=v_t(Z_t)\) almost everywhere. Their endpoints give a coupling \(\pi\). Since \(Z_1-Z_0=\int_0^1\dot Z_t\,dt\), time Jensen yields

\[ S_\pi\preceq\int_0^1\mathbb E[\dot Z_t\dot Z_t^\top]\,dt =\int_0^1 S_t\,dt. \]

Apply monotonicity, then convexity of \(\gamma\):

\[ \gamma(S_\pi) \le\gamma\!\left(\int_0^1 S_t\,dt\right) \le\int_0^1\gamma(S_t)\,dt. \]

Taking infima gives the reverse bound. Starting from an optimal coupling also supplies an attaining flow, proving the minimum statement. \(\square\)

Optimal couplings therefore generate constant-speed geodesics by the same straight-segment interpolation as classical transport. Indeed, the restriction to a time interval gives an upper bound proportional to its length, and the triangle inequality forces equality.

There is also a useful link with the perspective transform. Write momentum \(m=\rho v\). For positive densities,

\[ \gamma\!\left(\int\frac{mm^\top}{\rho}\,dx\right) =\max_{Q\in K_\gamma}\int\frac{m^\top Qm}{\rho}\,dx. \]

Each integrand is a quadratic perspective. Taking their supremum gives a convex action in mass and momentum, while continuity becomes \(\partial_t\rho+\nabla\cdot m=0\), linear in the unknowns. The closed perspective extends the formula to measures; for singular \(Q\), the local vacuum cost is zero on \(\ker Q\) and infinite in the other directions. The full action still forces finite-energy momentum to be absolutely continuous with respect to mass, because \(\gamma\) bounds Euclidean kinetic energy from below.

Can the Schatten exponent be smaller than one?

Yes, with a precise change of scope. The paper uses matrix norms, hence \(p\ge1\). For \(0\lt p\lt 1\), \(\gamma_p\) is a Schatten anti-norm on positive matrices: it is monotone and concave rather than convex [8]. The robust-cost theorem and convex coupling formulation above no longer apply. Nevertheless, the distance and dynamic statements extend further.

Theorem 3 (the half-Schatten threshold). Define \(\mathsf W_{\gamma_p}\) by the same static formula. For every \(p\in[1/2,\infty]\), it is a distance and satisfies the same Benamou–Brenier formula with action \(\mathcal E_{\gamma_p}\). For every \(0\lt p\lt 1/2\), the triangle inequality can fail.

Proof: the square root lifts to a Schatten-2p norm

The norm hidden inside the covariance. For a square-integrable random vector \(Z\) on a probability space \(\Omega\), define

\[ A_Z:L^2(\Omega)\longrightarrow\mathbb R^d, \qquad A_Zh=\mathbb E[Zh]. \]

Its adjoint maps \(u\) to \(Z\cdot u\), so \(A_ZA_Z^*=\mathbb E[ZZ^\top]\). Consequently

\[ N_p(Z):=\sqrt{\gamma_p(\mathbb E[ZZ^\top])} =\lVert A_Z\rVert_{S_{2p}}. \]

For \(p=\infty\), use the operator norm. Since \(Z\mapsto A_Z\) is linear and injective, this is a norm on random vectors as soon as \(2p\ge1\): the triangle inequality is the ordinary Schatten norm inequality.

Gluing proves the metric property. Glue a coupling of \((X,Y)\) to one of \((Y,Z)\) on a common probability space. Then

\[ N_p(Z-X)\le N_p(Y-X)+N_p(Z-Y). \]

Use optimal couplings on the right and minimize on the left. Symmetry is immediate, and separation follows from equivalence with the Euclidean \(L^2\) norm in fixed dimension.

The dynamic proof still works. The coupling-to-flow argument uses only Loewner monotonicity, which \(\gamma_p\) retains for every positive \(p\). For the converse, use the same superposition law on paths and Minkowski in \(N_p\), then scalar Cauchy–Schwarz:

\[ N_p(Z_1-Z_0) \le\int_0^1N_p(\dot Z_t)\,dt \le\left(\int_0^1N_p(\dot Z_t)^2\,dt\right)^{1/2}. \]

The last squared norm is exactly \(\gamma_p(S_t)\). This replaces the convexity step that was unavailable below one, and proves static–dynamic equality.

Why the lower endpoint 1/2 is sharp

In \(\mathbb R^2\), take

\[ \mu_0=\delta_0,\qquad \mu_1=\delta_{e_1},\qquad \mu_2=\tfrac12(\delta_{e_1+e_2}+\delta_{e_1-e_2}). \]

All relevant couplings are forced by the Dirac marginals. The two successive displacement matrices are \(e_1e_1^\top\) and \(e_2e_2^\top\); the direct displacement matrix is \(I_2\). Thus

\[ \mathsf W_{\gamma_p}(\mu_0,\mu_1) =\mathsf W_{\gamma_p}(\mu_1,\mu_2)=1, \] \[ \mathsf W_{\gamma_p}(\mu_0,\mu_2)=2^{1/(2p)}. \]

The direct distance exceeds two exactly when \(p\lt 1/2\), violating the triangle inequality. \(\square\)

This extension does not restore convexity of the static coupling objective below one. A small example shows why. Transport the uniform pair \(\pm e_1\) to \(\pm e_2\). The two bijective matchings give

\[ S_\pm=\begin{pmatrix}1&\pm1\\\pm1&1\end{pmatrix}, \qquad \gamma_p(S_\pm)=2. \]

Their equally weighted mixture gives \(I_2\), with \(\gamma_p(I_2)=2^{1/p}>2\) for \(p\lt 1\). Hence no supremum of linear quadratic costs can represent this gauge there. A metric can survive even though its static optimization problem is nonconvex.

Numerical experimentation: changing p changes the coupling

How does the spectral gauge change transport between more intricate distributions? The figure below uses two point clouds with several clusters and curved structure. Change the number of points in each cloud, then move \(p\) from \(1/2\) through \(1\) to the operator-norm endpoint. The same clouds are kept when \(p\) changes, so the differences come from the transport geometry.

For \(n\) points in each cloud, use the empirical measures

\[ \mu_n=\frac1n\sum_{i=1}^n\delta_{x_i},\qquad \nu_n=\frac1n\sum_{j=1}^n\delta_{y_j}, \]

A coupling has nonnegative weights \(\pi_{ij}\), with every row and column summing to \(1/n\). Its displacement matrix is

\[ S_\pi=\sum_{i,j}\pi_{ij}(y_j-x_i)(y_j-x_i)^\top. \]

Edges show the computed coupling, with thickness indicating transported mass. Green particles follow the straight-segment interpolation. The ellipse summarizes the directions of displacement, not the shape of either marginal; its scale stays fixed when \(p\) changes.

No entropy regularization is used. At \(p=1\), a linear assignment gives quadratic optimal transport. For finite \(p>1\), conditional-gradient steps combine assignments and report a convex optimality-gap estimate. At \(p=\infty\), a smooth spectral approximation is used internally, with its approximation error included in the bound. All displayed costs evaluate the actual requested gauge.

For \(p\lt 1\), concavity guarantees that some optimal plan is a matching, but finding the best one is a nonconvex problem. Here several initializations are refined and the best candidate is shown, without a global-optimality claim. Its displayed cost is an upper bound on the true optimum. Above one, mass splitting is allowed and can help balance the displacement spectrum. The slider stops at \(1/2\), the sharp metric threshold proved above.

References

  1. Gabriel Peyré. Muon Dynamics as a Spectral Wasserstein Flow. 2026. Sections 3–4 and Appendices A and E.
  2. François-Pierre Paty and Marco Cuturi. Subspace Robust Wasserstein Distances. ICML, PMLR 97, 5072–5081, 2019.
  3. Maurice Sion. On General Minimax Theorems. Pacific Journal of Mathematics 8(1), 171–176, 1958.
  4. François-Pierre Paty and Marco Cuturi. Regularized Optimal Transport is Ground Cost Adversarial. ICML, PMLR 119, 7532–7542, 2020.
  5. Tianyi Lin, Chenyou Fan, Nhat Ho, Marco Cuturi, and Michael I. Jordan. Projection Robust Wasserstein Distance and Riemannian Optimization. NeurIPS 33, 2020.
  6. Jean-David Benamou and Yann Brenier. A Computational Fluid Mechanics Solution to the Monge–Kantorovich Mass Transfer Problem. Numerische Mathematik 84, 375–393, 2000.
  7. Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient Flows in Metric Spaces and in the Space of Probability Measures. Second edition, Birkhäuser, 2008. Theorem 8.2.1.
  8. Jean-Christophe Bourin and Fumio Hiai. Norm and Anti-Norm Inequalities for Positive Semi-Definite Matrices. 2011 preprint revision.