Training Generative Networks with general Optimal Transport distances

TL;DR

This paper introduces a novel algorithm for training generative networks using flexible optimal transport distances, enabling explicit use of various cost functions like Wasserstein-2 and image-centered metrics for improved stability and results.

Contribution

The method allows explicit incorporation of any transportation cost function in training generative networks, extending beyond Euclidean distances used in prior approaches like WGANs.

Findings

Enables use of Wasserstein-2 metric for stable training.

Incorporates image-centered distances like SSIM.

Results show notable improvements in generative quality.

Abstract

We propose a new algorithm that uses an auxiliary neural network to express the potential of the optimal transport map between two data distributions. In the sequel, we use the aforementioned map to train generative networks. Unlike WGANs, where the Euclidean distance is ${\it implicitly}$ used, this new method allows to ${\it explicitly}$ use ${\it any}$ transportation cost function that can be chosen to match the problem at hand. For example, it allows to use the squared distance as a transportation cost function, giving rise to the Wasserstein-2 metric for probability distributions, which results in fast and stable gradient descends. It also allows to use image centered distances, like the structure similarity index, with notable differences in the results.