Wasserstein GANs and Lipschitz constraints

The GAN note ended with a problem. With an optimal discriminator the generator minimizes $2\,D_{\mathrm{JS}}(P\,\|\,Q) - 2\ln 2$, and when the real distribution $P$ and the generator’s distribution $Q$ have disjoint supports, $D_{\mathrm{JS}} = \ln 2$ wherever $Q$ sits. A (near-)optimal discriminator is then constant around the generated samples, so the minimax generator loss gets no gradient and the non-saturating one gets unstable updates (Arjovsky and Bottou, 2017, Theorems 2.1–2.2, 2.4 and 2.6; see the GAN note). The Wasserstein GAN (Arjovsky, Chintala and Bottou, 2017) replaces JS with the Wasserstein-1 distance, which keeps decreasing as $Q$ moves towards $P$ even when they do not overlap. The price is a constraint on the discriminator: it must be 1-Lipschitz. The rest of the note is about the losses and the two standard ways to meet that constraint.

Notation. $P$ is the real data distribution, $Q$ the generator’s distribution ($x = G(z)$ with noise $z \sim p_z$, as in the GAN note), and $f$ the critic: the WGAN name for the discriminator, a network with a real-valued output and no sigmoid. Its loss keeps the discriminator’s name, $\mathcal L_D$. Real samples are $x_r \sim P$ and fake ones $x_f = G(z)$. $\|\cdot\|_2$ is the Euclidean norm.

Wasserstein-1 distance#

Think of $P$ and $Q$ as two piles of earth of equal mass. A transport plan $\gamma(x, y)$ says how much mass moves from $x$ to $y$; its marginals must be $P$ and $Q$, and $\Pi(P, Q)$ is the set of all such plans (Peyré and Cuturi, 2019, Chapter 2). Moving mass $m$ over a distance $\|x - y\|$ costs $m\,\|x - y\|$, and $W_1$, also called the earth mover’s distance, is the cheapest total cost:

$$W_1(P, Q) = \inf_{\gamma \in \Pi(P, Q)} \mathbb{E}_{(x, y)\sim\gamma}\,\|x - y\| .$$

A network cannot optimize directly over all transport plans. For distributions on $\mathbb R^d$ with finite first moments, $\mathbb E\|x\| < \infty$ under both, $W_1$ is finite and the Kantorovich–Rubinstein duality (Villani, 2009, Remark 6.5) rewrites the infimum as a supremum over functions:

$$\begin{aligned} W_1(P, Q) = \sup_{\|f\|_L \le 1}\; \Big[\, &\mathbb{E}_{x\sim P}\,f(x)\\ &- \mathbb{E}_{x\sim Q}\,f(x) \Big], \end{aligned}$$

where $\|f\|_L \le 1$ ($f$ is 1-Lipschitz) means $|f(x) - f(y)| \le \|x - y\|$ for all $x, y$. This is where the critic and its constraint come from: the critic $f$ plays the role of the function inside the supremum. It scores real samples high and fake ones low, but it may not change faster than distance. Without the constraint, and with $P \ne Q$, the supremum would be $+\infty$, since any separating $f$ could be scaled up.

Why this helps. Take two point masses, $P = \delta_0$ (the data) and $Q_\theta = \delta_\theta$ (the model). For every $\theta \ne 0$ the supports are disjoint. The mixture $M = \tfrac12(P + Q_\theta)$ puts mass $\tfrac12$ on each atom, so each KL term in $D_{\mathrm{JS}}$ compares mass $1$ with mass $\tfrac12$ on a single atom:

$$\begin{aligned} D_{\mathrm{JS}}(P\,\|\,Q_\theta) &= \tfrac12 D_{\mathrm{KL}}(P\,\|\,M)\\ &\quad + \tfrac12 D_{\mathrm{KL}}(Q_\theta\,\|\,M)\\ &= \tfrac12\ln\frac{1}{1/2} + \tfrac12\ln\frac{1}{1/2}\\ &= \ln 2 \approx 0.693,\\ D_{\mathrm{KL}}(P\,\|\,Q_\theta) &= \infty,\\ W_1(P, Q_\theta) &= |\theta| . \end{aligned}$$

