Entropic Riemannian Neural Optimal Transport: Bridging Geometry and Scalability

TL;DR

Entropic RNOT is a new neural framework that solves the "curved space" optimal transport problem. By combining entropic regularization with neural pullback potentials, it enables fast, amortized transport on manifolds like spheres and rotation groups. It successfully refined protein-docking poses on SE(3), outperforming traditional physics-based minimization without needing retraining for new protein complexes.

The Motivation: Why Euclidean OT Fails in Curved Spaces

In the world of AI, we often treat data as points in flat Euclidean space. However, many real-world objects don't live there. Directions live on a sphere ($S^2$), orientations on a rotation group ($SO(3)$), and robot poses on $SE(3)$.

If you use standard "straight-line" distances to compare these, you get "distortion." An average of two rotations calculated as a linear midpoint might result in a transformation that isn't even a valid rotation. While the Sinkhorn algorithm helps solve these problems for finite sets of points, it doesn't give us a "machine" (a model) that can handle new, unseen points instantly.

The Core Method: Entropic RNOT

The authors propose Entropic RNOT, which avoids the heavy iterative solves of previous manifold OT methods.

1. Semidual Potential: Instead of learning a pair of potentials, they rely on a single target-side potential $g_{h}eta$.
2. Neural Pullback: They map non-Euclidean points into a Euclidean space using "landmark" distances or logarithmic coordinates, then pass them through a standard MLP.
3. Intrinsic Surrogates: Since entropic OT naturally produces a spread-out (probabilistic) plan, the authors use Heat-kernel smoothing and Barycentric projection to find the "center" of the transport, giving a deterministic map that respects the manifold's curvature.

Figure 1: While traditional Sinkhorn methods (dashed lines) explode in memory as data points increase, Entropic RNOT (solid lines) maintains a constant, lightweight memory footprint.

Experiments: From Synthetic Tests to Molecular Docking

Intrinsic Geometry Benchmarks

On manifolds like $H^2$ (Hyperbolic plane) and $SPD(3)$ (Symmetric Positive Definite matrices), Entropic RNOT recovered the ground truth plan with nearly zero KL divergence, while "Tangent-space" methods (which flatten the manifold at one point) failed as the data spread out.

Real-World Case: Protein-Ligand Docking

Molecular docking predicts how a drug (ligand) fits into a protein. The tool GNINA often produces a "cloud" of possible poses. Entropic RNOT was used to "denoise" this cloud.

• Success Rate: Poses refined within the 2 Å "near-native" region jumped from 10.3% to 75.9%.
• Amortization: The model learned from many proteins and can now refine poses for a totally new protein in milliseconds without retraining.

Figure 2: The refined pose (blue) captures the true crystal position (green), whereas the initial docking guess (pink) was far off in the pocket.

Critical Analysis & Future Outlook

The strength of Entropic RNOT is its independence from support size. Because it's a neural model, you can train it on a small batch and then apply it to millions of points at inference time.

Limitations:

• It requires the calculation of geodesic distances, which can be computationally expensive on complex, custom manifolds.
• Currently, the docking model is "unconditional," meaning it doesn't fully look at the protein sequence yet—it only looks at the geometry of the pose cloud.

The Takeaway: Entropic RNOT proves that we can have our cake and eat it too: the mathematical rigor of Riemannian geometry and the massive scalability of neural networks.

Conclusion

This work is a significant step toward making "geometric" AI practical for large-scale production, particularly in structural biology and robotics where the world is rarely flat.