📋 KEY INSIGHTS
- Optimal transport (OT) is a mathematical framework for comparing probability distributions by computing the minimum cost required to transport mass from one distribution to another. The Wasserstein distance W_p(P, Q) = (inf_{γ ∈ Γ(P,Q)} ∫ c(x,y)^p dγ(x,y))^{1/p} measures the distance between two distributions P and Q as the minimum expected transportation cost under the optimal coupling γ, where Γ(P,Q) is the set of joint distributions with marginals P and Q. Unlike KL divergence or Jensen-Shannon divergence, Wasserstein distance is a true metric, is defined even when distributions have disjoint supports, and metrises weak convergence — making it better-behaved for optimisation.
- The Sinkhorn algorithm (Cuturi, 2013) makes OT tractable at scale by adding an entropic regularisation term εH(γ) to the OT objective, converting the linear program into a strongly convex problem solvable by iterative matrix scaling. Regularised OT: W_ε(P, Q) = min_{γ ∈ Γ(P,Q)} [∫ c(x,y) dγ − ε H(γ)]. The Sinkhorn algorithm alternately normalises rows and columns of the cost matrix, converging in O(n² log n / ε) time — far cheaper than the O(n³) exact OT solvers. Smaller ε → closer to exact OT but slower convergence; larger ε → smoother, faster but more blurred transport plan.
- Wasserstein GANs (WGAN, Arjovsky et al., 2017) replaced the original GAN discriminator (which estimates JS divergence) with a critic that estimates the Earth Mover distance (Wasserstein-1 distance) between the real and generated distributions. The key advantage: Wasserstein distance provides meaningful gradients even when the generated distribution and real distribution have disjoint supports — the original GAN loss vanishes or becomes unstable in this case. WGAN training requires the critic to be 1-Lipschitz (originally enforced via weight clipping, later via gradient penalty in WGAN-GP), and uses the Kantorovich-Rubinstein duality: W_1(P_r, P_g) = sup_{||f||_L ≤ 1} [E_{x~P_r}[f(x)] − E_{x~P_g}[f(x)]].
- Domain adaptation with OT aligns feature distributions across source and target domains by learning a transport map T: X_S → X_T that minimises the Wasserstein distance between the transported source distribution T#P_S and the target distribution P_T. Discrete OT (computing the exact transport plan between finite sample sets) produces sample-level alignments used for dataset-level domain adaptation; continuous OT (learning a neural transport map, e.g. via input-convex neural networks) generalises to unseen source points. The connection to our transfer learning guide is direct: OT-based domain adaptation is a mathematically principled alternative to feature-space alignment methods like domain adversarial training.
- Sliced Wasserstein distance (SWD) makes high-dimensional OT tractable by projecting distributions onto random 1D lines and computing the average Wasserstein distance over projections: SW_p^p(P, Q) = E_{θ~S^{d-1}}[W_p^p(P_θ, Q_θ)], where P_θ and Q_θ are the 1D projections of P and Q onto direction θ. 1D Wasserstein distance has a closed-form solution (the L_p distance between sorted samples), making SWD O(n log n) in the number of samples — drastically cheaper than the O(n³) exact OT. SWD is used as a training loss for generative models (Sliced Wasserstein Autoencoders), point cloud processing, and texture synthesis.
- OT for data science practitioners is accessible via the Python Optimal Transport (POT) library:
ot.emd(a, b, M)computes exact OT,ot.sinkhorn(a, b, M, reg)computes regularised OT, andot.sliced_wasserstein_distance(X, Y)computes sliced Wasserstein distance. The cost matrix M is typically the squared Euclidean distance between feature vectors. Applications include: dataset comparison (is my training set representative of production data?), distribution shift monitoring (Wasserstein distance between reference and current feature distributions as a drift metric in our ML monitoring guide), colour transfer between images, and document comparison via word mover distance.
Optimal transport has emerged as one of the most practically impactful mathematical frameworks in modern machine learning, moving from theoretical geometry to deployed production systems in the span of a decade. The Wasserstein distance provides a geometrically meaningful way to compare distributions that respects the underlying space structure — two distributions of pixel intensities that are slightly shifted look very different under KL divergence but very similar under Wasserstein distance, correctly capturing perceptual similarity. This guide covers the full OT toolkit for machine learning practitioners: the Kantorovich formulation and Wasserstein distances, the Sinkhorn algorithm for scalable computation, Wasserstein GANs, optimal transport for domain adaptation and distribution alignment, sliced Wasserstein distance, word mover distance for text, and practical applications in distribution shift monitoring, colour transfer, and dataset comparison.
Wasserstein Distance and the Sinkhorn Algorithm
The Kantorovich formulation of optimal transport defines the Wasserstein-p distance between two probability distributions P (source) and Q (target) as the minimum expected cost of a transport plan: W_p^p(P, Q) = inf_{γ ∈ Γ(P,Q)} E_{(x,y)~γ}[c(x,y)^p], where c(x,y) is the cost of moving a unit of mass from x to y (typically the squared Euclidean distance for W_2 or the L1 distance for W_1) and Γ(P,Q) is the set of all joint distributions (couplings) whose marginals are P and Q. For discrete distributions P = Σ_i a_i δ_{x_i} and Q = Σ_j b_j δ_{y_j}, this becomes the linear program: min_{T ≥ 0} Σ_{ij} T_{ij} C_{ij} subject to T1 = a and T^T 1 = b, where C_{ij} = c(x_i, y_j) is the cost matrix and T is the transport matrix. This linear program has O(n³) complexity via standard solvers — prohibitive for large datasets. The Sinkhorn algorithm solves the entropically regularised OT problem W_ε(P, Q) = min_{γ ∈ Γ(P,Q)} [Σ_{ij} γ_{ij} C_{ij} + ε Σ_{ij} γ_{ij} log γ_{ij}] by iterative Bregman projections: initialise K = exp(−C/ε), then alternate u ← a / (Kv) and v ← b / (K^T u) until convergence, where the transport plan is T = diag(u) K diag(v). This GPU-parallelisable iteration converges in O(n²/ε) multiplications — dramatically faster than LP solvers for typical ε values. The connection to our linear algebra guide is direct: Sinkhorn iteration is a matrix scaling algorithm operating on the Gibbs kernel K = exp(−C/ε), and understanding its convergence properties requires familiarity with doubly stochastic matrices and Birkhoff polytopes.
| OT Method | Complexity | GPU? | High-Dim? | Primary Use |
|---|---|---|---|---|
| Exact OT (LP solver) | O(n³ log n) | No | Poor | Small n (<5000); ground truth |
| Sinkhorn (entropic OT) | O(n² / ε) | Yes | Moderate | Batch ML; GAN training; domain adaptation |
| Sliced Wasserstein | O(n log n × K projections) | Yes | Good | High-dim generative models; point clouds |
| Mini-batch OT | O(m² / ε) per mini-batch | Yes | Good | Stochastic GAN training; large datasets |
| Word Mover Distance | O(n³ log n) | No | Via embeddings | Document similarity; NLP |
Wasserstein GANs, Domain Adaptation and Distribution Shift Monitoring
Wasserstein GANs (WGAN, Arjovsky et al., 2017) replaced the original GAN value function V(D, G) = E[log D(x)] + E[log(1−D(G(z)))] with the Wasserstein-1 distance estimated via the Kantorovich-Rubinstein duality: W_1(P_r, P_g) = sup_{||f||_L ≤ 1} [E_{x~P_r}[f(x)] − E_{x~P_g}[f(x)]], training the critic f to maximise the difference in expectation between real and generated samples subject to 1-Lipschitz constraint. This formulation provides meaningful gradients even when P_r and P_g have disjoint support — the original JS-divergence-based GAN loss saturates to zero in this case, causing vanishing gradients and mode collapse. WGAN-GP (gradient penalty) enforces the Lipschitz constraint by adding λ E_{x~P_x}[(||∇_x f(x)||_2 − 1)²] to the critic loss, interpolating between real and generated samples. The connection to our diffusion models guide is direct: diffusion models have largely superseded GANs for image generation, but Wasserstein distance and score matching (the diffusion model training objective) are both distributional divergences for generative modelling, and OT provides theoretical grounding for understanding the geometry of learned generative model distributions. OT for domain adaptation computes the transport plan between labelled source samples and unlabelled target samples, then uses the transport plan to either reweight source samples (importance weighting) or directly map source features to the target distribution (feature-level alignment). The POT library implements both: ot.da.SinkhornTransport() for entropic transport-based domain adaptation and ot.da.EMDTransport() for exact transport. The connection to our ML monitoring guide is direct: Wasserstein distance between a reference feature distribution (from training or a recent stable window) and the current production distribution is the most geometrically meaningful drift detection metric — better than population stability index (PSI) or KS test for high-dimensional, correlated feature spaces. The nannyml and evidently libraries implement Wasserstein-based drift detection for production ML systems.
Word Mover Distance, Colour Transfer and Practical OT
Word Mover Distance (WMD, Kusner et al., 2015) uses OT to measure the semantic dissimilarity between two text documents by computing the minimum cumulative distance that embedded words from one document need to travel to match the words of the other. Each document is represented as a normalised bag of word embeddings (Word2Vec or GloVe from our word embeddings guide); the cost of transporting word i to word j is their Euclidean distance in embedding space; and WMD is the resulting Wasserstein-1 distance. WMD is particularly effective for short texts where TF-IDF fails (no word overlap between semantically similar phrases) and outperforms BOW baselines on k-NN text classification. Colour transfer between images is a classic OT application: given a source image and a target image, compute the OT transport plan between the source and target colour distributions (3D histograms in Lab colour space) and apply the transport map to transform source pixels to match the target colour palette — producing colour-graded images without manually adjusting colour curves. This application motivates the sliced Wasserstein approach: exact 3D OT is expensive, but computing 1D Wasserstein distances along multiple random projections of the colour space and averaging (sliced Wasserstein) achieves nearly identical visual results at a fraction of the cost. The connection to our synthetic data guide is direct: OT is used for distribution-matching-based synthetic data generation — the goal of synthetic data is to produce a dataset whose distribution matches the real data distribution, and OT provides both the objective function (minimise Wasserstein distance between synthetic and real distributions) and the generation mechanism (apply the learned transport map to noise samples).
✦ SUMMARIZE THIS ARTICLE WITH AI
ChatGPTDiffusion models (which use score matching as an alternative distributional divergence to Wasserstein) are in our Diffusion Models guide. ML system monitoring (Wasserstein drift detection) is in our ML Monitoring guide. Transfer learning and domain adaptation (feature-level alignment alternatives to OT) are in our Transfer Learning guide. Word embeddings (Word2Vec/GloVe used in Word Mover Distance) are in our Word Embeddings guide. Synthetic data generation (OT-based distribution matching) is in our Synthetic Data guide. Linear algebra (Sinkhorn matrix scaling) is in our Linear Algebra guide.
