On amortizing convex conjugates for optimal transport
This addresses a computational bottleneck in optimal transport for researchers and practitioners, though it is incremental as it builds on existing methods with fine-tuning.
The paper tackles the difficulty of computing convex conjugates in Euclidean Wasserstein-2 optimal transport by using amortized optimization to approximate the conjugate, which significantly improves the quality of transport maps on a benchmark and models various 2-dimensional couplings and flows.
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.