WGAN: Wasserstein GAN for Stable Training

A payments company is building a generator that synthesizes fake UPI transaction records to augment a fraud-detection dataset. Real transaction amounts cluster tightly: most payments fall in a handful of common bands (₹100–₹500 for small merchants, ₹1000–₹5000 for bills, and so on). Early in training, the generator's outputs occupy a completely different region of the amount-and-metadata space than real transactions do — wrong scale, wrong clustering, wrong correlations between amount and merchant category. The team watches the discriminator loss and sees something familiar: the discriminator's accuracy shoots to nearly 100% within a few hundred steps, its loss drops to almost zero, and the generator's gradient updates shrink to numbers so small the optimizer might as well not be running. The generator is still obviously bad, but the training signal has gone silent. This is not a bug in the code. It is a structural property of the loss function the original GAN uses, and it is the exact problem WGAN was designed to fix.

Why Jensen–Shannon divergence goes flat

The standard GAN's minimax objective is, at its optimum for a fixed generator, equivalent to minimizing the Jensen–Shannon (JS) divergence between the real distribution P_r and the generated distribution P_g. JS divergence is bounded: it ranges from 0 (identical distributions) to log 2 (completely disjoint supports), and critically, it is locally constant almost everywhere in that disjoint regime — nudging the generator's parameters a little does not change the divergence at all, so the gradient with respect to those parameters is exactly zero almost everywhere, not just small.

Here is why disjoint support is the normal case, not an edge case. Real transaction data lives on a low-dimensional manifold embedded in a much higher-dimensional space (a few underlying factors — merchant type, time of day, customer segment — determine most of the variation across dozens of raw fields). An untrained or partially trained generator's output also lives on some low-dimensional manifold. Two randomly-positioned low-dimensional manifolds inside a high-dimensional space generically do not intersect, in the same way two random lines in 3D space almost never cross. So for most of training, P_r and P_g have supports that don't overlap, and JS divergence sits at its ceiling of log 2 with zero local gradient — exactly the "confident, useless discriminator" the fraud-detection team observed.

Arjovsky, Chintala, and Bottou's toy example makes this precise and independently checkable. Let Z be a random variable uniform on [0,1]. Define P₀ as the distribution of the point (0, Z) — a vertical line segment sitting at x = 0 in the plane — and P_θ as the distribution of (θ, Z), the same segment shifted to x = θ. For any θ ≠ 0, the two supports are disjoint line segments. Work out what the standard divergences do:

KL divergence: (infinite) in both directions whenever θ ≠ 0, because P_θ assigns zero density everywhere P₀ has positive density (and vice versa), and KL's integrand p·log(p/q) blows up wherever q = 0 but p > 0.

JS divergence: let M = (P₀ + P_θ)/2 be the mixture. On the support of P₀, which has zero measure under P_θ, the mixture density is exactly half of P₀'s density. So KL(P₀ ‖ M) = ∫ p₀·log(p₀ / (p₀/2)) = ∫ p₀·log 2 = log 2, since p₀ integrates to 1. The symmetric term contributes the same log 2. JS = ½·log 2 + ½·log 2 = log 2, for every θ ≠ 0, all the way from θ = 0.00001 to θ = 1000. At θ = 0 it drops discontinuously to 0. There is no slope anywhere except at that single point — nothing for gradient descent to climb.

Wasserstein-1 (Earth-Mover) distance: the optimal way to transport the mass of P₀ onto P_θ is to match each point (0, z) to (θ, z) — same z, shifted x — which costs exactly |θ| per unit of mass moved, and no cheaper coupling exists since every point must move at least that far in the x-coordinate. So W(P₀, P_θ) = |θ| exactly: continuous, and differentiable everywhere except θ = 0, with gradient sign(θ) pointing the generator directly toward the real distribution. This is the entire case for WGAN in one example: same disjoint-support scenario, but one distance is flat and useless while the other is a ramp a gradient step can climb.

