Wasserstein GAN

The Wasserstein GAN (WGAN) is a generative adversarial network whose generator is trained to minimize an estimate of the Earth Mover’s distance, also called the Wasserstein-1 distance, between the distribution of its images and the data distribution, instead of the Jensen–Shannon divergence behind the original GAN. The discriminator is replaced by a critic, a network with an unbounded real output that is kept Lipschitz-continuous and trained to score real images high and generated ones low; the gap between its average scores is the distance estimate. Arjovsky, Chintala and Bottou, at the Courant Institute and Facebook AI Research, released it as an arXiv preprint (arXiv:1701.07875) and published it at ICML 2017. The original WGAN enforces the Lipschitz constraint by clipping the critic’s weights; WGAN-GP (Gulrajani et al., NIPS 2017) replaces the clipping with a penalty on the critic’s input gradient. Neither changes the generator: what changes is the loss and the constraint on the second network.

This article covers both papers. The GAN article summarizes WGAN among the remedies for unstable training, and its example measures how the minimax, non-saturating and WGAN-GP losses play the game and cover modes; the example here measures the distances and the critics themselves.

Problem Addressed

A GAN’s generator $g_\theta$ maps noise $z$ to an image, defining a distribution $\mathbb{P}_\theta$ that has no density to evaluate. When the generator’s input has fewer dimensions than an image, $\mathbb{P}_\theta$ lies on a low-dimensional set, and so, in the paper’s view, do natural images. Arjovsky and Bottou (ICLR 2017; read here in the authors’ copy) proved what this means for the original GAN. If the two distributions have disjoint supports, or lie on low-dimensional manifolds that do not perfectly align, a perfect discriminator exists whose gradient is zero almost everywhere on both (their Theorems 2.1 and 2.2); the Jensen–Shannon divergence is at its maximum and both Kullback–Leibler divergences are infinite however close the manifolds are (Theorem 2.3); and as the discriminator approaches the perfect one, the gradient of the minimax generator loss vanishes (Theorem 2.4). The non-saturating loss need not suffer from the vanishing gradient, but when the error of a nearly optimal discriminator is modelled as independent Gaussian noise, each coordinate of its expected generator gradient follows a centred Cauchy distribution, with infinite mean and variance (Theorem 2.6). Training therefore depends on how much the discriminator is trained, which is left to the user to decide.

The WGAN paper (§1–§2) frames the remedy as a choice of distance $\rho(\mathbb{P}_r, \mathbb{P}_\theta)$ between the real distribution $\mathbb{P}_r$ and the model. A weaker distance, one under which more sequences of distributions converge, makes it easier for the loss $\theta \mapsto \rho(\mathbb{P}_r, \mathbb{P}_\theta)$ to be continuous in the generator’s parameters, and continuity is what gradient descent on the generator needs.

Architecture

WGAN keeps the GAN’s two networks and changes the roles. The generator is any network that maps noise to images; the paper used the DCGAN generator of Radford et al., a DCGAN generator without batch normalization and with a constant number of filters, and a 4-layer ReLU multilayer perceptron with 512 units per layer. In most experiments the critic is the DCGAN discriminator without its final sigmoid, so that it outputs an unbounded score (WGAN, Figure 4 caption). Everything specific to WGAN lies in the distance the critic estimates and in how the critic is constrained.

The Earth Mover’s distance

For distributions $\mathbb{P}_r$ and $\mathbb{P}_g$ on a compact metric space such as the images $[0,1]^d$ , the paper (Eq. 1) defines

W(\mathbb{P}_r, \mathbb{P}_g) = \inf_{\gamma \in \Pi(\mathbb{P}_r, \mathbb{P}_g)} \mathbb{E}_{(x, y) \sim \gamma}\big[\lVert x - y \rVert\big],

where $\Pi(\mathbb{P}_r, \mathbb{P}_g)$ is the set of joint distributions (couplings) whose marginals are $\mathbb{P}_r$ and $\mathbb{P}_g$ . A coupling says how much probability mass to move from each $x$ to each $y$ , and $W$ is the cost of the cheapest plan, mass times distance. For histograms this is the earth mover’s distance of image retrieval, described with its transportation-problem solution in the similarity and distance measures article; here the points being moved are whole images.

The paper’s Example 1 shows the difference. Let $\mathbb{P}_0$ be uniform on the vertical segment $\{0\} \times [0, 1]$ in the plane and $\mathbb{P}_\theta$ the same segment shifted to $x = \theta$ . Then $W(\mathbb{P}_0, \mathbb{P}_\theta) = \lvert\theta\rvert$ , while the Jensen–Shannon divergence is $\log 2$ , both Kullback–Leibler divergences are $+\infty$ and the total variation distance is 1 for every $\theta \ne 0$ , and all four are 0 at $\theta = 0$ . Only the Earth Mover’s distance shrinks as $\theta \to 0$ , so only it gives the parameter a gradient. The paper’s definition of the Jensen–Shannon divergence in §2 omits the factors $\tfrac{1}{2}$ , which would make the value $2 \log 2$ , but its Example 1, its §4.3 and its proof of Theorem 2 (Appendix C) use the $\tfrac{1}{2}$ convention, under which the maximum is $\log 2$ ; this article uses the latter. Arjovsky and Bottou are inconsistent the other way: they define the divergence with the factors $\tfrac{1}{2}$ (§1) but give its value as $2 \log 2$ in Theorem 2.3.

