Technical Summary: Optimizing Computational-Statistical Runtime for Wasserstein Distance Estimation

Problem Statement

The squared Wasserstein distance, $W_2^2(P, Q)$, is a standard metric for measuring discrepancy between probability distributions $P$ and $Q$ in $\mathbb{R}^d$. In practical data analysis and machine learning tasks (e.g., parameter estimation, hypothesis testing), these distributions are unknown and must be estimated from empirical samples $\hat{P}_n$ and $\hat{Q}_n$ of size $n$.

Existing approaches face a fundamental bottleneck: computing $W_2^2(\hat{P}_n, \hat{Q}_n)$ exactly or with additive error $\epsilon$ typically requires runtime that scales poorly with $n$ and $\epsilon$ (e.g., $O(n^2/\epsilon^2)$ or worse). Conversely, purely statistical approaches that rely on the convergence rate of empirical measures often ignore the computational cost of solving the optimal transport problem on the resulting large datasets.

This paper addresses the Computational-Statistical Runtime (CSR). The goal is to design a randomized algorithm that estimates $W_2^2(P, Q)$ within $\epsilon$-additive error in expectation, where the runtime $O(f(\epsilon))$ accounts for both the cost of collecting samples (assumed $O(1)$ per sample) and the cost of computation. The authors argue that the target accuracy should align with the statistical error induced by the sample size, rather than treating the empirical measures as ground truth.

Methodology: Sample-Sketch-Solve Paradigm

The authors propose a three-phase paradigm to achieve fast CSRs, particularly for distributions with Hölder smooth densities on the domain $(0, 1)^d$:

1. Sample: Collect $n$ random samples from $P$ and $Q$ to form empirical measures $\hat{P}_n$ and $\hat{Q}_n$.
2. Sketch: Map these samples onto a regular Cartesian grid $G_h$ with cell width $h = \Theta(\epsilon^{-\frac{1}{1+\alpha}})$. All mass within a grid cell is collapsed to the cell center. This creates compressed discrete measures $G_h\#\hat{P}_n$ and $G_h\#\hat{Q}_n$.
3. Solve: Compute the exact squared Wasserstein distance between the sketched measures using a specialized algorithm that exploits the regular grid structure.

Key Technical Components

1. Grid-Based Discretization Error Bounds (Section 2)

The paper establishes a new upper bound on the approximation error introduced by collapsing mass to grid centers.

