Optimal Transport and EMD in Machine Learning
Optimal transport is useful when where probability moves matters. A shift of 0.1 in a prediction score should usually count as smaller than a shift of 0.8. The following examples show how that idea becomes a transport plan, a monitoring statistic, and a training signal.
Throughout, both distributions have total mass 1. We use ordinary distance as the cost per unit of moved mass, so the minimum total cost is Wasserstein-1 distance, $W_1$, also called EMD here.
1. Match Weighted Samples with a Cost Matrix
Imagine comparing embeddings from two datasets. Real embeddings have many coordinates; use one coordinate here so we can see every distance. The source distribution has half its probability at 0 and half at 2. The target needs one quarter at 1 and three quarters at 3:
| Distribution | Locations | Probability at each |
|---|---|---|
| Source $P$ | 0, 2 | 0.50, 0.50 |
| Target $Q$ | 1, 3 | 0.25, 0.75 |
The cost matrix $C$ lists the distance for every possible source–target pair: $C_{ij}=|x_i-y_j|$. Rows correspond to source locations 0 and 2; columns correspond to target locations 1 and 3.
$$C=\begin{pmatrix}|0-1|&|0-3|\\|2-1|&|2-3|\end{pmatrix} =\begin{pmatrix}1&3\\1&1\end{pmatrix}.$$We choose a transport plan $T$, where $T_{ij}$ is the amount of probability sent along that pair. Each row must sum to its source probability; each column must sum to its target probability. One optimal plan is:
$$T=\begin{pmatrix}0.25&0.25\\0&0.50\end{pmatrix}.$$The first row sends the source mass at 0 to two destinations. The second row sends all the mass at 2 to 3. The row sums are $(0.50,0.50)$ and the column sums are $(0.25,0.75)$, exactly as required.
| Move | Mass | Distance | Cost |
|---|---|---|---|
| 0 → 1 | 0.25 | 1 | 0.25 |
| 0 → 3 | 0.25 | 3 | 0.75 |
| 2 → 3 | 0.50 | 1 | 0.50 |
Why is this cheapest? Moving from 2 costs 1 either way. Moving from 0 costs 1 to reach target 1 but 3 to reach target 3. We therefore reserve all of target 1's capacity, 0.25, for mass from 0. The remaining 0.25 at 0 must take the longer trip.
For real embeddings, replace absolute distance with a meaningful distance between vectors, such as Euclidean distance. Feature scales affect the answer. A different cost, such as $1-\text{cosine similarity}$, can also define an OT problem, but does not automatically give the metric properties of $W_1$. The POT guide uses the same cost-matrix and mass-constraint formulation.
2. Detect a Drift in Prediction Scores
A model-monitoring dashboard compares three baseline scores, $[0.2,0.4,0.6]$, with today's scores, $[0.3,0.5,0.7]$. Give each sample probability $1/3$.
On a line, equally weighted samples can be matched in sorted order. Here the movements are $0.2\to0.3$, $0.4\to0.5$, and $0.6\to0.7$: each travels 0.1. Thus:
$$W_1=\frac{0.1+0.1+0.1}{3}=0.1.$$Interpretation: the cheapest matching moves probability by an average of 0.1 score units, or 10 percentage points. This measures a change in the score distribution; it does not by itself establish that prediction accuracy got worse. Sorted matching is the one-dimensional transport solution; see the 1D Wasserstein post for its cumulative-distribution form.
It detects more than a change in the mean. Compare $[0,1]$ with $[0.5,0.5]$. Both means are 0.5, but both samples must move by 0.5, so $W_1=0.5$. The distribution became more concentrated even though its average stayed the same.
3. Give a GAN a Signal When Distributions Are Separated
Consider a deliberately simple generator: every real sample is 0, and every generated sample is $\theta$, a location controlled by the generator. The whole unit of probability must travel from $\theta$ to 0, so:
$$W_1(P_{\mathrm{real}},P_\theta)=1\times|\theta-0|=|\theta|.$$| Generated location $\theta$ | $W_1$ | JS (bits) |
|---|---|---|
| 10 | 10 | 1 |
| 1 | 1 | 1 |
| 0.1 | 0.1 | 1 |
| 0 | 0 | 0 |
For every nonzero $\theta$, the distributions have no shared possible outcome, so equal-weight JS stays at its maximum of 1 bit. It cannot distinguish “almost there” from “very far away.” Wasserstein distance decreases as the generated point approaches 0. For $\theta>0$, its derivative is $+1$, so gradient descent decreases $\theta$; for $\theta<0$, the derivative is $-1$, so it increases $\theta$. Both updates move toward the real data at 0.
A WGAN uses a learned critic, a network that scores samples, to approximate a different formula for this transport distance. The critic is constrained to be 1-Lipschitz: its score difference between two inputs cannot exceed their distance. The generator then learns to raise the critic's scores on generated samples. The resulting signal depends on how well the critic and its constraint are learned; an arbitrary critic loss is not an exact EMD measurement. This is the motivation developed in the Wasserstein GAN paper.
When the Transport Problem Gets Large
The first example has only four source–target costs. With $n$ source samples and $m$ target samples, a dense cost matrix has $nm$ entries, and the plan must satisfy all the mass constraints together. This can be substantially more work than evaluating KL or JS on existing histogram bins.
Entropic regularization modifies the objective to favor more spread-out plans. Sinkhorn iterations solve this modified problem efficiently in many settings. The regularization strength controls how much spreading is encouraged; the result should not be treated as exact, unregularized $W_1$. The POT guide's regularized OT section explains this tradeoff.