Optimizing Computational-Statistical Runtime for Wasserstein Distance Estimation

TL;DR

This paper introduces a new approach for estimating Wasserstein distance efficiently from samples, especially for smooth distributions, by using a grid sketch and regularization to improve computational-statistical runtime.

Contribution

The authors develop a Sample-Sketch-Solve paradigm that compresses data via a grid sketch, enabling faster Wasserstein distance estimation with theoretical guarantees for smooth distributions.

Findings

Achieves $ ilde{O}( ext{error}^{- ext{exponent}})$ runtime for Wasserstein distance estimation.

Provides near-optimal runtime bounds for certain smoothness conditions in low dimensions.

Introduces a regular grid sketch that preserves accuracy while enabling faster algorithms.

Abstract

Squared Wasserstein distance is a frequently used tool to measure discrepancy between 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)$, algorithms for Wasserstein distance computation 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 is to estimate from samples the Wasserstein distance between potentially smooth measures up to $\epsilon$-additive error in expectation with respect to the sampling; we allow $O(1)$ computational cost for collecting a 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 $\epsilon$ error in $\epsilon^{-\max(2,\frac{d+1+o(1)}{1+\alpha})}$ time for $0 < \alpha < 1$ H\"older smooth distributions $P,Q$ on $(0,1)^{d}$; an optimal $\Theta(\epsilon^{-2})$ for $\alpha > 1/2$ when $d=2$ and nearly optimal as $\alpha \to 1$ when $d = 3$.