Alternative Divergences: Wasserstein Distance

As we saw, vanilla GAN training often struggles with instability. The generator might produce only a narrow range of outputs (mode collapse), or the learning process might diverge entirely due to vanishing or exploding gradients. A significant contributing factor to these issues lies in the very nature of the objective function used: the Jensen-Shannon (JS) divergence.

While the JS divergence is a valid way to measure the difference between probability distributions, it has properties that make it less than ideal for the adversarial training dynamics of GANs, especially in high-dimensional spaces like images.

The Problem with Jensen-Shannon Divergence

The original GAN paper showed that minimizing the GAN objective corresponds to minimizing the JS divergence between the real data distribution $P_r$ and the generated data distribution $P_g$. However, consider what happens when $P_r$ and $P_g$ have supports that do not overlap significantly, or when they lie on low-dimensional manifolds within a much higher-dimensional space (a common scenario for image data).

In such cases, it's often possible to find a discriminator that perfectly separates real samples from generated samples. When the discriminator achieves near-perfect accuracy:

1. The JS divergence between $P_r$ and $P_g$ effectively saturates at its maximum value (log 2).
2. More importantly, the gradients of the loss function with respect to the generator's parameters approach zero.

Think about it: if the discriminator can always tell real from fake, it provides very little useful information back to the generator about how to improve. The generator essentially hits a wall, receiving no meaningful signal to guide its updates. This vanishing gradient problem is a primary cause of training failure and mode collapse, as the generator stops learning effectively.

We need a way to measure the distance between $P_r$ and $P_g$ that provides smoother, more informative gradients even when the distributions are quite different or don't overlap perfectly.

Introducing the Wasserstein Distance (Earth Mover's Distance)

Enter the Wasserstein-1 distance, often referred to as the Earth Mover's Distance (EMD). Imagine $P_r$ and $P_g$ as two different piles of dirt (probability mass). The EMD represents the minimum "cost" required to transform one pile into the other. The cost is defined as the amount of dirt moved multiplied by the distance it is moved.

The Wasserstein distance calculates the minimum cost to transport mass from the generated distribution $P_g$ to match the real distribution $P_r$.

Crucially, even if the two piles of dirt (distributions) are located far apart (disjoint supports), the cost of moving the dirt provides a meaningful measure of how different they are. Unlike JS divergence, the Wasserstein distance doesn't saturate abruptly and generally provides a smoother metric.

The Kantorovich-Rubinstein Duality

Calculating the EMD directly using its primal definition involves finding an optimal "transport plan" (how much dirt to move from each point in $P_g$ to each point in $P_r$), which is computationally intractable in most interesting cases. Fortunately, the Kantorovich-Rubinstein duality provides an alternative formulation for the Wasserstein-1 distance:

$$W_1(P_r, P_g) = \sup_{\|f\|_L \le 1} \left( \mathbb{E}_{x \sim P_r}[f(x)] - \mathbb{E}_{\tilde{x} \sim P_g}[f(\tilde{x})] \right)$$

Let's break this down:

• $W_1(P_r, P_g)$ is the Wasserstein-1 distance between the real distribution $P_r$ and the generated distribution $P_g$.
• $\sup$ denotes the supremum, which is the least upper bound (similar to a maximum).
• $\|f\|_L \le 1$ means the supremum is taken over all functions $f$ that are 1-Lipschitz continuous. A function $f$ is K-Lipschitz if $|f(x_1) - f(x_2)| \le K \|x_1 - x_2\|$ for all $x_1, x_2$ in its domain. Essentially, a 1-Lipschitz function's rate of change is bounded by 1. This constraint is fundamental.
• $\mathbb{E}_{x \sim P_r}[f(x)]$ is the expected value of $f(x)$ when $x$ is sampled from the real data distribution $P_r$.
• $\mathbb{E}_{\tilde{x} \sim P_g}[f(\tilde{x})]$ is the expected value of $f(\tilde{x})$ when $\tilde{x}$ is sampled from the generator's distribution $P_g$ (where $\tilde{x} = g(z)$, with $z$ sampled from some prior noise distribution $p(z)$).

This dual formulation states that the Wasserstein distance is the maximum possible difference in expectations achievable by evaluating a 1-Lipschitz function $f$ on samples from the two distributions.

Wasserstein GANs (WGANs)

The insight of the WGAN paper (Arjovsky et al., 2017) was to use this dual formulation to define a new GAN objective. The core idea is:

1. Replace the Discriminator with a Critic: Instead of a discriminator outputting the probability of a sample being real (a value between 0 and 1, typically via a sigmoid), we use a network, called the critic (let's denote it $f_w$ with parameters $w$), that outputs an unbounded real number.
2. Train the Critic to Approximate the Supremum: The critic $f_w$ is trained to make the term $(\mathbb{E}_{x \sim P_r}[f_w(x)] - \mathbb{E}_{\tilde{x} \sim P_g}[f_w(\tilde{x})])$ as large as possible. By doing so, under the right conditions (specifically, if $f_w$ remains within the space of 1-Lipschitz functions), the critic's objective value approximates the Wasserstein distance $W_1(P_r, P_g)$.
3. Train the Generator to Minimize the Critic's Output: The generator $g_\theta$ (with parameters $\theta$) is trained to produce samples $\tilde{x} = g_\theta(z)$ that make the critic's output $f_w(\tilde{x})$ larger (closer to the output for real samples), thereby minimizing the difference the critic is trying to maximize. This effectively minimizes the estimated Wasserstein distance.

The objectives become:

• Critic Loss: Maximize $L_{\text{Critic}} = \mathbb{E}_{x \sim P_r}[f_w(x)] - \mathbb{E}_{z \sim p(z)}[f_w(g_\theta(z))]$. In practice, neural networks are trained via gradient descent (minimization), so we typically minimize $-L_{\text{Critic}}$.
• Generator Loss: Minimize $L_{\text{Generator}} = - \mathbb{E}_{z \sim p(z)}[f_w(g_\theta(z))]$. The generator tries to fool the critic by producing samples that score highly (as if they were real).

Why This Helps Stability

Using the Wasserstein distance via the critic provides several advantages:

1. Meaningful Gradients: The Wasserstein distance provides gradients that are non-zero and informative even when $P_r$ and $P_g$ have disjoint supports. This drastically reduces the vanishing gradient problem that plagued original GANs.
2. Improved Correlation with Sample Quality: Empirically, the WGAN loss (the estimated Wasserstein distance) often correlates much better with the perceptual quality of the generated samples during training. Lower loss generally means better-looking images, which wasn't always true for the original GAN loss.
3. Reduced Risk of Mode Collapse: Theoretical arguments and empirical results suggest that WGANs are less prone to mode collapse, likely due to the more stable gradients allowing the generator to explore the data distribution more effectively.

The Lipschitz Constraint: A Practical Hurdle

The theory behind WGAN relies heavily on the critic $f_w$ being (or approximating) a 1-Lipschitz function. If the critic's weights grow too large, it can violate this constraint, making the approximation of the Wasserstein distance invalid and potentially leading back to unstable training.

Therefore, an important part of implementing WGANs is enforcing this Lipschitz constraint on the critic network. The original WGAN paper proposed a simple, albeit sometimes problematic, method called weight clipping. Subsequent research introduced more sophisticated and often preferred techniques like the gradient penalty (WGAN-GP) and spectral normalization. We will examine these methods for enforcing the Lipschitz constraint in the following sections.