Optimal Transport, Visually

← text notes · Foundations › Probability. How to move one distribution onto another as cheaply as possible, and why the answer gives a distance between distributions that deep learning keeps reaching for (Wasserstein distance, FID, OT flow matching, domain adaptation).

1. The idea: sand into holes

Picture a probability distribution as a pile of sand: where the pile is tall, there's a lot of probability. You have one pile, μ (the source), and a set of holes shaped like a second distribution, ν (the target). Both hold the same total amount (probability 1).

The optimal transport question: what's the cheapest way to shovel all the sand into the holes? Moving a grain a distance d costs d (or d², depending on the cost we pick). The minimum total cost is a number that says how far apart the two distributions are. That number is the Wasserstein distance, also called the earth mover's distance.

What to notice. In 1-D, the optimal rule is simple: sort the grains left to right, sort the holes left to right, and match the 1st with the 1st, the 2nd with the 2nd, and so on. Nothing ever crosses over. Switch to random matching and the grains criss-cross, so the average distance goes up. Optimal transport is the matching with the smallest possible average distance.

2. A worked example with numbers

Swap sand for bread. Three bakeries produce 30, 20 and 50 loaves (total 100). Three cafés need 40, 35 and 25 loaves (total 100). Shipping one loaf costs its distance on the map. Divide everything by 100 and it's a probability problem: μ = (0.3, 0.2, 0.5), ν = (0.4, 0.35, 0.25).

A transport plan is a table: “ship γ_ij loaves from bakery i to café j.” It has to obey two rules:

Row sums = supply. Each bakery ships out exactly what it baked (row i adds up to μ_i).
Column sums = demand. Each café receives exactly what it needs (column j adds up to ν_j).

Among all tables obeying both rules, OT picks the one with the smallest total cost Σ_ij cost_ij × γ_ij. Drag the shops around and watch the plan re-route.

The naive plan fills café 1 first, then café 2, then café 3, without looking at distances. It still obeys both rules (it's a valid plan), but it costs more. Notice that the optimal plan sometimes splits one bakery's output across two cafés. That splitting is the key difference between the two classic formulations, coming up next.

3. Maps vs. plans (Monge vs. Kantorovich)

There are two ways to write down “a way of moving the sand.”

The problem with maps (and a reminder about “pushforward”)

Monge's constraint is written $T_\#\mu = \nu$, read “T pushes μ forward to ν.” In plain words: take every grain of μ, move it to T(grain), and the pile you end up with must look exactly like ν. In deep learning terms: if X ~ μ is an input and T is your network, then $T_\#\mu$ is simply the distribution of the outputs T(X).

Now look at the right-hand picture above. If μ is a single spike of mass 1 and ν is two spikes of mass ½, no function T can do it: a function sends one point to one point, so the whole spike lands in one place. Kantorovich fixed this by letting mass split. That makes the problem always solvable, and it turns it into a linear program (a table of unknowns, linear constraints, linear cost), which is why it's the version everyone uses.

Good to know: when μ is a smooth density (no spikes) and the cost is squared distance, the best Kantorovich plan doesn't split anything. It turns out to be a Monge map after all (Brenier's theorem). §4 lets you see this: the optimal plan collapses onto a thin curve, which is the graph of the map y = T(x).

4. A plan is a joint distribution

Here is the view that makes OT click for probability. Divide the plan table by the total mass and γ is a joint distribution of a pair (X, Y):

its row sums (marginal over y) give back μ, the distribution of X;
its column sums (marginal over x) give back ν, the distribution of Y.

A joint distribution with prescribed marginals is called a coupling. So OT means choosing the coupling of μ and ν that makes X and Y as close as possible on average, i.e. the one with the smallest 𝔼[|X − Y|²].

Below, the square is γ(x, y): source position x runs up the left side (with μ drawn there), target position y runs along the top (with ν). Dark cells mean a lot of mass travels from that x to that y. The dashed diagonal is “don't move” (y = x).

Play with the blur ε (the “entropic regularisation” behind the Sinkhorn algorithm):

ε = 0: the exact OT plan, a thin increasing curve. Each x goes to one y (a Monge map, as Brenier promised), and it's the “sort and match” rule from §1.
small ε: the curve gets a soft halo. It costs slightly more, but it's much faster to compute on GPUs and it's differentiable. This is what libraries like POT and geomloss usually compute.
ε → ∞: the plan becomes $\mu(x)\,\nu(y)$, the independent coupling: X and Y ignore each other and sand gets sent everywhere. It's still a valid plan (the marginals are right), just a terrible one.

