Wasserstein GAN: Stable Training

A payments company processing UPI transactions wants to train a fraud classifier, but confirmed fraud cases make up under 0.05% of the transaction log — nowhere near enough positive examples to teach a model the shape of fraudulent behaviour. The standard fix is synthetic data: train a GAN on the few thousand confirmed fraud transactions and use its generator to manufacture thousands more with the same statistical fingerprints — unusual merchant-category combinations, odd time-of-day clustering, rapid small-value probing before one large withdrawal. The team builds a textbook discriminator-generator pair, trains it overnight, and wakes up to a discriminator reporting 99.8% accuracy and a generator whose loss curve has been flat for six hours. The samples it produces are useless — three transaction patterns repeated with cosmetic variation. Nothing crashed. No exception was thrown. The training loop ran exactly as coded and still failed, because the loss function itself stopped carrying any signal the moment the generator's output distribution and the real fraud distribution stopped overlapping. That specific failure — a discriminator so confident it goes flat, and a generator that receives essentially zero useful gradient — is precisely the pathology Wasserstein GAN (WGAN) was built to fix. This chapter derives why it happens and why replacing the loss function with an optimal-transport distance fixes it, with every number checked.

Why vanilla GANs stall: the geometry problem

In a standard GAN, the discriminator D is trained with binary cross-entropy to separate real samples (label 1) from generated samples (label 0). For a fixed generator, the discriminator that minimises this loss is known in closed form: D*(x) = p_r(x) / (p_r(x) + p_g(x)), where p_r and p_g are the densities of the real and generated distributions. Substituting this optimal D* back into the generator's objective reduces it — after some algebra — to (a constant plus) twice the Jensen–Shannon (JS) divergence between P_r and P_g. So a GAN trained to convergence on the inner discriminator loop is, in effect, minimising JS(P_r, P_g).

JS divergence is defined as JS(P, Q) = ½·KL(P‖M) + ½·KL(Q‖M), where M = (P+Q)/2 is the average distribution and KL is the Kullback–Leibler divergence. JS is always finite and lies in [0, log 2]. That boundedness is exactly the problem. Whenever P and Q have supports that don't overlap — or overlap only on a set of measure zero, which is the generic case whenever both distributions live on thin, low-dimensional manifolds inside a huge ambient space (a 64×64 RGB image lives in a 12,288-dimensional pixel space, but "plausible faces," or "plausible fraud transactions," occupy a far lower-dimensional manifold within it) — JS divergence doesn't shrink gradually as the manifolds get closer. It sits at exactly log 2 no matter how close or far apart the two disjoint supports are, and only drops to 0 once they coincide. A locally constant function has zero derivative. That is the flat generator loss the fraud-detection team watched for six hours: the discriminator wasn't wrong to reach 99.8% accuracy — with disjoint supports, a good discriminator can classify perfectly, and JS-divergence theory says its output should saturate. The KL term is worse: if P_r assigns zero density anywhere P_g doesn't, or vice versa, KL diverges to +∞, which is even less useful as a training signal. Neither classical divergence measures "how far" two non-overlapping distributions are from each other — both effectively answer only "do they overlap: yes or no."

Earth mover's distance: measuring the cost of transport

