On amortizing convex conjugates for optimal transport

TL;DR

This paper introduces an amortized optimization approach to efficiently compute convex conjugates in Wasserstein-2 optimal transport, improving transport map quality and modeling capabilities in continuous spaces.

Contribution

It proposes a novel method combining amortized approximation with fine-tuning to enhance conjugate computation in optimal transport, enabling better transport maps and modeling in continuous settings.

Findings

Improved transport map quality on Wasserstein-2 benchmarks

Able to model complex 2D couplings and flows

Method and code are publicly available

Abstract

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.