Riemannian Neural Optimal Transport

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.

💡 Research Summary

Paper Overview
The manuscript addresses the problem of learning optimal transport (OT) maps on compact Riemannian manifolds using neural networks. While neural OT has become a powerful tool for generative modeling in Euclidean spaces, extending it to manifolds faces two major obstacles: (1) existing methods rely on discretizing the manifold (e.g., via meshes or a finite set of landmark points), and (2) such discretizations inevitably suffer from the curse of dimensionality (CoD), meaning that the number of parameters required to achieve a fixed accuracy grows exponentially with the intrinsic dimension of the manifold.

Negative Result – CoD for Discrete‑Output Maps
The authors formalize “discrete‑output” transport maps as those whose push‑forward of the source measure is supported on at most (m) points. They prove (Theorem 3.1) that for any absolutely continuous source and target measures (\mu,\nu) on a compact manifold, the root‑mean‑square error between any map in this class and the true optimal map (T^\star) is bounded below by (C,m^{-1/p}), where (p) is the manifold dimension. Consequently, to achieve an error (\delta) one needs at least (m\ge (C/\delta)^p) support points, i.e., exponential growth in (p). This barrier applies to all discretization‑based approaches, including the recently proposed Riemannian Convex Potential Maps (RCPM), which the authors show (Corollary 3.2) are a special case of discrete‑output maps and therefore inherit the same CoD.

Positive Contribution – Riemannian Neural OT (RNOT)
To bypass the CoD, the paper introduces RNOT, a continuous neural‑network parameterization of OT potentials that respects the geometric constraints of the problem by construction. The key steps are:

1. Feature Embedding – Choose a continuous injective map (\phi:M\to\mathbb{R}^n). A practical construction uses distances to a maximal (\delta)-separated set of landmarks (Gromov’s distance‑to‑landmarks embedding), guaranteeing injectivity on compact manifolds.

2. Universal Approximation Class – Let (\mathcal{F}) be a dense family of functions on (\mathbb{R}^n) (e.g., ReLU deep networks). The pull‑back class (\phi^\ast\mathcal{F}={f\circ\phi:f\in\mathcal{F}}) inherits universality on (M) thanks to the injectivity of (\phi).

3. Implicit c‑concavity – For any (\psi\in\phi^\ast\mathcal{F}), define the potential (\varphi=\psi^{c}) via the c‑transform. This operation automatically enforces the required c‑concavity without explicit constraints.

4. Transport Map Construction – The learned OT map is obtained by the Riemannian exponential: (T(x)=\exp_x(-\nabla\varphi(x))). Because (\varphi) is continuous, (T) is a continuous map whose image is not limited to a finite set.

Theoretical Guarantees
The authors prove a universality theorem (Theorem 4.1) stating that approximating the optimal potential within the implicit class (\mathcal{C}(\phi^\ast\mathcal{F})) yields pointwise convergence of the induced transport maps to the true optimal map. Moreover, they derive explicit bounds on the number of network parameters (W) and depth (L) needed to achieve a prescribed uniform error (\varepsilon). Under mild smoothness assumptions (the optimal potential belongs to a Hölder class of order (r)), they show that (W = O(\varepsilon^{-k})) and (L = O(\log \varepsilon^{-1})) for some constant (k) depending on (r) but independent of the manifold dimension. This sub‑exponential (indeed polynomial) complexity directly contrasts with the exponential requirement of discrete methods.

Empirical Validation
Experiments are conducted on both synthetic and real manifold‑valued datasets:

• Synthetic manifolds – Spheres (S^2), tori (T^2), and the rotation group (SO(3)). The authors compare RNOT against RCPM and mesh‑based OT. For a target Wasserstein error of (10^{-2}), RNOT uses roughly 10–15× fewer parameters and converges faster.

• Real data – Spherical image data (e.g., Earth observation) and 3D human pose data represented on (SO(3)). RNOT achieves comparable or better Fréchet Inception Distance (FID) and Wasserstein scores while reducing training time by ~30 % and memory consumption by ~40 %.

Importantly, because the learned map is continuous, sampling at test time is amortized: a single forward pass of the neural network followed by the exponential map yields new data points without any additional optimization.

Conclusions and Future Directions
The paper makes two central contributions: (1) a rigorous proof that any discretization‑based OT method on manifolds suffers from an exponential parameter blow‑up with dimension, and (2) the introduction of RNOT, a theoretically grounded, continuous neural OT framework that avoids this blow‑up and enjoys polynomial‑scale complexity. The work opens the door to scalable generative modeling, domain adaptation, and physics‑informed simulation on non‑Euclidean domains. Future research avenues suggested include extending the approach to non‑quadratic costs, handling non‑compact manifolds, and developing more efficient automatic‑differentiation schemes for the Riemannian exponential and gradient operations.