Optimal Transportation by Orthogonal Coupling Dynamics

TL;DR

This work introduces Orthogonal Coupling Dynamics (OCD), a projection-based evolution that computes optimal transport maps by sliding samples along a marginal-preserving tangent space, effectively turning the constrained OT problem into a sequence of regression problems via conditional expectations. It proves key structural properties—marginal preservation, monotone cost descent, and instability of sub-optimal couplings—and derives a McKean–Vlasov density evolution for the joint distribution, with a detailed Gaussian analysis showing OT recovery under commuting covariances and exponential convergence. The authors present a nonparametric Monte-Carlo algorithm inspired by opinion dynamics to estimate conditional expectations and evolve particle pairs, along with a RK4 time integrator and complexity assessments. Numerical experiments demonstrate OCD’s capacity to recover nonlinear Monge maps, learn distributions, and interpolate colors, while highlighting favorable scalability relative to traditional OT solvers and potential as a building block for data-driven transport models. Overall, OCD provides a scalable, regression-oriented pathway to compute OT maps and Wasserstein distances, with broad implications for numerical OT, distribution learning, and generative modeling.

Abstract

Many numerical and learning algorithms rely on the solution of the Monge-Kantorovich problem and Wasserstein distances, which provide appropriate distributional metrics. While the natural approach is to treat the problem as an infinite-dimensional linear programming, such a methodology limits the computational performance due to the polynomial scaling with respect to the sample size along with intensive memory requirements. We propose a novel alternative framework to address the Monge-Kantorovich problem based on a projection type gradient descent scheme. The dynamics builds on the notion of the conditional expectation, where the connection with the opinion dynamics is leveraged to devise efficient numerical schemes. We demonstrate that the resulting dynamics recovers random maps with favourable computational performance. Along with the theoretical insight, the proposed dynamics paves the way for innovative approaches to construct numerical schemes for computing optimal transport maps as well as Wasserstein distances.

Figures (21)

• Figure 1: Geometric illustration of OCD; (a) The gradient descent, which leaves the polytope $\Pi(\mu,\nu)$, is projected onto the tangent space $\ker(\mathsf{P}_{p_t,X})\times\ker(\mathsf{P}_{p_t,Y})$; (b) Initial condition as product Gaussian probability density; (c) Evolved by $v^{\textrm{ocd}}$, the initial density is updated by projected gradient descent in $L^2$-cost; (d) Evolved by $v^{\textrm{gd}}$, the initial density is updated by gradient descent in $L^2$-cost, leaving $\Pi(\mu,\nu)$.
• Figure 2: Execution time (top) and memory consumption (down) of OCD-piecewise linear algorithm with Runge Kuta 4th order solver with $\Delta t=0.1$ against Earth Mover's Distance (EMD) method JMLR:v22:20-451 in learning the map between samples of $\mu=\mathcal{N}( 0, I)$ and $\nu=\mathcal{N}( 1, I)$ in (a) $n=3$ for a range of $N_p$ and (b) $n=1,...,128$ with $N_p=1000$. Here we scale $\epsilon$ with a rule of thumb $\varepsilon = 3 n {N_p}^{-1/4}/4$.
• Figure 3: Evolution of point particles under the coupling created by OCD.
• Figure 4: Evolution of cross-correlation matrix, $J_t$, in the cone of symmetric positive semi-definite matrices.
• Figure 5: Recovery of nonlinear Monge map between (a) normal distribution and (b) its push-forward by softmax map $T_i=\nabla_{x_i} \log \left(\exp(x_1)+\exp(x_2)\right)$; (c) Starting from independent samples, the pair $(X_t,Y_t)$ is evolved by OCD and $\mathbb{E}\lVert X_t-Y_t \rVert_2^2$ converges towards the Wasserstein distance $d^2$; (d) First component of the map generated by OCD; (e) First component of the exact Monge map.

Theorems & Definitions (18)

• Remark 3.2: Stability of the orthogonal complement operator $\mathsf{S}_p$
• Remark 3.3: On the tangent space $\textrm{Tan}_{p_t,2}\Pi(\mu,\nu)$
• proof : Proof of Proposition \ref{['prop:marg-pres']} (Marginal preservation)
• proof : Proof of Proposition \ref{['prop:pos']} (Descent in Cost)
• proof : Proof of Proposition \ref{['prop:sub-opt']} (Instability of Sub-Optimal Couplings)
• proof : Proof of Theorem \ref{['theorem:vlasov']} (McKean--Vlasov system)
• Lemma 3.4: Preparatory for Proposition \ref{['prop:spd']}
• proof
• proof : Proof of Proposition \ref{['prop:spd']} (Symmetric positive semidefinite correlation)
• Lemma 3.5: Orthogonal projections onto marginal‚Äìpreserving subspaces