# Universal Neural Optimal Transport

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

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 55 references that are in the map.
- First posted on arXiv 30 Nov 2022.

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

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [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)
- [Neural Operator: Graph Kernel Network for Partial Differential Equations](https://mlsanity.com/paper/2003.03485/) (2020)
- [On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport](https://mlsanity.com/paper/1805.09545/) (2018)
- [Invertible Residual Networks](https://mlsanity.com/paper/1811.00995/) (2018)
- [Regularisation of Neural Networks by Enforcing Lipschitz Continuity](https://mlsanity.com/paper/1804.04368/) (2018)
- [Joint Distribution Optimal Transportation for Domain Adaptation](https://mlsanity.com/paper/1705.08848/) (2017)
- [Generalized Sliced Wasserstein Distances](https://mlsanity.com/paper/1902.00434/) (2019)
- [Learning with a Wasserstein Loss](https://mlsanity.com/paper/1506.05439/) (2015)
- [Primal Wasserstein Imitation Learning](https://mlsanity.com/paper/2006.04678/) (2020)
- [Proximal Optimal Transport Modeling of Population Dynamics](https://mlsanity.com/paper/2106.06345/) (2021)
- [Reversible Architectures for Arbitrarily Deep Residual Neural Networks](https://mlsanity.com/paper/1709.03698/) (2017)
- [Sliced Wasserstein Kernels for Probability Distributions](https://mlsanity.com/paper/1511.03198/) (2015)
- [Supervised Training of Conditional Monge Maps](https://mlsanity.com/paper/2206.14262/) (2022)
- [Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networks](https://mlsanity.com/paper/2106.00774/) (2021)
- [Meta Optimal Transport](https://mlsanity.com/paper/2206.05262/) (2022)
- [Distilled Wasserstein Learning for Word Embedding and Topic Modeling](https://mlsanity.com/paper/1809.04705/) (2018)
- [Dataset Dynamics via Gradient Flows in Probability Space](https://mlsanity.com/paper/2010.12760/) (2020)
- [Rethinking Initialization of the Sinkhorn Algorithm](https://mlsanity.com/paper/2206.07630/) (2022)
- [On the explainable properties of 1-Lipschitz Neural Networks: An Optimal Transport Perspective](https://mlsanity.com/paper/2206.06854/) (2022)

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

- [Slicing Wasserstein Over Wasserstein Via Functional Optimal Transport](https://mlsanity.com/paper/2509.22138/) (2025)
- [DDEQs: Distributional Deep Equilibrium Models through Wasserstein Gradient Flows](https://mlsanity.com/paper/2503.01140/) (2025)
- [Measure-to-measure Regression with Transformers](https://mlsanity.com/paper/2605.28075/) (2026)
- [Experts-Guided Unbalanced Optimal Transport for ISP Learning from Unpaired and/or Paired Data](https://mlsanity.com/paper/2512.05635/) (2025)
- [C-kNN-LSH: A Nearest-Neighbor Algorithm for Sequential Counterfactual Inference](https://mlsanity.com/paper/2602.02371/) (2026)
- [Generative Modeling on Metric Graphs via Neural Optimal Transport](https://mlsanity.com/paper/2606.16273/) (2026)

## Abstract (arXiv)

Optimal Transport (OT) problems are a cornerstone of many applications, but solving them is computationally expensive. To address this problem, we propose UNOT (Universal Neural Optimal Transport), a novel framework capable of accurately predicting (entropic) OT distances and plans between discrete measures for a given cost function. UNOT builds on Fourier Neural Operators, a universal class of neural networks that map between function spaces and that are discretization-invariant, which enables our network to process measures of variable resolutions. The network is trained adversarially using a second, generating network and a self-supervised bootstrapping loss. We ground UNOT in an extensive theoretical framework. Through experiments on Euclidean and non-Euclidean domains, we show that our network not only accurately predicts OT distances and plans across a wide range of datasets, but also captures the geometry of the Wasserstein space correctly. Furthermore, we show that our network can be used as a state-of-the-art initialization for the Sinkhorn algorithm with speedups of up to 7.4×, significantly outperforming existing approaches.