Optimal Transportation — 1. Relationship with Deep Learning
By: Fei Sun
Tags: Convex Optimization, Probability Distributions, Deep Learning, Computational Geometry, Optimal Transport

Relationship with Deep Learning

Introduction

Recently, I have followed professor David Gu of Stony Brook university on some computational geometry stuff. It’s enlightening that fundamental mathematics have such profound impact on computer science and machine learning. Since my background is algebraic geometry/topology, I find this field smooth and exciting to follow. I decide to write a series of posts based on what I learned and show some cool stuff regarding optimal transportation (first).

I will try not to get too technical in the first post cause I don’t want to scare people away. I will briefly show how optimal transportation can improve deep learning, especially in GAN case.

Deep Learning and model collapse

GAN stands for Generative Adversarial Network. There are plenty of articles online introducing GAN model, so I will not go into details of it. In short, GAN is a generative model consisting of two neural networks, one is discriminator and another is generator. The discriminator network is used to judge whether an image is “fake” or “real”, which is usually pre-trained. The generator model will generate images taking input points from latent space and output an image. The latent space is a vector space and need to be trained. The GAN model stacks generator and discriminator together and a random point flows into generator then to discriminator. The output from generator model is manually labeled as “real”, the discriminator however, at the beginning, will deem it to be “fake” and the large error will force the back propagation to update generator weights to output better image. In this process, the latent space will be assigned meaning and eventually represent the output distribution.

Let us see more broadly into deep learning. Often time when we train a DL model, e.g, NLP models, we use embedding. A word or an object is embedded into a high dimensional matrix or space, so that each point is represented by a high dimensional vector in some space. The model is then trained to endow the space a meaning so that our training sample data will be distributed accordingly in this space. Now, it seems we can answer the question: what does deep learning really learn? It learns the probability distribution on a manifold.

However, the distribution of the latents points in latent space may not be continuous. This may cause problem, for example, what if our input point lies in the gap area of latent space? Then the model has no way of telling which distribution it comes from and will generate a very poor result. Let’s show this issue with a concrete example: the MNIST digit data set. Use GAN, we can generate digit image from a large latent space, say, 100 dimensional vector space. By choosing appropriate learning rate and number of epochs, the discriminator can do a pretty good job telling “real” and “fake” images

Accuracy of Discriminator Model

but the GAN model loss (binary cross-entropy) stays at a pretty high level even after 100 epochs. This suggests the generator may not be able to fool the discriminator even after long time training.

Let us plot the latent space after projecting it to a lower dimensional space (2-dimensional) using UMAP.

Immediately we see that the distribution is highly non-continuous and has large gaps. What would happen if we just use this 2-dimensional space as latent space to generate digit image?

There are clearly many ambiguous outputs and some of them are not recognizeable. No wonder the discriminator is not fooled and easily catches the “fake” ones.

Optimal Transportation

We can resolve this issue by using optimal transportation. Optimal transportation theory tries to solve the following problem: we have two measurable spaces

with probability distribution functions

Probability Distribution Function

We say a map T from X to Y measure preserving if

Measure Preserving Map

Then given a cost function c

Cost Function

we want to calculate the measure preserving transportation T (which is called the optimal transportation) that minimizes the cost

Monge Problem

The integral we try to minimize is often called the Wasserstein distance and the above expression is called the Monge problem. Via some mathematical manipulation, we can transfer Monge problem to Brenier problem (I’m hoping to give more maths details in this series!):

Brenier Problem

It can be shown that Brenier problem has a solution if we take quadratic cost function

Brenier Cost Function

The function u in the Brenier problem is called the Brenier potential (see diagrams below) whose gradient gives the transportation map T.

We apply optimal transportation theory from Brenier point of view to the 2-dimensional latent space given above. First, we endow a uniform distribution to the 2-dim latent space in the sense that each point has equal weight, this will be our Y space. Then we build the Brenier potential as follows:

We do a triangulation (Delaunay) of Y and the corresponding Voronoi diagram, and calculate each Voronoi cell’s area. We build Brenier potential by lift the Voronoi cells into 3-dimensional space as hyperplanes and calculate their upper envelope. The projection to 2D space of the Legendre dual of this upper envelope gives a weighted Delaunay triangulation of Y space. Iteratively adjust the heights of the liftings and upper envelope so that its projection to 2D space gives the Voronoi decomposition such that cells’ areas equal to the given weights. This gives the X space.

The Brenier potential obtained this way would be a piece-wise linear space, and the projection will give the Voronoi diagram of the X space so that each cell has measure the target weight, i.e, in our case, a uniform distributed space.

Brenier Potential

We see that the projected space obtained this way has no gap between distributions, and singular points have measure zero. If we run GAN from this space instead of original 2-dim latent space, the model collapsing issue would be drastically improved.

GAN after OT

It is worth noticing that majority of the generated digit is clean and unambiguous, but still some of them are fuzzy, e.g, the bottom row third to the right. Reason is the data point in the OT latent space is on the edge of two distributions, a.k.a, a singular point. Although singularities have zero measure (it’s codimension 1 in our space), it can happen that some points lie in this zero measure set. However, this is already a huge improvement of the original latent space.