Module 5 — Mode collapse and training stabilization

The DCGAN of module 4 works on MNIST for most of a training run, then quietly stops covering the ten digit classes. Or the losses go wild and the generator produces uniform grey. Or nothing visible changes for a hundred epochs. These pathologies made GANs famous for being hard to train. This module explains what is happening, and applies the three fixes that keep them alive.

Symptoms and their common cause

There are three canonical GAN failure modes. Being able to name them from a sample sheet is half the battle.

Mode collapse. The generator produces a small subset of the data distribution. On MNIST, only three or four digit classes appear across a batch of 64 samples from independent noise. On faces, all generated faces look like near-duplicates of one prototype. Diversity has collapsed, quality per sample may look fine.

Vanishing gradients. The discriminator has become too strong. It classifies every fake with probability near zero, its logits saturate, and the generator's gradient dies. Training curves flatline, samples stop changing.

Oscillation and divergence. The two networks chase each other in a cycle: $D$ learns to spot the current fakes, $G$ shifts to fool the current $D$, and neither improves. Loss curves swing violently. In the worst case the model diverges to nonsense.

The three symptoms have a common root: the objective is a minimax game with no monotonic quality signal. Nothing in ordinary SGD guarantees that alternating descent on a saddle point converges. If we want reliable training, we have to redesign either the loss or the model's regularity.

Wasserstein: replace the metric

The Wasserstein GAN (WGAN, Arjovsky 2017) is the redesign that opened the door to stable GAN training. Instead of asking the discriminator to output a probability, we ask a critic to output a real number, and we train it to give large scores to reals and small scores to fakes.

The generator's loss becomes:

$\min_G \; -\mathbb{E}_{z}[D(G(z))]$

and the critic's:

$\max_D \; \mathbb{E}_{x \sim p_{\text{data}}}[D(x)] - \mathbb{E}_{z}[D(G(z))]$

Under the constraint that $D$ is 1-Lipschitz — its output changes by at most one when its input changes by one, in some norm — this objective is the Wasserstein-1 distance between the data and the generator's distributions. A meaningful distance means the critic's loss is finally a quality signal: it should decrease as training progresses.

The Lipschitz constraint has to be enforced somehow. WGAN-GP (gradient penalty, Gulrajani 2017) does it by adding a penalty on the gradient norm at interpolation points:

def gradient_penalty(D, x_real, x_fake, device):
    batch = x_real.size(0)
    eps = torch.rand(batch, 1, 1, 1, device=device)
    x_hat = eps * x_real + (1 - eps) * x_fake
    x_hat.requires_grad_(True)
    d_hat = D(x_hat)
    grads = torch.autograd.grad(
        outputs=d_hat, inputs=x_hat,
        grad_outputs=torch.ones_like(d_hat),
        create_graph=True, retain_graph=True,
    )[0]
    grads = grads.view(batch, -1)
    return ((grads.norm(2, dim=1) - 1) ** 2).mean()

The critic step then combines the Wasserstein term with 10.0 * gradient_penalty(...). The generator step is unchanged.

Spectral normalisation: another route to Lipschitz

Spectral normalisation (Miyato 2018) enforces the same 1-Lipschitz constraint by normalising each layer's weight matrix by its largest singular value. It requires no penalty, no extra hyperparameter, no interpolation samples, and is applied by wrapping the layers:

import torch.nn.utils.parametrizations as P

layer = P.spectral_norm(nn.Conv2d(64, 128, 4, 2, 1))

In practice spectral normalisation on the discriminator alone often stabilises training as well as gradient penalty, at a fraction of the compute. It is now the default choice in most codebases that were still training with vanilla BCE.

Do not stack every fix at once

Mode collapse tempts practitioners to combine Wasserstein loss, gradient penalty, spectral normalisation, R1 regularisation, exponential moving averages of the generator, and instance noise, all at once. That combination often works but also often breaks in ways that are very hard to diagnose. Add one stabiliser at a time and keep the sample sheet from module 4 open — the visual is your quality signal.

