Formal Riemannian Structure of the Wasserstein metric
Given a closed and convex space , two probability measures on the same space, , the 2-Wasserstein metric is defined as
where is the set of transport plans from to . The Wasserstein metric is indeed a metric in the sense that it satisfies the desired properties of a distance function between probability measures on . Moreover, the Wasserstein metric can be used to define a formal Riemannian metric on . 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 would be the time derivative at 0 of a curve, , where . 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 . 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 . This choice of tangent vectors is justified by the following lemma
- Lemma A vector belongs to the tangent cone at iff
where we are taking the norm. Divergence condition implies that our tangent vectors are equivalent up to a vector field with zero divergence. This implies that is in fact a gradient of some function , in which case our continuity equation becomes
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 in our space, , we can define the Riemannian metric as follows
Here, are tangent vectors at , and are solutions to the modified continuity equation from the previous section. This metric defines an inner product at every point in our space, . 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.