5. Why ML likes Wasserstein distance

KL divergence and Jensen–Shannon compare distributions point by point: “at this location, how different are the heights?” They don't care how far apart the mass is. Wasserstein compares them by how far the mass has to travel. Slide the target away from the source:

Why this matters for training. Once the two distributions stop overlapping (|θ| > width), JS is stuck at log 2 and KL is ∞. Neither changes as you move, so their gradient is zero and a generator gets no signal about which way to go. W₁ = |θ| keeps shrinking as you move closer, so its gradient always points the right way. That's the core argument of the Wasserstein GAN paper (Arjovsky et al., 2017). Real data and early generator samples live on thin, non-overlapping manifolds, so this is the normal situation, not a corner case.

6. Point clouds & straight-line flows

In practice you rarely have densities. You have samples, e.g. a batch of embeddings from domain A and a batch from domain B. With n equally-weighted points on each side, OT becomes an assignment problem: pair every source point with exactly one target point so that the average squared distance is smallest. Slide t to move every point along a straight line to its partner.

Link to flow matching. Sliding t moves each point along $x_t = (1-t)\,x_0 + t\,x_1$. With the optimal pairing, the paths are short and orderly, and (for squared cost) two moving points never occupy the same spot at the same time. With random pairing, paths are long and tangled. Flow-matching models learn a velocity field along exactly these paths, and tangled paths are harder to learn and need more integration steps. That's why OT-coupled flow matching (Tong et al., 2023; Pooladian et al., 2023) pairs noise and data inside each minibatch with OT before training. The Path-Decoupled Hyperbolic Flow Matching paper (Li et al., 2026) builds on this family.

7. The formal definitions, piece by piece

Now the symbols, each one tied back to a picture above.

Ingredients

$\mu$ on a space $\mathcal X$: the source distribution (the sand pile, the bakeries, the batch from domain A).
$\nu$ on $\mathcal Y$: the target distribution (the holes, the cafés, domain B).
$c(x,y)$: the cost of moving one unit of mass from x to y. Usually $\|x-y\|$ or $\|x-y\|^2$.

Monge's problem (maps, §3 left)

$$\inf_{T}\ \int_{\mathcal X} c\big(x, T(x)\big)\, d\mu(x) \quad\text{subject to}\quad T_\#\mu = \nu .$$

$\int c(x,T(x))\,d\mu(x)$ is just “the average cost per grain when grains are drawn from μ”, i.e. $\mathbb E_{X\sim\mu}[c(X,T(X))]$. The constraint $T_\#\mu=\nu$ says “the outputs T(X) must be distributed as ν”. Formally $(T_\#\mu)(B) = \mu(T^{-1}(B))$: the probability of landing in a region B equals the probability of starting somewhere that T sends into B.

Kantorovich's problem (plans, §2 and §4)

$$\mathrm{OT}_c(\mu,\nu) \;=\; \min_{\gamma\in\Pi(\mu,\nu)} \int_{\mathcal X\times\mathcal Y} c(x,y)\, d\gamma(x,y) \;=\; \min_{\gamma\in\Pi(\mu,\nu)} \mathbb E_{(X,Y)\sim\gamma}\big[c(X,Y)\big]$$

$\Pi(\mu,\nu)$ is the set of couplings: joint distributions of (X, Y) whose first marginal is μ and second marginal is ν (the “row sums / column sums” rule). Formally, $\gamma(A\times\mathcal Y)=\mu(A)$ and $\gamma(\mathcal X\times B)=\nu(B)$.

The discrete version (what code actually solves)

With weights $a\in\mathbb R^n$ on source points, $b\in\mathbb R^m$ on target points, and cost matrix $C_{ij}=c(x_i,y_j)$:

$$\min_{P\in\mathbb R_{\ge 0}^{n\times m}} \ \sum_{i,j} C_{ij}P_{ij} \quad\text{s.t.}\quad P\mathbf 1_m = a,\;\; P^\top\mathbf 1_n = b .$$

That's §2 exactly: $P$ is the shipping table, $P\mathbf 1 = a$ is “row sums = supply”, $P^\top\mathbf 1 = b$ is “column sums = demand”. It's a linear program with $nm$ unknowns.

