Light Entropic Optimal Transport on Riemannian Manifolds

Xavier Aramayo-Carrasco; Petr Mokrov; Alexander Korotin

cs.LGOct 2, 2026

Light Entropic Optimal Transport on Riemannian Manifolds

Authors: Xavier Aramayo-Carrasco, Petr Mokrov, Alexander Korotin

Organizations: Applied AI Institute, Moscow, Russia · Applied AI Institute AXXX, Moscow, Russia

Abstract

Entropic Optimal Transport (EOT) has become a practical framework for learning stochastic couplings between complex distributions, with applications in generative modeling and domain adaptation. However, most EOT solvers are designed for Euclidean spaces, while manifold extensions remain limited and often rely on costly iterative methods, simulated dynamics, or generic neural models that do not fully exploit the underlying geometry. We introduce ManifoldLightOT, a light approach for learning kernel-induced EOT couplings directly on common manifolds. Using the kernel form of the EOT solution, we construct geometry-specific Gibbs kernels together with compatible potential parameterizations for spheres, tori, $\mathrm{SO}(3)$ , and $\mathrm{SE}(3)$ . These choices yield closed-form normalization and directly sampleable conditional distributions. Our formulation naturally extends to products of manifolds, making it applicable to more complex geometries. The parameters of the potentials are optimized directly from samples using Monte Carlo estimates of the learning objective. Through synthetic and real-world experiments, we show that ManifoldLightOT often outperforms existing manifold OT methods while retaining direct sampling.

Figures & tables

Method	${\mathbb{S}}^{60}$	${\mathbb{S}}^{100}$	${\mathbb{T}}^{60}$	${\mathbb{T}}^{100}$	$\mathrm{SO}(3)$	$\mathrm{SE}(3)$
ManifoldLightOT (ours)	$\mathbf{0.07\pm 0.00}$	$0.45\pm 0.05$	$\mathbf{0.08\pm 0.00}$	$\mathbf{0.08\pm 0.00}$	$\mathbf{0.03\pm 0.01}$	$\mathbf{0.18\pm 0.01}$
RNOT	$0.18\pm 0.04$	$1.47\pm 0.21$	$2.15\pm 0.20$	$23.15\pm 0.57$	$0.61\pm 0.27$	$0.23\pm 0.04$
ERNOT	$0.22\pm 0.01$	$\mathbf{0.32\pm 0.09}$	$0.27\pm 0.02$	$0.27\pm 0.03$	$0.47\pm 0.09$	$5.54\pm 0.72$
NeuralOT	$0.27\pm 0.06$	$4.38\pm 8.24$	$61.59\pm 0.06$	$61.61\pm 0.07$	$1.27\pm 2.58$	$0.27\pm 0.05$

Table 2: MMD 2 results across different manifolds, reported in units of

10^{-2}

. We report mean

\pm

standard deviation over five seeds. Lower is better.

Figure 1: Prediction on continental drift from 150 Ma to the present day for ManifoldLightOT.

Figure 2: Pole-density plots for the $\{0001\}$ orientation family in the EBSD texture-transfer task, shown for the RD and TD conditions. For each orientation quaternion $q$ , we take the fixed direction $h=(0,0,1)$ , rotate it into the sample frame, project the resulting unit vector onto the two-dimensional disk using stereographic projection, and estimate the density of the projected points. Colors indicate orientation concentration in multiples of random distribution (MRD).

Method	Direction	MMD SO(3)	Texture Index Error	Normalized ODF Difference	M-Index Error
ManifoldLightOT (ours)	RD	0.000592	0.01168	0.00270	0.00425
ManifoldLightOT (ours)	TD	0.000830	0.00112	0.00265	0.00883
RNOT	RD	0.000635	0.09022	0.01298	0.01304
RNOT	TD	0.000680	0.00441	0.01936	0.02033
ERNOT	RD	0.000901	0.11089	0.12242	0.01650
ERNOT	TD	0.001056	0.13406	0.18600	0.01554