The Earth-Mover distance, computed by hand

The general definition of the Wasserstein-1 distance is an infimum over all joint distributions (couplings) γ with marginals P_r and P_g: W(P_r, P_g) = inf_{γ∈Π(Pr,Pg)} E_(x,y)~γ[‖x − y‖]. That infimum over all possible transport plans is intractable to compute directly for anything beyond toy cases — which is exactly why WGAN needs the duality trick in the next section. But for one-dimensional distributions there is a closed form that is easy to verify by hand: W₁(P, Q) equals the sum of the absolute differences between the two cumulative distribution functions.

Take a simplified snapshot of the fraud-detection generator early in training. Bucket transaction amounts into five bands (₹100, ₹200, ₹300, ₹400, ₹500) and compare the real distribution to what the generator currently produces:

Summing the last column gives W₁ = 0.30 + 0.30 + 0.10 + 0.05 + 0.00 = 0.75 bucket-units. With buckets spaced ₹100 apart, that scales to W₁ = ₹75: on average, each unit of probability mass in the generator's output would need to move about ₹75 along the amount axis to match the real distribution. Notice this number moves smoothly and means something concrete even though the two distributions overlap at every bucket — nothing here is the disjoint-support pathology of the toy example, and yet Wasserstein distance is still the more informative summary because it weights how far mass sits from where it should be, not merely whether the two histograms differ. A discriminator-style loss (or JS divergence, which is well-defined and non-constant when supports overlap like this) would still register that the two histograms differ, but it would not tell you whether the generator overshot by one bucket or five — W₁ does, because it is built directly out of a transport cost.

Kantorovich–Rubinstein duality: from an intractable infimum to a trainable network

Computing W₁ by searching over every possible coupling γ is only feasible in the 1D case above. In the high-dimensional space of real transaction records or images, that search space is unmanageable. The Kantorovich–Rubinstein duality theorem rewrites the same quantity as a supremum over a much friendlier object — scalar functions f with a bounded rate of change:

W(P_r, P_g) = sup_‖f‖L≤1 ( E_x~Pr[f(x)] − E_x~Pg[f(x)] )

where the supremum ranges over all 1-Lipschitz functions f: a function is 1-Lipschitz when |f(a) − f(b)| ≤ ‖a − b‖ for every pair of points a, b — it never changes faster than the distance between its inputs. This is precisely the form a neural network critic can approximate: train a network D_w to maximize E[D_w(x)] − E[D_w(x̃)] over real batches x and generated batches x̃, subject to D_w staying 1-Lipschitz, and the value it converges to is (an estimate of) W(P_r, P_g). This is why WGAN's second network is called a critic rather than a discriminator: it has no sigmoid, no probability to output, and no "real vs. fake" classification task. It outputs an unconstrained real number that scores how far a sample sits from the real distribution on the axis the duality theorem hands it, and the entire remaining engineering problem in WGAN is how to keep that network 1-Lipschitz while it trains.

Weight clipping: the crude first fix

The original WGAN paper's answer was blunt: after every critic optimizer step, clip every weight into a fixed range [−c, c] (commonly c = 0.01). Clamping the weights bounds how much the critic's output can change per unit of input change, which is a (loose) way to bound its Lipschitz constant. It works well enough to demonstrate that the Wasserstein objective fixes vanishing gradients — loss curves in the original paper correlate with sample quality for the first time in GAN training — but the mechanism is crude in a way that compounds with depth.

