Optimizing Computational-Statistical Runtime for Wasserstein Distance Estimation

arXiv:2605.20122v1 Announce Type: cross 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-Holder 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 1/2 when d=2 and nearly optimal as alpha o 1 when d = 3.

Source: arXiv cs.LG | 2026-05-20