Table 3: Results on the crystallographic texture transfer task for the rolling-direction (RD) and transverse-direction (TD) conditions. Lower is better for all metrics.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Sphere
Method	${\mathbb{S}}^{2}$	${\mathbb{S}}^{10}$	${\mathbb{S}}^{20}$	${\mathbb{S}}^{40}$	${\mathbb{S}}^{60}$	${\mathbb{S}}^{80}$	${\mathbb{S}}^{100}$
ManifoldLightOT	$0.08\pm 0.01$	$0.07\pm 0.01$	$0.08\pm 0.01$	$0.06\pm 0.00$	$0.07\pm 0.00$	$0.13\pm 0.01$	$0.45\pm 0.05$
RNOT	$0.13\pm 0.06$	$0.10\pm 0.02$	$0.10\pm 0.01$	$0.11\pm 0.02$	$0.18\pm 0.04$	$0.43\pm 0.09$	$1.47\pm 0.21$
ERNOT	$0.12\pm 0.03$	$0.15\pm 0.01$	$0.20\pm 0.03$	$0.20\pm 0.04$	$0.22\pm 0.01$	$0.25\pm 0.02$	$0.32\pm 0.09$
NeuralOT	$0.19\pm 0.07$	$0.12\pm 0.04$	$0.17\pm 0.02$	$0.17\pm 0.02$	$0.27\pm 0.06$	$0.29\pm 0.04$	$4.38\pm 8.24$
Torus

Appendix

Table 4: MMD 2 results on high-dimensional spheres and tori, reported in units of

10^{-2}

. We report mean

\pm

standard deviation over five seeds. Lower is better.

Optimal Transport

Optimal transport compares and transforms probability distributions by finding low-cost ways to move mass between them.

See all 287 papers

Entropic Riemannian Neural Optimal Transport
May 5, 2026Alessandro Micheli, Silvia Sapora, Anthea Monod +1Riemannian Geometry Optimal Transport
Imperial College London London, UK · University of Oxford Oxford, UK · Statens Serum Institut Copenhagen, Denmark +1
Sample complexity of unbalanced entropic OT
Jun 23, 2026Francisco Andrade, Gabriel Peyré, Clarice PoonSample Complexity Entropy Regularization
INRIA (PreMeDICaL & HeKA) · INRIA · CNRS and ENS, PSL Universit´e +1
Density-Reweighted Entropic Optimal Transport: Decoupling Geometry from Sampling Density
Aug 17, 2026Keyi Li, Yuval Kluger, Boris LandaOptimal Transport Entropy Regularization
Program of Applied & Computational Mathematics, Yale University · Department of Pathology, Yale University · Department of Electrical and Computer Engineering, Yale University

May 5, 2026J/K move · Enter open · S save

Entropic Riemannian Neural Optimal Transport

Alessandro Micheli, Silvia Sapora, Anthea Monod, Samir Bhatt

Imperial College London London, UK · University of Oxford Oxford, UK · Statens Serum Institut Copenhagen, Denmark+1

Riemannian Geometry Optimal Transport

Open paper AskPDF arXiv

Many machine learning problems involve data supported on curved spaces such as spheres, rotation groups, hyperbolic spaces, and general Riemannian manifolds, where Euclidean geometry can distort distances, averages, and the resulting optimal transport (OT) problem. Existing manifold OT methods have pursued amortized out-of-sample maps, while entropic regularization has made discrete OT more scalable, but these advantages have remained largely disjoint. We propose Entropic Riemannian Neural Optimal Transport (Entropic RNOT), a unified framework that combines intrinsic entropic OT with amortized out-of-sample evaluation on Riemannian manifolds. Our method learns a single target-side Schrödinger potential through a neural pullback parameterization, recovers the induced Gibbs coupling, and uses the resulting conditional laws to construct intrinsic transport surrogates. These include barycentric projections on Cartan-Hadamard manifolds and heat-smoothed conditional surrogates on stochastically complete manifolds, the latter turning possibly atomic target laws into absolutely continuous ones. For fixed regularization

\varepsilon>0

, we prove that the proposed hypothesis class recovers the entropic optimal coupling in strong probabilistic metrics. As consequences, barycentric surrogates converge in

L^2

, while heat-smoothed surrogates are stable at fixed heat time and asymptotically unbiased as the heat time vanishes. The guarantees hold for compactly supported data on possibly noncompact manifolds. Empirically, our method matches or improves over Euclidean, tangent-space, and log-Euclidean baselines on benchmarks over

\mathbb{S}^2

\mathrm{SO}(3)

\mathrm{SPD}(3)

\mathrm{SE}(3)

, and

\mathbb{H}^2

, scales favorably relative to discrete manifold Sinkhorn, and in a protein-ligand docking application, refines poses on

\mathrm{SE}(3)

without retraining or per-instance optimization.

Abstract

Optimal Transport

Explore similar work

Entropic Riemannian Neural Optimal Transport

Sample complexity of unbalanced entropic OT

Density-Reweighted Entropic Optimal Transport: Decoupling Geometry from Sampling Density