[WGF 1.1] Basics of Optimal Transport

This post is about the basics of optimal transport (OT) problems and the definition and properties of Wasserstein distances as a special OT problem instance.

Welcome to my long post series on Optimal Transport and Wasserstein Gradient Flows! Today we will start with the very basics by understanding what optimal transport (OT) problems exactly are, and how we can define Wasserstein distances as a special type of OT problem. The problem that defines the Wasserstein distance enjoys some nice properties and interesting characteristics, which are what the first two posts will all be about. For a brief overview about what this whole series will be about, you can take a look at the introduction. Without further adieu, let’s dive right in!

Notations

Throughout the series, we will mostly use the notations \(\mu\) and \(\nu\) for probability measures over the spaces \(\mathcal{X}\) and \(\mathcal{Y}\), respectively. In most cases, we will consider the simplest case of Euclidean spaces, \(\mathcal{X} = \mathcal{Y} = \mathbb{R}^{d}\). If not, we assume that \(\mathcal{X}\) and \(\mathcal{Y}\) are (at least) Polish spaces, which one can think of as a complete, separable metric space.

We will denote by \(\mathcal{P} (\mathcal{X})\) the set of probability measures over \(\mathcal{X}\), and by \(\mathcal{P}_{p} (\mathcal{X})\) the set of such probability measures with finite \(p\)-th moment with respect to the metric \(\| \cdot \|\):

$$ \mathcal{P}_{p} (\mathcal{X}) := \left\{ \mu \in \mathcal{P} (\mathcal{X}) : \int \lVert \boldsymbol{x} \rVert^{p} \mathrm{d} \mu < \infty \right\}. $$

Monge Formulation

The very first, original version of the optimal transport (OT) problem was formulated by Monge [Mon81] as

$$ \inf_{\substack{T: \mathcal{X} \rightarrow \mathcal{Y} \\ T_{\sharp} \mu = \nu}} \int c(\boldsymbol{x}, T(\boldsymbol{x})) \mathrm{d} \mu (\boldsymbol{x}), $$

where \(c(\boldsymbol{x}, \boldsymbol{y})\) is the cost function and \(T: \mathcal{X} \rightarrow \mathcal{Y}\) is a map such that \(T_{\sharp} \mu = \nu\), which is a mathematical formulation of the idea of mass transportation through the notion of push-forwards of measures.

Intuitively, \(T_{\sharp} \mu = \nu\) means that \(T\) moves the mass spread on \(\mathcal{X}\) w.r.t. the measure \(\mu\) to \(\mathcal{Y}\) w.r.t. the measure \(\nu\), which is exactly what transportation problems are designed for. For the simplest case, we could choose the quadratic cost \(c(\boldsymbol{x}, \boldsymbol{y}) = \lVert T(\boldsymbol{x}) {}-{} \boldsymbol{x} \rVert^{2}\). Then we can write

$$ \begin{align} \inf_{\substack{T: \mathcal{X} \rightarrow \mathcal{Y} \\ T_{\sharp} \mu = \nu}} \int \lVert T(\boldsymbol{x}) {}-{} \boldsymbol{x} \rVert^{2} \mathrm{d} \mu (\boldsymbol{x}). \end{align} $$

The solution to this problem will find a transportation map \(T\) from \(\mu\) to \(\nu\) so that we suffer the minimum cost of transporting the masses among all possible ways of transporting all the mass from the initial distribution \(\mu\) to the target \(\nu\), and thus it makes sense why we call this an optimal transport problem.

Despite being a natural way to write down the optimal transport problem via mathematics, one caveat of the Monge formulation is that the solution of the problem might not exist at all for some cases. As a counterexample, let us think of the following measures defined in a 1D space (or on any axis line on a \(d\)-dimensional space):

$$ \begin{align} \mu = \delta_0, \quad \nu = \frac{1}{2} \delta_{1} + \frac{1}{2} \delta_{-1}. \end{align} $$

