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$