• Assumption: $P$ and $Q$ possess $\alpha$-Hölder smooth densities ($0 < \alpha < 1$) and are bounded away from zero and infinity.
• Result: The authors prove that $|W_2^2(G_h\#P, G_h\#Q) - W_2^2(P, Q)| \leq K h^{1+\alpha}$.
• Mechanism: Unlike standard triangle inequality arguments which yield $O(h)$, this result utilizes a novel coupling argument based on the regularity of the Brenier map (the optimal transport map). By leveraging recent results from Collins and Tong [2025] regarding the global Hölder continuity of the Brenier map, the authors show that first-order error terms cancel out (similar to a midpoint rule), resulting in a tighter $O(h^{1+\alpha})$ bound.
• Non-Smooth Case: For distributions without smoothness assumptions, the bound reverts to $O(h)$, consistent with standard triangle inequality arguments.

2. Efficient Exact Computation on Regular Grids (Section 3)

Once the data is sketched onto a regular grid, the problem structure allows for significantly faster exact computation.

• Structure: The sketched measures are supported on a grid with $L$ divisions per axis, where each atom has a weight that is a multiple of $1/n$. The ground cost (squared Euclidean distance) is separable across dimensions.
• Reduction: The authors utilize a reduction by Auricchio et al. [2018] to transform the optimal transport problem between $d$-dimensional histograms with separable costs into a minimum-cost flow problem on a $(d+1)$-partite graph.
• Algorithm: By scaling the problem to ensure integral demands, capacities, and costs, they apply the deterministic exact minimum-cost flow algorithm of van den Brand et al. [2023].
• Complexity: The exact computation on the sketch runs in $\tilde{O}(L^{d+1+o(1)})$ time.

3. Synthesis: The CSR Algorithm (Section 4)

The final algorithm combines the statistical sampling, the grid sketching, and the fast solver.

• Sample Size: The number of samples $n$ is chosen to balance statistical error and discretization error, typically $n = \tilde{O}(\epsilon^{-\max(2, d/2)})$.
• Grid Resolution: The grid size is set to $L = \Theta(\epsilon^{-\frac{1}{1+\alpha}})$.
• Total Runtime: The combined runtime is dominated by the solver step, yielding a Computational-Statistical Runtime of:
  $$\tilde{O}\left(\epsilon^{-\max\left(2, \frac{d+1+o(1)}{1+\alpha}\right)}\right)$$

Key Results and Contributions

1. Improved Runtime for Smooth Distributions

The paper demonstrates that for $\alpha$-Hölder smooth distributions, the CSR can be nearly optimal, matching the statistical lower bound of $\Omega(\epsilon^{-2})$ in specific regimes:

• Dimension $d=2$: For $\alpha > 1/2$, the algorithm achieves $\tilde{O}(\epsilon^{-2})$. This matches the optimal statistical rate, meaning the computational cost does not exceed the cost of gathering sufficient samples.
• Dimension $d=3$: As $\alpha \to 1$ (Lipschitz smoothness), the runtime approaches $\tilde{O}(\epsilon^{-2})$. For general $0 < \alpha < 1$, the runtime is $\tilde{O}(\epsilon^{-\frac{4}{1+\alpha}})$.
• General Dimension $d$: The runtime is $\tilde{O}(\epsilon^{-\max(2, \frac{d+1}{1+\alpha})})$.

2. Comparison with Prior Work

The authors compare their results against existing methods in Table 2:

• Empirical Plug-in + Lahn et al. [2019]: Achieves $\tilde{O}(\epsilon^{-4})$ in $d=2$ and $\tilde{O}(\epsilon^{-5})$ in $d=3$.
• Entropic OT / Sinkhorn (Chizat et al. [2020]): Achieves $\tilde{O}(\epsilon^{-6})$ in $d=2$ and $d=3$.
• This Work: Achieves $\tilde{O}(\epsilon^{-2})$ in $d=2$ (for $\alpha > 1/2$) and $\tilde{O}(\epsilon^{-2})$ in $d=3$ (as $\alpha \to 1$).

3. Non-Smooth Case

In the absence of smoothness assumptions ($\alpha \to 0$), the method yields a CSR of $\tilde{O}(\epsilon^{-\max(2, d+1)})$. For $d=2$, this is $\tilde{O}(\epsilon^{-3})$, and for $d=3$, it is $\tilde{O}(\epsilon^{-4})$.

Significance and Claims

The paper claims that its primary contribution is the Sample-Sketch-Solve paradigm, which successfully decouples the statistical estimation error from the computational complexity of the transport problem.

• Compression without Error: The authors show that mapping samples to a regular grid compresses the data without increasing the asymptotic error, provided the underlying distributions are smooth.
• Regularization for Speed: The grid structure regularizes the problem, enabling the use of fast exact flow algorithms that are intractable for general discrete point sets.
• Optimality: The work argues that for sufficiently smooth distributions in low dimensions ($d=2, 3$), it is possible to achieve the optimal statistical rate of $\epsilon^{-2}$ computationally. This challenges the notion that computing Wasserstein distances inherently requires super-linear or super-quadratic runtime relative to the desired precision.

The authors explicitly state that their results rely on the assumption that $P$ and $Q$ are supported on a bounded domain (specifically $(0, 1)^d$) and possess Hölder smooth densities. They do not claim these results extend to non-smooth distributions or unbounded domains without modification. The significance lies in bridging the gap between statistical estimation theory and computational optimal transport, demonstrating that under realistic smoothness assumptions, the "computational bottleneck" can be effectively removed.