Optimal Transport in WGANs: A Geometric Perspective

An important result of the use of optimal transport lies in the branch of applied mathematics, specifically in computer image generation. The emergence of WGAN (Wasserstein Generative Adversarial Network) as a Machine Learning algorithm has brought a substantial change to the traditional GAN (Generative Adversarial Network) algorithm.

The Space of Images and the Natural Image Manifold

The essential idea is the following: Let’s take as the space of all possible images, where the dimension is .

An arbitrary image will generally represent an image with noise; that is, it will be a random sample of different pixels without any meaningful correlation between them. Those images containing identifiable objects, which we will call natural images, are a tiny minority within the vast space of .

We assume as a hypothesis that the set of natural images forms a submanifold of . In other words, they lie on a lower-dimensional geometric surface that can be mapped locally with a function where . Furthermore, the probability measure associated with the data distribution of natural images is singular with respect to the Lebesgue measure in , meaning that .

The Problem with Traditional GANs

Within this context, GANs train a model by pitting two sets of data against each other: the real, natural images and the generated images. The model attempts to identify whether the generated images belong to the set of natural images.

The problem occurs when, after starting to train the model, the supports of the data distributions—those of the generated images and those of the natural images—lie on two completely different, disjoint manifolds.

The training model seeks the infimum of a loss function, which we will call . Traditional loss functions compare data point by point. With these standard loss functions, as we modify the model’s parameters, ends up converging to a constant, providing a gradient of zero. As a result, the model becomes stuck and stops learning.

The Wasserstein Solution

The Wasserstein distance (derived from optimal transport) gives us a fundamentally different interpretation.

While traditional cost functions look at the data vertically (comparing probabilities at the exact same point in space), the Wasserstein distance looks at it horizontally. It measures the “work” required to transport the probability mass from the generated distribution to the natural image distribution.

By introducing the Wasserstein distance into the training algorithm, we allow that, even if these two manifolds are separated, we can always find a meaningful gradient—a direction in which to train the model so that the generated images move closer to the natural images.