Wasserstein distance

Choose the cost $c(x,y)=\|x-y\|^p$ and take a p-th root:

$$W_p(\mu,\nu) \;=\; \Big(\min_{\gamma\in\Pi(\mu,\nu)} \mathbb E_{(X,Y)\sim\gamma}\,\|X-Y\|^p\Big)^{1/p}.$$

This is a genuine metric on distributions (symmetric, zero only if μ = ν, triangle inequality). $W_1$ is the earth mover's distance (§1, §5) and $W_2$ is the one with the nice geometry (§4, §6).

Closed forms worth memorising

1-D: sort and match. With $F^{-1}, G^{-1}$ the quantile functions of μ and ν,
$$W_p^p(\mu,\nu)=\int_0^1 \big|F^{-1}(u)-G^{-1}(u)\big|^p\,du .$$
“Match the u-th quantile of μ with the u-th quantile of ν”, which is the §1 animation.
Gaussians: for $\mathcal N(m_1,\Sigma_1)$ and $\mathcal N(m_2,\Sigma_2)$,
$$W_2^2 = \|m_1-m_2\|^2 + \operatorname{tr}\!\Big(\Sigma_1+\Sigma_2-2\big(\Sigma_2^{1/2}\Sigma_1\Sigma_2^{1/2}\big)^{1/2}\Big).$$
This is exactly the FID score: fit a Gaussian to Inception features of real images and one to generated images, then compute this. It's also useful for the Gaussian-embedding papers (Betser et al. 2026, Gaussianity in practice): if embeddings are Gaussian, the W₂ distance between two domains has this formula.

Entropic OT and Sinkhorn (the ε slider in §4)

$$\min_{P\in\Pi(a,b)} \sum_{ij} C_{ij}P_{ij} \;-\; \varepsilon\, H(P), \qquad H(P) = -\sum_{ij} P_{ij}\log P_{ij}.$$

The entropy term $H(P)$ rewards spread-out plans, which is the blur. The solution has the form $P = \operatorname{diag}(u)\,K\,\operatorname{diag}(v)$ with $K_{ij}=e^{-C_{ij}/\varepsilon}$. The Sinkhorn algorithm finds it by alternately rescaling rows to match a and columns to match b. It uses only matrix–vector products, so it's GPU-friendly and differentiable (Cuturi, 2013).

Duality: the shipping company's view

$$\mathrm{OT}_c(\mu,\nu) = \max_{f,g}\ \Big\{\ \mathbb E_{\mu}[f(X)] + \mathbb E_{\nu}[g(Y)] \ :\ f(x)+g(y)\le c(x,y)\ \text{for all } x,y \Big\}.$$

Story: a shipping company offers to do the job. It charges f(x) per loaf to pick up at x and g(y) per loaf to drop off at y. To stay competitive, pickup + dropoff can never exceed what you'd pay to ship it yourself, c(x, y). The most the company can earn equals your cheapest self-shipping cost. For $W_1$ this becomes the Kantorovich–Rubinstein form

$$W_1(\mu,\nu)=\sup_{\|f\|_{\text{Lip}}\le 1}\ \mathbb E_{\mu}[f(X)]-\mathbb E_{\nu}[f(Y)],$$

which is where the WGAN “critic” comes from: f is a neural network kept 1-Lipschitz (by weight clipping or a gradient penalty).

8. Reading a paper's OT definition (data-selection example)

Here is a definition from a data-selection paper ("Optimal Transport between Discrete Distributions", which reads like TSDS, Liu et al., 2024, the OT-based baseline in the NTK-Selector paper). It's the discrete Kantorovich problem from §7, i.e. the bakery problem from §2, with different letters:

$$\begin{aligned}&\min_{\gamma\in\mathbb R_{\ge 0}^{|U|\times|V|}}\ \sum_{i=1}^{|U|}\sum_{j=1}^{|V|}\gamma_{ij}\,f(u_i,v_j)\\ &\text{s.t.}\ \ \sum_{j=1}^{|V|}\gamma_{ij}=\mu_i\ \ \forall i,\qquad \sum_{i=1}^{|U|}\gamma_{ij}=\nu_j\ \ \forall j .\end{aligned}$$

