Riemannian Neural Optimal Transport

TL;DR

This paper extends neural optimal transport methods to Riemannian manifolds, introducing Riemannian Neural OT maps that improve scalability and approximation accuracy in high-dimensional geometric settings.

Contribution

The paper develops Riemannian Neural OT maps that avoid discretization, providing a scalable neural network approach for optimal transport on manifolds.

Findings

RNOT maps approximate Riemannian OT with sub-exponential complexity.

Experiments show improved scalability over discretization methods.

Competitive performance demonstrated on synthetic and real data.

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.