Riemannian Neural Optimal Transport

TL;DR

We address learning optimal transport on compact Riemannian manifolds and identify a fundamental CoD barrier for discrete-output OT methods that discretize the manifold.We introduce Riemannian Neural OT (RNOT), a continuous, geometry-aware framework that parameterizes OT potentials with neural nets and recovers transport maps via the exponential map, thereby avoiding discretization and preserving manifold structure.Under mild regularity, we prove that RNOT potentials and maps admit polynomial-$\varepsilon^{-1}$ complexity and that the induced maps converge to the true Riemannian OT map, with a stability guarantee yielding pointwise convergence on full-$\mu$-measure sets.Empirical results on continental drift and synthetic spheres/ tori demonstrate improved scalability and competitive performance against discretization-based baselines, validating the dimension-friendly guarantees in practice.

Abstract

Computational optimal transport (OT) offers a principled framework for generative modeling. Neural OT methods, which use neural networks to learn an OT map (or potential) from data in an amortized way, can be evaluated out of sample after training, but existing approaches are tailored to Euclidean geometry. Extending neural OT to high-dimensional Riemannian manifolds remains an open challenge. In this paper, we prove that any method for OT on manifolds that produces discrete approximations of transport maps necessarily suffers from the curse of dimensionality: achieving a fixed accuracy requires a number of parameters that grows exponentially with the manifold dimension. Motivated by this limitation, we introduce Riemannian Neural OT (RNOT) maps, which are continuous neural-network parameterizations of OT maps on manifolds that avoid discretization and incorporate geometric structure by construction. Under mild regularity assumptions, we prove that RNOT maps approximate Riemannian OT maps with sub-exponential complexity in the dimension. Experiments on synthetic and real datasets demonstrate improved scalability and competitive performance relative to discretization-based baselines.

Figures (5)

• Figure 1: Continental drift optimal transport on $\mathbb{S}^{2}$. Left to right: source mass distribution $\mu$ ($\sim 150$ million years ago), target distribution $\nu$ (present day), transported distribution $T_{\#}\mu$, and geodesic trajectories induced by the learned transport map $T$.
• Figure 2: Optimal transport on $\mathbb{S}^2$ (top) and $\mathbb{T}^2$ (bottom) from a representative run corresponding to Tables \ref{['tab:S2']} and \ref{['tab:T2']}, respectively. From left to right: uniform source $\mu$, wrapped normal target $\nu$, geodesic trajectories induced by the learned transport $T$, and pushforward $T_{\#}\mu$.
• Figure 3: KL divergence (log scale) versus dimension $p\in \{2,\ldots,10\}$ for transport from a uniform source to a wrapped normal target on $\mathbb{S}^{p}$ (left panel) and $\mathbb{T}^{p}$ (right panel). Curves show RNOT (ours) and RCPM under a sweep of regularization values $\gamma$. RCPM performance degrades as $p$ increases --- most notably for small $\gamma$, where the map is sharp --- whereas RNOT remains stable across dimensions.
• Figure 4: Embedding diagnostics on $\mathbb{S}^2$: coverage radius $R_M^{\mathrm{val}}$ (left), minimum pairwise separation $s_M$ (center), and near-collision fraction $\rho_M(\varepsilon)$ (right) for random vs. farthest-point sampling (FPS). FPS achieves lower coverage radius while maintaining comparable separation, with $\rho_M \approx 0$ for all $M$ tested.
• Figure 5: KL divergence scaling with dimension on high-dimensional manifolds. We compare our approach against RCPMs with varying regularization parameters $\gamma \in \{0.001, 0.005, 0.01, 0.05, 0.1, 1.0\}$. RCPM results are shown for $p \in \{2, \ldots, 10\}$ only due to computational intractability at higher dimensions, while our method extends to $p = 40$. Our method achieves consistently lower KL divergence across all tested dimensions and scales favorably to high-dimensional settings, empirically supporting our theoretical results.

Theorems & Definitions (58)

• Definition 2.1: Def. 3.1 from CorderoErausquin2001
• Proposition 2.3: Distance-to-Landmarks Embedding (Gromov1983)
• Theorem 3.1: CoD Barrier for Discrete-Output Maps
• Corollary 3.2: CoD Barrier for RCPMs
• Theorem 4.1: Universality of Implicit $c$-Concave Potentials.
• Definition 4.2: Riemannian Neural OT Potentials and Maps
• Theorem 5.1: Polynomial $\varepsilon$-Complexity for Functions on Manifolds
• Corollary 5.2: Polynomial $\varepsilon^{-1}$-Complexity for RNOT Potentials
• Theorem 5.3: Pointwise Stability of RNOT Maps
• Proposition 1.1: Normal neighborhoods