Paper symbol	Plain meaning	In the bakery widget (§2)
$(A, f)$, a metric space	A finite set of examples $A$ with a distance $f(a,a')$ between any two, e.g. distance between embeddings	The map; $f$ = straight-line distance
$U \subseteq A$, points $u_i$	The source examples (e.g. a candidate pool)	The bakeries
$V \subseteq A$, points $v_j$	The target examples (e.g. a small target-task set)	The cafés
$\mu_i = \mu(u_i)$	Probability mass sitting on example $u_i$. Uniform would be $1/\|U\|$	Supply 30, 20, 50 ÷ 100
$\nu_j = \nu(v_j)$	Probability mass on target example $v_j$	Demand 40, 35, 25 ÷ 100
$\gamma\in\mathbb R_{\ge0}^{\|U\|\times\|V\|}$	A non-negative $\|U\|\times\|V\|$ table; $\gamma_{ij}$ = mass moved from $u_i$ to $v_j$	The “Plan γ” table
$\sum_{i,j}\gamma_{ij}f(u_i,v_j)$	Total cost = mass × distance, summed over every pair	“Total cost” readout
$\sum_j \gamma_{ij}=\mu_i$	Everything at $u_i$ gets shipped out (row sums)	“row sum = μᵢ” column
$\sum_i \gamma_{ij}=\nu_j$	Each $v_j$ receives exactly what it should (column sums)	“col sum = νⱼ” row

Two remarks help when reading such papers:

$U$ and $V$ living in one space $A$ just means one distance function $f$ compares any source example with any target example. In practice $A$ is “all texts or images” and $f$ is a distance between their embeddings (often cosine or Euclidean in a pretrained model's feature space). The OT answer is only as meaningful as that embedding space.
Why OT is useful for selection. The optimal value (with $f$ a metric, this is $W_1(\mu,\nu)$) measures how far the source distribution is from the target distribution. If you get to choose the weights $\mu$ on the candidate pool, e.g. by putting mass only on selected examples, then making this value small means choosing candidates whose distribution covers the target distribution. That's a different goal from just grabbing each target point's nearest neighbours: the column constraints force every target example to be served, so selections can't all pile up near one popular target region. Look at §2: with the column-sum rule, every café gets its full demand even if it's far from all bakeries.

9. Where it shows up in papers

Generative models: Wasserstein GAN (Arjovsky, Chintala & Bottou, 2017) trains against W₁ via the dual (§5, §7). FID is W₂ between Gaussian fits (§7).
Flow matching / diffusion: OT-conditional flow matching (Tong et al., 2023) and multisample flow matching (Pooladian et al., 2023) use minibatch OT couplings to straighten paths (§6).
Domain adaptation: Courty et al. (2017) align source and target feature distributions with an OT plan.
Transferability / dataset similarity: OTDD (Alvarez-Melis & Fusi, 2020) measures distance between labelled datasets with OT. OTCE (Tan et al., 2021) uses an OT coupling between source and target to predict transfer performance. Both sit next to LEEP and LogME in the transferability notes.
Computation: Sinkhorn (Cuturi, 2013). Libraries: POT (Python Optimal Transport), geomloss, ott-jax.
Textbooks: Peyré & Cuturi, Computational Optimal Transport (2019) is readable and ML-oriented and has most of these pictures. Villani, Optimal Transport: Old and New (2009) is the deep-theory reference.

10. Cheat sheet

Optimal transport (OT): Cheapest way to rearrange distribution μ into distribution ν, given a per-unit moving cost c(x, y).
Transport plan / coupling γ: A joint distribution of (X, Y) with marginals μ and ν. “How much mass goes from x to y.” In code: an n×m matrix with fixed row and column sums.
Transport map T: A function sending each x to one T(x). It can't split mass. Exists for nice (smooth) μ with squared cost (Brenier).
Pushforward T_#μ: The distribution of T(X) when X ~ μ.
Wasserstein distance W_p: (min over couplings of 𝔼‖X − Y‖^p)^1/p. It measures how far mass must travel, and it stays informative when supports don't overlap.
Earth mover's distance: Another name for W₁.
Entropic OT / Sinkhorn: OT with a blur term −εH(P). Fast, differentiable, GPU-friendly. ε → 0 gives exact OT, ε → ∞ gives the independent coupling μ⊗ν.
Dual potentials f, g: Pickup/dropoff “prices” whose sum never exceeds the cost. In W₁ this is a single 1-Lipschitz function (the WGAN critic).