Diagnosing mode collapse concretely

The instinct is to look at loss curves, and they will not tell you. The reliable diagnostic uses either a classifier or a coverage metric.

On MNIST, take a pretrained digit classifier — a small ConvNet trained separately — and count the class distribution over 10,000 samples from the generator. A healthy generator produces roughly 10% of each class. A collapsed generator produces 40% of one class and 0% of another:

with torch.no_grad():
    z = torch.randn(10000, 100, device=device)
    samples = G(z)
    predicted = classifier(samples).argmax(dim=1).cpu()
    counts = torch.bincount(predicted, minlength=10)
print(counts.float() / counts.sum())

On datasets without labels, the FID metric of module 8 catches diversity loss — but not always well, which is one of its limitations we come back to there.

Practical tricks that help without redesigning the loss

Even inside a plain DCGAN, small habits reduce the chance of collapse.

Update ratio. Train the discriminator more than the generator (say five critic steps per generator step in WGAN-GP) so the critic never falls too far behind. The exact ratio matters less than the principle that a lagging discriminator gives the generator a free ride.

Feature matching. Instead of purely fooling $D$, the generator matches the mean of $D$'s intermediate features on fakes to that on reals. This provides a smoother signal than the final classification loss.

Two time-scale update rule (TTUR). Give the discriminator a higher learning rate than the generator. Heusel 2017 showed convergence to a local Nash equilibrium under this asymmetry. In code: opt_d = Adam(D.parameters(), lr=4e-4) and opt_g = Adam(G.parameters(), lr=1e-4).

Exponential moving average of the generator's weights. Evaluate and sample from an EMA copy of $G$ rather than the live weights. The samples are smoother and more consistent across runs.

A glimpse at StyleGAN

The state of the art of GANs on high-resolution images — before diffusion took over — is StyleGAN and its descendants (Karras 2018 onwards). Two design ideas dominate.

Progressive growing. The generator and discriminator start with low-resolution outputs (4 by 4 pixels), and new layers are faded in over training to double the resolution up to 1024 by 1024. Training is drastically easier when the network first learns coarse structure and only later gets to worry about hair strands and skin pores.

Style modulation. Instead of injecting the noise $z$ only at the input, StyleGAN passes it through an eight-layer MLP to produce a style vector $w$, and injects $w$ at every layer via adaptive normalisation. This decouples high-level structure (pose, face shape) from fine detail (hair, texture), and enables the famous latent-space edits — smiling, ageing, changing gender — because the disentangled space of $w$ supports them.

StyleGAN produces near-photographic faces, and its second version was the reference generative model for faces until diffusion caught up. Its main limitation is what all GANs share: no easy conditioning on text, no easy per-class control, and a training recipe that takes days on multiple GPUs.

When a GAN is still the right choice

Latency at inference. A GAN generates in a single forward pass — often under a millisecond on a GPU — while diffusion (module 6) needs tens or hundreds of steps. For real-time avatars, live video synthesis or on-device generation, GANs remain the pragmatic choice even in 2026.

In summary

• Mode collapse, vanishing gradients and oscillation are three canonical GAN pathologies with a common root: the minimax objective has no monotonic quality signal, so ordinary SGD offers no convergence guarantee.
• Wasserstein loss with gradient penalty and spectral normalisation are the two dominant fixes, both aiming to keep the discriminator 1-Lipschitz so its gradient stays meaningful.
• Diagnose collapse with a class distribution from a pretrained classifier, or with FID over 10,000 samples, not with the loss curves.
• StyleGAN with progressive growing and style modulation was the reference for face generation before diffusion overtook it; GANs are still the right choice when generation latency dominates.

Next module: diffusion models, a different training paradigm entirely that avoids most of these pathologies at the cost of slow sampling.