Wasserstein Distance

Piles of dirt

Imagine two distributions as piles of dirt on a number line. One has 100 kg piled at position 0; the other has 100 kg piled at position 10. To turn the first pile into the second, you must move 100 kg through a distance of 10 — total work 1000 kg-meters.

More generally: μ is a probability distribution, ν is another. The Wasserstein distance W_p(μ, ν) is the minimum work to transport mass from μ to ν, where work for a mass moving distance d is d^p per unit mass. The infimum is taken over all transport plans:

W_p(μ, ν)^p = inf_π ∫ |x - y|^p dπ(x, y)

inf over joint distributions π with marginals μ and ν

A coupling π is any joint distribution on (x, y) whose marginals are μ and ν. Picking the optimal π is the optimal transport problem.

For 1D distributions, the optimal coupling is the quantile coupling: pair the u-th quantile of μ with the u-th quantile of ν. The Wasserstein distance reduces to integrating |F⁻¹_μ(u) − F⁻¹_ν(u)|^p over u ∈ [0, 1]. This makes 1D Wasserstein extremely fast — sort the samples, then compute.

Wasserstein vs KL — why it matters

The KL divergence is the standard way to compare distributions in statistics. But KL has three problems:

Property	KL divergence	Wasserstein
Symmetric	No — KL(μ‖ν) ≠ KL(ν‖μ)	Yes — true metric
Disjoint supports	Infinity	Finite (depends on geometry)
Triangle inequality	No	Yes
Sensitive to geometry	No — ignores distances	Yes — distance is the cost
Closed form for Gaussians	Yes	Yes (W₂)
Computable	Closed form when densities exist	Linear program / Sinkhorn

The third row makes Wasserstein indispensable for generative models. Early in GAN training, the generator distribution P_g and target P_data have nearly disjoint supports — KL is infinite, gives no useful gradient. Wasserstein is finite and decreases as P_g moves closer to P_data, providing a meaningful loss signal throughout training.

Kantorovich-Rubinstein duality

The infimum-over-couplings definition is hard to compute directly. The W₁ case has an equivalent dual representation:

W_1(μ, ν) = sup_{f: 1-Lipschitz} |∫ f dμ - ∫ f dν|

A function f is 1-Lipschitz if |f(x) − f(y)| ≤ |x − y| for all x, y. The supremum is over all such f.

This dual is striking: instead of computing optimal transport plans, you compute the worst-case 1-Lipschitz function distinguishing μ from ν. The maximizer f* is the "Kantorovich potential" — geometrically, it's the function pointing in the direction of mass flow.

WGAN — Wasserstein in deep learning

Wasserstein GAN (Arjovsky, Chintala, Bottou, 2017) replaced the JS divergence of vanilla GANs with W₁. The critic network f (replacing the discriminator) is trained to maximize:

L_critic = E_(x~P_data)[f(x)] - E_(x~P_g)[f(x)]    subject to f being 1-Lipschitz

This is exactly the Kantorovich-Rubinstein dual, and the optimum equals W₁(P_data, P_g). The generator G then minimizes the same quantity — pulling P_g toward P_data along the W₁ gradient.

Enforcing the 1-Lipschitz constraint requires care:

• Weight clipping (original WGAN). Clip each weight to [−c, c]. Crude but effective; can cause vanishing gradients.
• Gradient penalty (WGAN-GP, Gulrajani et al.). Add λ·E[(‖∇f(x̂)‖ − 1)²] to the loss, with x̂ a random interpolation between real and fake samples. Cleaner and more stable.
• Spectral normalization (Miyato et al.). Normalize each weight matrix by its largest singular value. Strict bound, computationally cheap.

WGAN's training stability over vanilla GAN was a major step in generative modeling. It directly motivated diffusion models, normalizing flows, and continuous-time generative models — all of which lean on Wasserstein geometry.

Worked example — two discrete distributions

μ has mass at x = 0, 1, 2 with weights (1/3, 1/3, 1/3). ν has mass at y = 5, 6, 7 with weights (1/3, 1/3, 1/3). Find W₁.

