Optimizing Computational-Statistical Runtime for Wasserstein Distance Estimation

Squared Wasserstein distance is a frequently used tool to compare probability distributions. This distance is typically computed between empirical measures of size $n$ from two underlying random samples. Unfortunately, even in lower dimensional Euclidean space problems $\left( d \in \{2,3\} \right)$, Wasserstein distance algorithms with approximate or exact precision guarantees scale poorly in the runtime as a function of $n$ and the desired precision. In response, we consider the computational-statistical runtime, where the goal to estimate from samples the Wasserstein distance up to the $\varepsilon$-additive error which is achievable under a sample; we allow $O(1)$ time per sample. Towards this, we develop a Sample-Sketch-Solve paradigm where we introduce a regular cartesian grid sketch of the samples. We show that (especially under $\alpha$-H\"older smooth distributions) this can compress the data without increasing asymptotic error, and also regularizes the structure which enables faster exact algorithms. Ultimately, we approximate $W_2^2(P,Q)$ within $\varepsilon$ error in $\varepsilon^{-\max(2,\frac{d+1+o(1)}{2})}$ time for Lipschitz $P,Q$ on $(0,1)^d$; nearly an optimal $\Theta(\varepsilon^{-2})$ for $d \leq 3$.