Wasserstein Critics and the Lipschitz Constraint

Two distributions on a line: one a point mass at 0, the other a point mass at \(\theta\). Jensen-Shannon divergence between them is \(\log 2\) for every \(\theta \neq 0\) and 0 at \(\theta = 0\). It is a step function, so its gradient with respect to \(\theta\) is zero everywhere useful and undefined at the one place it changes. Wasserstein distance between the same two distributions is \(|\theta|\), which has a gradient pointing in the right direction from anywhere. That example, from the WGAN paper (Arjovsky et al., 2017, arXiv:1701.07875), is the whole argument.

Why it matters for images

The situation in the toy example is not exotic; it is the normal case for image GANs. Natural images lie on a low-dimensional manifold in pixel space, and a generator's output lies on another low-dimensional manifold. Two such manifolds in a high-dimensional space almost surely have measure-zero intersection, so the supports are effectively disjoint, JS divergence is pinned at its maximum, and its gradient carries no directional information. The discriminator can achieve perfect separation and hand the generator nothing.

Wasserstein distance, the cost of the cheapest plan for transporting mass from one distribution to the other, degrades gracefully. It measures how far the mass has to move, which is defined and informative regardless of overlap.

The dual form, and why the critic must be Lipschitz

The primal Wasserstein distance is an optimisation over transport plans and is intractable directly. Kantorovich-Rubinstein duality converts it into

The supremum is over all 1-Lipschitz functions \(f\). Parameterise \(f\) as a network, maximise the difference of expectations, and you have a critic. It is called a critic rather than a discriminator because it does not classify: it outputs an unbounded real score whose difference between the two distributions estimates the distance. There is no sigmoid and no probability.

The 1-Lipschitz constraint is not a regularisation nicety. Without it the supremum is unbounded and the objective is meaningless: any critic can scale its output to infinity. Everything difficult about WGANs is the enforcement of this constraint.

Three ways to enforce it, none clean

Weight clipping, the original proposal, clamps every weight into \([-c, c]\). It works and the paper says plainly that it is a terrible way to do it. Small \(c\) causes vanishing gradients through a deep critic; large \(c\) takes a long time to reach the constraint. Clipped critics also tend to learn very simple functions, wasting capacity.

Gradient penalty (WGAN-GP) adds a penalty on \((\lVert \nabla_{\hat{x}} f(\hat{x}) \rVert_2 - 1)^2\) evaluated at points \(\hat{x}\) interpolated between real and generated samples (Gulrajani et al., 2017, arXiv:1704.00028). It uses the fact that a differentiable function is 1-Lipschitz exactly when its gradient norm is at most 1 everywhere, and it enforces this on a sampled subset of the space rather than everywhere. It is the most widely used variant. It costs a double backward pass per step, and it is incompatible with batch normalisation in the critic, since the penalty is defined per example and batch norm makes the critic's output depend on the rest of the batch.

Spectral normalisation divides each weight matrix by its largest singular value, estimated with one power-iteration step per forward pass. This bounds the Lipschitz constant of each layer, and therefore of the composition, by construction rather than by penalty. It is cheap, needs no extra loss term, and it is conservative: the product of per-layer bounds is an over-estimate of the network's true Lipschitz constant, so the critic is more constrained than necessary.

When it breaks

The reported loss is not the Wasserstein distance. It is a lower-bound estimate produced by a critic that is neither fully optimised nor exactly 1-Lipschitz. The WGAN literature's claim that the loss correlates with sample quality is a useful empirical observation, not a guarantee, and the correlation is much weaker across models than within a single training run.

Critic iterations are a real cost. The duality requires the critic to be near its supremum for the generator's gradient to be meaningful, which is why WGAN recipes specify five critic steps per generator step. That is roughly a five-fold increase in discriminator compute compared with a standard GAN.

Gradient penalty samples a line, not the space. The penalty is enforced only along interpolations between real and fake samples, on the assumption that this is where the constraint matters. It is a heuristic, and the constraint can be violated badly elsewhere without the penalty noticing.

The advantage narrowed. With spectral normalisation, hinge losses, and modern regularisation, standard non-saturating GANs became competitive with WGANs on image benchmarks. The durable contribution was not the specific loss but the diagnosis: adversarial training fails when the divergence has no gradient under disjoint support, and constraining the discriminator's Lipschitz constant fixes more than it was originally introduced to fix.