It is easy to see that a deterministic map can never push-forward \(\mu\) to \(\nu\). Instead, it would make more sense to think of a randomized assignment which sends the mass concentrated in \(0\) to either \(1\) or \(-1\) with equal probability. In a similar spirit, a convex relaxation of the Monge formulation (called the Monge-Kantorovich formulation or just the Kantorovich formulation), designed to find a probability measure instead of a deterministic map like \(T\), can admit the existence, uniqueness, and many useful characteristics of the solution. This would require us to think of finding a coupling measure on the product space \(\mathcal{X} \times \mathcal{Y}\), defined as below, instead of a deterministic map \(T: \mathcal{X} \rightarrow \mathcal{Y}\) with \(T_{\sharp} \mu = \nu\).

Example. As a light example, we can immediately think of the coupling measure for the counterexample of Monge’s formulation:

$$ \begin{align} \pi = \frac{1}{2} \delta_{(0, 1)} + \frac{1}{2} \delta_{(0, -1)}. \end{align} $$

In general, there usually are (infinitely) many different coupling measures for a given tuple of probability measures \(\mu\) and \(\nu\). For instance, when \(\mu = \nu = \frac{1}{2} \delta_{1} + \frac{1}{2} \delta_{-1}\), all of the following measures for \(\alpha \in [-1, 1]\):

$$ \begin{align} \pi = \frac{1 + \alpha}{4} (\delta_{(1, 1)} + \delta_{(-1, -1)}) + \frac{1 {}-{} \alpha}{4} (\delta_{(1, 1)} + \delta_{(-1, -1)}). \end{align} $$

can be a coupling. In fact, the case when \(\alpha = 0\) is when the two are independent and \(\pi\) is just the product measure, while the cases \(\alpha = \pm 1\) can be thought of as the two measures being completely correlated to each other.

Measure Theory Bits: Properties of \(\Pi(\mu, \nu)\)

The notion of compactness of sets in \(\mathcal{P} (\mathcal{X})\) (for Polish spaces \(\mathcal{X}\)) can be characterized as closed and tight sets by Prokhorov’s theorem, which is quite similar to the Heine-Borel theorem that shows compact sets on \(\mathbb{R}^{d}\) are equivalent to closed and bounded sets. (We omit the proof of Prokhorov’s theorem in this article.) In fact, Prokhorov’s theorem defines compact sets of measures in \(\mathcal{P}(\mathcal{X})\) as those we can concentrate into compact sets in the space \(\mathcal{X}\) where the measures live on.

Remark. Here, compactness and closedness are all with respect to the topology of weak convergence.

Going back to the constraint \(\Pi(\mu, \nu)\), we can see that this set is nonempty because it always contains the product measure of \(\mu\) and \(\nu\). Also, \(\Pi(\mu, \nu)\) is a convex set because, by definition, any convex combination of \(\pi \in \Pi(\mu, \nu)\) must also have the same marginal distributions. Moreover, we can show that \(\Pi(\mu, \nu)\) must be compact, which we demonstrate in the following theorem.

Proof. Note that it suffices to show that \(\Pi(\mu, \nu)\) is compact, or equivalently closed and tight by Prokhorov’s theorem. To show closedness, note that for any sequence \((\pi_{n})_{n \in \mathbb{N}}\) in \(\Pi(\mu, \nu)\) with \(\pi_{n} \rightarrow \pi\) weakly, we have

$$ \begin{align}
\int_{\mathcal{X} \times \mathcal{Y}} f(\boldsymbol{x}) \mathrm{d} \pi &= \lim_{n \rightarrow \infty} \int_{\mathcal{X} \times \mathcal{Y}} f(\boldsymbol{x}) \mathrm{d} \pi_{n} = \int_{\mathcal{X}} f \ \mathrm{d} \mu, \\
\int_{\mathcal{X} \times \mathcal{Y}} g(\boldsymbol{y}) \mathrm{d} \pi &= \lim_{n \rightarrow \infty} \int_{\mathcal{X} \times \mathcal{Y}} g(\boldsymbol{y}) \mathrm{d} \pi_{n} = \int_{\mathcal{Y}} g \ \mathrm{d} \nu
\end{align} $$

for any bounded, continuous \(f: \mathcal{X} \rightarrow \mathbb{R}\) and \(g: \mathcal{Y} \rightarrow \mathbb{R}\), which means that \(\pi \in \Pi(\mu, \nu)\) and thus \(\Pi(\mu, \nu)\) is closed.

