Wasserstein-2 Generative Networks

TL;DR

This paper introduces a new end-to-end algorithm for training Wasserstein-2 generative models that avoids bias and scales well, using input convex neural networks and cycle-consistency regularization.

Contribution

It presents a novel non-minimax algorithm for Wasserstein-2 optimal transport that leverages cycle-consistency and input convex neural networks, improving scalability and bias reduction.

Findings

The algorithm accurately estimates Wasserstein-2 distances.

It performs well on image-to-image translation tasks.

Theoretical analysis of the generative mapping properties.

Abstract

We propose a novel end-to-end non-minimax algorithm for training optimal transport mappings for the quadratic cost (Wasserstein-2 distance). The algorithm uses input convex neural networks and a cycle-consistency regularization to approximate Wasserstein-2 distance. In contrast to popular entropic and quadratic regularizers, cycle-consistency does not introduce bias and scales well to high dimensions. From the theoretical side, we estimate the properties of the generative mapping fitted by our algorithm. From the practical side, we evaluate our algorithm on a wide range of tasks: image-to-image color transfer, latent space optimal transport, image-to-image style transfer, and domain adaptation.