# GeONet: a neural operator for learning the Wasserstein geodesic

arXiv: https://arxiv.org/abs/2209.14440 · HTML: https://mlsanity.com/paper/2209.14440/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 25 of its 64 references that are in the map.
- First posted on arXiv 28 Sep 2022.

## Builds on (most cited first, 25 of 25)

- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators](https://mlsanity.com/paper/1910.03193/) (2019)
- [Computational Optimal Transport](https://mlsanity.com/paper/1803.00567/) (2018)
- [Neural Operator: Graph Kernel Network for Partial Differential Equations](https://mlsanity.com/paper/2003.03485/) (2020)
- [Learning the solution operator of parametric partial differential equations with physics-informed DeepOnets](https://mlsanity.com/paper/2103.10974/) (2021)
- [Physics-Informed Neural Operator for Learning Partial Differential Equations](https://mlsanity.com/paper/2111.03794/) (2021)
- [Optimal Transport for Domain Adaptation](https://mlsanity.com/paper/1507.00504/) (2015)
- [Neural Operator: Learning Maps Between Function Spaces](https://mlsanity.com/paper/2108.08481/) (2021)
- ["Zero-Shot" Super-Resolution using Deep Internal Learning](https://mlsanity.com/paper/1712.06087/) (2017)
- [On the eigenvector bias of Fourier feature networks: From regression to solving multi-scale PDEs with physics-informed neural networks](https://mlsanity.com/paper/2012.10047/) (2020)
- [How to train your neural ODE: the world of Jacobian and kinetic regularization](https://mlsanity.com/paper/2002.02798/) (2020)
- [Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration](https://mlsanity.com/paper/1705.09634/) (2017)
- [Gradient-enhanced physics-informed neural networks for forward and inverse PDE problems](https://mlsanity.com/paper/2111.02801/) (2021)
- [Stochastic Optimization for Large-scale Optimal Transport](https://mlsanity.com/paper/1605.08527/) (2016)
- [Computational Optimal Transport: Complexity by Accelerated Gradient Descent Is Better Than by Sinkhorn's Algorithm](https://mlsanity.com/paper/1802.04367/) (2018)
- [On the relation between optimal transport and Schr"odinger bridges: A stochastic control viewpoint](https://mlsanity.com/paper/1412.4430/) (2014)
- [Vocabulary Learning via Optimal Transport for Neural Machine Translation](https://mlsanity.com/paper/2012.15671/) (2020)
- [Meta Optimal Transport](https://mlsanity.com/paper/2206.05262/) (2022)
- [Cutoff for exact recovery of Gaussian mixture models](https://mlsanity.com/paper/2001.01194/) (2020)
- [Enhanced DeepONet for Modeling Partial Differential Operators Considering Multiple Input Functions](https://mlsanity.com/paper/2202.08942/) (2022)
- [Learning High Dimensional Wasserstein Geodesics](https://mlsanity.com/paper/2102.02992/) (2021)
- [Optimal Transport-based Coverage Control for Swarm Robot Systems: Generalization of the Voronoi Tessellation-based Method](https://mlsanity.com/paper/2011.08337/) (2020)
- [Learning to Generate Wasserstein Barycenters](https://mlsanity.com/paper/2102.12178/) (2021)

## Built on by (most cited first, 3 of 3)

- [New universal operator approximation theorem for encoder-decoder architectures](https://mlsanity.com/paper/2503.24092/) (2025)
- [Embedding Empirical Distributions for Computing Optimal Transport Maps](https://mlsanity.com/paper/2504.17740/) (2025)
- [Neural solver for Wasserstein Geodesics and optimal transport dynamics](https://mlsanity.com/paper/2602.22003/) (2026)

## Abstract (arXiv)

Optimal transport (OT) offers a versatile framework to compare complex data distributions in a geometrically meaningful way. Traditional methods for computing the Wasserstein distance and geodesic between probability measures require mesh-specific domain discretization and suffer from the curse-of-dimensionality. We present GeONet, a mesh-invariant deep neural operator network that learns the non-linear mapping from the input pair of initial and terminal distributions to the Wasserstein geodesic connecting the two endpoint distributions. In the offline training stage, GeONet learns the saddle point optimality conditions for the dynamic formulation of the OT problem in the primal and dual spaces that are characterized by a coupled PDE system. The subsequent inference stage is instantaneous and can be deployed for real-time predictions in the online learning setting. We demonstrate that GeONet achieves comparable testing accuracy to the standard OT solvers on simulation examples and the MNIST dataset with considerably reduced inference-stage computational cost by orders of magnitude.