A critic is a composition of L layers, and the Lipschitz constant of a composition is bounded by the product of the per-layer Lipschitz constants. Clipping controls each layer only loosely and indirectly (it bounds the weight magnitudes, not the operator norm exactly), so the effective per-layer scale drifts to whatever value the optimizer finds convenient subject to the box constraint. If that effective scale sits even slightly off 1.0, the error compounds multiplicatively across depth: a critic whose layers each carry an effective scale of about 0.9 has its gradient signal shrink to roughly 0.9¹⁰ ≈ 0.35 of its original magnitude by the tenth layer; at 1.1 per layer it grows to roughly 1.1¹⁰ ≈ 2.59. Neither number looks catastrophic in isolation, but stack it across the 20–50 layers of a realistic critic and the gradient either vanishes or explodes long before it reaches the generator. Gulrajani et al. (2017) additionally documented that clipping pushes the critic's weights toward the two extremes ±c — the optimizer, fighting the box constraint, saturates almost every weight at the boundary rather than using the full range, which caps the critic to something close to a piecewise-linear function and starves it of the capacity needed to approximate a good 1-Lipschitz function on complex data.

WGAN-GP: enforcing Lipschitz continuity with a gradient penalty

Gulrajani, Ahmed, Arjovsky, Dumoulin, and Courville's WGAN-GP (NeurIPS 2017, arXiv:1704.00028) replaces the hard clip with a soft penalty derived from a direct property of the optimal critic: for the true 1-Lipschitz function that achieves the Kantorovich–Rubinstein supremum, the gradient norm equals exactly 1 almost everywhere along straight-line paths between paired real and generated points. Instead of restricting weights, WGAN-GP samples a point x̂ on the line between a real sample x and a generated sample x̃ (x̂ = εx + (1−ε)x̃ for ε ~ U(0,1)), computes the critic's gradient at that interpolated point via autograd, and penalizes the squared distance of that gradient's norm from 1:

import torch

LAMBDA_GP = 10
N_CRITIC = 5

def gradient_penalty(critic, real, fake, device):
    batch_size = real.size(0)
    # shape-agnostic: works for image tensors (N,C,H,W) and for
    # tabular feature vectors like the UPI transaction records (N, num_features)
    eps = torch.rand(batch_size, *([1] * (real.dim() - 1)), device=device)
    eps = eps.expand_as(real)
    interpolated = eps * real + (1 - eps) * fake
    interpolated.requires_grad_(True)

    scores = critic(interpolated)
    grads = torch.autograd.grad(
        outputs=scores,
        inputs=interpolated,
        grad_outputs=torch.ones_like(scores),
        create_graph=True,
        retain_graph=True,
    )[0]

    grads = grads.view(batch_size, -1)
    grad_norm = grads.norm(2, dim=1)
    penalty = ((grad_norm - 1) ** 2).mean()
    return penalty

Walking through it: eps is one random interpolation weight per sample in the batch, shaped with a singleton dimension for every axis after the batch dimension (so it works whether real is a 4-D image tensor or, as in the UPI example, a 2-D tensor of tabular feature vectors) and then broadcast across every remaining dimension via expand_as so the same ε applies to every pixel or feature of a given sample. interpolated.requires_grad_(True) tells autograd to track this tensor as a leaf so gradients can be taken with respect to it (not with respect to the critic's weights — this penalty flows into the critic's input, which is a different computational path than the usual backward pass through weights). torch.autograd.grad with grad_outputs=torch.ones_like(scores) computes ∂(Σ_i scores_i)/∂interpolated, i.e. each sample's own gradient, because summing before differentiating and using per-sample outputs on a batch dimension is algebraically equivalent to differentiating each sample independently. create_graph=True is essential: it keeps this gradient computation itself part of the graph, so that when penalty is later backpropagated, gradients can flow through the gradient-norm computation into the critic's weights — a second-order derivative, which is exactly what "penalizing the gradient" requires. Flattening to (batch_size, -1) and taking an L2 norm per row gives one scalar gradient-norm per sample; squaring the deviation from 1 and averaging gives the penalty term.

Two more changes come bundled with this fix: WGAN-GP critics use layer normalization instead of batch normalization, because batch norm makes each sample's critic output depend on the other samples in its batch (through the batch statistics), which breaks the per-sample gradient penalty derivation above — the penalty needs D_w(x̂) to be a clean function of x̂ alone. And weight clipping is removed entirely; nothing constrains the weight magnitudes directly anymore, only the penalty shapes the function's local slope.

