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.

This paper is about teaching neural networks to move one cloud of data into another in the most efficient way, when the data lives on a curved space rather than on a flat grid. This “most efficient moving” problem is called optimal transport. It is useful for generative modelling: instead of generating data directly, a model starts with a simple distribution and learns how to transform it into the target data distribution. Existing neural optimal transport methods mostly work in ordinary flat, Euclidean space, but many real datasets naturally live on curved spaces such as spheres, tori, or other geometric surfaces. The authors show that a common strategy for optimal transport on curved spaces has a fundamental scaling problem. Many methods approximate the transport map using a finite set of target points, rather like replacing a smooth movement with a lookup table of possible destinations. The paper proves that any such discrete-output approach runs into the curse of dimensionality: to keep the same accuracy as the manifold dimension grows, the number of parameters must grow exponentially. In plain terms, these methods become impractical very quickly on high-dimensional curved data. To address this, the paper introduces Riemannian Neural Optimal Transport (RNOT). Instead of discretising the curved space, RNOT represents the transport rule continuously using a neural network, while building in the geometry of the manifold from the start. It learns a mathematical object called a transport potential, enforces the right optimal-transport structure through a transform, and then converts this into a map that moves points along the manifold’s natural geometry. The main theoretical result is that, under regularity assumptions, RNOT can approximate the true optimal transport map without the exponential blow-up suffered by discrete methods. The authors prove polynomial-style complexity bounds for approximating the relevant potentials and show that the induced transport maps converge to the correct map on almost all points of the manifold. They also test the method empirically. One case study models continental drift on the sphere, learning a transport map from Earth’s land-mass distribution 150 million years ago to the present day. The learned paths give interpretable movement trajectories consistent with plate-tectonic motion. The authors also test synthetic data on spheres and tori, including higher-dimensional settings, and find that RNOT stays stable as dimension increases, while a discretisation-based baseline degrades. Overall, the contribution is a new way to do neural optimal transport on curved spaces that is more scalable in dimension and better aligned with the geometry of the data.