Optimizing Computational-Statistical Runtime for Wasserstein Distance Estimation

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 $ε$-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 $α$-Hölder 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 $ε$ error in $ε^{-\max(2,\frac{d+1+o(1)}{1+α})}$ time for $0 < α< 1$ Hölder smooth distributions $P,Q$ on $(0,1)^{d}$; an optimal $Θ(ε^{-2})$ for $α> 1/2$ when $d=2$ and nearly optimal as $α\to 1$ when $d = 3$.

• The paper introduces the Sample-Sketch-Solve framework to minimize computational cost while achieving a prescribed statistical accuracy in Wasserstein distance estimation.
• It derives tight grid-based discretization error bounds, demonstrating that the additional error scales as O(h^(1+α)) for smooth density functions.
• The method establishes minimax optimal runtime in low dimensions, outperforming conventional Sinkhorn-based and plug-in estimation approaches.

Optimizing Computational-Statistical Runtime for Wasserstein Distance Estimation

Overview

The paper "Optimizing Computational-Statistical Runtime for Wasserstein Distance Estimation" (2605.20122) addresses the efficient estimation of the squared $2$-Wasserstein distance $W_2^2(P, Q)$ between two probability distributions with access to random sampling. The core focus is minimizing the total computational cost necessary to achieve a given statistical accuracy—an objective formalized as the computational-statistical runtime (CSR). The authors develop a Sample-Sketch-Solve framework that leverages regular gridding and smoothness assumptions on the underlying distributions. Their main contributions include novel discretization error bounds, new algorithmic reductions for efficiently computing Wasserstein distances over grids, and refined CSR bounds in low dimensions under density regularity assumptions.

Motivation and Problem Setting

Estimating Wasserstein distances is a fundamental task in statistics, especially for comparing distributions in generative modeling, hypothesis testing, and parameter estimation. While several computational methods exist for $W_2^2$ between empirical measures, high computational cost remains a bottleneck even in moderate dimensions. Existing algorithms optimize computational complexity for finite sample sets, largely ignoring the statistical accuracy dictated by the finite-sample plug-in estimation error. This work advocates aligning computational strategies with statistical objectives, embracing the perspective that samples are proxies for unknown underlying continuous distributions, and targeting overall additive estimation error in $W_2^2(P, Q)$ up to a prescribed tolerance $\epsilon$.

The CSR is then defined as the total time required to produce an estimate $\widehat{A}(\epsilon)$ satisfying

$$\mathbb{E} |\widehat{A}(\epsilon) - W_2^2(P, Q)| \leq \epsilon,$$

where the cost of drawing each additional sample is $O(1)$, and the algorithm adaptively selects sample size and computational procedure.

Main Contributions

Sample-Sketch-Solve Paradigm

The algorithmic approach involves three steps:

1. Sample: Draw $n$ samples from $P$ and $W_2^2(P, Q)$0 (with $W_2^2(P, Q)$1 determined by desired accuracy and statistical considerations).
2. Sketch: Quantize sample points onto a regular Cartesian grid of suitable resolution. This “sketching” step compresses data with negligible additional error for smooth distributions.
3. Solve: Exactly solve the optimal transport problem on the compressed grid-aggregated empirical measures.

Grid-Based Discretization Error Bounds

A major theoretical innovation lies in the tight analysis of error introduced by grid quantization. The authors show that for ($W_2^2(P, Q)$2)-controlled (i.e., $W_2^2(P, Q)$3-Hölder smooth and uniformly bounded) densities, the additional error due to gridding scales as $W_2^2(P, Q)$4 in grid width $W_2^2(P, Q)$5. This leverages recent regularity results for the Brenier map and intricate cancellation properties in the quantization process. For rough distributions ($W_2^2(P, Q)$6), a universal $W_2^2(P, Q)$7 bound is provided.

Algorithmic Reductions for Grid-Supported Measures

For empirical measures supported on grids, the computation of $W_2^2(P, Q)$8 with squared Euclidean cost is reduced, via [Auricchio et al., 2018], to minimum-cost flow on a $W_2^2(P, Q)$9-partite graph—enabling exact solutions with improved arithmetic runtime. With modern deterministic solvers [Brand et al., 2023], the complexity is $W_2^2$0, where $W_2^2$1 is the number of grid cells per axis and $W_2^2$2 the inverse of the atom mass granularity.

Tight Computational-Statistical Runtime Bounds

Combining the above, the minimal CSR for additive $W_2^2$3 error is shown to be

$W_2^2$4

for distributions with $W_2^2$5-Hölder densities on $W_2^2$6. Notably:

• In $W_2^2$7, for $W_2^2$8, this matches the minimax optimal $W_2^2$9—the rate in one dimension.
• In $W_2^2(P, Q)$0, the exponent approaches $W_2^2(P, Q)$1 as $W_2^2(P, Q)$2 (Lipschitz densities), nearly achieving the conjectural optimum.
• Without smoothness, the best possible is $W_2^2(P, Q)$3.

These rates strictly outperform the previously best-known SOTA algorithms relying on Sinkhorn regularization [Chizat et al., 2020] or direct plug-in estimation, both computationally and statistically, especially in low dimensions and for smooth input measures.

Implications and Future Directions

The practical import is twofold. First, the proposed methods allow high-accuracy Wasserstein distance estimation in settings, such as imaging or spatial data analysis, where $W_2^2(P, Q)$4 or $W_2^2(P, Q)$5 and smoothness is plausible, at computational cost matching the information-theoretic lower bounds. Second, the paradigm highlights the power of adapting algorithmic designs to the smoothness of the input and leveraging discretization structure—a methodology transferable to other transport-type inference problems.

Theoretically, the grid discretization bounds, backed by Brenier regularity under $W_2^2(P, Q)$6-Hölder densities, nearly close the gap between computational and statistical lower limits for Wasserstein estimation. Future research directions include:

• Extending these results to higher dimensions and lower smoothness regimes, identifying whether further algorithmic speedups are obtainable.
• Generalization to other optimal transport costs and divergences, including those with non-Euclidean or learned ground metrics.
• Adapting these ideas to related inference settings, such as integral probability metrics, MMD, or barycenter computation.

Conclusion

This work establishes minimax-rate computational-statistical runtimes for Wasserstein distance estimation in low-dimensional, smooth settings, via a rigorous blend of discretization analysis and algorithmic reduction. The Sample-Sketch-Solve framework delineates the fundamental role of input regularity in optimal transport estimation, providing both actionable algorithms and theoretical understanding that sharpen the boundaries of tractable statistical inference with Wasserstein distances.