Two theorems generalize the example. Theorem 1: if the generator $g_\theta(z)$ is continuous in $\theta$ , so is $W(\mathbb{P}_r, \mathbb{P}_\theta)$ ; if $g$ is locally Lipschitz with local constants whose expectation over $z$ is finite (their Assumption 1), $W$ is continuous everywhere and differentiable almost everywhere; neither holds for the Jensen–Shannon divergence or the KL divergences. Corollary 1 shows that any feedforward network with Lipschitz nonlinearities and a prior with $\mathbb{E}\lVert z \rVert < \infty$ satisfies the assumption (the paper says ReLUs need a similar but more technical proof). Theorem 2 orders the topologies: convergence in KL implies convergence in Jensen–Shannon and total variation, which are equivalent, and these imply convergence in $W$ , which is equivalent to convergence in distribution. The Earth Mover’s distance is the weakest of them, so it is continuous in the most cases.

The Kantorovich–Rubinstein duality and the critic

The infimum over couplings is, in the paper’s words, highly intractable for images. The Kantorovich–Rubinstein duality, which the paper takes from Villani (2009), rewrites it as a supremum over functions (WGAN, Eq. 2):

W(\mathbb{P}_r, \mathbb{P}_\theta) = \sup_{\lVert f \rVert_L \le 1} \; \mathbb{E}_{x \sim \mathbb{P}_r}\big[f(x)\big] - \mathbb{E}_{x \sim \mathbb{P}_\theta}\big[f(x)\big],

where the supremum runs over all 1-Lipschitz functions $f$ , those with $\lvert f(x) - f(y) \rvert \le \lVert x - y \rVert$ . Taking $K$ -Lipschitz functions instead gives $K \cdot W$ . A network $f_w$ whose weights lie in a compact set is $K$ -Lipschitz for some $K$ that depends on that set and not on the individual weights, so maximizing (Eq. 3)

\max_{w \in \mathcal{W}} \; \mathbb{E}_{x \sim \mathbb{P}_r}\big[f_w(x)\big] - \mathbb{E}_{z \sim p(z)}\big[f_w(g_\theta(z))\big]

estimates $W$ up to the factor $K$ and up to the network’s capacity, if the supremum is attained within the family. Theorem 3 gives the generator its gradient: for an optimal 1-Lipschitz $f$ , $\nabla_\theta W(\mathbb{P}_r, \mathbb{P}_\theta) = -\mathbb{E}_{z \sim p(z)}\big[\nabla_\theta f(g_\theta(z))\big]$ when both sides are defined. The generator therefore minimizes $-\mathbb{E}_z[f_w(g_\theta(z))]$ , backpropagating through the critic as a GAN generator backpropagates through its discriminator. The critic is not a classifier, which is why the paper renames it.

Appendix D places this among integral probability metrics, $\sup_{f \in \mathcal{F}} \mathbb{E}_{\mathbb{P}_r} f - \mathbb{E}_{\mathbb{P}_\theta} f$ for a function class $\mathcal{F}$ : 1-Lipschitz functions give $W$ , functions bounded in $[-1, 1]$ give total variation, and the unit ball of a reproducing-kernel Hilbert space gives the maximum mean discrepancy. Appendix E shows that an energy-based GAN with an optimal discriminator minimizes total variation, which has the regularity of the Jensen–Shannon divergence, not of $W$ .

Because $W$ is continuous and differentiable almost everywhere, the authors argue that the critic can and should be trained to optimality: a better critic gives a more reliable gradient, where a better GAN discriminator gives a vanishing one (§3). Their Figure 3 trains both to optimality on two Gaussians: the discriminator saturates, while the critic converges to a linear function with clean gradients everywhere. They also argue that training the critic to optimality makes mode collapse impossible, since collapse comes from the generator’s optimum against a fixed discriminator being a sum of point masses where the discriminator is highest.

The WGAN-GP paper (Proposition 1, proved in its Appendix A) adds a property of the optimal critic that its method relies on. If $f^\ast$ is optimal and differentiable, $\pi$ is an optimal coupling with $\pi(x = y) = 0$ , and $x_t = t x + (1 - t) y$ for a coupled pair of a generated $x$ and a real $y$ , then with probability 1 the gradient at $x_t$ is the unit vector $(y - x_t) / \lVert y - x_t \rVert$ : the optimal critic rises with slope exactly 1 along the straight lines from generated points to the real points they are coupled with. Its Corollary 1: $f^\ast$ has gradient norm 1 almost everywhere under $\mathbb{P}_r$ and $\mathbb{P}_g$ .

Weight clipping

The original WGAN keeps the critic’s weights in a compact set by clamping every weight to $[-c, c]$ after each update, with $c = 0.01$ in all experiments (Algorithm 1); the ICML version presents it as a simple way to keep the weights in a compact space. Appendix G.1 argues that with a small clipping value the critic behaves like a WGAN critic even with a sigmoid output, because the saturating regime is never reached, and Appendix G.2 that batch normalization still behaves as a Lipschitz function in practice.

Gulrajani et al. (2017, §3) identified the failure modes of clipping. They also tried L2 norm clipping, weight normalization and L1 and L2 weight decay and found similar problems, which batch normalization in the critic mitigates only partly:

Capacity underuse (§3.1). A critic constrained by clipping and trying to reach its largest gradient norm learns very simple functions. Trained to optimality on toy 2D distributions against a fixed “generator” (the real data plus unit-variance Gaussian noise), clipped critics ignored the higher moments of the data and modelled simple approximations of the optimal critic (their Figure 1a), whereas penalized critics did not.
Exploding and vanishing gradients (§3.2). On the Swiss Roll with 12-layer ReLU MLPs for both networks and no batch normalization, the gradient of the critic loss with respect to successive layers either grew or decayed exponentially going back through the network for each of $c \in \{0.1, 0.01, 0.001\}$ ; only careful tuning of $c$ avoids one or the other.
Weights at the bounds. Under clipping, the weights were pushed toward the two extremes of the clipping range, unlike under the gradient penalty (Figure 1b).

