Optimal Transport — The Mathematics of Moving Distributions

Optimal transport measures the distance between two probability distributions as the cheapest way to haul sand from one pile into the shape of another. From the definition of the Wasserstein distance to the Sinkhorn algorithm that makes it practical, and on to WGAN, FID and Flow Matching — with numpy code and the failure modes that bite in production.

Sinkhorn Distances: Lightspeed Computation of Optimal Transport

Where an ordinary distance quietly fails

Most of machine learning boils down to pulling a model's distribution toward the data's distribution. To pull, you first need a single number saying how far apart they currently are. The usual choice is KL divergence — and KL has one situation it handles badly: when the two distributions do not overlap at all.

Say the real data is a thin spike at $x=0$ and the model is a thin spike at $x=\theta$. Whether $\theta$ is $0.01$ or $100$, KL is infinite on the sole grounds that the overlap is zero (Jensen–Shannon would sit flat at $\log 2$). So moving $\theta$ changes nothing — the gradient is zero, even though any human would say $\theta=0.01$ is the closer one. That intuition uses where the spike stands, while KL only compares probabilities vertically, point by point, never sideways along the axis.

Optimal transport fills exactly that gap. It defines distance as the minimum cost of hauling one distribution until it takes the shape of the other. Because you are hauling, horizontal distance ends up inside the number.

The metaphor: a moving company's quote

You run an earthworks outfit. Sand is piled at several depots, and several construction sites each need a fixed amount. The bill is "amount moved × distance moved." There are countless ways to finish the job, but exactly one cheapest total, and that cheapest total is the "distance" between where the sand is and where it is wanted.

This is not only a metaphor but the literal origin. In 1781 Gaspard Monge posed it as the problem of cuttings and embankments — moving earth.

Intuition: you choose a transport table, not a matching

Monge's version had a flaw. It looks for a one-to-one assignment — "all the sand at depot $i$ goes to site $\sigma(i)$" — so if you have three depots and two sites, no such assignment exists. Leonid Kantorovich's 1942 fix is almost anticlimactic: let the sand be split. What you choose then is a table.

Entry $(i,j)$ of the table $P$ is how much moves from depot $i$ to site $j$. Only three constraints apply: every entry non-negative, each row sums to that depot's stock, each column sums to that site's demand. The set of tables satisfying this is called the transport polytope.

The total bill is "amount × unit price" summed over every cell, so both the objective and the constraints are linear in the entries. Optimal transport is therefore a linear program. The moment splitting is allowed, a solution is guaranteed to exist and an entire existing toolbox becomes available — that is why Kantorovich's relaxation was the turning point.

Mechanism: defining the Wasserstein distance

All of the above condenses into one line.

$\mu$ and $\nu$ are the two distributions being compared, $x_i$ and $y_j$ are the locations of their points, and $a_i$ and $b_j$ are the probability mass sitting there (the sand). $U(a,b)$ is the transport polytope from above, and $d$ is distance on the ground, usually Euclidean. The exponent $p$ sets how the bill scales: $p=1$ charges proportionally to distance, $p=2$ charges the square, which strongly punishes long hauls. The outer $1/p$ power converts the result back into units of length.

In plain words the formula says: consider every possible transport table, take the one with the smallest total bill, and call that bill the distance. The $p=1$ case has its own name, the Earth Mover's Distance — a name that is already the explanation.

Why brute force is hopeless

The naive plan is to try every way of moving. But just assigning $n$ points to $n$ points one-to-one gives $n!$ options — about $2.4 \times 10^{18}$ at $n=20$.

Even the dedicated exact solvers (network simplex, auction algorithms) cost roughly $n^3$ with a log factor for $n$ points. Fine up to a few thousand points — far too slow for a generative model comparing two minibatches at every training step.

In one dimension, sorting is all you need

There is one exception. If the points lie on a line, the optimal plan requires no thought: match them in sorted order. If two hauling routes ever crossed, untangling the crossing would strictly shorten the total distance.

$$W_p^p(\mu,\nu)=\int_0^1 \left|F^{-1}(u)-G^{-1}(u)\right|^p du$$

Here $F^{-1}$ and $G^{-1}$ are the quantile functions — the value below which a fraction $u$ of the mass sits. The formula says: line up matching quantiles, measure the gap, average. With equal sample counts you just sort both arrays and subtract elementwise, which is $O(n\log n)$. In Python, scipy.stats.wasserstein_distance is the one-liner.

This "one dimension is dirt cheap" fact pays off later. Project a high-dimensional cloud onto many random directions, take the 1-D Wasserstein distance along each, and average — that is the Sliced Wasserstein distance, an approximation people genuinely ship.