Optimal Transport: Why Wasserstein Distance Sees Geometry That KL Divergence Misses

Two probability distributions can be very different in a way that KL divergence barely notices, or notices too harshly. Move a narrow bump of probability a little to the right, and KL divergence changes smoothly. Move it far away so the two bumps no longer overlap, and KL divergence can blow up or become undefined, even though the intuitive answer is simple: the distance is just how far the mass had to travel.

Optimal transport makes that intuition precise. It asks: what is the cheapest way to reshape one distribution into another, when moving a unit of mass across a distance incurs a cost? The answer, the Wasserstein distance, measures difference in the geometry of the underlying space, not only overlap of densities. That one change explains why optimal transport shows up in generative modelling, domain adaptation, shape comparison, and distribution matching, and also why it can be expensive and fragile in high dimensions.

The earth mover intuition

Imagine each distribution as a pile of earth spread over a landscape. To turn the first pile into the second, you may dig earth from some locations and deposit it elsewhere. Moving earth a short distance is cheap. Moving it across the landscape is expensive. The total cost depends on both how much earth moves and how far it moves.

KL divergence asks a different question. It compares densities point by point: at each location, how surprising is one density under the other? If the two distributions put mass in disjoint places, that pointwise comparison breaks down. Optimal transport instead compares mass, and allows mass to travel.

A simple example makes the contrast sharp. Let P put all its mass at 0 and let Q put all its mass at t, where t is a positive number. Intuitively, the distance between P and Q should grow with t. The Wasserstein distance does exactly that: it equals t. KL divergence, by contrast, is not a useful guide here, because the supports do not overlap. This is not a pathological curiosity. Real data often live on low dimensional manifolds, and two empirical distributions rarely share support in a meaningful way.

From Monge to Kantorovich

The classical formulation is due to Monge. Given a source distribution and a target distribution, find a map T that pushes the source onto the target while minimising total transport cost. The difficulty is that such a map may not exist: a single source point cannot split its mass between two target points under a deterministic map.

Kantorovich relaxed the problem by allowing mass to split. Instead of a map, we look for a coupling, a joint distribution whose marginals are the two distributions of interest. The coupling says how much mass travels from each source location to each target location. The optimal transport cost is the minimum, over all couplings, of the expected transport cost.

For distributions P and Q on a space with a ground cost c, the Kantorovich problem is:

$$\inf_{\pi \in \Pi(P, Q)} \int c(x, y) \, d\pi(x, y)$$

Here Pi(P, Q) is the set of all couplings of P and Q. When the ground cost is a power of a metric, for example c(x, y) = ||x - y||^p with p at least 1, the p-th root of this minimum is the p-Wasserstein distance:

$$W_p(P, Q) = \left( \inf_{\pi \in \Pi(P, Q)} \int ||x - y||^p \, d\pi(x, y) \right)^{1/p}$$

The two most used cases are p = 1 and p = 2. W1 is often called the earth mover distance. W2 has extra structure: under regularity conditions, the optimal map is the gradient of a convex function, a result known as Brenier's theorem. That structure is why W2 appears so often in theory and in algorithms that need a genuine map rather than only a distance.

A useful way to read W1 is through its dual form. For distributions with finite first moments:

$$W_1(P, Q) = \sup_{||f||L \le 1} \left( \mathbb{E}{X \sim P}[f(X)] - \mathbb{E}_{Y \sim Q}[f(Y)] \right)$$

The supremum is over all 1-Lipschitz functions f. In words, W1 is the largest difference in expectation that a function with slope at most 1 can detect. This dual view is the bridge to Wasserstein GANs: a critic network, constrained to be approximately Lipschitz, tries to maximise exactly this gap, while the generator tries to minimise it.

The discrete case: a linear program

In practice we usually have samples, not densities. Suppose P is supported on points x1 to xm with weights a1 to am, and Q is supported on points y1 to yn with weights b1 to bn. A transport plan is an m by n matrix Gamma with nonnegative entries, where row i sums to ai and column j sums to bj. Entry Gamma(i, j) is the mass moved from xi to yj. If C(i, j) is the cost of moving mass from xi to yj, the discrete problem is:

$$\min_{\Gamma \ge 0} \sum_{i=1}^{m} \sum_{j=1}^{n} \Gamma_{i,j} C_{i,j}$$

subject to:

$$\sum_{j=1}^{n} \Gamma_{i,j} = a_i, \quad \sum_{i=1}^{m} \Gamma_{i,j} = b_j$$

This is a linear program. That is good news for correctness and bad news for scale: a naive solution costs far more than most machine learning routines can afford when m and n are large. Much of modern optimal transport research is about approximations that keep the geometry while making this problem tractable.

One dimension is special

In one dimension, optimal transport has a closed form that is worth knowing, because it also gives a fast sanity check in higher dimensional work. For distributions on the real line with cumulative distribution functions F_P and F_Q, and quantile functions F_P^{-1} and F_Q^{-1}, the p-Wasserstein distance is:

$$W_p(P, Q) = \left( \int_0^1 |F_P^{-1}(u) - F_Q^{-1}(u)|^p \, du \right)^{1/p}$$

For p = 1, this simplifies to the area between the two cumulative distribution functions:

$$W_1(P, Q) = \int_{-\infty}^{\infty} |F_P(x) - F_Q(x)| \, dx$$

The intuition is sorting. In one dimension, the optimal plan never crosses itself: the smallest source mass goes to the smallest target mass, the second smallest to the second smallest, and so on. That monotone matching is both optimal and easy to compute.

The code below computes W1 between two equal sized samples in one dimension using exactly this sorted matching. It needs no solver and no density estimate.