Arjovsky, Chintala and Bottou had themselves reported a negative result (§4.3): WGAN training sometimes became unstable with a momentum-based optimizer such as Adam ( $\beta_1 > 0$ ) on the critic or with high learning rates, since the critic’s loss is nonstationary; the cosine between the Adam step and the gradient usually turned negative where the loss blew up. They switched to RMSProp.

Gradient penalty (WGAN-GP)

A differentiable function is 1-Lipschitz if and only if its gradient has norm at most 1 everywhere. Gulrajani et al. (2017, §4) therefore constrain the gradient directly, softly and only where it matters most. The critic minimizes (their Eq. 3)

L = \underbrace{\mathbb{E}_{\tilde x \sim \mathbb{P}_g}\big[D(\tilde x)\big] - \mathbb{E}_{x \sim \mathbb{P}_r}\big[D(x)\big]}_{\text{original critic loss}} + \lambda \, \underbrace{\mathbb{E}_{\hat x \sim \mathbb{P}_{\hat x}}\Big[\big(\lVert \nabla_{\hat x} D(\hat x) \rVert_2 - 1\big)^2\Big]}_{\text{gradient penalty}},

with the critic written $D$ , as in their paper. The design choices, as they give them:

Sampling distribution. $\hat x = \epsilon x + (1 - \epsilon)\tilde x$ with $x$ a real image, $\tilde x$ a generated one and $\epsilon \sim U[0, 1]$ , so the penalty acts on straight lines between real and generated samples. Proposition 1 motivates this: the optimal critic has slope 1 along lines between coupled points, and enforcing the constraint everywhere is intractable. The pairs are random, not the optimal coupling; the authors found this sufficient in experiments.
Penalty coefficient. $\lambda = 10$ in all the paper’s experiments, from toy tasks to large ImageNet CNNs.
No batch normalization in the critic. The penalty acts on each input’s gradient separately, which batch normalization invalidates by mixing examples; layer normalization is their recommended drop-in replacement.
Two-sided penalty. The penalty pulls the norm toward 1 rather than only below it. In Appendix C, over 174 random architectures, the two-sided penalty gave the higher Inception Score in 100 trials and a one-sided $\max(0, \lVert \nabla D \rVert - 1)^2$ in 77 (as printed; the two counts add up to 177, more than the 174 architectures), a difference the authors note is not significant at $p < 0.05$ .
Smooth activations. The penalty’s gradient involves second derivatives of the activations. Piecewise-linear activations trained fine, but they could not train networks with ELU, and a smooth substitute, $\mathrm{softplus}(2x + 2)/2 - 1$ , fixed it (Appendix D).

Their Algorithm 1 uses $\lambda = 10$ , $n_{\text{critic}} = 5$ critic iterations per generator iteration, and Adam with learning rate $\alpha = 0.0001$ , $\beta_1 = 0$ and $\beta_2 = 0.9$ .

Training

WGAN’s Algorithm 1 alternates critic and generator updates. For each generator iteration, the critic takes $n_{\text{critic}}$ steps on minibatches of $m$ real images and $m$ noise vectors, each a gradient ascent step on the critic objective with RMSProp at learning rate $\alpha$ , followed by clipping; then the generator takes one RMSProp step on $-\frac{1}{m}\sum_i f_w(g_\theta(z^{(i)}))$ . The defaults used in all experiments were $\alpha = 0.00005$ , $c = 0.01$ , $m = 64$ and $n_{\text{critic}} = 5$ . As printed, the critic loop runs “for $t = 0, \dots, n_{\text{critic}}$ ”, which would be $n_{\text{critic}} + 1$ steps; the text and the official code use $n_{\text{critic}}$ . The official code also trains the critic 100 times per generator step for the first 25 generator iterations and at every 500th, a detail its README says the paper omitted, so that the critic starts near its optimum and the loss curve is not distorted early on.

WGAN-GP’s Algorithm 1 has the same structure with the penalty added to each critic loss and Adam for both networks. The paper’s own settings vary by experiment: Appendix B uses the Algorithm 1 settings for the random-architecture study, while Appendix E lists Adam with $\beta_1 = 0.5$ and $\beta_2 = 0.9$ for the LSUN robustness experiments, and the CIFAR-10 ResNet of Appendix F uses a learning rate of $2 \times 10^{-4}$ decayed linearly to 0 over 100,000 generator iterations.

The critic’s loss as a training signal

Because the critic is trained relatively well before each generator update, its objective (Eq. 3) at that point estimates the Earth Mover’s distance, up to constant factors set by the way the Lipschitz constant is constrained (WGAN, §4.3); the critic’s loss is its negative. The paper’s Figure 4 plots this estimate during training on LSUN Bedrooms at 64 × 64 pixels for a DCGAN generator, an MLP generator, and MLPs for both networks, after a median filter, and the authors describe a clear correlation between lower values and better samples; when training failed, the curve stayed flat. A standard GAN’s discriminator loss, converted to a lower bound on the Jensen–Shannon divergence, did not track sample quality and often stayed near $\log 2$ (Figure 5), and the same held for its generator loss (Appendix F, Figure 9). They presented this as, to their knowledge, the first GAN loss showing properties of convergence. They did not claim a new evaluation measure: the scale depends on the critic’s architecture, which makes models with different critics hard to compare, and a critic of finite capacity gives an estimate of unknown accuracy.

