NEURAL OPTIMAL TRANSPORT WITH GENERAL COST FUNCTIONALS

Abstract

We present a novel neural-networks-based algorithm to compute optimal transport (OT) plans and maps for general cost functionals. The algorithm is based on a saddle point reformulation of the OT problem and generalizes prior OT methods for weak and strong cost functionals. As an application, we construct a functional to map data distributions with preserving the class-wise structure of data.

1. INTRODUCTION

Optimal transport (OT) is a powerful framework to solve mass-moving and generative modeling problems for data distributions. Recent works (Korotin et al., 2022c; Rout et al., 2022; Korotin et al., 2021b; Fan et al., 2021a; Daniels et al., 2021) propose scalable neural methods to compute OT plans (or maps). They show that the learned transport plan (or map) can be used directly as the generative model in data synthesis (Rout et al., 2022) and unpaired learning (Korotin et al., 2022c; Rout et al., 2022; Daniels et al., 2021; Gazdieva et al., 2022) . Compared to WGANs (Arjovsky et al., 2017) which employ OT cost as the loss for generator (Rout et al., 2022, 3) , these methods provide better flexibility: the properties of the learned model can be controlled by the transport cost function. Existing neural OT plan (or map) methods consider distance-based cost functions, e.g., weak or strong quadratic costs (Korotin et al., 2022c; Fan et al., 2021a; Gazdieva et al., 2022) . Such costs are suitable for the tasks of unpaired image-to-image style translation (Zhu et al., 2017, Figures 1, 2 ) and image restoration (Lugmayr et al., 2020) . However, they do not take into account the class-wise structure of data or available side information, e.g., the class labels. As a result, such costs are hardly applicable to certain tasks such as the dataset transfer where the preservation the class-wise structure is needed (Figure 1 ). We tackle this issue. Contributions. We propose the extension of neural OT which allows to apply it to previously unreleased problems. For this, we develop a novel neural-networks-based algorithm to compute optimal transport plans for general cost functionals ( 3). As an example, we construct ( 4) and test ( 6) the functional for mapping data distributions with the preservation the class-wise structure. The learner has the access to labeled input data ∼ P and only partially labeled target data ∼ Q. Notation. The notation of our paper is based on that of (Paty & Cuturi, 2020; Korotin et al., 2022c) . For a compact Hausdorf space S we use P(S) to denote the set of Borel probability distributions on S. We denote the space of continuous R-valued functions on S endowed with the supremum norm by C(S). Its dual space is the space M(S) ⊃ P(S) of finite signed Borel measures over S. For a



Figure 1: The setup of class-guided dataset transfer. Input P = n α n P n , target Q = n β n Q n distributions are mixtures of N classes. The task is to learn a transport map T preserving the class.The learner has the access to labeled input data ∼ P and only partially labeled target data ∼ Q.

