Beating Mode Collapse in GANs: From DCGAN to WGAN‑GP

Subtitle: A practical guide to generating diverse anime faces with Wasserstein GAN and gradient penalty

If you’ve ever trained a Generative Adversarial Network (GAN), you’ve likely faced the dreaded mode collapse — the generator finds a few “safe” outputs and repeats them endlessly. Your loss curves look great, but the generated images are all the same.

In this article, I’ll walk through a head‑to‑head comparison of two GAN variants:

• DCGAN — a strong convolutional baseline
• WGAN‑GP — an improved version that virtually eliminates mode collapse

We’ll train both on 64×64 anime faces, analyse the results, and deploy an interactive Gradio app so you can test them live.

The Problem: Mode Collapse

Standard GANs minimise the Jensen‑Shannon divergence between real and fake distributions. This often leads to vanishing gradients or mode collapse — the generator covers only one or a few modes of the data distribution (e.g., only faces with blonde hair and blue eyes).

We need a more stable training dynamic and a loss function that encourages covering the whole distribution.

Enter WGAN‑GP

Wasserstein GAN with Gradient Penalty (WGAN‑GP) makes three key changes:

1. Discriminator → Critic — no sigmoid at the output. The critic scores realness without bounding it to [0,1].
2. Wasserstein loss — instead of BCE, we use critic(fake).mean() - critic(real).mean(). This provides meaningful gradients everywhere.
3. Gradient penalty — enforces the 1‑Lipschitz constraint by penalising the gradient norm on interpolated samples (λ = 10).
4. Critic updates per generator update = 5 — the critic is trained more often to provide reliable gradients.

Implementation Overview (PyTorch)

Data Preparation

• Dataset: Anime Faces (64×64) / Pokemon Sprites
• Resize to 64×64, normalise to [-1, 1]
• Batch size = 32 (adjust for GPU memory)
• Mixed precision (torch.cuda.amp) for speed

DCGAN Generator

python

nn.ConvTranspose2d(nz, ngf*8, 4,1,0),
nn.BatchNorm2d(ngf*8), nn.ReLU(True),
...  # up to 64×64 RGB
nn.Tanh()

DCGAN Discriminator

python

nn.Conv2d(nc, ndf, 4,2,1), nn.LeakyReLU(0.2),
...  # down to 1×1
nn.Sigmoid()

WGAN‑GP Critic (no sigmoid, InstanceNorm)

python

nn.Conv2d(nc, ndf, 4,2,1), nn.LeakyReLU(0.2),
nn.InstanceNorm2d(ndf*2, affine=True),
...  # final layer no activation

Gradient Penalty Function

python

def compute_gradient_penalty(critic, real, fake, device):
    alpha = torch.rand(batch_size,1,1,1, device=device)
    interpolated = (alpha*real + (1-alpha)*fake).requires_grad_(True)
    critic_interp = critic(interpolated)
    grad = torch.autograd.grad(outputs=critic_interp, inputs=interpolated,
                               grad_outputs=torch.ones_like(critic_interp),
                               create_graph=True, retain_graph=True)[0]
    grad = grad.view(batch_size, -1)
    return ((grad.norm(2, dim=1) - 1) ** 2).mean()

Training Loop (WGAN‑GP)

python

for epoch in range(epochs):
    for real_imgs in dataloader:
        # Train critic 5 times
        for _ in range(CRITIC_ITERS):
            fake = generator(noise)
            gp = gradient_penalty(critic, real, fake)
            loss_critic = critic(fake).mean() - critic(real).mean() + LAMBDA_GP * gp
            loss_critic.backward()
            optimizer_critic.step()
        # Train generator once
        fake = generator(noise)
        loss_generator = -critic(fake).mean()
        loss_generator.backward()
        optimizer_generator.step()

Results & Observations

I trained both models for 10 epochs on 15k anime faces (Kaggle T4×2 GPU).

• DCGAN — generated plausible faces but showed clear mode collapse. Many outputs looked nearly identical (same hair style, same eye colour).
• WGAN‑GP — produced much more diverse samples: different hair colours, accessories, facial expressions. The critic loss remained stable, and the gradient penalty converged around 10.

Loss Curves

• DCGAN losses oscillate widely, typical of standard GANs.
• WGAN‑GP losses are smooth and correlate with visual quality.

Diversity Grid (64 images)

• DCGAN 64‑grid: many repeated faces.
• WGAN‑GP 64‑grid: high variety, no obvious repetitions.

Live Demo: Gradio App

I built an interactive web app using Gradio. You can:

• Choose DCGAN or WGAN‑GP
• Adjust number of images, random seed, and temperature (noise scale)
• Compare both models side‑by‑side using the same latent vector
• Generate a 64‑image grid to visually assess mode collapse
• View training loss curves

👉 Try it yourself: https://a2aa16b9b2232fc631.gradio.live

Key Takeaways

1. Wasserstein loss + gradient penalty effectively eliminates mode collapse.
2. Critic updates per generator update (5:1) are crucial — never skip this.
3. InstanceNorm works better than BatchNorm in the critic.
4. Mixed precision (FP16) saves memory without hurting quality.
5. A simple Gradio app makes your model accessible for real‑world testing.

What’s Next?

• Train for more epochs (50–100) to further improve quality.
• Compute FID (Fréchet Inception Distance) for quantitative comparison.
• Extend the same architecture to other unpaired translation tasks (CycleGAN).

References

• Radford et al. (2015) — DCGAN
• Gulrajani et al. (2017) — WGAN‑GP
• PyTorch DCGAN Tutorial

Have you faced mode collapse in your own GAN projects? What tricks worked for you? Let’s discuss in the comments!

GAN #WGAN #DeepLearning #PyTorch #GenerativeAI #ModeCollapse