By the 1D quantile coupling, pair x = 0 with y = 5, x = 1 with y = 6, x = 2 with y = 7. Each transfers 1/3 of mass through distance 5:

W_1 = (1/3)·5 + (1/3)·5 + (1/3)·5 = 5

The two distributions have identical shapes shifted by 5, so the Wasserstein-1 distance is exactly the shift. Makes intuitive sense: minimum effort to morph one into the other is to translate every grain of dirt 5 units.

For W₂ same example:

W_2^2 = (1/3)·25 + (1/3)·25 + (1/3)·25 = 25
W_2 = 5

For pure translations W₁ = W₂. They diverge when transport involves spreading or compressing mass.

Gaussian closed form (Bures-Wasserstein)

For Gaussian distributions N(μ₁, Σ₁) and N(μ₂, Σ₂):

W_2(N(μ_1, Σ_1), N(μ_2, Σ_2))^2 = ‖μ_1 - μ_2‖^2 + tr(Σ_1 + Σ_2 - 2(Σ_1^(1/2) Σ_2 Σ_1^(1/2))^(1/2))

This is the Bures-Wasserstein formula. It combines the mean displacement with a covariance term that's symmetric in Σ₁ and Σ₂.

Special cases:

• Σ₁ = Σ₂ = Σ: W₂² = ‖μ₁ − μ₂‖² (pure translation).
• μ₁ = μ₂ = 0, diagonal Σᵢ: W₂² = Σⱼ (σ_(1j) − σ_(2j))² (sum of standard deviation differences).
• 1D Gaussians: W₂² = (μ₁ − μ₂)² + (σ₁ − σ₂)².

This closed form is used in covariance matrix interpolation, generative modeling with Gaussian priors, and shape analysis.

Sinkhorn — fast computation

For discrete distributions with n and m points, exact W_p is a linear program — solvable but slow for large n, m. Cuturi (2013) introduced the Sinkhorn algorithm, which adds an entropy regularizer to the transport plan:

min_π   <π, C>  -  (1/λ) H(π)        subject to marginals

This entropic problem has a unique solution computable by matrix scaling (Sinkhorn-Knopp iterations). Each iteration is O(nm) instead of O((n + m)³) for exact transport. The Sinkhorn distance S_λ converges to W_p as λ → ∞. With λ moderate, Sinkhorn is differentiable (used for end-to-end training of neural networks with transport loss).

For 2D image distributions with n = m = 10⁴ pixels, exact transport takes hours; Sinkhorn finishes in seconds. This is what made optimal transport practical at scale.

Where Wasserstein appears

• Wasserstein GAN. Stable generative model training. Foundation for image synthesis and style transfer.
• Diffusion models. Score-based generative models can be viewed as W₂ gradient flow toward the data distribution.
• Domain adaptation. Align distributions across source and target domains by minimizing Wasserstein distance.
• Word Mover's Distance. Document similarity using W₁ on bags of word embeddings (Kusner et al. 2015).
• Distributional reinforcement learning. Algorithms like QR-DQN (Quantile Regression DQN) and IQN use Wasserstein-aware value-distribution updates.
• Color transfer / image processing. Match color distributions between images via optimal transport.
• Shape analysis. Compare anatomical shapes by treating them as point cloud distributions; W₂ gives Bures-Wasserstein-style mean shapes.
• Climate science. Compare modeled vs observed weather distributions, satellite versus reanalysis fields.
• PDE analysis. Heat equation, Fokker-Planck, porous medium equation are all gradient flows in W₂ (Jordan-Kinderlehrer-Otto 1998).
• Topological data analysis. Comparing persistence diagrams uses a Wasserstein-like distance.

Python — discrete optimal transport

import numpy as np
from scipy.stats import wasserstein_distance
import ot   # Python Optimal Transport library

