Implementing WGANs for Stable Training

Hello! Welcome to our next lesson on generative models.

In our last session, we implemented a Deep Convolutional GAN (DCGAN). We saw how a set of specific architectural guidelines—like using strided convolutions, batch normalization, and specific activation functions—dramatically improved the stability and output quality compared to a simple MLP-based GAN. However, even with these improvements, the fundamental training process can still be brittle. The loss function we used, Binary Cross-Entropy, doesn't always correlate with image quality, and training can still suffer from vanishing gradients or mode collapse.

Today, we will address these core issues by changing the loss function itself. Your learning outcome is to implement Wasserstein GANs (WGANs) for improved training stability. This is a significant theoretical step forward, introducing a loss function grounded in a more robust measure of distance between probability distributions.

We'll cover:

• The theoretical flaws of the standard GAN loss function (which approximates Jensen-Shannon divergence).
• The concept of Wasserstein distance (or Earth Mover's distance) and why it's a better metric for training GANs.
• The architectural and loss function changes required to build a WGAN, including the role of the "critic" and the Lipschitz constraint.
• Implementing the original WGAN using weight clipping.
• Identifying the drawbacks of weight clipping and implementing the improved WGAN with Gradient Penalty (WGAN-GP).

By the end of this lesson, you'll understand and be able to implement a more robust and principled GAN training framework that has become a cornerstone of modern generative modeling.

1. The Trouble with Divergence: Why GANs are Hard to Train

The core problem with vanilla GAN training lies in its loss function. The minimax game with an optimal discriminator aims to minimize the Jensen-Shannon (JS) divergence between the real data distribution $p_r$ and the generated data distribution $p_g$.

This sounds reasonable, but it has a major practical flaw. As your previous experience with programming might suggest, a good cost function should provide a useful signal (a non-zero gradient) to guide optimization. The JS divergence fails to do this when the two distributions are "disjoint"—that is, when they don't have significant overlap.

In high-dimensional spaces like images, it's almost certain that the manifold of real images and the manifold of generated images will be disjoint, especially early in training.

As the article explains, when the distributions are disjoint, the discriminator can become "perfect" very quickly, classifying real and fake images with 100% accuracy. For the generator, this is a dead end: the discriminator's gradient becomes zero, providing no information on how to improve. It's like a student getting a score of 0, but with no feedback on what they did wrong.

2. A Better Metric: The Wasserstein (Earth Mover's) Distance

To solve the vanishing gradient problem, the WGAN paper proposed using a different metric: the Wasserstein-1 distance, also known as the Earth Mover's (EM) distance.

The intuition is simple: imagine the two probability distributions are two piles of dirt. The EM distance is the minimum "cost" to transform one pile into the shape of the other. The cost is defined as the amount of dirt moved multiplied by the distance it is moved.

Unlike JS divergence, which gives a binary "they are different" or "they are the same" signal for disjoint distributions, the Wasserstein distance provides a smooth, continuous measure of how far apart they are. This gives the generator a meaningful gradient to follow, even when its outputs are easily distinguishable from real data.

The WGAN Loss Function

Calculating the EM distance directly is intractable. However, the Kantorovich-Rubinstein duality allows us to re-express it in a more usable form:

$W(p_r,p_g)={\sup }_{\Vert f{\Vert }_L\le 1}{\mathbb{E}}_{x\sim p_r}[f(x)]-{\mathbb{E}}_{x\sim p_g}[f(x)]$

This looks complex, but the idea is powerful:

1. We need to find a function $f$ that has a Lipschitz constant of at most 1 (it is 1-Lipschitz). This essentially means the function's gradient norm is bounded by 1 everywhere.
2. The GAN's discriminator is repurposed into a "critic". Its job is no longer to output a probability, but to learn the function $f$ that maximizes the difference between the scores for real images (${\mathbb{E}}_{x\sim p_r}[f(x)]$) and fake images (${\mathbb{E}}_{x\sim p_g}[f(x)]$).
3. The value of this maximization is the Wasserstein distance, which becomes the critic's loss.
4. The generator's loss is simply $-{\mathbb{E}}_{x\sim p_g}[f(x)]$, as it tries to produce images that get a high score from the critic.

3. Implementing the Original WGAN with Weight Clipping

Let's put this theory into practice. Our starting point will be the DCGAN code from the previous lesson. The changes are surprisingly minimal. The first WGAN paper proposed a simple (though flawed) way to enforce the 1-Lipschitz constraint: weight clipping. After each critic update, we simply clamp all its weights to a small range, like [-0.01, 0.01].

This WGAN with weight clipping is already a huge improvement. The loss of the critic now meaningfully correlates with the quality of the generated images—as the loss goes down, the image quality goes up. This gives us a reliable metric to track training progress, something we lacked with standard GANs.

4. An Improvement: WGAN with Gradient Penalty (WGAN-GP)

Weight clipping is a simple hack, but it has significant problems.

• Reduced Model Capacity: It forces the network to learn a much simpler function than it might need to by pushing weights to the boundaries (-c and +c).
• Vanishing/Exploding Gradients: If the clipping value c is too small, gradients can vanish. If it's too large, it takes a long time to have any effect, making training unstable.

A much better way to enforce the Lipschitz constraint is with a gradient penalty. The theory states that a differentiable function is 1-Lipschitz if and only if its gradients have a norm of at most 1 everywhere.

The WGAN-GP paper proposes adding a "penalty" term to the critic's loss. This term punishes the critic if the norm of its gradient (with respect to its input) deviates from 1.

The penalty is calculated on images $\hat{x}$ that are interpolated between real ($x_r$) and fake ($x_g$) samples: $\hat{x}=ϵx_r+(1-ϵ)x_g$.

The new critic loss becomes:
$L_{critic}=\underset{\text{Original WGAN Loss}}{\underbrace{\mathbb{E}[f(x_g)]-\mathbb{E}[f(x_r)]}}+\underset{\text{Gradient Penalty}}{\underbrace{\lambda \mathbb{E}[(\Vert {\nabla }_{\hat{x}}f(\hat{x}){\Vert }_2-1{)}^2]}}$

5. Implementing WGAN-GP

This brings us to the modern, standard implementation of WGAN. The main change is to replace the weight clipping loop with a function that calculates this gradient penalty and adds it to the loss. Your CS background will be useful here, as this involves using PyTorch's autograd capabilities more directly.

Conclusion

Congratulations on completing this lesson! You've tackled one of the most important theoretical and practical advancements in GANs. Moving from a standard GAN to a WGAN-GP fundamentally changes the training dynamic, making it more stable and interpretable.

Key Takeaways:

• Standard GANs suffer from training instability because their loss function (approximating JS divergence) provides vanishing gradients for disjoint distributions.
• Wasserstein distance offers a smoother metric that provides useful gradients even when the critic can easily distinguish real from fake.
• WGANs repurpose the discriminator into a critic that outputs a real-valued score. The training goal is to estimate the Wasserstein distance.
• This requires enforcing a 1-Lipschitz constraint on the critic.
• Weight clipping was the first method proposed but is flawed, as it reduces model capacity and is sensitive to hyperparameters.
• WGAN-GP is the superior, modern approach, using a gradient penalty to enforce the constraint, leading to much more stable training and higher-quality results.

Preview of the next lesson:
Now that we have a stable foundation for GAN training with WGAN-GP, we can explore more advanced architectures that offer greater control over the generation process and can produce even higher-fidelity images. In our next lesson, we will dive into advanced GAN architectures, focusing on Conditional GANs (cGANs), which allow us to guide image generation with labels, and getting a first look at the ideas behind StyleGAN, a model famous for its photorealistic outputs.