What Is Optimal Transport?

Moving mass at minimum cost

Optimal transport (OT) answers a simple question: given two probability distributions, what is the cheapest way to rearrange the mass of one to match the other?

The canonical image is a pile of sand. Flatten it into the shape of a source distribution, then imagine rearranging the grains into the shape of a target distribution. Each grain travels some distance; OT finds the rearrangement that minimises the total work done.

The same idea applies to any pair of discrete distributions — histograms, point clouds, or probability vectors. Two sensors with different noise profiles, two images with different pixel distributions, two sets of detections at different time steps: OT gives a geometrically meaningful measure of how far apart they are.

Two histograms

Start with a concrete 1-D example. Suppose mass is concentrated on the left (source a) and needs to move right (target b):

import numpy as np
import torch
import torchmatch

N_BINS = 8
grid = np.arange(N_BINS, dtype=np.float32)

# Source: mass on the left.
a = np.array([0.30, 0.25, 0.20, 0.15, 0.05, 0.03, 0.01, 0.01], dtype=np.float32)
# Target: mass on the right.
b = np.array([0.01, 0.01, 0.03, 0.05, 0.15, 0.20, 0.25, 0.30], dtype=np.float32)

# Cost matrix: squared Euclidean distance between grid positions.
# C[i, j] = (i - j)^2
C = (grid[:, None] - grid[None, :]) ** 2   # (8, 8)

The cost matrix encodes the price of moving mass between any pair of bins. Moving between adjacent bins is cheap; moving across the histogram is expensive.

The transport plan

A transport plan P is an N×M matrix where P[i, j] is the amount of mass moved from source bin i to target bin j. Two constraints must hold:

• Row sums equal the source weights: P.sum(axis=1) == a
• Column sums equal the target weights: P.sum(axis=0) == b

These are the marginal constraints. Any P satisfying them is a valid transport plan; the optimal plan minimises the total cost ⟨P, C⟩ = ∑_ij P[i,j] · C[i,j].

torchmatch.transport.matrix.solve finds this plan. It expects a 3-D input (B, N, M) to support batches; add an unsqueeze for single problems:

cost_t = torch.tensor(C).unsqueeze(0)    # (1, 8, 8)
a_t    = torch.tensor(a).unsqueeze(0)    # (1, 8)
b_t    = torch.tensor(b).unsqueeze(0)    # (1, 8)

log_plan = torchmatch.transport.matrix.solve(
    cost_t, a=a_t, b=b_t, reg=0.01, n_iter=500
)
P = log_plan.exp().squeeze(0).numpy()    # (8, 8)

print(f"Row sums ≈ a: {P.sum(axis=1).round(3)}")
print(f"Col sums ≈ b: {P.sum(axis=0).round(3)}")
print(f"Transport cost: {(P * C).sum():.4f}")

The plan is returned in log space (log_plan) for numerical stability. Call .exp() to recover probabilities. The diagonal-dominant structure of P reflects the nature of the problem: nearby bins exchange the most mass, distant bins exchange very little.

Why log space?

With small regularisation (e.g. reg=0.01), many entries of P are extremely small — values like 1e-200 that underflow to zero in float32. Working in log space avoids this: log_plan[b, i, j] stores log P[i, j], which stays in a numerically stable range even when Pi, j itself would vanish.

Wasserstein distance

The Wasserstein distance (or earth-mover distance) between two distributions is the minimum total transport cost over all valid plans:

W(a, b) = min_P ⟨P, C⟩   subject to marginal constraints

It is not a raw per-element comparison. Two distributions with identical histograms shifted by one bin have a small Wasserstein distance; two distributions with the same mean but swapped peaks have a larger one. The geometry of the underlying space enters the computation through C.

import numpy as np
import torch
import torchmatch

def make_gaussian_hist(mean, std, n=8):
    grid = np.arange(n, dtype=np.float32)
    w = np.exp(-0.5 * ((grid - mean) / std) ** 2)
    return (w / w.sum()).astype(np.float32)

N_BINS = 8
C = (np.arange(N_BINS, dtype=np.float32)[:, None] -
     np.arange(N_BINS, dtype=np.float32)[None, :]) ** 2

pairs = [
    (make_gaussian_hist(3.0, 0.8), make_gaussian_hist(4.0, 0.8), "nearby"),
    (make_gaussian_hist(1.5, 0.8), make_gaussian_hist(6.5, 0.8), "far apart"),
]
cost_t = torch.tensor(C).unsqueeze(0)

for a_i, b_i, label in pairs:
    log_p = torchmatch.transport.matrix.solve(
        cost_t,
        a=torch.tensor(a_i).unsqueeze(0),
        b=torch.tensor(b_i).unsqueeze(0),
        reg=0.02, n_iter=300,
    )
    P_i = log_p.exp().squeeze(0).numpy()
    w = (P_i * C).sum()
    print(f"{label:12s}  W = {w:.4f}")

The "far apart" pair has a larger Wasserstein distance — the mass must travel further. Unlike, say, cross-entropy or L2, Wasserstein distance gives a meaningful answer even when the distributions have non-overlapping support.

Connection to the assignment problem

When both distributions are uniform over N points, the optimal transport plan is a permutation matrix — exactly the output of an assignment solver. OT is a continuous generalisation of the linear assignment problem (LAP): it handles arbitrary distributions, not just unit-mass matchings.

import numpy as np
import torch
import torchmatch

rng = np.random.default_rng(42)
N = 5
a_unif = np.ones(N, dtype=np.float32) / N
b_unif = np.ones(N, dtype=np.float32) / N
C_lap  = rng.random((N, N)).astype(np.float32)

log_plan_unif = torchmatch.transport.matrix.solve(
    torch.tensor(C_lap).unsqueeze(0),
    a=torch.tensor(a_unif).unsqueeze(0),
    b=torch.tensor(b_unif).unsqueeze(0),
    reg=0.001, n_iter=1000,
)
P_unif = log_plan_unif.exp().squeeze(0)

# With very small regularisation, P ≈ a permutation matrix.
# The assignment solver gives the same optimal cost.
assignment = torchmatch.assignment.solve(torch.tensor(C_lap))
lap_cost  = C_lap[np.arange(N), assignment.numpy()].sum()
ot_cost   = (P_unif.numpy() * C_lap).sum()
print(f"LAP cost: {lap_cost:.4f}   OT cost: {ot_cost:.4f}")

With reg=0.001 the regularised OT cost closely tracks the exact LAP cost. Increasing reg smears the plan and raises the cost — this tradeoff is the subject of the next tutorial.

See also

• Quickstart: matrix.solve and samples.loss in one page.
• Sinkhorn algorithm: why regularisation is necessary and how Sinkhorn iteration works.
• Algorithms: the full derivation of the OT problem and solver families.