The only plan moves all the mass a distance $|\theta|$. $D_{\mathrm{JS}}$ is constant in $\theta$, so its derivative is $0$; $W_1$ has slope $\pm 1$ and always points the model towards the data. This is a point-mass version of the parallel-lines example of the WGAN paper (Arjovsky, Chintala and Bottou, 2017, Example 1). More generally, write $Q_\theta$ for the distribution of $G_\theta(z)$. If the data live in a compact set, as the paper assumes throughout, and $G_\theta$ is continuous in $\theta$, so is $W_1(P, Q_\theta)$. On an unbounded space continuity of $G_\theta$ alone is not enough, and an integrability condition is needed as well. If $G_\theta$ is a feedforward network and the prior has $\mathbb E\|z\| < \infty$, which supplies such a condition, $W_1(P, Q_\theta)$ is continuous and also differentiable almost everywhere in $\theta$. Neither statement holds for $D_{\mathrm{JS}}$ or KL, which can jump (Arjovsky, Chintala and Bottou, 2017, Theorem 1 and Corollary 1).

Figure

D_JS is flat for disjoint supports; W₁ keeps a slope

WGAN losses#

The critic maximizes the dual objective, so its loss is the negative of it. The generator only appears in the second expectation and wants the critic’s score on fakes to be high:

$$\begin{aligned} \mathcal L_D &= -\mathbb{E}_{x\sim P}\big[f(x)\big] + \mathbb{E}_{z\sim p_z}\big[f(G(z))\big],\\ \mathcal L_G &= -\mathbb{E}_{z\sim p_z}\big[f(G(z))\big]. \end{aligned}$$

What does $-\mathcal L_D$ measure? A network family contains only some of the $K$-Lipschitz functions, so even its best member gives $-\mathcal L_D \le K\,W_1(P, Q)$: a lower bound, tight when the family contains $K$ times an optimal 1-Lipschitz critic ($K = 1$ for a 1-Lipschitz family). In training, $-\mathcal L_D$ is a minibatch estimate of that value. Unlike a GAN discriminator, the critic should be trained close to optimality: the better the critic, the more reliable the gradient of $W_1$, whereas a better GAN discriminator saturates. It is therefore updated several times per generator step, $n_{\text{critic}} = 5$ in the original algorithm (Arjovsky, Chintala and Bottou, 2017, Section 3 and Algorithm 1).

With a constrained critic trained close to optimality, $-\mathcal L_D$ correlates with sample quality during a run, which the GAN discriminator’s loss does not (Section 4.2, Figures 3–4). The constant $K$ depends on the critic’s architecture, so the number is hard to compare between models with different critics, and the authors do not propose it as an evaluation metric.

The losses alone do not keep $f$ Lipschitz. The original WGAN clipped every critic weight to $[-c, c]$ after each step, which makes the critic $K$-Lipschitz for some $K$ that depends on $c$ and the architecture. Its authors called this “a clearly terrible way” to enforce the constraint: a large $c$ makes the critic slow to train to optimality, and a small $c$ can make gradients vanish (Arjovsky, Chintala and Bottou, 2017, Section 3). WGAN-GP showed further that clipping biases the critic towards overly simple functions, and that without careful tuning of $c$ the gradients explode or vanish (Gulrajani et al., 2017, Sections 3.1–3.2). The two methods below replace clipping.

WGAN-GP: gradient penalty#

WGAN-GP (Gulrajani et al., 2017) keeps the WGAN losses and adds a penalty to the critic:

$$\begin{aligned} \mathcal L_D = {}& \mathbb{E}_{z\sim p_z}\big[f(G(z))\big] - \mathbb{E}_{x\sim P}\big[f(x)\big]\\ &+ \lambda\,\mathbb{E}_{\hat x}\Big[\big(\|\nabla_{\hat x} f(\hat x)\|_2 - 1\big)^2\Big]. \end{aligned}$$