To show tightness, observe that the singletons \( \{\mu\} \) and \( \{\nu\} \) each trivially form a compact set. Therefore they are also tight, i.e., for any \(\epsilon > 0\), there exists compact sets \(K_{1} \subset \mathcal{X}\) and \(K_{2} \subset \mathcal{Y}\) such that \(\mu(K_{1}^{c}) < \frac{\epsilon}{2}\) and \(\mu(K_{2}^{c}) < \frac{\epsilon}{2}\). By Tychonoff’s theorem, the product \(K_{1} \times K_{2}\) is also compact, and for all \(\pi \in \Pi(\mu, \nu)\) we have

$$ \begin{align}
\pi \left((K_{1} \times K_{2})^{c}\right) \le \pi(K_{1}^{c} \times \mathcal{Y}) + \pi(\mathcal{X} \times K_{2}^{c}) = \mu (K_{1}^{c}) + \mu(K_{2}^{c}) < \epsilon
\end{align} $$

and thus \(\Pi(\mu, \nu)\) is tight. \(\square\)


Monge-Kantorovich Formulation

We are now ready to write down the Monge-Kantorovich formulation of the optimal transport problem [Kan42], which we will simply refer to as the optimal transport problem from now on.

One thing to note is that the Monge-Kantorovich formulation is a convex optimization problem in the sense that the constraint \( \Pi (\mu, \nu) \) is a convex set and the integral is a linear functional on the measure \(\pi\).

Remark. If we further assume that \( \mathcal{X} \) and \( \mathcal{Y} \) are finite sets, then \( (\mathrm{P}) \) becomes a Linear Programming (LP) problem, which we can solve using methods like the Sinkhorn algorithm.

We can show that the infimum is always attained by some \(\pi \in \Pi(\mu, \nu)\), i.e., we can think of the \(\inf\) as a \(\min\) instead, whenever the cost function \(c(\boldsymbol{x}, \boldsymbol{y})\) is lower semicontinuous.

Proof. Let us define

$$I(\pi) := \int_{\mathcal{X} \times \mathcal{Y}} c(\boldsymbol{x}, \boldsymbol{y}) \mathrm{d} \pi.$$

Since \(\mathrm{P} (\mu, \nu) = \inf_{\pi \in \Pi(\mu, \nu)} I(\pi)\), there exists a sequence \((\pi_{n})_{n \in \mathbb{N}}\) such that \(\lim_{ n \to \infty } I(\pi_{n}) = \mathrm{P} (\mu, \nu)\). Also, since \(\Pi(\mu, \nu)\) is a compact space (Theorem 1.6), we can find a subsequence \((\pi_{i_{n}})_{n \in \mathbb{N}}\) such that \(\pi_{i_{n}} \rightarrow \pi^{\star}\) weakly for some \(\pi^{\star} \in \Pi(\mu, \nu)\). Now, the direction \(\mathrm{P} (\mu, \nu) \le I(\pi^{\star})\) is obvious by definition. To show equality, we can observe that \(c(\boldsymbol{x}, \boldsymbol{y})\) is lower semicontinuous, and by the Portmanteau theorem the weak convergence \(\pi_{i_{n}} \rightarrow \pi^{\star}\) implies

$$\begin{align}
I(\pi^{\star}) = \int_{\mathcal{X} \times \mathcal{Y}} f \ \mathrm{d} \pi^{\star} \le \lim \inf_{n \rightarrow \infty} \int_{\mathcal{X} \times \mathcal{Y}} f \ \mathrm{d} \pi_{i_{n}} \le \lim_{n \rightarrow \infty} \int_{\mathcal{X} \times \mathcal{Y}} f \ \mathrm{d} \pi_{n} = \mathrm{P} (\mu, \nu).
\end{align}$$

Therefore we have \(\mathrm{P} (\mu, \nu) \le I(\pi^{\star})\). \(\square\)


Wasserstein distances

Perhaps the most popular instance among optimal transport problems would be the Wasserstein distance between measures. Specifically, choosing \(\mathcal{X} = \mathcal{Y} = \mathbb{R}^{d}\) and \(c(\boldsymbol{x}, \boldsymbol{y}) = \lVert \boldsymbol{x} {}-{} \boldsymbol{y} \rVert^{p}\) for each \(p \ge 1\) corresponds to the \(p\)-Wasserstein distance, defined as below.

