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 a transport plan 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 Riemannian metric on ${\mathcal {P}}_{2}(X)$. Such a metric 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)\;\mathrm {suchthat} \;\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$