Neural Optimal Transport with General Cost Functionals

We introduce a novel neural network-based algorithm to compute optimal transport (OT) plans for general cost functionals. In contrast to common Euclidean costs, i.e., ℓ 1 or ℓ 2 , such functionals provide more flexibility and allow using auxiliary information, such as class labels, to construct the required transport map. Existing methods for general cost functionals are discrete and do not provide an out-of-sample estimation. We address the challenge of designing a continuous OT approach for general cost functionals in high-dimensional spaces, such as images. We construct two example functionals: one to map distributions while preserving the class-wise structure and the other one to preserve the given data pairs. Additionally, we provide the theoretical error analysis for our recovered transport plans. Our implementation is available at https://github.com/machinestein/gnot Figure 1: Results of our method with the pair-guided cost functional ( 6.2) applied to the supervised image-to-image translation task (Celeba-MaskHQ dataset, 256 × 256 images).