For the optimization problem that defines the Wasserstein distance, we can show that the infimum is always attained by some \(\pi \in \Pi(\mu, \nu)\), i.e., we can think of the \(\inf\) as a \(\min\) instead.

An important property of \(p\)-Wasserstein distances is that they turn out to be a metric over the space \(\mathcal{P}_{p} (\mathbb{R}^{d})\), as shown in Theorem 1.10.

Proof. The first two are trivial by the definition of \(\mathcal{W}_{p}\). For the third property, we can easily see that \(\mathcal{W}_{p} (\mu, \mu) = 0\) because the measure \(\pi (\mathrm{d} x, \mathrm{d} y) = \mu (\mathrm{d} x) \delta_{x} (\mathrm{d} y)\), which is just the measure \(\mu\) fully concentrated on \(\boldsymbol{x} = \boldsymbol{y}\), can attain the infimum value of \(0\). For the converse direction, \(\mathcal{W}_{p} (\mu, \nu) = 0\) is equivalent to

$$ \begin{align} \inf_{\pi \in \Pi(\mu, \nu)} \left( \int \lVert \boldsymbol{x} {}-{} \boldsymbol{y} \rVert^{p} \mathrm{d} \pi \right) = 0, \end{align} $$

which implies that there exists some \(\pi \in \Pi (\mu, \nu)\) such that, for \((\boldsymbol{X}, \boldsymbol{Y}) \sim \pi\), we have \(\lVert \boldsymbol{X} {}-{} \boldsymbol{Y} \rVert^{p} = 0\) or equivalently \(\boldsymbol{X} = \boldsymbol{Y}\) almost surely, which concludes that \(\mu = \nu\).

For the last property, we use the following lemma, also known as the gluing lemma.


Informal Proof of Lemma 1.10. We can think of a triple of random variables \((\boldsymbol{X}, \boldsymbol{Y}, \boldsymbol{Z})\) constructed by first drawing \(\boldsymbol{Y} \sim \nu\), and then drawing \(\boldsymbol{X}\) and \(\boldsymbol{Z}\) from the conditional distribution of \(\pi_{1, 2}\) and \(\pi_{2, 3}\) given \(\boldsymbol{Y}\), respectively. Then we can set \(\pi_{1, 2, 3}\) to be the law of the triple \((\boldsymbol{X}, \boldsymbol{Y}, \boldsymbol{Z})\). \(\square\)

Remark. A more formal proof of the gluing lemma will involve disintegration of measures, as shown in this article written by Prof. Filippo Santambrogio.


Now suppose that \(\pi_{1, 2}^{\star} \in \mathcal{P} (\mathbb{R}^{d} \times \mathbb{R}^{d})\) and \(\pi_{2, 3}^{\star} \in \mathcal{P} (\mathbb{R}^{d} \times \mathbb{R}^{d})\) are the optimal coupling measures that attain the infimum for \(\mathcal{W}_{p} (\mu, \nu)\) and \(\mathcal{W}_{p} (\nu, \rho)\), respectively. By the gluing lemma, there exists \(\pi_{1, 2, 3} \in \mathcal{P}(\mathcal{X} \times \mathcal{Y} \times \mathcal{Z})\) such that the marginal on \(\mathcal{X} \times \mathcal{Y}\) and \(\mathcal{Y} \times \mathcal{Z}\) is \(\pi_{1, 2}^{\star}\) and \(\pi_{2, 3}^{\star}\), respectively. Letting \(\pi_{1, 3}\) be the marginal of \(\pi_{1, 2, 3}\) on \(\mathcal{X} \times \mathcal{Z}\), we have