The first two terms are the WGAN critic loss written in the other order, and the generator loss $\mathcal L_G$ is unchanged. The symbol $\mathcal L_D$ now includes the penalty, so $-\mathcal L_D$ is no longer the estimate of the lower bound on $K\,W_1$ from the previous section: it is that estimate minus $\lambda$ times the penalty. To monitor training, log the real-minus-fake score difference $\mathbb E_P f - \mathbb E_Q f$ and the penalty separately. Read the difference with the earlier caveats, and one more: a soft constraint keeps the critic only approximately 1-Lipschitz, so the difference is not guaranteed to be a lower bound on $W_1$.

Why the gradient norm. For a differentiable $f$ on $\mathbb R^d$, being 1-Lipschitz is the same as $\|\nabla f(x)\|_2 \le 1$ everywhere (one direction from the definition of the derivative, the other from the mean value theorem along a segment). Pushing $\|\nabla_{\hat x} f(\hat x)\|_2$ towards $1$ therefore keeps $f$ approximately 1-Lipschitz. The literal constraint would be a one-sided penalty, $\max(0, \|\nabla f\|_2 - 1)^2$. WGAN-GP uses the two-sided $(\|\nabla f\|_2 - 1)^2$, which does not exclude the optimal critic. On a compact space, if the optimal critic $f^*$ is differentiable and the optimal coupling moves every point ($\pi(x = y) = 0$), then $f^*$ has gradient norm $1$ along the segments between optimally coupled points, and almost everywhere under $P$ and $Q$ (Gulrajani et al., 2017, Proposition 1 and Corollary 1). In their experiments the two-sided penalty also worked slightly better (Section 4 and Appendix C). The constraint is soft: it is only checked at the sampled points $\hat x$.

Where $\hat x$ comes from. Pair each real sample with a fake one, draw $\alpha \sim U(0, 1)$ and take a random point on the segment between them:

$$\hat x = \alpha\,x_r + (1 - \alpha)\,x_f .$$

Proposition 1 concerns pairs drawn from the optimal coupling, which is unknown. WGAN-GP pairs real and fake samples at random instead. This is a heuristic motivated by the proposition: enforcing the unit norm everywhere is intractable, and enforcing it along these segments works well in practice (Gulrajani et al., 2017, Section 4). In code, $\alpha$ is drawn once per image and broadcast over all pixels. The gradient of the critic’s output with respect to $\hat x$ is computed with torch.autograd.grad(..., create_graph=True), so that the penalty itself can be backpropagated. It is then flattened per sample into a vector $g$, and its norm enters the penalty ${(\|g\|_2 - 1)^2}$. The paper uses $\lambda = 10$ in all experiments (Section 4).

Figure

The gradient penalty is evaluated between real and fake samples

No batch normalization in the critic. The penalty treats the critic as a function of one image: it constrains $\|\nabla_{\hat x} f(\hat x)\|_2$ for each $\hat x$ separately. With batch normalization the critic maps a whole minibatch of inputs to a minibatch of outputs, and a penalty on the gradient with respect to each input no longer describes the function that is being constrained. WGAN-GP therefore leaves batch normalization out of the critic and allows only normalizations that do not mix examples; the authors recommend layer normalization (Gulrajani et al., 2017, Section 4). Instance normalization, which normalizes each channel of each image over its spatial positions, also qualifies.

SN-GAN: spectral normalization#

The second approach constrains the weights instead of the gradients (Miyato et al., 2018). Miyato et al. trained their SN-GANs with the standard GAN loss and the hinge loss rather than the Wasserstein loss, with WGAN-GP as a baseline (Section 4, Eqs. 15–17). The normalization itself is a Lipschitz constraint and applies equally to a WGAN critic. Each layer’s weight matrix is divided by its spectral norm, the largest singular value:

$$\begin{gathered} \hat W = \frac{W}{\sigma_{\max}(W)},\\ \sigma_{\max}(W) = \max_{h \ne 0}\frac{\|W h\|_2}{\|h\|_2} . \end{gathered}$$

