ON AMORTIZING CONVEX CONJUGATES FOR OPTIMAL TRANSPORT

Abstract

This paper focuses on computing the convex conjugate operation that arises when solving Euclidean Wasserstein-2 optimal transport problems. This conjugation, which is also referred to as the Legendre-Fenchel conjugate or c-transform, is considered difficult to compute and in practice, Wasserstein-2 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 of the baselines, methods, and solvers in this paper are available at http://github. com/facebookresearch/w2ot.

1. INTRODUCTION

Optimal transportation (Villani, 2009; Ambrosio, 2003; Santambrogio, 2015; Peyré et al., 2019) is a thriving area of research that provides a way of connecting and transporting between probability measures. While optimal transport between discrete measures is well-understood, e.g. with Sinkhorn distances (Cuturi, 2013) , optimal transport between continuous measures is an open research topic actively being investigated (Genevay et al., 2016; Seguy et al., 2017; Taghvaei and Jalali, 2019; Korotin et al., 2019; Makkuva et al., 2020; Fan et al., 2021; Asadulaev et al., 2022) . Continuous OT has applications in generative modeling (Arjovsky et al., 2017; Petzka et al., 2017; Wu et al., 2018; Liu et al., 2019; Cao et al., 2019; Leygonie et al., 2019) , domain adaptation (Luo et al., 2018; Shen et al., 2018; Xie et al., 2019 ), barycenter computation (Li et al., 2020; Fan et al., 2020; Korotin et al., 2021b), and biology (Bunne et al., 2021; 2022; Lübeck et al., 2022) . This paper focuses on estimating the Wasserstein-2 transport map between measures α and β in Euclidean space, i.e. supp(α) = supp(β) = R n with the Euclidean distance as the transport cost. The Wasserstein-2 transport map, T : R n → R n , is the solution to Monge's primal formulation: T ∈ arg inf T ∈T (α,β) E x∼α x -T (x) 2 2 , where T (α, β) := {T : T # α = β} is the set of admissible couplings and the push-forward operator # is defined by T # α(B) := α(T -1 (B)) for a measure α, measurable map T , and all measurable sets B. T exists and is unique under general settings, e.g. as in Santambrogio (2015, Theorem 1.17), and is often difficult to solve because of the coupling constraints T . Almost every computational method instead solves the Kantorovich dual, e.g. as formulated in Villani (2009, §5) and Peyré et al. (2019, §2.5 ). This paper focuses on the dual associated with the negative inner product cost (Villani, 2009, eq. 5.12) , which introduces a dual potential function f : R n → R and solves: " f ∈ arg sup f ∈L 1 (α) -E x∼α [f (x)] -E y∼β [f (y)] where L 1 (α) is the space of measurable functions that are Lebesgue-integrable over α and f is the convex conjugate, or Legendre-Fenchel transform, of a function f defined by: f (y) := -inf x∈X J f (x; y) with objective J f (x; y) := f (x) -x, y . (3)