WGAN-GP keeps the property: the negative critic loss of a WGAN-GP on LSUN Bedrooms converged as the generator trained (Gulrajani et al., Figure 5a). On a random subset of 1,000 MNIST images, large unregularized models overfit, and the training and validation critic losses diverged for both WGAN and WGAN-GP, a sign that the critic overfits and no longer estimates $W$ ; with WGAN-GP the training loss even rose while the validation loss fell (Figure 5b). The authors present the gap as a way to detect overfitting in the critic against the same loss the networks are trained on.

Example: Distances and Critics on a Shifted Shape

The example computes the Jensen–Shannon divergence and the Earth Mover’s distance exactly for two image distributions with disjoint supports, then trains a small critic on the same pair with weight clipping and with the gradient penalty and compares its estimate with the exact distance. The data are 16 × 24 grayscale images of one soft round shape, a Gaussian blob with $\sigma = 1.5$ pixels and values in $[0, 1]$ . The real distribution $\mathbb{P}_r$ is uniform over 32 images with the blob’s centre at column 5 and rows evenly spaced from 4 to 11; the model distribution $\mathbb{P}_\theta$ is the same 32 images with the shape shifted right by $\theta$ pixels, as in the paper’s Example 1. Images are compared as vectors in $\mathbb{R}^{384}$ with the Euclidean distance, the ground metric of the paper’s definition.

With equal weights on equally many points, an optimal coupling is a permutation (Birkhoff’s theorem on doubly stochastic matrices), so the primal is an assignment problem. The dual is the Kantorovich–Rubinstein form restricted to the 64 support points: a value $f$ at each point, maximizing the difference of means subject to $f(u) - f(v) \le \lVert u - v \rVert$ for every ordered pair. The script also prints the slack of the Lipschitz constraint on the coupled pairs, which complementary slackness of the linear program makes zero, the discrete counterpart of WGAN-GP’s Proposition 1 (NumPy 2.3.5, SciPy 1.16.3; under a second).

import numpy as np
import scipy
from scipy.optimize import linear_sum_assignment, linprog

H, W, SIGMA, K = 16, 24, 1.5, 32
yy, xx = np.mgrid[0:H, 0:W]

def blob(cx, cy):  # a soft round shape centred at (cx, cy), values in [0, 1]
    return np.exp(-((xx - cx) ** 2 + (yy - cy) ** 2) / (2 * SIGMA ** 2)).ravel()

CY = np.linspace(4, 11, K)
real = np.stack([blob(5.0, cy) for cy in CY])          # P_r: 32 equally likely images

def jsd(p_pts, q_pts):  # exact JSD (natural log, 1/2 convention) of two uniform discrete distributions
    pts = np.unique(np.vstack([p_pts, q_pts]), axis=0)
    idx = lambda a: np.array([np.flatnonzero((pts == r).all(1))[0] for r in a])
    p = np.bincount(idx(p_pts), minlength=len(pts)) / len(p_pts)
    q = np.bincount(idx(q_pts), minlength=len(pts)) / len(q_pts)
    m = (p + q) / 2
    kl = lambda a: np.sum(a[a > 0] * np.log(a[a > 0] / m[a > 0]))
    return 0.5 * kl(p) + 0.5 * kl(q)

def w1_primal(a, b):  # equal weights: an optimal coupling is a permutation
    C = np.linalg.norm(a[:, None] - b[None], axis=2)
    r, c = linear_sum_assignment(C)
    return C[r, c].mean(), c

def w1_dual(a, b):  # Kantorovich-Rubinstein: max E_r f - E_g f over f 1-Lipschitz on the points
    pts = np.vstack([a, b]); n = len(pts)
    D = np.linalg.norm(pts[:, None] - pts[None], axis=2)
    i, j = np.where(~np.eye(n, dtype=bool))
    A = np.zeros((len(i), n)); A[np.arange(len(i)), i] = 1; A[np.arange(len(i)), j] = -1
    w = np.r_[np.full(len(a), 1 / len(a)), np.full(len(b), -1 / len(b))]
    res = linprog(-w, A_ub=A, b_ub=D[i, j], bounds=[(None, None)] * n, method="highs")
    return -res.fun, res.x

print(f"NumPy {np.__version__}, SciPy {scipy.__version__}; {K} images of {H}x{W}, sigma {SIGMA}")
print(f"non-overlap ceiling sqrt(2)*||g|| = {np.sqrt(2) * np.linalg.norm(real[0]):.3f}")
print(f"{'theta':>5} {'min dist':>9} {'JSD':>7} {'W1 primal':>10} {'W1 dual':>8} {'max slack on coupled pairs':>27}")
for theta in [0, 0.5, 1, 2, 3, 4, 6, 8, 12]:
    fake = np.stack([blob(5.0 + theta, cy) for cy in CY])   # P_theta: the same shapes shifted right
    dmin = np.linalg.norm(real[:, None] - fake[None], axis=2).min()
    wp, perm = w1_primal(real, fake)
    wd, f = w1_dual(real, fake)
    slack = np.abs(f[:K] - f[K + perm] - np.linalg.norm(real - fake[perm], axis=1)).max()
    print(f"{theta:5.1f} {dmin:9.3f} {jsd(real, fake):7.4f} {wp:10.4f} {wd:8.4f} {slack:27.1e}")

Output:

