Formal Riemannian Structure of the Wasserstein metric

Given a closed and convex space $X\subseteq R^{d}$, two probability measures on the same space, $\mu ,\nu \in {\mathcal {P}}_{2}(X)$, the 2-Wasserstein metric is defined as

$W_{2}(\mu ,\nu ):=\min _{\gamma \in \Gamma (\mu ,\nu )}\left(\int |x_{1}-x_{2}|^{2}\,d\gamma (x_{1},x_{2})\right)^{1/2}$

where $\Gamma (\mu ,\nu )$ is the set of transport plans from $\mu$ to $\nu$. The Wasserstein metric is indeed a metric in the sense that it satisfies the desired properties of a distance function between probability measures on ${\mathcal {P}}_{2}(X)$. Moreover, the Wasserstein metric can be used to define a formal Riemannian metric on ${\mathcal {P}}_{2}(X)$. Such a formal metric structure allows one to define angles and lengths of vectors at each point in our ambient space.

Tangent Space Induced by the Wasserstein Metric

A convenient way to formalize tangent vectors in this setting is to consider time derivatives of curves on the manifold. A tangent vector at a point $\rho$ would be the time derivative at 0 of a curve, $\rho (t)$, where $\rho (0)=\rho$. Since we are dealing with a space of probability measures, additional restrictions need to be added in order to make our tangent space well-defined. For example, we would like our trajectory to satisfy the continuity equation ${\frac {\partial \rho }{\partial t}}+\nabla \cdot (\rho v)=0$. There are many such vector fields that solve the continuity equation, so we can restrict to a vector field that minimizes kinetic energy, which is defined as $\int \rho |v|^{2}$. This choice of tangent vectors is justified by the following lemma

Lemma A vector $v\in L^{2}(\rho ;X)$ belongs to the tangent cone at $\rho$ iff

$\lVert v+w\rVert \geq \lVert v\rVert \;\forall w\in L^{2}(\rho ;X)\;{\mbox{such that}}\;\nabla \cdot (w\rho )=0$

where we are taking the $L^{2}(\rho ,X)$ norm. Divergence condition implies that our tangent vectors are equivalent up to a vector field with zero divergence. This implies that $v$ is in fact a gradient of some function $u$, in which case our continuity equation becomes

${\frac {\partial \rho }{\partial t}}+\nabla \cdot (\rho \nabla u)=0$

This is an elliptic partial differential equation, so one can apply tools used for that class of PDEs in order to determine existence and uniqueness of the tangent vectors.

Riemannian Metric Induced by the Wasserstein Metric

Given two tangent vectors at a point $\rho$ in our space, ${\mathcal {P}}_{2}(X)$, we can define the Riemannian metric as follows

$\left\langle {\frac {\partial \rho }{\partial t_{1}}},{\frac {\partial \rho }{\partial t_{2}}}\right\rangle _{\rho }=\int \rho \langle \nabla u_{1},\nabla u_{2}\rangle$

Here, ${\frac {\partial \rho }{\partial t_{1}}},{\frac {\partial \rho }{\partial t_{2}}}$ are tangent vectors at $\rho$, and $u_{1},u_{2}$ are solutions to the modified continuity equation from the previous section. This metric defines an inner product at every point in our space, ${\mathcal {P}}_{2}(X)$. This not only allows one to define geodesics in this space, but the metric can be used to define calculus operators such as gradients and Hessians. These operators can be applied to in a similar manner to the same operators in finite dimensional Riemannian manifolds.