Optimal Transport and Wasserstein Distances
Why KL divergence is useless between distributions that do not overlap, how the cost of moving mass gives a metric that is not, and what entropic regularisation trades away to make it computable.
Put one distribution's mass at \(x = 0\) and another's at \(x = \theta\), both as point masses. KL divergence between them is infinite for every \(\theta \ne 0\) and zero at \(\theta = 0\): no gradient, no notion that \(\theta = 0.01\) is nearly right. Every divergence built on likelihood ratios has this defect, and it is not an edge case. Distributions on a low-dimensional manifold inside a high-dimensional space, which is what images and activations are, generically fail to overlap.
Optimal transport measures the cost of moving mass instead of comparing densities pointwise, so it gives \(\lvert \theta \rvert\), which is exactly the answer wanted.
From Monge to Kantorovich
Monge's original problem asks for a map \(T\) pushing \(\mu\) onto \(\nu\) minimising \(\int c(x, T(x))\, \mathrm{d}\mu(x)\). It can be infeasible, since a deterministic map cannot split an atom of mass. Kantorovich's relaxation optimises over couplings instead, joint distributions \(\gamma\) with the right marginals:
This is a linear program: linear objective, linear marginal constraints, always feasible because the product coupling exists. That structure is why duality is available, and for \(p = 1\) the Kantorovich-Rubinstein dual is the form most people meet first,
a supremum over 1-Lipschitz witness functions. Wasserstein GANs are this dual with \(f\) parameterised by a network and the Lipschitz constraint enforced approximately, which is the entire reason the critic's output is a score rather than a probability.
The cost, and the standard escape
Solving the exact linear program for \(n\) support points costs roughly \(O(n^3 \log n)\), which rules it out inside a training loop. Entropic regularisation adds \(\varepsilon H(\gamma)\) to the objective, and the optimal coupling then takes the form \(\mathrm{diag}(u)\, K\, \mathrm{diag}(v)\) with \(K = e^{-C/\varepsilon}\), so the problem reduces to alternately rescaling rows and columns. That is the Sinkhorn-Knopp algorithm: cheap, differentiable, GPU-friendly, matrix-multiplication-shaped, and several orders of magnitude faster than transportation solvers (Cuturi, NeurIPS 2013, arXiv:1306.0895).
The connection to generative modelling runs deeper than the GAN dual. Displacement interpolation, the geodesic between two distributions under \(W_2\), moves mass along straight lines at constant speed. That is precisely the conditional path flow matching regresses against, and it is why flow-matching papers describe their probability paths as optimal-transport paths rather than as a heuristic (Lipman et al., ICLR 2023, arXiv:2210.02747).
When it breaks
Sample complexity is exponential in dimension. Estimating \(W_p\) between distributions from \(n\) samples converges at roughly \(n^{-1/d}\) for dimension \(d\). In hundreds of dimensions the empirical Wasserstein distance is dominated by sampling error, so a reported "Wasserstein distance" between two sets of high-dimensional embeddings is usually measuring sample size. Sliced variants, which average one-dimensional projections, exist because of this.
Entropic regularisation biases. The Sinkhorn objective is not zero when the two distributions are identical, since the entropy term prefers a blurred coupling. Sinkhorn divergences subtract the self-transport terms to correct this, and code that omits the correction has a systematic offset that grows with \(\varepsilon\).
The regulariser is a numerical dial too. Small \(\varepsilon\) approaches true OT and makes \(K = e^{-C/\varepsilon}\) underflow; large \(\varepsilon\) is stable and answers a different question. Log-domain Sinkhorn implementations exist entirely to push the usable range of \(\varepsilon\) lower.
The Lipschitz constraint is rarely enforced. In practice weight clipping and gradient penalties approximate the constraint in the WGAN dual, so the quantity being optimised is not \(W_1\). It is a useful objective; calling it a Wasserstein distance is a convention rather than a fact.
14 flashcards for this concept
Click a card to reveal the answer.