What a singular value is. A linear map stretches space. It sends the unit circle to an ellipse, and the singular values say how much it stretches in each direction: the right singular vectors $v_i$ are mapped to $\sigma_i u_i$, the semi-axes of the ellipse. The largest one, $\sigma_{\max}$, is the most any vector can be stretched, so it is the Lipschitz constant of $h \mapsto W h$. After dividing by it, no direction is stretched by more than $1$.

Figure

A matrix stretches the unit circle into an ellipse; its largest singular value is the longest semi-axis

For a network that alternates linear layers and 1-Lipschitz activations (ReLU, LeakyReLU), Lipschitz constants multiply under composition, so (Miyato et al., 2018, Eq. 7)

$$\|f\|_L \le \prod_{l} \sigma_{\max}(W_l),$$

which is at most $1$ after normalization.

Unlike the gradient penalty, the constraint is imposed on the weights rather than checked at sampled points, and there is no $\lambda$ to tune. In PyTorch it is a wrapper, torch.nn.utils.spectral_norm(nn.Conv2d(...)) (the newer API is torch.nn.utils.parametrizations.spectral_norm), which estimates $\sigma_{\max}$ with one step of power iteration per forward pass in training mode instead of a full SVD. In practice the bound is approximate, for three reasons. Power iteration only estimates $\sigma_{\max}$. For a convolution the wrapper normalizes the kernel reshaped into a matrix, which is a proxy for the convolution’s operator norm, not equal to it; the true norm also depends on stride and padding (Gouk et al., 2021, Section 3.2; Miyato et al., 2018, footnote 3). And normalization layers placed between the convolutions are not 1-Lipschitz in general. Spectral normalization therefore keeps the critic’s Lipschitz constant under control rather than guaranteeing that it is at most $1$. The product bound also need not be tight, so the normalized critic is $K$-Lipschitz for some $K$ that can be well below $1$; a WGAN critic with spectral normalization then estimates a lower bound on $K\,W_1$ rather than $W_1$ itself.

Two dual views#

Both f-GAN, from the GAN note, and WGAN train the discriminator as the function inside a supremum. f-GAN uses the convex-conjugate bound on an f-divergence, which compares density ratios and therefore inherits the support problem of $D_{\mathrm{JS}}$ and KL. WGAN uses the Kantorovich–Rubinstein dual of a transport cost, which compares locations and stays informative when supports are disjoint. The Lipschitz constraint is the price of the second view: it is what ties the critic to the geometry of the data space.

Notebook#

The WGAN notebook of the HSE-LAMBDA Deep Generative Models repository (pinned commit; cell numbers count every cell, markdown included, from 1) implements all three variants on CelebA at $64 \times 64$. All three training functions use Adam with learning rate $2\cdot10^{-4}$ and betas $(0.5, 0.9)$, and $n_{\text{critic}} = 5$ critic steps per generator step; the WGAN-GP paper uses learning rate $10^{-4}$ and betas $(0, 0.9)$ (Gulrajani et al., 2017, Algorithm 1). The basic and WGAN-GP runs are set to 2 epochs, the spectral-normalization run to 20.

