WGAN's success may have nothing to do with the Wasserstein distance

WGAN的成功，可能跟Wasserstein距离没啥关系 苏剑林 · 2021-03-15 · Machine Learning optimization GAN generative model

WGAN, short for Wasserstein GAN, is regarded as one of the more important theoretical breakthroughs in the history of GANs. It replaced the f-divergence used to measure the discrepancy between two probability distributions in GANs with the Wasserstein distance, which made WGAN's training process much more stable and generally produced better generation quality as well. The Wasserstein distance is related to optimal transport and belongs to the family of Integral Probability Metrics (IPMs), a class of probability metrics that typically enjoy nicer theoretical properties. As a result, WGAN's emergence attracted many people to study and understand GANs from the perspective of optimal transport and IPMs.

However, a recent paper on Arxiv, Wasserstein GANs Work Because They Fail (to Approximate the Wasserstein Distance), points out that although WGAN was derived from the Wasserstein GAN formulation, the WGANs that actually work well today do not approximate the Wasserstein distance particularly well — and in fact, if we make WGAN approximate the Wasserstein distance more accurately, the results actually get worse. In truth, this is something I've long wondered about myself: the Wasserstein distance itself doesn't inherently explain why it should improve GAN performance. This paper's conclusions confirm that suspicion, so the real reason why GANs succeed remains rather mysterious. More below the fold.

Background and review

This post is a discussion of the WGAN training process, and isn't really an introductory piece. For those new to GANs, I recommend The Art of Mutual Sparring: Straight to WGAN-GP; for the connection between f-divergences and GANs, see A Brief Introduction to f-GAN: The GAN Model Factory and Designing GANs: Another GAN Factory; for the theoretical derivation of WGAN, see From the Wasserstein Distance and Duality Theory to WGAN; and for an analysis of the GAN training process, see A Dynamical-Systems View of Optimization Algorithms (Part 4): The Third Stage of GAN Training.

Generally speaking, a GAN corresponds to a min-max process:

\begin{equation}\min_G \max_D \mathcal{L}(D, G)\end{equation}

Of course, in general the loss functions of the discriminator and generator may differ, but the form above is already representative enough. The original, most basic GAN is usually called the vanilla GAN, and takes the form

\begin{equation}\min_G \max_D \mathbb{E}_{x\sim p(x)}\left[\log D(x)\right] + \mathbb{E}_{z\sim q(z)}\left[\log (1 - D(G(z)))\right]\label{eq:v-gan}\end{equation}

As can be seen in Towards Principled Methods for Training Generative Adversarial Networks, A Stunning Explanation of Wasserstein GAN, or various GAN-related posts on this blog, the vanilla GAN is effectively minimizing the JS divergence between the two distributions. The JS divergence is one particular f-divergence, and all f-divergences share a common problem: when two distributions have almost no overlap, the divergence becomes a constant, which means the gradient is zero. Since we optimize via gradient descent, this means we can't optimize properly in that regime. This is exactly why WGAN came about — it uses the Wasserstein distance to design a new kind of GAN:

\begin{equation}\min_G \max_{\Vert D\Vert_{L}\leq 1} \mathbb{E}_{x\sim p(x)}\left[D(x)\right] - \mathbb{E}_{z\sim q(z)}\left[D(G(z))\right]\label{eq:w-gan}\end{equation}

A clear distinction from earlier GANs is that WGAN explicitly imposes a Lipschitz constraint $\Vert D\Vert_{L}\leq 1$ on the discriminator $D$. Since the Wasserstein distance is fairly well-defined for almost any pair of distributions (even ones with no overlap), WGAN in theory resolves the vanishing-gradient and unstable-training issues that plague traditional f-divergence-based GANs.

There are two main approaches for imposing a Lipschitz constraint on the discriminator: one is Spectral Normalization (SN), which you can read about in Lipschitz Constraints in Deep Learning: Generalization and Generative Models — nowadays many GANs (not just WGAN variants) apply spectral normalization to the discriminator or even the generator to stabilize training. The other is Gradient Penalty (GP), which includes both a penalty centered at 1 (WGAN-GP) and one centered at 0 (WGAN-div); see WGAN-div: An Unsung Hero That Fills in WGAN's Gaps. Current results suggest the zero-centered penalty has better theoretical properties and empirical performance.

Performance ≠ approximation

In fact, the observation that "WGAN doesn't approximate the Wasserstein distance very well" isn't new — a 2019 paper, How Well Do WGANs Estimate the Wasserstein Metric?, already discussed this systematically. The paper we're covering here instead sets up fairly rigorous experiments to pin down the relationship between how well WGAN performs and how accurately it approximates the Wasserstein distance.