$$ \begin{align} \mathcal{W}_{p} (\mu, \rho) &\le \left( \int_{\mathcal{X} \times \mathcal{Z}} \lVert \boldsymbol{x} {}-{} \boldsymbol{z} \rVert^{p} \mathrm{d} \pi_{1, 3} \right)^{1/p} \\
&= \left( \int_{\mathcal{X} \times \mathcal{Y} \times \mathcal{Z}} \lVert \boldsymbol{x} {}-{} \boldsymbol{y} + \boldsymbol{y} {}-{} \boldsymbol{z} \rVert^{p} \mathrm{d} \pi_{1, 2, 3} \right)^{1/p} \\
&\le \left( \int_{\mathcal{X} \times \mathcal{Y}} \lVert \boldsymbol{x} {}-{} \boldsymbol{y} \rVert^{p} \mathrm{d} \pi_{1, 2}^{\star} \right)^{1/p} + \left( \int_{\mathcal{Y} \times \mathcal{Z}} \lVert \boldsymbol{y} {}-{} \boldsymbol{z} \rVert^{p} \mathrm{d} \pi_{2, 3}^{\star} \right)^{1/p} \\
&= \mathcal{W}_{p} (\mu, \nu) + \mathcal{W}_{p} (\nu, \rho) \end{align}$$

as desired, which completes the proof. \(\square\)


By Theorem 1.10, we can equip the space \(\mathcal{P}_{p} (\mathbb{R}^{d})\) with the Wasserstein distance \(\mathcal{W}_{p}\) to obtain a metric space \((\mathcal{P}_{p} (\mathbb{R}^{d}), \mathcal{W}_{p})\). If not specified otherwise, we will simply write \(\mathcal{P}_{p} (\mathbb{R}^{d})\) as the metric space endowed with the Wasserstein distance to save space. Theorem 1.10 will reappear in the very beginning of the second part of the series, which will be about how \(\mathcal{P}_{p} (\mathbb{R}^{d})\) and functions defined on this space must look like, thus leading to a natural way to define the notion of gradients on this complicated, infinite-dimensional space. (For the first part, our main interest will be on the optimal transport problem itself rather than properties of the metric space \(\mathcal{P}_{p} (\mathbb{R}^{d})\).)

Summary

To summarize what we’ve seen in post, for Polish spaces \(\mathcal{X}\) and \(\mathcal{Y}\), the optimal transport (OT) problem with respect to measures \(\mu \in \mathcal{P} (\mathcal{X})\), \(\nu \in \mathcal{P} (\mathcal{Y})\) is defined as

$$\begin{align}
\mathrm{P} (\mu, \nu) &= \inf_{\pi \in \Pi(\mu, \nu)} \int_{\mathcal{X} \times \mathcal{Y}} c(\boldsymbol{x}, \boldsymbol{y}) \ \mathrm{d} \pi,
\end{align}$$

where the constraint below is a nonempty, convex, and compact set:

$$\begin{align}
\Pi(\mu, \nu) &:= \left\{ \pi \in \mathcal{P} (\mathcal{X} \times \mathcal{Y}): \int_{\mathcal{Y}} \mathrm{d} \pi = \mu, \int_{\mathcal{X}} \mathrm{d} \pi = \nu, \ \pi \ge 0 \right\}.
\end{align}$$

The OT problem always admits a solution for lower semicontinuous costs \(c(\boldsymbol{x}, \boldsymbol{y})\), and when \(\mathcal{X} = \mathcal{Y} = \mathbb{R}^{d}\), we can use cost functions of the form \(c(\boldsymbol{x}, \boldsymbol{y}) = \| \boldsymbol{x} {}-{} \boldsymbol{y} \|^{p}\) to define \(p\)-Wasserstein distances (for \(p \ge 1\)), which we can equip as a metric to the space \(\mathcal{P}_{p} (\mathbb{R}^{d})\).

In the remaining posts for the first part, we will explore how we can formalize the dual of the optimal transport problem, prove strong duality, how to characterize the optimal solutions using gradients of convex functions, and some other (possibly less relevant but) interesting topics related to optimal transport problems. Stay tuned!

References

[CNR25] S. Chewi, J. Niles-Weed, and P. Rigollet. Statistical optimal transport. École d’Été de Probabilités de Saint-Flour XLIX, 2025.

[Kan42] L. V. Kantorovich. On the translocation of masses. Dokl. Akad. Nauk SSSR 37 (1942), 227–229.

[Mon81] Gaspard Monge. Mémoire sur la théorie des déblais et des remblais. De l’Imprimerie Royale, 1781.

Jaewook

I am a first-year PhD student at Stanford EE. I am interested in various topics in optimization theory and algorithms, applications to ML/AI settings, and other random math-related stuff.

Leave a Reply

Your email address will not be published. Required fields are marked *