Entropic Riemannian Neural Optimal Transport

This paper introduces Entropic RNOT, a unified framework that combines intrinsic entropic regularization with amortized neural learning to efficiently solve optimal transport problems on Riemannian manifolds, offering strong theoretical convergence guarantees and superior empirical performance across diverse curved spaces compared to existing baselines.

The Gist

Imagine you are trying to move a pile of sand from one spot to another, but the ground isn't flat. Maybe it's a sphere, a twisted knot, or a curved surface like a saddle. In the real world, data often lives on these curved surfaces (like the rotation of a robot arm or the shape of a molecule), not on flat, grid-like paper.

This paper introduces a new tool called Entropic RNOT to solve the problem of moving "data sand" across these curved landscapes efficiently and accurately.

Here is the breakdown of what they did, using simple analogies:

1. The Problem: The Flat Map vs. The Curved Earth

Most computer programs assume the world is flat (Euclidean). If you try to draw a straight line between two points on a globe using a flat map, the distance and direction get distorted.

• The Issue: When data lives on curved shapes (like a sphere or a rotation group), standard math tricks break down. They either get the distances wrong or require so much computing power to solve that they become useless for large datasets.
• The Old Solutions:
  • Method A: Flatten the curve, do the math, then fold it back. This introduces errors (like trying to flatten an orange peel without tearing it).
  • Method B: Calculate the perfect path for every single grain of sand individually. This is incredibly accurate but takes forever (like calculating a route for every single car in a city traffic jam).

2. The Solution: Entropic RNOT

The authors created a "smart guide" (a neural network) that learns how to move data on these curved surfaces without flattening them or calculating every single path individually.

Think of it like this:

• The "Entropic" Part (The Foggy Lens): Instead of demanding a single, perfect, rigid path for every grain of sand, the method allows for a little bit of "fog" or randomness. Imagine you are trying to get from point A to point B, but instead of one strict road, you have a cloud of possible paths. This "fog" makes the math much easier and faster to solve, similar to how a blurry photo is easier to process than a high-definition one.
• The "Neural" Part (The Learning Guide): Instead of solving the math problem from scratch every time you have new data, they train a neural network (a type of AI) to learn the "shape" of the solution. Once trained, this network can instantly tell you where to move any new piece of data, even ones it has never seen before. This is called amortization—you pay the computing cost once during training, and then the "guide" works for free later.

3. How It Works: The "Heat" and the "Center"

The paper describes two clever ways to turn the "fuzzy cloud" of possible paths into a concrete answer:

• The "Center of Gravity" (Barycentric Projection): If you are on a curved surface like a sphere (Cartan-Hadamard manifolds), the method finds the "center of gravity" of the fuzzy cloud. It's like asking, "If all these possible paths were people, where would they stand if they held hands and found their average spot?" This gives a single, clear destination.
• The "Heat Smoothing" (Heat-Smoothed Surrogates): For more complex shapes, they use a concept called "heat." Imagine dropping a drop of ink (the data) into water. At first, it's a sharp dot. As time passes (heat time), it spreads out into a smooth cloud. The method uses this spreading effect to turn sharp, jagged data points into smooth, flowing distributions. This makes the data easier to handle and prevents the math from getting stuck on tiny, noisy details.

4. What They Proved

The authors didn't just guess; they proved mathematically that:

• Their "smart guide" can learn the perfect solution if given enough training.
• The "center of gravity" method gets closer and closer to the true answer as the training improves.
• The "heat smoothing" method is stable and doesn't introduce weird biases, even as the "heat" (randomness) is turned down.

5. Real-World Test: Fixing Protein Docking

To show it works, they tested it on a very specific, real-world problem: Protein-Ligand Docking.

• The Scenario: Imagine a key (a drug molecule) trying to fit into a lock (a protein). Computers try to guess how the key fits, but they often get the orientation slightly wrong.
• The Test: They took thousands of "wrong" guesses generated by other software and used their Entropic RNOT to "refine" them.
• The Result: The method successfully nudged the drug molecules into the correct position much better than previous methods. It reduced the error from a large distance (11.24 Å) to a very small, accurate distance (3.47 Å). Crucially, it did this without needing to re-calculate the math for every single drug molecule individually; the trained "guide" just applied the rules it learned.

Summary

This paper presents a new way to move data on curved surfaces that is:

1. Accurate: It respects the true geometry of the data (no flattening).
2. Fast: It learns a reusable model so it doesn't have to re-solve the math for every new piece of data.
3. Stable: It uses "fog" and "heat" concepts to make the math robust and easy to compute.

They proved it works mathematically and showed it works in practice by fixing the orientation of drug molecules, making it a powerful tool for machine learning on complex, curved data.

Technical Summary

Technical Summary: Entropic Riemannian Neural Optimal Transport

1. Problem Statement

Many machine learning applications involve data supported on curved spaces (Riemannian manifolds) such as spheres ($S^2$), rotation groups ($SO(3)$), rigid poses ($SE(3)$), and symmetric positive definite matrices ($SPD$). In these settings, standard Euclidean approximations distort distances, averages, and the resulting Optimal Transport (OT) problems.

Existing approaches face a trade-off:

• Manifold OT methods often pursue amortized, out-of-sample transport maps but suffer from computational bottlenecks, frequently requiring iterative inner optimizations for each new instance.
• Entropic Regularization (e.g., Sinkhorn iterations) makes discrete OT scalable and numerically stable but does not inherently provide an amortized model; each new pair of distributions typically requires solving a new optimization problem.

The paper addresses the gap of combining intrinsic geometric OT with amortized out-of-sample evaluation and entropic regularization on possibly non-compact Riemannian manifolds.