The Wasserstein-1 distance (also called Earth Mover's distance) fixes this by measuring transport cost instead of density overlap. Formally:

W(P, Q) = inf_γ∈Π(P,Q) 𝔼_(x,y)~γ[‖x − y‖]

where Π(P, Q) is the set of all joint distributions ("couplings") γ whose two marginals are P and Q. Think of P as a pile of earth and Q as the shape you want to reshape it into: γ is a transport plan specifying how much mass moves from each point of P to each point of Q, and the cost of a plan is the total mass moved times the distance moved. W(P, Q) is the cost of the cheapest possible plan. Unlike JS, this distance grows continuously with how far apart the distributions are, even when their supports never touch — moving a unit mass 1000 units away costs 1000, not "log 2 regardless."

The infimum over the space of all couplings is not something you can optimise directly with backpropagation — it's an intractable optimisation over an enormous, difficult-to-parameterise object. This is where the Kantorovich–Rubinstein duality result becomes the engineering trick that makes WGAN possible: for the Wasserstein-1 distance specifically,

W(P, Q) = sup_{‖f‖L ≤ 1} 𝔼_x~P[f(x)] − 𝔼_x~Q[f(x)]

Here ‖f‖_L ≤ 1 means f is 1-Lipschitz: |f(x) − f(y)| ≤ |x − y| for every pair of inputs, i.e. f's output cannot change faster than its input does. This converts an intractable search over transport plans into a search over a constrained family of scalar functions — which is exactly what a neural network can be trained to approximate. WGAN replaces the discriminator with a critic f_w, a network with no final sigmoid, trained to maximise 𝔼[f_w(x_real)] − 𝔼[f_w(x_fake)] subject to the Lipschitz constraint — approximating the dual supremum — while the generator is trained to shrink that same gap. Because W doesn't saturate the way JS does, this gap keeps providing gradient no matter how far apart the two distributions currently are.

Worked example: exact numbers on a two-point-mass toy

To see the vanishing-gradient failure and the fix in numbers instead of hand-waving, strip the problem down to the simplest possible generator. Let the real distribution P_r be a point mass at 0 — all probability concentrated at x = 0. Let the generator ignore its noise input entirely and output a single learnable number θ, so its distribution P_θ is also a point mass, located at x = θ. This degenerate setup is deliberate: it isolates the gradient-flow question from every other complication.

JS divergence, case θ ≠ 0. The two point masses sit at different locations, so their supports are disjoint. The mixture M = (P_r + P_θ)/2 puts probability ½ on x = 0 and ½ on x = θ. At x = 0: P_r = 1, M = ½, so this term of KL(P_r‖M) contributes 1·log(1/(½)) = log 2, and there is no other point where P_r has mass, so KL(P_r‖M) = log 2 exactly. By the same argument KL(P_θ‖M) = log 2. So JS(P_r, P_θ) = ½·log2 + ½·log2 = log 2 ≈ 0.6931 nats — for every θ ≠ 0, whether θ = 0.0001 or θ = 1000. At θ = 0 exactly, P_θ = P_r and JS drops to 0. JS(θ) is therefore a step function: constant at log 2 everywhere except a single discontinuous drop to 0 at the origin. Its derivative is 0 at every point where it's differentiable. A generator starting at θ = 3 sees dJS/dθ = 0 — no signal, ever, to move toward 0, under the (theoretically correct) assumption of a converged, optimal discriminator. KL(P_r‖P_θ) is worse still: P_r(0)/P_θ(0) = 1/0 = +∞ for any θ ≠ 0.

Wasserstein-1 distance, same setup. Two point masses admit exactly one feasible coupling — the deterministic one that moves the entire unit mass from θ to 0 — so the infimum in the definition of W is trivially that single plan's cost: W(P_r, P_θ) = |θ − 0| = |θ|. This is continuous everywhere and, for θ ≠ 0, differentiable with dW/dθ = θ/|θ| = sign(θ).

Tracing gradient descent. Start at θ₀ = 3 with learning rate α = 0.1, minimising W(θ) = |θ|:

θ₁ = θ₀ − α·sign(θ₀) = 3 − 0.1(1) = 2.9
θ₂ = 2.9 − 0.1(1) = 2.8
…
θ₂₉ = 3 − 0.1×29 = 0.1
θ₃₀ = 0.1 − 0.1(1) = 0.0 — converged exactly, generator distribution now matches the real one.

Under the JS objective, θ never leaves 3 across all 30 steps, because dJS/dθ = 0 identically once the supports separate. That is the whole content of "Wasserstein GANs train more stably": the gradient magnitude under W is constant (sign(θ), magnitude 1) regardless of how far apart the distributions currently sit, while JS's gradient vanishes identically the instant the supports separate — which is true almost everywhere near initialisation for real data, since a freshly initialised generator produces something close to noise, nowhere near the real manifold.

Check the dual side too: for θ = 3, the function f(x) = −x is 1-Lipschitz (|f(x) − f(y)| = |x − y| exactly), and 𝔼[f(x_real)] − 𝔼[f(x_fake)] = f(0) − f(3) = 0 − (−3) = 3 = W(P_r, P₃). This f is the optimal critic for this toy problem — it achieves the Kantorovich–Rubinstein supremum exactly, confirming the duality holds numerically here, not just in the abstract statement.

The full loop, end to end

Common misconception: "just remove the sigmoid"

A student who has just met WGAN often assumes the whole change is cosmetic: take the existing discriminator, delete its final sigmoid so it outputs a raw score instead of a probability, and call the result a critic. Deleting the sigmoid is necessary but is nowhere near sufficient, and treating it as the whole story is the single most common misconception about this topic. Without an enforced Lipschitz constraint, an unconstrained critic has no reason to stay bounded: because 𝔼[f(x_real)] − 𝔼[f(x_fake)] scales linearly if you scale f by any constant c (multiply every weight by 10 and the "distance estimate" multiplies by 10 too, with no penalty), gradient ascent on this unconstrained objective simply drives the critic's weights toward infinity. The reported "distance" would grow without bound and the gradients backpropagating into the generator would explode rather than provide a stable, informative signal. The 1-Lipschitz constraint is not a minor technical footnote — it is the condition that makes 𝔼[f(x_real)] − 𝔼[f(x_fake)] actually equal a real, meaningful distance (via Kantorovich–Rubinstein duality) instead of an arbitrarily inflatable number. Everything that makes WGAN training stable — bounded, non-saturating, non-exploding gradients — depends on that constraint being enforced, not on the sigmoid being absent.

Weight clipping in practice, and why WGAN-GP exists

The original WGAN paper enforces the Lipschitz constraint the crudest way that still works: after every critic update, clip every individual weight into [−c, c] (the paper uses c = 0.01). A linear layer with all weights bounded by c has a bounded operator norm, and composing bounded-slope layers keeps the whole network's global Lipschitz constant bounded — a sufficient, though far from tight, condition. In practice this has known problems. Because the critic's unconstrained loss almost always wants to push |𝔼[f(x_real)] − 𝔼[f(x_fake)]| larger, gradient ascent drives most weights straight to one of the two clip boundaries, −c or c, wasting most of the network's representational capacity on what is effectively a near-linear function. Set c too small and the clipped network's gradients vanish across many stacked layers; set c too large and gradients can still explode in the interval between one clipping step and the next. Gulrajani et al.'s WGAN-GP replaces the hard clip with a soft penalty added directly to the critic loss:

L_C = −(𝔼[f(x_real)] − 𝔼[f(x_fake)]) + λ·𝔼_x̂[(‖∇_x̂f(x̂)‖₂ − 1)²]

where x̂ = ε·x_real + (1 − ε)·x_fake for ε drawn uniformly from [0, 1] — a random point on the straight-line segment between a paired real and fake sample — and λ is typically 10. This term directly pushes the critic's gradient norm toward exactly 1 (the tightest case of the 1-Lipschitz condition) along the interpolation, instead of only indirectly bounding it through weight magnitude, and it removes the capacity-wasting saturation that hard clipping causes. The original paper's own recipe for the clipped version is also worth remembering precisely: RMSProp rather than Adam, learning rate 5×10⁻⁵, n_critic = 5 critic updates per generator update, and c = 0.01. The Adam choice matters — momentum accumulated across a loss landscape that shifts abruptly whenever weights hit or leave the clip boundary tends to amplify oscillation rather than damp it, which is precisely the instability WGAN was designed to remove.

The training step, traced

import torch
import torch.nn as nn
import torch.optim as optim

# Critic: outputs a single unbounded real score, NOT a probability
critic = nn.Sequential(
    nn.Linear(784, 256), nn.LeakyReLU(0.2),
    nn.Linear(256, 1)          # no sigmoid
)
generator = nn.Sequential(
    nn.Linear(100, 256), nn.ReLU(),
    nn.Linear(256, 784), nn.Tanh()
)

opt_critic = optim.RMSprop(critic.parameters(), lr=5e-5)
opt_gen    = optim.RMSprop(generator.parameters(), lr=5e-5)

CLIP = 0.01
N_CRITIC = 5

for step in range(10000):
    # ---- 1. Train the critic n_critic times ----
    for _ in range(N_CRITIC):
        real = sample_real_batch(64)              # (64, 784)
        z = torch.randn(64, 100)
        fake = generator(z).detach()               # freeze G

        critic_loss = -(critic(real).mean() - critic(fake).mean())
        opt_critic.zero_grad()
        critic_loss.backward()
        opt_critic.step()

        # enforce 1-Lipschitz via weight clipping
        for p in critic.parameters():
            p.data.clamp_(-CLIP, CLIP)

    # ---- 2. Train the generator once ----
    z = torch.randn(64, 100)
    fake = generator(z)
    gen_loss = -critic(fake).mean()
    opt_gen.zero_grad()
    gen_loss.backward()
    opt_gen.step()

Trace the sign conventions carefully, because they're the part students most often get backwards. critic_loss is the negative of the K–R dual gap, so minimising it — which is what opt_critic.step() does — maximises 𝔼[f(real)] − 𝔼[f(fake)], exactly matching the "sup" in the duality formula. fake = generator(z).detach() stops gradients from flowing into the generator's parameters during the critic phase, so only the critic's weights change across those five inner iterations, matching the diagram's "w frozen here" / "critic frozen" separation. gen_loss = -critic(fake).mean() is the negative of 𝔼[f(x̃)]; minimising it pushes the generator to make the critic score its fake batch as highly as possible — i.e. to shrink the estimated Wasserstein distance from the real distribution. Because fake here is not detached in step 2, gradients flow generator → critic → loss and back, but opt_gen.step() only updates generator.parameters(), so the critic's weights are read but never written during this phase. The clamp_ call after every critic step is the (crude but functional) stand-in for the Lipschitz constraint derived above — remove it, and the critic loss will trend toward −∞ or +∞ instead of tracking a bounded distance.

Active recall

Q1. Write the Kantorovich–Rubinstein dual for W(P, Q) and state precisely what constraint is placed on f.

Q2. For the toy model P_r = δ₀, P_θ = δ_θ, compute JS(P_r, P_θ) at θ = 5 and at θ = 0.0001. Why are both answers the same number even though 0.0001 is "almost" 0?

Q3. Compute W(P_r, P_θ) and its gradient at θ = −2, then take one gradient-descent step with learning rate 0.05.

Q4. A student claims: "I can turn any GAN into a WGAN just by removing the sigmoid from the discriminator's last layer." What is missing from this claim?

Q5. WGAN-GP replaces weight clipping with a gradient-penalty term. Write it, and state in one sentence what quantity it pushes toward 1 and where it's evaluated.

Q6. The original WGAN paper uses RMSProp rather than Adam for both networks. What property of the clipped critic's loss landscape makes Adam's momentum term risky here?

Answers.

A1. W(P, Q) = sup_{‖f‖L ≤ 1} 𝔼_x~P[f(x)] − 𝔼_x~Q[f(x)]. The constraint is that f must be 1-Lipschitz: |f(x) − f(y)| ≤ |x − y| for every x, y — f's output cannot change faster than its input.

A2. For any θ ≠ 0, the two point masses have disjoint support, so KL(P_r‖M) = KL(P_θ‖M) = log 2 exactly (derived above), giving JS = log 2 ≈ 0.6931 nats at both θ = 5 and θ = 0.0001 — the same number, because JS only measures whether the supports overlap at all, not how far apart they are. JS is a step function of θ: constant at log 2 for every nonzero θ, dropping to 0 only in the limit θ = 0 exactly. It gives no indication that 0.0001 is "closer" than 5.

A3. W(P_r, P₋₂) = |−2| = 2. dW/dθ = sign(θ) = sign(−2) = −1. One gradient-descent step: θ_new = θ − α·dW/dθ = −2 − 0.05×(−1) = −2 + 0.05 = −1.95. |θ_new| = 1.95 < 2, confirming the update moved the generator's point mass toward the real one at 0, as expected.

A4. Missing is the Lipschitz constraint. Without it, scaling every weight of the critic by a constant c scales 𝔼[f(real)] − 𝔼[f(fake)] by the same c with no penalty, so unconstrained gradient ascent drives the critic's weights toward infinity instead of settling on a value that equals the true Wasserstein distance. The constraint (via clipping or a gradient penalty) is what makes the critic's score actually correspond to the Kantorovich–Rubinstein dual, not just an arbitrarily inflatable number; removing the sigmoid alone is necessary but does nothing to prevent this divergence.

A5. λ·𝔼_x̂[(‖∇_x̂f(x̂)‖₂ − 1)²], added to the critic loss. It pushes the norm of the critic's gradient (with respect to its input, not its weights) toward 1, evaluated at random points x̂ = ε·x_real + (1 − ε)·x_fake on the straight-line segment between paired real and fake samples.

A6. Hard clipping makes the effective loss landscape non-smooth: whenever a weight hits or leaves the clip boundary between steps, the function the critic represents can change abruptly, so the loss surface has sharp, non-stationary shifts in curvature. Adam's per-parameter momentum and adaptive step sizes accumulate biased estimates across exactly these abrupt shifts, amplifying oscillation instead of damping it; RMSProp (and plain SGD without momentum) proved more robust to this non-stationarity in the paper's experiments, which is why it's the recommended optimiser for the clipped-critic variant.

Think About It

Think about this: How would you explain wasserstein gan: stable training to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind wasserstein gan: stable training, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.