$L^2$ over Wasserstein: Statistical Analysis for Optimal Transport

TL;DR

This paper develops a statistical framework for optimal transport in the space of random probability measures, extending Wasserstein geometry to account for uncertainty and enabling new inference methods.

Contribution

It introduces the $L^2$ over Wasserstein space, establishing its geometric structure and applying it to statistical convergence, Bayesian consistency, and transformer models.

Findings

Established the Riemannian structure of $L^2$ over Wasserstein space.

Proved convergence results for empirical measures within this framework.

Refined Bayesian posterior convergence in Wasserstein topology.

Abstract

Optimal transport provides an inherently geometric and highly structured framework for studying spaces of probability measures, supplying a rich theoretical toolkit for contemporary statistics, machine learning, and generative modelling. In applications, however, the measures of interest are almost never known precisely, calling for a theory of optimal transport that accounts for statistical uncertainty. We construct such a framework, lifting the classical theory to the setting of random probability measures. We introduce the $L^2$ over Wasserstein space establishing that it inherits the formal Riemannian structure of the Wasserstein space by characterising distances and geodesic geometry. The structure induces random flows with Wasserstein gradient flow sample paths, making it the natural extension of the Wasserstein space which allows for random gradient flow dynamics. We ensemble statistical convergence results of the optimal transport machinery using the empirical measure within the $L^2$ over Wasserstein framework. Moreover, in the setting of Bayesian non-parametrics, we refine Schwartz's consistency theorem to the Wasserstein topology and deduce posterior convergence of the same machinery in the $L^2$ over Wasserstein space. We demonstrate that the growing theory of random token sampling for transformer models using self-attention flow paths can be embedded into the our framework. The results provide a unified treatment of random optimal transport and its consequences for principled inference and generative modelling under the statistical uncertainty of random sampling.