import numpy as np

def wasserstein_1d_sample(x, y):
    x = np.sort(np.asarray(x, dtype=float))
    y = np.sort(np.asarray(y, dtype=float))
    if x.shape != y.shape:
        raise ValueError("This shortcut needs equal sized samples.")
    return float(np.mean(np.abs(x - y)))

# Two samples from similar but shifted distributions
rng = np.random.default_rng(0)
source = rng.normal(loc=0.0, scale=1.0, size=1000)
target = rng.normal(loc=2.0, scale=1.0, size=1000)
print(wasserstein_1d_sample(source, target))

The output is close to 2, the size of the shift, as intuition suggests. If the two samples came from the same distribution, the value would be small but not zero, because finite samples differ by chance. That sampling noise matters: always compare an observed Wasserstein distance against a null value obtained by splitting one sample in two, or by permuting labels, before claiming that two distributions differ.

Entropic regularisation and the Sinkhorn algorithm

Exact optimal transport is expensive, and its solution can be sparse and unstable: a tiny change in the data can switch the optimal plan. Entropic regularisation adds a penalty that favours diffuse, smooth plans. The regularised problem is:

$$\min_{\Gamma \in \Pi(a, b)} \sum_{i,j} \Gamma_{i,j} C_{i,j} + \varepsilon \sum_{i,j} \Gamma_{i,j} \log \Gamma_{i,j}$$

Here epsilon is a positive regularisation parameter. Larger epsilon gives a smoother plan and a faster algorithm, but a more biased distance. Smaller epsilon approaches exact transport but needs more iterations and more careful numerics.

The regularised problem can be solved with Sinkhorn iterations, which alternate between rescaling rows and columns of a kernel matrix derived from the cost matrix. Each iteration is a matrix-vector product, so the method is fast on GPUs and differentiable, which is why it appears inside neural networks and in differentiable sorting, matching, and clustering layers.

The trade-off must be stated plainly. Sinkhorn distance is not the Wasserstein distance. It is a regularised approximation, and its value depends on epsilon. Reporting a Sinkhorn number without epsilon, and without the ground cost, is like reporting a loss without saying which loss. If the downstream decision depends on the exact value, validate against exact transport on a subsample.

Where Wasserstein helps, and where it misleads

Wasserstein distance earns its place when geometry matters. Three settings are typical.

First, generative modelling. A generator that produces samples slightly off the data manifold should be penalised in proportion to how far off it is, not infinitely because densities do not overlap. The WGAN idea uses the W1 dual to build a critic with that property. In practice the Lipschitz constraint is enforced only approximately, by weight clipping or gradient penalty, so a WGAN loss is an approximation to an approximation. Treat it as a training signal, not as a certified distance.

Second, domain adaptation and distribution matching. If source and target domains differ by a shift, rotation, or deformation of features, a transport plan can align them in a way that a pointwise density ratio cannot. The plan itself is often more useful than the scalar distance: it tells you which source examples correspond to which target examples.

Third, comparing structured objects. Distributions over images, point clouds, word embeddings, or histograms can be compared with a ground cost that respects the structure of the objects, for example Euclidean distance between points in a cloud, or cosine distance between embeddings. The choice of ground cost is the modelling choice. Wasserstein distance inherits its meaning from that cost.

The limitations are equally important.

Sample complexity in high dimensions is the main one. Estimating a Wasserstein distance from samples gets harder as dimension grows, roughly because the space has more directions in which two samples can differ. In high dimensions, use sliced Wasserstein distances, which average one dimensional Wasserstein distances over random projections, or use embeddings and dimensionality reduction first, and report that choice.

The ground metric can dominate the result. Change the feature scaling and you change the transport cost, even if the data are unchanged. Standardise features with a reason, not by habit, and check sensitivity to that scaling before drawing conclusions.

Mass must be conserved in the standard formulation. If one distribution has extra mass that the other does not, for example outliers or a new class, standard transport is forced to move that mass somewhere, possibly at high cost. Unbalanced optimal transport relaxes exact marginal constraints and is the right tool when creation or destruction of mass is part of the problem.

Finally, a small Wasserstein distance does not mean two distributions are the same in every respect. It means they are close under one chosen ground cost. Two distributions can have the same mean, the same Wasserstein distance to a reference, and very different tails, modes, or conditional structure. Use Wasserstein as one diagnostic among several, alongside moment checks, conditional checks, and task performance.

A practical checklist

Before using a Wasserstein distance in a project, answer these questions in order:

1. What is the ground cost, and why does it reflect what "far" means in this problem?
2. Are the features scaled so that the ground cost is meaningful, or is one feature silently dominating?
3. Is exact transport feasible, or is entropic regularisation needed, and if so, what is epsilon?
4. What is the null value of the distance between two samples from the same distribution at this sample size?
5. Does mass need to be conserved, or is unbalanced transport more honest?
6. Is the transport plan itself useful, or only the scalar distance?

If any answer is unclear, the number that comes out of the solver will be precise and uninformative. Optimal transport is powerful precisely because it forces you to say what distance means before you measure it.

Takeaway

KL divergence compares densities at the same point. Wasserstein distance asks what it costs to move mass from one distribution to another. When distributions are separated, live on different manifolds, or differ by a geometric shift, that transport view gives a smooth, interpretable answer where a pointwise comparison fails. The price is computation, sample complexity in high dimensions, and a strong dependence on the ground cost. Choose the ground cost carefully, report regularisation honestly, and use the transport plan, not only the distance, whenever the correspondence between samples is what you actually need.

---

Connect with me on LinkedIn: linkedin.com/in/umeshtharukaofficial