First, the paper compares the performance of gradient penalty (GP) versus a method called $c\text{-transform}$ for implementing WGAN. $c\text{-transform}$ was likewise proposed in How Well Do WGANs Estimate the Wasserstein Metric?, and it approximates the Wasserstein distance better than gradient penalty does. The two figures below illustrate this:

Approximation quality of WGAN-GP vs. c-transform relative to the Wasserstein distance in a static test
Approximation quality of WGAN-GP vs. c-transform relative to the Wasserstein distance during training

However, the generation quality of $c\text{-transform}$ turns out to be worse than that of gradient penalty:

Comparison of generation quality between WGAN-GP and c-transform

To be honest, the figure the original paper chose here is a bit unfortunate — in practice WGAN-GP's results can be far better than what's shown on the right. Still, we can tentatively conclude:

1. WGANs that perform well in practice do not approximate the Wasserstein distance well during training;
2. Approximating the Wasserstein distance more accurately does not actually help improve generation quality.

Theory ≠ experiment

Now let's think about where the problem lies. Whether it's the original GAN $\eqref{eq:v-gan}$, WGAN $\eqref{eq:w-gan}$, or any other GAN variant, in practice they all share two common features:

1. $\min$ and $\max$ are trained alternately;
2. Each training step uses only a randomly sampled batch.

What's problematic about these two features?

First, almost every GAN formulation is written as $\min\limits_G \max\limits_D$, which in theory requires first solving $\max\limits_D$ exactly, and only then moving on to $\min\limits_G$, in order to actually be optimizing the probability metric that the GAN corresponds to. If instead we only alternate the optimization, then in theory there's no way to approximate the probability metric very precisely. Even though WGAN, thanks to using the Wasserstein distance, doesn't suffer from vanishing gradients and so can afford to train the discriminator for several extra steps $D$ per alternation (or use a larger learning rate $D$ for it), it still can't precisely approximate the Wasserstein distance. This is one source of the discrepancy.

Second, training on a randomly sampled batch rather than the full training set leads to a curious consequence: the Wasserstein distance between two randomly chosen batches from the training set can actually be larger than the Wasserstein distance between a training batch and the average of its samples, as shown below:

Left: a batch of real samples. Middle: the average sample. Right: sample cluster centers. Judged by Wasserstein distance, the real samples actually score worse than the two blurry ones on the right

This tells us that, under batch-based training, if you want more realistic-looking samples, you necessarily cannot be optimizing the Wasserstein distance precisely — and conversely, if you are optimizing the Wasserstein distance very precisely, you won't get more realistic samples, because the blurry, averaged samples actually achieve a smaller Wasserstein distance.

Mathematics ≠ vision

Mathematically speaking, the properties of the Wasserstein distance really are quite elegant; in some sense it's the best possible way to measure the discrepancy between any two distributions. But mathematics is one thing, and the Achilles' heel of the Wasserstein distance is that it depends on a specific choice of ground metric:

\begin{equation}\mathcal{W}[p,q]=\inf_{\gamma\in \Pi[p,q]} \iint \gamma(x,y) d(x,y) dxdy\end{equation}

That is, we need to specify beforehand a function $d(x,y)$ that measures the discrepancy between two samples. For many practical scenarios — say, comparing two images — designing such a metric function is itself one of the hardest problems imaginable. WGAN simply uses the Euclidean distance $d(x,y)=\Vert x - y\Vert_2$, which, while mathematically sound, is not sound from the standpoint of visual perception: two images that we perceive as very similar need not have a small Euclidean distance. So the more precisely we try to approximate the Wasserstein distance, the more the visual quality can actually suffer. The original paper also ran an experiment: using $c\text{-transform}$ to get a better approximation of the Wasserstein distance, the resulting generated samples end up looking similar to K-Means cluster centers — and K-Means, of course, also uses Euclidean distance as its metric:

Similarity between c-transform's results and K-Means

So now the reason for WGAN's success becomes rather mysterious: WGAN was derived from the Wasserstein distance, yet its actual implementation deviates somewhat from the Wasserstein distance — and it's quite possible that this very deviation is the key to WGAN's success. The original paper argues that the crucial ingredient in WGAN is the introduction of the Lipschitz constraint: adding a Lipschitz constraint (via spectral normalization or gradient penalty) to more or less any GAN variant tends to improve both performance and stability to some degree. So it's the Lipschitz constraint that's doing the heavy lifting, rather than the Wasserstein distance as one might have imagined.

But this remains more of a conclusion than a theoretical analysis. It seems there's still a long way to go before we truly understand GANs at a deep level.

Brief summary

This post has mainly shared a recent paper which points out that how well WGAN approximates the Wasserstein distance has no necessary bearing on how well WGAN actually performs. Reaching a better understanding of GAN theory and practice remains, it seems, a genuinely difficult task.