Training loop and what n_critic actually buys you

The full critic loss combines the (negated, since optimizers minimize) Wasserstein estimate with the penalty, and the critic is trained several steps for every one generator step:

# critic, generator: pre-built nn.Module instances (assumed helper, not shown)
# opt_critic, opt_gen: torch.optim.Adam instances (assumed helper, not shown)
# get_real_batch(batch_size): loads a batch of real samples (assumed helper, not shown)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

for step in range(1, 10001):
    for _ in range(N_CRITIC):
        real = get_real_batch(64).to(device)
        z = torch.randn(64, 100, device=device)
        fake = generator(z)

        critic_real = critic(real)
        critic_fake = critic(fake.detach())
        gp = gradient_penalty(critic, real, fake.detach(), device)

        critic_loss = critic_fake.mean() - critic_real.mean() + LAMBDA_GP * gp

        opt_critic.zero_grad()
        critic_loss.backward()
        opt_critic.step()

    z = torch.randn(64, 100, device=device)
    fake = generator(z)
    gen_loss = -critic(fake).mean()

    opt_gen.zero_grad()
    gen_loss.backward()
    opt_gen.step()

Minimizing critic_loss = critic_fake.mean() − critic_real.mean() + λ·gp is the same as maximizing critic_real.mean() − critic_fake.mean() (the Wasserstein estimate) while pulling the gradient norm toward 1, matching the duality objective derived earlier. fake.detach() in the critic step stops gradients from flowing into the generator while only the critic is being updated. N_CRITIC = 5 means the critic sees five optimizer steps for every one generator step, which is standard for both weight-clipped and gradient-penalized WGAN and is not just a tuning knob — it is the direct consequence of the theory. In the original GAN, an unusually well-trained discriminator saturates and starves the generator of gradient (that is precisely the vanishing-gradient failure this chapter opened with). In WGAN, D_w's loss is an estimate of an actual distance, so a more converged critic gives a more accurate distance estimate and a better gradient direction, never a saturated one. Training the critic close to optimality before each generator update is not just safe under the Wasserstein objective — it is what the duality theorem asks for.

Comparing the three training regimes

Common misconception

Students who have just learned weight clipping often assume it is a minor implementation detail — "clip to some small range and the Lipschitz constraint is satisfied, so any reasonably small c should work about the same." The depth calculation above is the direct correction: c does not act in isolation, it interacts multiplicatively with network depth through the composition-of-Lipschitz-constants bound. A c that keeps a 4-layer critic well-behaved can leave a 30-layer critic's gradients vanishing or exploding, and there is no principled formula for choosing c that works across architectures — the WGAN paper's own authors described tuning it as delicate and dataset-dependent. That fragility, not a stylistic preference, is why WGAN-GP replaced clipping outright rather than simply publishing better guidance on choosing c.

The WGAN-GP training mechanism

Active recall

Attempt each question before reading its answer.

1. Using the UPI amount-bucket example, suppose training progresses and the generator's distribution improves to P_g = [0.15, 0.35, 0.35, 0.10, 0.05] over the same five buckets, while the real distribution stays P_r = [0.10, 0.30, 0.40, 0.15, 0.05]. Recompute W₁ and state it in rupees.

Answer: CDF_r is unchanged: [0.10, 0.40, 0.80, 0.95, 1.00]. CDF_g for the new distribution: [0.15, 0.50, 0.85, 0.95, 1.00]. Absolute differences: |0.10−0.15|=0.05, |0.40−0.50|=0.10, |0.80−0.85|=0.05, |0.95−0.95|=0.00, |1.00−1.00|=0.00. Sum = 0.20 bucket-units, i.e. ₹20 at ₹100 spacing — down from ₹75. The distance shrank smoothly and proportionally as the generator improved, which is exactly the informative, non-saturating signal Section 2 showed JS divergence cannot provide in the disjoint-support regime.