NumPy 2.3.5, SciPy 1.16.3; 32 images of 16x24, sigma 1.5
non-overlap ceiling sqrt(2)*||g|| = 3.760
theta  min dist     JSD  W1 primal  W1 dual  max slack on coupled pairs
  0.0     0.000  0.0000     0.0000  -0.0000                     0.0e+00
  0.5     0.622  0.6931     0.6223   0.6223                     0.0e+00
  1.0     1.219  0.6931     1.2193   1.2193                     0.0e+00
  2.0     2.252  0.6931     2.2523   2.2523                     0.0e+00
  3.0     2.989  0.6931     2.9894   2.9894                     0.0e+00
  4.0     3.427  0.6931     3.4275   3.4275                     0.0e+00
  6.0     3.725  0.6931     3.7253   3.7253                     0.0e+00
  8.0     3.758  0.6931     3.7584   3.7584                     0.0e+00
 12.0     3.760  0.6931     3.7599   3.7599                     7.8e-08

The Jensen–Shannon divergence behaves as in Example 1: it jumps from 0 to $\log 2 \approx 0.6931$ as soon as the shape moves half a pixel and stays there, so it says nothing about how far the shape has moved. The Earth Mover’s distance grows from 0 with $\theta$ , and the primal and dual agree to every printed digit, which is the Kantorovich–Rubinstein duality holding exactly on this finite problem. The slack on coupled pairs is zero up to solver tolerance (at most $7.8 \times 10^{-8}$ ): the optimal potential rises by exactly the distance between each real image and its partner. In every row $W$ equals the smallest distance between any real and any shifted image, so every unit of mass travels the shortest distance available.

The example also shows a limit that Example 1 hides. Here $W$ is measured in pixel space, and the pixel distance between two copies of a shape stops growing once they no longer overlap appreciably: it comes within 1% of the ceiling $\sqrt{2}\,\lVert g \rVert$ of two disjoint blobs, 3.760, by $\theta = 6$ , and from $\theta = 8$ to $\theta = 12$ it barely changes. In our reading, the Wasserstein distance is only as informative as its ground metric: it gives a usable signal while generated and real images still overlap in pixel space, and a signal that flattens, as the JSD’s does, when they are far apart in that metric. A critic constrained to be 1-Lipschitz in pixel space works with the same metric.

The second script trains a critic on the same pairs with the generator held fixed, as Gulrajani et al. did for their Figure 1a, so that the critic’s estimate $\hat W = \mathbb{E}_{\mathbb{P}_r} f - \mathbb{E}_{\mathbb{P}_\theta} f$ can be compared with the exact $W$ . The critic is a ReLU MLP, 384 → 128 → 128 → 1, without normalization. Each method uses its paper’s defaults: clipping at $c = 0.01$ with RMSProp at learning rate $5 \times 10^{-5}$ , and the gradient penalty with $\lambda = 10$ and Adam at $10^{-4}$ with $\beta = (0, 0.9)$ . Clipping is applied once at initialization as well as after each update, as the official WGAN code clamps the weights before every critic step. Each critic takes 2,000 steps on batches of 64 real and 64 shifted images; the script prints $\hat W$ after 1,000 and 2,000 steps, the gradient norm $\lVert \nabla_x f \rVert$ at the real and shifted images and at 1,024 random interpolates, and for clipping the fraction of weights at the bound. It loops over $\theta \in \{1, 2, 4, 8\}$ and seeds 0, 1 and 2 (PyTorch 2.14.1 with SciPy 1.18.1 for the exact $W$ , CPU, one thread, 25 to 40 s in our runs; two runs give identical output).

import sys
import numpy as np
import scipy
import torch
import torch.nn as nn
from scipy.optimize import linear_sum_assignment

torch.set_num_threads(1)
H, W, SIGMA, K = 16, 24, 1.5, 32
yy, xx = torch.meshgrid(torch.arange(H, dtype=torch.float32), torch.arange(W, dtype=torch.float32), indexing="ij")

def blob(cx, cy):  # same soft round shape as in the first script
    return torch.exp(-((xx - cx) ** 2 + (yy - cy) ** 2) / (2 * SIGMA ** 2)).flatten()

CY = torch.linspace(4, 11, K)
real = torch.stack([blob(5.0, float(cy)) for cy in CY])

def exact_w1(a, b):  # optimal assignment = exact W1 for two uniform 32-point distributions
    C = torch.cdist(a.double(), b.double()).numpy()
    r, c = linear_sum_assignment(C)
    return C[r, c].mean()

def critic():
    return nn.Sequential(nn.Linear(H * W, 128), nn.ReLU(), nn.Linear(128, 128), nn.ReLU(), nn.Linear(128, 1))

def grad_norm(f, x):
    x = x.clone().requires_grad_(True)
    g, = torch.autograd.grad(f(x).sum(), x)
    return g.norm(dim=1)