# 1D Wasserstein-1 via scipy (closed-form via sorted samples)
mu = np.random.normal(0, 1, 1000)
nu = np.random.normal(3, 1, 1000)
w1 = wasserstein_distance(mu, nu)
print(f'W_1 (scipy 1D) = {w1:.3f}')   # ≈ 3.0 (pure shift)

# Multi-dim: linear program via POT
# Two distributions in 2D
xs = np.random.randn(100, 2)
xt = np.random.randn(100, 2) + np.array([5.0, 0.0])
a = np.ones(100) / 100   # uniform weights
b = np.ones(100) / 100
M = ot.dist(xs, xt, metric='euclidean')   # 100×100 cost matrix
T = ot.emd(a, b, M)   # optimal transport plan
w1_exact = (T * M).sum()
print(f'W_1 (POT, exact) = {w1_exact:.3f}')

# Sinkhorn — much faster for large n
T_sinkhorn = ot.sinkhorn(a, b, M, reg=0.1)
w1_sinkhorn = (T_sinkhorn * M).sum()
print(f'Sinkhorn approximation = {w1_sinkhorn:.3f}')

# Gaussian closed form for W_2
def w2_gaussian(mu1, Sigma1, mu2, Sigma2):
    S1_sqrt = np.linalg.cholesky(Sigma1) if Sigma1.ndim > 0 else np.sqrt(Sigma1)
    inner = S1_sqrt @ Sigma2 @ S1_sqrt
    inner_sqrt = np.linalg.cholesky(inner + 1e-12 * np.eye(len(inner)))
    cov_term = np.trace(Sigma1 + Sigma2 - 2 * inner_sqrt)
    mean_term = np.sum((mu1 - mu2) ** 2)
    return np.sqrt(mean_term + cov_term)

w2 = w2_gaussian(np.zeros(2), np.eye(2), np.array([5., 0.]), np.eye(2))
print(f'W_2 (Gaussian, pure shift) = {w2:.3f}')   # ≈ 5.0

Common pitfalls

• Confusing W_p with KL. W_p is symmetric and finite; KL is asymmetric and can be infinite. They measure different things — W_p is geometric, KL is informational.
• Choosing the wrong cost function. Euclidean for spatial data, embedding-distance for text, Manhattan for histograms with grid structure. The cost shape changes the optimal transport entirely.
• Forgetting normalization. Wasserstein assumes both inputs are probability measures (total mass = 1). For unnormalized signals, use unbalanced optimal transport or normalize first.
• Sinkhorn regularization bias. Entropic Sinkhorn underestimates true Wasserstein for moderate λ. The difference vanishes as λ → ∞ but at λ = 0.1 it can be substantial.
• Curse of dimensionality. Empirical W_p between n samples in d dimensions has error O(n^(−1/d)) — slow convergence in high d. Estimators like sliced Wasserstein and entropic Sinkhorn mitigate this.
• Lipschitz violations in WGAN. Forgetting to enforce 1-Lipschitz on the critic (or using wrong λ in gradient penalty) breaks the dual formulation. The reported "Wasserstein loss" is no longer a true Wasserstein distance.

History

Gaspard Monge posed the original optimal transport problem in 1781 — "How to move a pile of earth to fill a hole using least work?" He couldn't solve it generally. Leonid Kantorovich relaxed the problem from deterministic maps to probability couplings (1942), proving existence and duality — this earned him the 1975 Nobel in economics for the closely related work on resource allocation. Leonid Vasershtein (1969) studied the metric properties of what later became known as the Wasserstein distance (the spelling "Wasserstein" is the German transliteration adopted in the Western literature). Cédric Villani's 2003 and 2009 books became the canonical references; Villani won the 2010 Fields Medal partly for this work. The deep learning revolution came with Wasserstein GAN (Arjovsky et al. 2017), which showed Wasserstein-1 was the right loss for generative modeling. Cuturi's Sinkhorn algorithm (2013) made large-scale computation practical. Marco Cuturi and Gabriel Peyré's 2019 book "Computational Optimal Transport" is the modern reference for ML practitioners.