2. Ripple check: if the same two original distributions (P_r = [0.10, 0.30, 0.40, 0.15, 0.05], P_g = [0.40, 0.30, 0.20, 0.10, 0.00]) were measured with buckets spaced ₹250 apart instead of ₹100, what is W₁ in rupees now, and why doesn't the bucket-unit value of 0.75 change?

Answer: The bucket-unit sum stays 0.75 because it depends only on the shape of the two distributions (the probability masses), which haven't changed — only the physical spacing between buckets changed. W₁ in rupees is the bucket-unit distance times the per-bucket ground distance: 0.75 × ₹250 = ₹187.50. This is the general rule: Wasserstein distance scales linearly with whatever metric defines "distance" in the underlying space, so changing units or the cost function changes the numeric answer even when the probability distributions themselves are identical.

3. Why does JS divergence give exactly zero gradient to the generator when P_r and P_g have disjoint support, while Wasserstein distance does not?

Answer: In the disjoint-support toy example, JS divergence equals log 2 for every nonzero separation θ between the two distributions — it is a flat plateau with a single discontinuous drop to 0 exactly at θ = 0, so its derivative with respect to θ is zero everywhere except that one non-differentiable point. Wasserstein distance equals |θ| in the same setup: a straight ramp with slope ±1 everywhere except θ = 0, giving the optimizer a well-defined direction to move in from any starting separation.

4. In the gradient penalty code, why is the penalty enforced at interpolated points x̂ = εx + (1−ε)x̃ rather than only at the real samples x or only at the fake samples x̃?

Answer: The 1-Lipschitz constraint needs to hold everywhere the critic will be evaluated, but the theoretical justification for WGAN-GP specifically shows that the optimal critic's gradient has norm exactly 1 along straight lines connecting paired real and fake points — that is where enforcing the constraint is both theoretically grounded and where it most directly shapes the region the critic actually needs to be well-behaved in during training, rather than spending penalty budget on regions of space neither distribution nor the transport path passes through.

5. A student sets a weight-clipping WGAN's clip value to c = 0.5 on a fairly deep critic (20+ layers). The critic loss curve looks smooth, but the generator's outputs stay close to noise for thousands of iterations. What's the likely cause, and what's the fix?

Answer: A smooth loss curve with a clipped critic can be misleading: clipping tends to push most weights toward the boundary ±c, which caps the critic's function class to something close to piecewise-linear regardless of how "smooth" its loss trajectory looks, and the composition-of-Lipschitz-constants effect means the effective per-layer scale compounding across 20+ layers can drift far from a well-behaved regime even while the scalar loss value looks stable. The fix is to drop weight clipping for the gradient penalty (WGAN-GP): it does not restrict the weight values directly, so the critic keeps its full capacity while the interpolated-gradient penalty keeps it close to 1-Lipschitz where it matters.

6. True or false, with justification: "Because the WGAN critic has no sigmoid, it should be trained to near-convergence (many steps) before every single generator update, unlike the vanilla GAN discriminator, which should be kept weak so it doesn't saturate."

Answer: True. In the vanilla GAN, an overly strong discriminator saturates against a mediocre generator and its JS-divergence-based loss plateaus at log 2 with zero local gradient — training it to convergence actively hurts the generator's signal. In WGAN, the critic's loss approximates an actual distance (not a divergence with a fixed ceiling), so a more converged critic gives a more accurate distance estimate and a cleaner gradient direction for the generator; it never saturates the way a JS-based discriminator does. This is exactly why n_critic ≈ 5 is standard practice in both weight-clipped WGAN and WGAN-GP, rather than the roughly 1:1 update ratio common in vanilla GAN training.

Think About It

Think about this: How would you explain wgan: wasserstein gan for 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.