# ENOT: Expectile Regularization for Fast and Accurate Training of Neural Optimal Transport

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

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 35 of its 57 references that are in the map.
- First posted on arXiv 6 Mar 2024.

## Builds on (most cited first, 30 of 35)

- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [A Style-Based Generator Architecture for Generative Adversarial Networks](https://mlsanity.com/paper/1812.04948/) (2018)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Improved Training of Wasserstein GANs](https://mlsanity.com/paper/1704.00028/) (2017)
- [StarGAN v2: Diverse Image Synthesis for Multiple Domains](https://mlsanity.com/paper/1912.01865/) (2019)
- [Fast and Accurate Deep Network Learning by Exponential Linear Units (ELUs)](https://mlsanity.com/paper/1511.07289/) (2015)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [Optimal transport mapping via input convex neural networks](https://mlsanity.com/paper/1908.10962/) (2019)
- [Neural Optimal Transport](https://mlsanity.com/paper/2201.12220/) (2022)
- [Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein](https://mlsanity.com/paper/2201.12324/) (2022)
- [Wasserstein-2 Generative Networks](https://mlsanity.com/paper/1909.13082/) (2019)
- [Do Neural Optimal Transport Solvers Work? A Continuous Wasserstein-2 Benchmark](https://mlsanity.com/paper/2106.01954/) (2021)
- [Proximal Optimal Transport Modeling of Population Dynamics](https://mlsanity.com/paper/2106.06345/) (2021)
- [Generative Modeling with Optimal Transport Maps](https://mlsanity.com/paper/2110.02999/) (2021)
- [Reinforcement Learning from Passive Data via Latent Intentions](https://mlsanity.com/paper/2304.04782/) (2023)
- [Watch and Match: Supercharging Imitation with Regularized Optimal Transport](https://mlsanity.com/paper/2206.15469/) (2022)
- [Recent Advances in Optimal Transport for Machine Learning](https://mlsanity.com/paper/2306.16156/) (2023)
- [Scalable Computations of Wasserstein Barycenter via Input Convex Neural Networks](https://mlsanity.com/paper/2007.04462/) (2020)
- [Cross-Domain Imitation Learning via Optimal Transport](https://mlsanity.com/paper/2110.03684/) (2021)
- [Optimal Transport for Offline Imitation Learning](https://mlsanity.com/paper/2303.13971/) (2023)
- [2-Wasserstein Approximation via Restricted Convex Potentials with Application to Improved Training for GANs](https://mlsanity.com/paper/1902.07197/) (2019)
- [The Monge Gap: A Regularizer to Learn All Transport Maps](https://mlsanity.com/paper/2302.04953/) (2023)
- [Kernel Neural Optimal Transport](https://mlsanity.com/paper/2205.15269/) (2022)
- [On amortizing convex conjugates for optimal transport](https://mlsanity.com/paper/2210.12153/) (2022)
- [Imitation Learning with Sinkhorn Distances](https://mlsanity.com/paper/2008.09167/) (2020)
- [On Scalable and Efficient Computation of Large Scale Optimal Transport](https://mlsanity.com/paper/1905.00158/) (2019)
- [Scalable Optimal Transport Methods in Machine Learning: A Contemporary Survey](https://mlsanity.com/paper/2305.05080/) (2023)
- [A gradual, semi-discrete approach to generative network training via explicit Wasserstein minimization](https://mlsanity.com/paper/1906.03471/) (2019)
- [Adversarial Computation of Optimal Transport Maps](https://mlsanity.com/paper/1906.09691/) (2019)

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

- [Hierarchical Refinement: Optimal Transport to Infinity and Beyond](https://mlsanity.com/paper/2503.03025/) (2025)
- [A Statistical Learning Perspective on Semi-dual Adversarial Neural Optimal Transport Solvers](https://mlsanity.com/paper/2502.01310/) (2025)
- [Displacement-Sparse Neural Optimal Transport](https://mlsanity.com/paper/2502.01889/) (2025)
- [HOTA: Hamiltonian framework for Optimal Transport Advection](https://mlsanity.com/paper/2507.17513/) (2025)
- [Quadratically Regularized Optimal Transport: Localization Bounds and Affine Case Analysis](https://mlsanity.com/paper/2605.24644/) (2026)
- [A Biconvex Formulation for Stable Transport of Mixture Models with a Unique Solution](https://mlsanity.com/paper/2606.02515/) (2026)

## Abstract (arXiv)

We present a new approach for Neural Optimal Transport (NOT) training procedure, capable of accurately and efficiently estimating optimal transportation plan via specific regularization on dual Kantorovich potentials. The main bottleneck of existing NOT solvers is associated with the procedure of finding a near-exact approximation of the conjugate operator (i.e., the c-transform), which is done either by optimizing over non-convex max-min objectives or by the computationally intensive fine-tuning of the initial approximated prediction. We resolve both issues by proposing a new, theoretically justified loss in the form of expectile regularisation which enforces binding conditions on the learning process of dual potentials. Such a regularization provides the upper bound estimation over the distribution of possible conjugate potentials and makes the learning stable, completely eliminating the need for additional extensive fine-tuning. Proposed method, called Expectile-Regularised Neural Optimal Transport (ENOT), outperforms previous state-of-the-art approaches on the established Wasserstein-2 benchmark tasks by a large margin (up to a 3-fold improvement in quality and up to a 10-fold improvement in runtime). Moreover, we showcase performance of ENOT for varying cost functions on different tasks such as image generation, showing robustness of proposed algorithm. OTT-JAX library includes our implementation of ENOT algorithm https://ott-jax.readthedocs.io/en/latest/tutorials/ENOT.html