2. Methodology: Entropic RNOT

The authors propose Entropic Riemannian Neural Optimal Transport (Entropic RNOT), a unified framework that learns a reusable, manifold-aware transport model.

Core Formulation

The method is based on the semidual formulation of entropic OT. Instead of learning a transport map directly, the model learns a target-side Schrödinger potential $g_\theta$.

• Parameterization: The potential is parameterized via a neural pullback. A continuous feature map $\phi: K_\nu \to \mathbb{R}^n$ (where $K_\nu$ is the support of the target distribution) maps manifold points to Euclidean space. A Euclidean neural network $a_\theta$ is composed with $\phi$ to form the hypothesis class.
• Centering: Since Schrödinger potentials are identifiable only up to an additive constant, the model uses a centered pullback class $C_\nu(\phi^* \mathcal{F})$ to ensure uniqueness.
• Optimization: The model is trained by maximizing the semidual objective $J_\varepsilon(g_\theta)$ using stochastic gradient ascent on minibatches. The source-side potential $f^\varepsilon_\theta$ is recovered via the soft $c$-transform (a log-sum-exp operation) of the learned target potential.

Intrinsic Transport Surrogates

Once the Gibbs coupling $\pi^\varepsilon_\theta$ is induced by the learned potentials, the paper extracts deterministic transport surrogates suitable for different manifold geometries:

1. Barycentric Projections: On Cartan–Hadamard manifolds (complete, simply connected, non-positive curvature), the conditional laws define a deterministic transport map via the Riemannian barycenter (Fréchet mean).
2. Heat-Smoothed Surrogates: On complete stochastically complete manifolds (a broader class including compact manifolds, Euclidean spaces, and products like $SE(3)$), the method applies heat smoothing to the conditional target laws. This converts potentially atomic conditional distributions (from finite samples) into absolutely continuous distributions. A point prediction (mode) is then derived from this smoothed density.

3. Key Contributions

The paper makes three primary contributions:

1. Framework Introduction: Entropic RNOT is the first intrinsic neural framework for static entropic OT on Riemannian manifolds that combines the semidual formulation with amortized out-of-sample evaluation.
2. Theoretical Guarantees: For a fixed regularization parameter $\varepsilon > 0$, the authors prove that the proposed hypothesis class can recover the entropic optimal coupling in strong probabilistic metrics (KL divergence, Total Variation, weak convergence). Consequently:
   • Barycentric surrogates converge in $L^2(\mu)$ on Cartan–Hadamard manifolds.
   • Heat-smoothed surrogates are stable at any fixed heat time $t > 0$ and are asymptotically unbiased as $t \to 0$.
   • These guarantees hold for compactly supported data on possibly non-compact manifolds.
3. Empirical Validation: The method demonstrates strong transport quality across diverse geometries ($S^2, SO(3), SPD(3), SE(3), H^2$), outperforming ambient Euclidean, tangent-space, and log-Euclidean baselines. It scales favorably in memory and time compared to discrete manifold Sinkhorn and achieves significant improvements in a real-world protein–ligand docking application.

4. Experimental Results

Synthetic Benchmarks

Evaluated on $S^2, SO(3), SPD(3), SE(3),$ and $H^2$ with wrapped normal distributions.

• Accuracy: Entropic RNOT consistently recovers the discrete manifold Sinkhorn reference plan more accurately than all baselines, with the largest gains observed on $SPD(3)$, $SE(3)$, and $H^2$ where intrinsic geometry is most critical.
• Metrics: It achieves significantly lower Plan KL divergence and endpoint geodesic errors compared to ambient Euclidean and tangent-space linearization methods.

Scalability

• Complexity: Discrete manifold Sinkhorn requires an $O(N^2)$ memory footprint for the cost matrix, becoming infeasible for large support sizes (e.g., $N=32,768$).
• Performance: Entropic RNOT training time and memory usage remain constant with respect to support size $N$ (dependent only on batch size). Inference throughput scales linearly with $N$, enabling the processing of millions of samples per second.

Real-World Application: Protein–Ligand Docking

The method was applied to refine rigid poses on $SE(3)$ for protein–ligand docking using the CrossDocked2020 dataset.

• Setup: A single model was trained on pooled complexes to refine held-out docking poses toward the docking engine's top-ranked binding basin. No crystallographic structures were used during training or inference.
• Results:
  • Reduced top-1 RMSD from 11.24 Å (no refinement) to 3.47 Å.
  • Improved success rate within 2 Å from 10.3% to 75.9%.
  • Outperformed both physics-based minimization (GNINA) and per-instance discrete Sinkhorn (which failed due to small target sets per complex).

5. Significance and Limitations

Significance:
The paper claims to provide the first intrinsic neural framework that unifies the scalability of entropic regularization with the generalization capabilities of amortized neural OT on manifolds. It offers a practical solution for high-dimensional, non-Euclidean transport tasks where discrete methods are computationally prohibitive.

Limitations (as stated by the authors):

• Theoretical Scope: Theoretical guarantees are established for fixed $\varepsilon > 0$ and compact supports; the vanishing-regularization regime ($\varepsilon \to 0$) is not addressed.
• Geometric Constraints: Barycentric map recovery guarantees require the Cartan–Hadamard setting; outside this, barycenters may be non-unique or unstable.
• Application Specifics: In the docking experiment, the method acts as a refinement/denoising procedure for existing pose ensembles rather than a de novo generative model. It currently ignores receptor pocket context and treats ligands as rigid bodies, omitting torsional flexibility.
• Computational Dependencies: Performance relies on efficient geodesic distance evaluation and stable log-sum-exp computations.