def train(method, fake, seed, steps=2000, batch=64, c=0.01, lam=10.0):
    torch.manual_seed(seed)
    gen = torch.Generator().manual_seed(seed + 100)
    f = critic()
    if method == "clip":   # WGAN, Algorithm 1: RMSProp, lr 5e-5, weights clipped to [-c, c]
        opt = torch.optim.RMSprop(f.parameters(), lr=5e-5)
        with torch.no_grad():
            for p in f.parameters(): p.clamp_(-c, c)
    else:                  # WGAN-GP, Algorithm 1: Adam, lr 1e-4, betas (0, 0.9), lambda 10
        opt = torch.optim.Adam(f.parameters(), lr=1e-4, betas=(0.0, 0.9))
    est_half = None
    for step in range(steps):
        x = real[torch.randint(0, K, (batch,), generator=gen)]
        xt = fake[torch.randint(0, K, (batch,), generator=gen)]
        loss = f(xt).mean() - f(x).mean()                     # critic loss: E_g f - E_r f
        if method == "gp":
            eps = torch.rand(batch, 1, generator=gen)
            xh = (eps * x + (1 - eps) * xt).requires_grad_(True)
            g, = torch.autograd.grad(f(xh).sum(), xh, create_graph=True)
            loss = loss + lam * ((g.norm(dim=1) - 1) ** 2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
        if method == "clip":
            with torch.no_grad():
                for p in f.parameters(): p.clamp_(-c, c)
        if step == steps // 2 - 1:
            with torch.no_grad(): est_half = (f(real).mean() - f(fake).mean()).item()
    with torch.no_grad():
        est = (f(real).mean() - f(fake).mean()).item()
    e = torch.rand(1024, 1, generator=gen)
    xi = e * real[torch.randint(0, K, (1024,), generator=gen)] + (1 - e) * fake[torch.randint(0, K, (1024,), generator=gen)]
    gi = grad_norm(f, xi)
    at_bound = None
    if method == "clip":
        w = torch.cat([p.detach().flatten() for p in f.parameters()])
        at_bound = (w.abs() >= 0.99 * c).float().mean().item()
    return est_half, est, grad_norm(f, real).mean().item(), grad_norm(f, fake).mean().item(), gi, at_bound

thetas = [float(t) for t in sys.argv[1].split(",")] if len(sys.argv) > 1 else [1, 2, 4, 8]
seeds = [int(s) for s in sys.argv[2].split(",")] if len(sys.argv) > 2 else [0, 1, 2]
print(f"PyTorch {torch.__version__}, SciPy {scipy.__version__}; critic 384-128-128-1 ReLU, 2000 critic steps")
print(f"{'theta':>5} {'W1':>6} {'seed':>4} {'method':>6} {'est@1000':>9} {'est@2000':>9} {'est/W1':>7} "
      f"{'|grad| real':>11} {'fake':>7} {'interp mean':>11} {'interp min-max':>15} {'at +-c':>7}")
for theta in thetas:
    fake = torch.stack([blob(5.0 + theta, float(cy)) for cy in CY])
    w1 = exact_w1(real, fake)
    for seed in seeds:
        for method in ["clip", "gp"]:
            eh, est, gr, gf, gi, ab = train(method, fake, seed)
            ab = f"{ab:7.0%}" if ab is not None else f"{'-':>7}"
            print(f"{theta:5.1f} {w1:6.3f} {seed:4d} {method:>6} {eh:9.4f} {est:9.4f} {est / w1:7.3f} "
                  f"{gr:11.4f} {gf:7.4f} {gi.mean():11.4f} {gi.min():7.4f}-{gi.max():.4f} {ab}")

Output:

PyTorch 2.14.1, SciPy 1.18.1; critic 384-128-128-1 ReLU, 2000 critic steps
theta     W1 seed method  est@1000  est@2000  est/W1 |grad| real    fake interp mean  interp min-max  at +-c
  1.0  1.219    0   clip    0.0402    0.0407   0.033      0.0803  0.0458      0.0839  0.0419-0.0839     86%
  1.0  1.219    0     gp    1.2704    1.3330   1.093      1.1182  1.1167      1.0600  0.9370-1.1560       -
  1.0  1.219    1   clip    0.0371    0.0384   0.031      0.0799  0.0445      0.0799  0.0566-0.0799     85%
  1.0  1.219    1     gp    1.2608    1.3252   1.087      1.1073  1.1108      1.0561  0.9147-1.1394       -
  1.0  1.219    2   clip    0.0427    0.0426   0.035      0.0878  0.0434      0.0874  0.0414-0.0880     86%
  1.0  1.219    2     gp    1.2968    1.3570   1.113      1.1281  1.1287      1.0640  0.9442-1.1754       -
  2.0  2.252    0   clip    0.0689    0.0689   0.031      0.0319  0.0359      0.0822  0.0393-0.0830     84%
  2.0  2.252    0     gp    2.5397    2.7340   1.214      1.2169  1.2142      1.1211  0.9677-1.3056       -
  2.0  2.252    1   clip    0.0741    0.0747   0.033      0.0368  0.0401      0.0889  0.0432-0.0898     84%
  2.0  2.252    1     gp    2.4750    2.6217   1.164      1.1956  1.1961      1.1093  0.9763-1.2590       -
  2.0  2.252    2   clip    0.0758    0.0758   0.034      0.0342  0.0403      0.0905  0.0404-0.0918     83%
  2.0  2.252    2     gp    2.5826    2.7149   1.205      1.2271  1.1862      1.1171  0.9693-1.2966       -
  4.0  3.428    0   clip    0.0812    0.0830   0.024      0.0308  0.0387      0.0766  0.0343-0.0779     83%
  4.0  3.428    0     gp    3.7800    4.0266   1.175      1.1484  1.1454      1.1782  1.1235-1.2531       -
  4.0  3.428    1   clip    0.0829    0.0863   0.025      0.0352  0.0365      0.0795  0.0372-0.0805     85%
  4.0  3.428    1     gp    3.8471    4.0730   1.188      1.1064  1.0875      1.1680  1.0715-1.2474       -
  4.0  3.428    2   clip    0.0829    0.0839   0.024      0.0334  0.0372      0.0774  0.0376-0.0788     84%
  4.0  3.428    2     gp    3.8160    4.0240   1.174      1.1703  1.1074      1.1724  1.1139-1.2441       -
  8.0  3.758    0   clip    0.0702    0.0702   0.019      0.0277  0.0332      0.0462  0.0277-0.0517     78%
  8.0  3.758    0     gp    4.0664    4.2244   1.124      1.0764  1.1682      1.1699  1.1563-1.1801       -
  8.0  3.758    1   clip    0.0748    0.0748   0.020      0.0328  0.0321      0.0483  0.0322-0.0545     77%
  8.0  3.758    1     gp    4.2422    4.4080   1.173      1.1183  1.0909      1.1895  1.0598-1.2009       -
  8.0  3.758    2   clip    0.0759    0.0759   0.020      0.0344  0.0317      0.0494  0.0317-0.0576     77%
  8.0  3.758    2     gp    4.1073    4.2836   1.140      1.0729  1.1694      1.1758  1.1147-1.1883       -

The clipped critic has converged (its estimate barely moves between steps 1,000 and 2,000) to a function whose gradient norm is 0.03 to 0.09, so in this network clipping at 0.01 acts as a Lipschitz bound far below 1, and its estimate is 2% to 3.5% of $W$ . That would be harmless if the factor were constant, since the paper’s Eq. 3 only claims $W$ up to a multiplicative constant. It is not: the ratio falls from about 0.033 at $\theta = 1$ to about 0.020 at $\theta = 8$ , and in all three seeds the estimate at $\theta = 8$ is lower than at $\theta = 4$ , although $W$ rose from 3.428 to 3.758. In this run the clipped critic orders the two larger shifts the wrong way round. More than three quarters of its weights sit at $\pm 0.01$ , the pile-up at the bounds that Gulrajani et al. show in their Figure 1b.

The penalized critic’s gradient norms sit near 1, with means from 1.06 to 1.23, and its estimate rises with $\theta$ and flattens where $W$ does. It overshoots, by 9% to 21%, and it has not converged: in every run it was still rising between steps 1,000 and 2,000. In our reading, the penalty is a soft constraint and acts only on the sampled interpolates, so the critic can buy a larger difference of means with gradient norms somewhat above 1, which is what the printed norms at the data points show. Neither critic recovers $W$ exactly: clipping gives an estimate with the wrong and varying scale, the penalty one with roughly the right scale that slightly exceeds it. These are statements about a toy critic trained on fixed distributions, not about the relative merit of the methods on images.

Inference

At sampling time only the generator is used: a noise vector goes through one forward pass to produce an image, as in any GAN, and the critic is discarded. Its estimate can still be computed on held-out real images and generated samples to monitor training, which is how Gulrajani et al. detected overfitting.

Performance

WGAN (2017). The paper reports no quantitative sample-quality score. Its experiments are qualitative: on a mixture of 8 Gaussians on a circle, WGAN learned the distribution without mode collapse and appeared to match the circle before the individual bumps (Figure 2). On LSUN Bedrooms at 64 × 64, with the hyperparameters of Algorithm 1 throughout, a WGAN and a standard GAN with the non-saturating loss (“the −log D trick”) both produced high-quality samples with DCGAN networks (Figure 6). With a generator without batch normalization and with a constant number of filters, the standard GAN failed to learn while WGAN still produced samples (Figure 7), and with a 4-layer 512-unit MLP generator WGAN produced lower-quality samples than with DCGAN but better than the standard GAN, which showed a significant degree of mode collapse (Figure 8). The authors saw no evidence of mode collapse in any WGAN experiment, and state that samples were not cherry-picked.

WGAN-GP (2017). Gulrajani et al. report:

Random architectures (§5.1). 200 architectures sampled from a DCGAN-based set (depth, nonlinearity, normalization and filter counts in each network), each trained on 32 × 32 ImageNet with WGAN-GP and with the standard GAN objective for 100,000 iterations, with success defined as an Inception Score above a threshold. At a threshold of 3.0, 88 architectures succeeded only with WGAN-GP, 1 only with the standard GAN, 110 with both and 1 with neither; at 5.0 the counts were 147, 0, 42 and 11 (Table 2, extended in Table 5 of the appendix). Their standard DCGAN scored 7.24.
Varied architectures on LSUN Bedrooms (§5.2). The DCGAN baseline and six other architectures, among them tanh nonlinearities, gated multiplicative nonlinearities and a 101-layer ResNet for both networks, each trained for 200,000 iterations with WGAN-GP, WGAN with clipping, DCGAN and LSGAN. Only WGAN-GP trained every architecture with one shared set of hyperparameters; with every other method some were unstable or collapsed modes (Figure 2). The authors believe this was the first successful training of very deep residual networks in a GAN, without claiming it is impossible otherwise.
Against clipping on CIFAR-10 (§5.3). On CIFAR-10, WGAN-GP with RMSProp at the clipping run’s learning rate converged faster and to a higher Inception Score than weight clipping, and Adam with a higher learning rate improved it further; WGAN-GP converged more slowly than DCGAN in wall-clock time but its score was more stable at convergence (Figure 3).
Inception Scores on CIFAR-10 (§5.4, Table 3). A ResNet trained with WGAN-GP scored $7.86 \pm 0.07$ unsupervised, the best in their table, above DFM ( $7.72 \pm 0.13$ ) and DCGAN ( $6.16 \pm 0.07$ ). With label information added by the AC-GAN method, it scored $8.42 \pm 0.10$ , above every model in their supervised list except SGAN ( $8.59 \pm 0.12$ ). They also show 128 × 128 LSUN Bedroom samples from a deep ResNet, which they believe are at least competitive with the best reported at any resolution.
Text with a continuous generator (§5.5). A character-level language model on the Google Billion Word dataset, whose 1D-CNN generator outputs softmax vectors that are fed to the critic without sampling, produced recognizably language-like 32-character samples with frequent spelling errors. They could not get comparable results with the standard GAN objective. Their explanation: a distribution on one-hot sequences and a continuous generator distribution have infinite KL divergences and a saturated JSD, while the conditions of WGAN’s Theorem 1 still hold, so $W$ stays continuous.

Fréchet Inception Distance was introduced in the same year, and the FID paper evaluates WGAN-GP with it.

Strengths

A loss that is continuous in the generator under mild assumptions, even when the generated and real distributions do not overlap (WGAN, Theorem 1), where the Jensen–Shannon divergence is constant.
No balancing act. The critic can be trained close to optimality without its gradient vanishing, so the relative capacity and training of the two networks need less tuning (WGAN, §4.4).
A meaningful training curve. The critic’s loss estimates the distance and tracked sample quality in both papers’ experiments, and its gap between training and validation data can reveal overfitting.
Architectural robustness. With the gradient penalty, a single set of hyperparameters trained a wide range of architectures, including ones that standard GAN training could not (Gulrajani et al., §5.1–§5.2).

Limitations

The constraint is approximate. Clipping bounds the Lipschitz constant only through an unknown, architecture-dependent $K$ , underuses capacity and, without careful tuning of $c$ , makes gradients vanish or explode (Gulrajani et al., §3). The gradient penalty constrains the gradient only softly and only on sampled interpolates. The example above shows both effects on a small critic.
The estimate is not a calibrated distance. Its scale depends on the critic’s architecture and capacity, so values from different critics are not comparable (WGAN, §4.3).
Cost. Each generator step needs several critic steps, and the penalty needs a second backward pass through the critic. WGAN-GP converged more slowly than DCGAN in wall-clock time on CIFAR-10 (Gulrajani et al., §5.3 and Figure 3).
Optimizer sensitivity. The original WGAN was unstable with momentum on the critic or high learning rates (WGAN, §4.3), and WGAN-GP rules out batch normalization in the critic.
Convergence is not guaranteed. Mescheder, Geiger and Nowozin (2018) showed on a simple example with point-mass data that neither WGAN nor WGAN-GP with a fixed number of critic updates per generator update always converges to the equilibrium, while noting that WGAN-GP works in practice.
The ground metric is pixel distance. In our reading of the example, the Earth Mover’s distance in pixel space flattens once shapes stop overlapping, so its advantage over the JSD is largest when the distributions are already close in pixel space.

Historical Importance

Arjovsky and Bottou’s companion paper explained the saturation and instability of GANs by the geometry of distributions on low-dimensional manifolds, and WGAN turned the analysis into an algorithm: choose a distance under which the generator’s loss is continuous, and approximate it with a constrained network. Its authors also presented its critic loss as, to their knowledge, the first GAN loss to show properties of convergence, one they could use to monitor training.

WGAN-GP made the idea practical. Its penalty trained deep residual GANs and a GAN language model with almost no tuning, and set the then-best unsupervised Inception Score on CIFAR-10 in its table. Gradient penalties are also used with other losses: StyleGAN’s R1 regularization, described in the StyleGAN article, is a zero-centred gradient penalty on real images, one of the simplified penalties for which Mescheder et al. (2018) proved local convergence under suitable assumptions. Spectral normalization, another way of bounding the discriminator’s Lipschitz constant, is summarized in the GAN article.

Variants

WGAN with weight clipping (Arjovsky, Chintala & Bottou, 2017): the original critic constraint, $c = 0.01$ , RMSProp.
WGAN-GP (Gulrajani et al., 2017): the two-sided penalty with $\lambda = 10$ , no batch normalization in the critic (layer normalization instead), Adam. Its authors also tested a one-sided penalty, which performed slightly worse in their comparison (Appendix C), and conditioned the CIFAR-10 model on class labels with the AC-GAN method (Table 3). In their conclusion they suggest adapting the penalty to the standard GAN objective as a direction for future work.
WGAN-GP with two time-scale updates. Heusel et al. (2017), the paper that introduced FID, trained WGAN-GP on CIFAR-10 and LSUN Bedrooms with one critic update per generator update and a critic learning rate of $3 \times 10^{-4}$ against the generator’s $10^{-4}$ , instead of the original code’s five critic updates (§3, Table 1); the FID paper article gives the settings and results.

Implementations

The authors’ code, as of November 2026:

martinarjovsky/WassersteinGAN on GitHub (PyTorch, BSD 3-Clause licence; latest commit f7a01e8, December 2018). main.py defaults to RMSProp at $5 \times 10^{-5}$ for both networks, clipping to $[-0.01, 0.01]$ before each critic step, and 5 critic iterations (100 for the first 25 generator iterations and every 500th); an --adam option switches to Adam. The README notes that plotting the negative critic loss reproduces the paper’s curves with a median filter.
igul222/improved_wgan_training on GitHub (TensorFlow, Python 2, MIT licence; latest commit fa66c57, June 2017), the code the WGAN-GP paper links. It has scripts for the toy datasets, MNIST, CIFAR-10, the 64 × 64 architectures and the language model; the README says the 64 × 64 script trains on ImageNet instead of the paper’s LSUN Bedrooms. Settings differ by script: the toy script uses $\lambda = 0.1$ , with a comment that a smaller $\lambda$ is faster on toy tasks but unnecessary with more critic iterations, and Adam with $\beta_1 = 0.5$ , while the CIFAR-10 ResNet script uses $\lambda = 10$ and Adam with $\beta_1 = 0$ , $\beta_2 = 0.9$ .

The critic loss and gradient penalty are a few lines in any automatic-differentiation framework, as in the example above.