Entropic Riemannian Neural Optimal Transport

Optimal transport (OT) is a fundamental tool in machine learning, but it breaks down when data lives on curved spaces like spheres, rotation groups, or hyperbolic manifolds—common in robotics, molecular modeling, and shape analysis. Euclidean distances distort the geometry, and existing manifold OT methods either lack scalability or fail to provide out-of-sample maps. In a new preprint, Alessandro Micheli, Silvia Sapora, Anthea Monod, and Samir Bhatt introduce Entropic Riemannian Neural Optimal Transport (Entropic RNOT), a unified framework that combines entropic regularization with neural amortization on Riemannian manifolds. The model learns a single target-side Schrödinger potential through a neural pullback parameterization, recovers the induced Gibbs coupling, and uses conditional laws to construct intrinsic transport surrogates—barycentric projections on Cartan-Hadamard manifolds and heat-smoothed conditional surrogates on stochastically complete manifolds.

For a fixed regularization ε>0, the authors prove that their hypothesis class recovers the entropic optimal coupling in strong probabilistic metrics. Barycentric surrogates converge in L², while heat-smoothed surrogates are stable and asymptotically unbiased as heat time vanishes—all for compactly supported data on possibly noncompact manifolds. Empirically, Entropic RNOT matches or improves over Euclidean, tangent-space, and log-Euclidean baselines on benchmarks covering S², SO(3), SPD(3), SE(3), and hyperbolic space ℍ². It scales favorably relative to discrete manifold Sinkhorn, and in a protein-ligand docking application, it refines poses on SE(3) without retraining or per-instance optimization.