# On amortizing convex conjugates for optimal transport

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

## Key facts

- Cited by 31 papers in the map (data through 14 Oct 2026).
- Builds on 42 of its 97 references that are in the map.
- First posted on arXiv 21 Oct 2022.

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [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)
- [Identity Mappings in Deep Residual Networks](https://mlsanity.com/paper/1603.05027/) (2016)
- [Glow: Generative Flow with Invertible 1x1 Convolutions](https://mlsanity.com/paper/1807.03039/) (2018)
- [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)
- [Identifying and attacking the saddle point problem in high-dimensional non-convex optimization](https://mlsanity.com/paper/1406.2572/) (2014)
- [Learning to Optimize: A Primer and A Benchmark](https://mlsanity.com/paper/2103.12828/) (2021)
- [Optimal transport mapping via input convex neural networks](https://mlsanity.com/paper/1908.10962/) (2019)
- [Stochastic Optimization for Large-scale Optimal Transport](https://mlsanity.com/paper/1605.08527/) (2016)
- [Large-Scale Optimal Transport and Mapping Estimation](https://mlsanity.com/paper/1711.02283/) (2017)
- [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)
- [Convex Potential Flows: Universal Probability Distributions with Optimal Transport and Convex Optimization](https://mlsanity.com/paper/2012.05942/) (2020)
- [Generative Modeling with Optimal Transport Maps](https://mlsanity.com/paper/2110.02999/) (2021)
- [Large-Scale Wasserstein Gradient Flows](https://mlsanity.com/paper/2106.00736/) (2021)
- [Continuously Differentiable Exponential Linear Units](https://mlsanity.com/paper/1704.07483/) (2017)
- [On the regularization of Wasserstein GANs](https://mlsanity.com/paper/1709.08894/) (2017)
- [Supervised Training of Conditional Monge Maps](https://mlsanity.com/paper/2206.14262/) (2022)
- [Faster Matchings via Learned Duals](https://mlsanity.com/paper/2107.09770/) (2021)
- [Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networks](https://mlsanity.com/paper/2106.00774/) (2021)
- [Continuous Wasserstein-2 Barycenter Estimation without Minimax Optimization](https://mlsanity.com/paper/2102.01752/) (2021)
- [Scalable Computations of Wasserstein Barycenter via Input Convex Neural Networks](https://mlsanity.com/paper/2007.04462/) (2020)
- [Neural Optimal Transport with General Cost Functionals](https://mlsanity.com/paper/2205.15403/) (2022)
- [2-Wasserstein Approximation via Restricted Convex Potentials with Application to Improved Training for GANs](https://mlsanity.com/paper/1902.07197/) (2019)
- [Multi-marginal Wasserstein GAN](https://mlsanity.com/paper/1911.00888/) (2019)
- [Continuous Regularized Wasserstein Barycenters](https://mlsanity.com/paper/2008.12534/) (2020)

## Built on by (most cited first, 30 of 31)

- [Multisample Flow Matching: Straightening Flows with Minibatch Couplings](https://mlsanity.com/paper/2304.14772/) (2023)
- [Optimal Flow Matching: Learning Straight Trajectories in Just One Step](https://mlsanity.com/paper/2403.13117/) (2024)
- [Entropic Neural Optimal Transport via Diffusion Processes](https://mlsanity.com/paper/2211.01156/) (2022)
- [Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation](https://mlsanity.com/paper/2311.15100/) (2023)
- [The Monge Gap: A Regularizer to Learn All Transport Maps](https://mlsanity.com/paper/2302.04953/) (2023)
- [Conditional Optimal Transport on Function Spaces](https://mlsanity.com/paper/2311.05672/) (2023)
- [Light Schr"odinger Bridge](https://mlsanity.com/paper/2310.01174/) (2023)
- [Neural Optimal Transport with Lagrangian Costs](https://mlsanity.com/paper/2406.00288/) (2024)
- [Energy-Guided Continuous Entropic Barycenter Estimation for General Costs](https://mlsanity.com/paper/2310.01105/) (2023)
- [Progressive Entropic Optimal Transport Solvers](https://mlsanity.com/paper/2406.05061/) (2024)
- [Robust Barycenter Estimation using Semi-Unbalanced Neural Optimal Transport](https://mlsanity.com/paper/2410.03974/) (2024)
- [Conditional simulation via entropic optimal transport: Toward non-parametric estimation of conditional Brenier maps](https://mlsanity.com/paper/2411.07154/) (2024)
- [Computing high-dimensional optimal transport by flow neural networks](https://mlsanity.com/paper/2305.11857/) (2023)
- [ENOT: Expectile Regularization for Fast and Accurate Training of Neural Optimal Transport](https://mlsanity.com/paper/2403.03777/) (2024)
- [Convexity in ReLU Neural Networks: beyond ICNNs?](https://mlsanity.com/paper/2501.03017/) (2025)
- [Fast and scalable Wasserstein-1 neural optimal transport solver for single-cell perturbation prediction](https://mlsanity.com/paper/2411.00614/) (2024)
- [Interpolation for Robust Learning: Data Augmentation on Wasserstein Geodesics](https://mlsanity.com/paper/2302.02092/) (2023)
- [On a Neural Implementation of Brenier's Polar Factorization](https://mlsanity.com/paper/2403.03071/) (2024)
- [Scalable unsupervised alignment of general metric and non-metric structures](https://mlsanity.com/paper/2406.13507/) (2024)
- [A Statistical Learning Perspective on Semi-dual Adversarial Neural Optimal Transport Solvers](https://mlsanity.com/paper/2502.01310/) (2025)
- [Sample and Map from a Single Convex Potential: Generation using Conjugate Moment Measures](https://mlsanity.com/paper/2503.10576/) (2025)
- [Neural Optimal Transport Meets Multivariate Conformal Prediction](https://mlsanity.com/paper/2509.25444/) (2025)
- [Multidimensional Uncertainty Quantification via Optimal Transport](https://mlsanity.com/paper/2509.22380/) (2025)
- [Deep Legendre Transform](https://mlsanity.com/paper/2512.19649/) (2025)
- [Amortized Maximum Inner Product Search with Learned Support Functions](https://mlsanity.com/paper/2603.08001/) (2026)
- [Hyperparameter Trajectory Inference with Conditional Lagrangian Optimal Transport](https://mlsanity.com/paper/2603.01771/) (2026)
- [Implicit Neural Optimal Transport via Fixed-Point Optimization](https://mlsanity.com/paper/2605.10792/) (2026)
- [A lift for input-convex neural net training](https://mlsanity.com/paper/2605.24274/) (2026)
- [A Biconvex Formulation for Stable Transport of Mixture Models with a Unique Solution](https://mlsanity.com/paper/2606.02515/) (2026)
- [Conditional Neural Optimal Transport for Predicting Cellular Phenotypes from Molecular Structure](https://mlsanity.com/paper/2608.14293/) (2026)

## Abstract (arXiv)

This paper focuses on computing the convex conjugate (also known as the Legendre-Fenchel conjugate or c-transform) that appears in Euclidean Wasserstein-2 optimal transport. This conjugation is considered difficult to compute and in practice, methods are limited by not being able to exactly conjugate the dual potentials in continuous space. To overcome this, the computation of the conjugate can be approximated with amortized optimization, which learns a model to predict the conjugate. I show that combining amortized approximations to the conjugate with a solver for fine-tuning significantly improves the quality of transport maps learned for the Wasserstein-2 benchmark by Korotin et al. (2021a) and is able to model many 2-dimensional couplings and flows considered in the literature. All baselines, methods, and solvers are publicly available at http://github.com/facebookresearch/w2ot.