• Cells 1–5: introduction, the earth mover’s distance with its Kantorovich–Rubinstein dual, and the WGAN objective with a real-valued critic.
• Cells 6–9: data. CelebA: 202,599 aligned and cropped face images of 178 × 218 pixels (width × height) of 10,177 people, each annotated with 40 binary attributes such as smiling or eyeglasses. The course’s dataset class returns the attributes, and this notebook ignores them. The images are resized to $64 \times 64$ and scaled to $[-1, 1]$, batch size 1024 with drop_last=True (197 batches per epoch), and cell 9 shows a grid of real images. Cell 7 selects the device (CUDA, Apple MPS or CPU). Resize((64, 64)) with a pair of sizes scales the two sides independently, so the aspect ratio changes from 178 : 218 to 1 : 1 and the faces become about 22% wider relative to their height; Resize(64) followed by CenterCrop(64) keeps the proportions.
• Cells 10–16: basic WGAN. Cell 12 defines the networks. The generator follows the DCGAN guidelines: $z \in \mathbb R^{100}$ passes through a bias-free linear layer to $512 \cdot 4 \cdot 4 = 8192$ units, BatchNorm1d and ReLU, is reshaped to $512 \times 4 \times 4$, and four transposed convolutions (kernel 5, stride 2, padding 2, output padding 1) take it through 256, 128 and 64 channels to $3 \times 64 \times 64$, with batch norm and ReLU after the first three and Tanh after the last. A kernel of 5 is not divisible by the stride 2, so these layers can produce checkerboard artifacts. The critic applies four stride-2 convolutions with kernel 5, $3 \to 64 \to 128 \to 256 \to 512$ channels, each followed by LeakyReLU(0.2), with InstanceNorm2d(affine=True) after the last three and no normalization at the input, then a $4 \times 4$ convolution to one score without a sigmoid. train_wgan (cell 14) uses d_loss = -mean(real_score) + mean(fake_score) and g_loss = -mean(fake_score). The baseline applies no Lipschitz constraint, neither clipping nor penalty (cell 11). Its critic objective is therefore unbounded (the supremum above is $+\infty$), its scores and D_loss can drift without limit, and the loss does not estimate $W_1$. Cell 13 notes that weight clipping can be added by clamping the critic’s weights after each update. Each training function draws 16 latent vectors once (fixed_noise, cells 14, 19 and 24) and decodes them after every epoch with the generator in evaluation mode, so the grids follow the same codes through training.
• Cells 17–21: WGAN-GP. gradient_penalty (cell 18) draws $\alpha \sim U(0,1)$ per sample, interpolates, differentiates with create_graph=True and penalizes the per-sample $L_2$ norm; lambda_gp = 10. The critic is the same network, and its InstanceNorm2d normalizes each sample separately, so the per-sample penalty remains valid (no batch normalization in the critic). In train_wgan_gp (cell 19) the penalty receives fake_imgs without detach(), so each critic step also backpropagates through the generator. Those gradients are cleared before the generator step, so the results are unaffected, but the computation is wasted.
• Cells 22–27: spectral normalization. DiscriminatorSN wraps every convolution in spectral_norm; in its inner blocks the normalized convolution is followed by InstanceNorm2d(affine=True) and LeakyReLU, so the bound above holds only approximately. Training uses the plain WGAN losses.
• Cell 28: summary.

Only cells 7, 9, 26 and 27 have saved outputs: the device (mps), the grid of real images, progress bars of the spectral-normalization run, which was interrupted during epoch 11 of 20, and samples from that partly trained generator. No loss values are saved. The basic and WGAN-GP training and sampling cells have no saved outputs, so for those variants the notebook documents procedures only.

Next#

The WGAN gave the generator a useful gradient, and spectral normalization kept the critic in check. The GANs so far still produce small images of one kind at a time: faces at 64×64, digits, bedrooms. The next note follows what it took to reach large, varied, high-resolution images: growing the networks progressively, self-attention for structure that spans the whole image, and BigGAN, which scales the same recipe to all of ImageNet.

References#

• M. Arjovsky, S. Chintala, L. Bottou. Wasserstein GAN. arXiv:1701.07875, 2017; published at ICML 2017 as “Wasserstein generative adversarial networks”. Section numbers refer to the arXiv version.
• M. Arjovsky, L. Bottou. Towards principled methods for training generative adversarial networks. ICLR, 2017.
• I. Gulrajani, F. Ahmed, M. Arjovsky, V. Dumoulin, A. Courville. Improved training of Wasserstein GANs. NeurIPS, 2017.
• T. Miyato, T. Kataoka, M. Koyama, Y. Yoshida. Spectral normalization for generative adversarial networks. ICLR, 2018.
• C. Villani. Optimal Transport: Old and New. Springer, 2009. Chapters 5–6: the Kantorovich duality and the Wasserstein distances.
• G. Peyré, M. Cuturi. Computational optimal transport. Foundations and Trends in Machine Learning, 11(5–6), 2019.
• H. Gouk, E. Frank, B. Pfahringer, M. J. Cree. Regularisation of neural networks by enforcing Lipschitz continuity. Machine Learning, 110, 2021.
• HSE-LAMBDA, Deep Generative Models course repository: WGAN.ipynb.