Scalable Approximation Algorithms for $p$-Wasserstein Distance and Its Variants

We thank the reviewer for their thorough review and constructive feedback. The two main contributions of our paper are as follows:

• A $O(\log n)$ relative approximation algorithm for $p=1$ has been known for almost three decades and has had a significant impact. Despite significant effort, no such algorithm was known for $p=2$. We present the first $O(\log n)$ relative approximation algorithm for the $p$-Wasserstein distance, for any fixed value $p > 1$.

• Under reasonable assumptions, we show that any algorithm that additively approximates $p$-Wasserstein distance requires $\Omega(n^2/\delta^p)$ time. This hardness also extends to other robust variants including $\lambda$-ROBOT and partial $p$-Wasserstein distance. In contrast, the $p$-RPW distance, due to its robustness, admits an $O(n^2/\delta^3)$ time algorithm for any value $p \ge 1$.

We note that our contributions are significant from a theoretical standpoint. The primary reason to implement and test our relative approximation algorithm is to understand the gap between its worst-case theoretical guarantees and its practical performance. Having said that, we can include a comparison of our algorithm with the standard Hungarian algorithm as well as an implementation of our approximation algorithm of $p$-RPW with that of [1].

Comparison with OT.emd: There are two reasons for not making a direct comparison between our algorithm and OT.emd.

• OT.emd requires $O(n^2)$ space which means we can only execute them on small instances. This restricts our ability to compare our algorithm with OT.emd and Sinkhorn on larger inputs and understand how their performance scales with input size.

• Our prototype implementation is written in Python where as OT.emd is based on a highly optimized C++ implementation. Comparing the performance of codes written in different languages, especially only in small instances, can mislead us which is why we have avoided a direct comparison.

Maybe I missed it, but I am not sure what "$an$" and "$an^{(3/2)}$" refer to in Figure 1.b and 1.d.

Response: The plot uses a $\log$–$\log$ scale, where any polynomial function appears as a straight line, with the slope indicating the polynomial’s exponent. In Plot 1.b, we included the functions $f(n) = an$ and $f(n) = an^{3/2}$ (for some constant $a$) to illustrate that, in our experiments, the running time of the Hungarian algorithm with our data structure is bounded by $O(n^{3/2})$.

Did you implement the algorithm to compute RPW?

Response: We show that additive approximation requires $\Omega(n^2/\delta^p)$ time due to their sensitivity to noise and the robustness of $p$-RPW allows for significantly faster additive approximation with an execution time of $O(n^2/\delta^3)$ (for all values of $p$). Our intention was to highlight this gap, which can be accomplished without the need for an implementation. Nonetheless, we are in the process of implementing our algorithm and commit to comparing it with the algorithm of [1] in the next version of our paper.

How does your algorithm introduced to compute RPW differ from the ones in [1]?

Response: The algorithm introduced in [1] computes the OT profile and uses it to approximate the $p$-RPW, which takes $O(n^2/\delta^p+ n/\delta^{2p})$ time. We use a guessing procedure along with an early stopping criteria to reduce the execution time to $O(n^2/\delta^3)$. More precisely, our algorithm picks a guess $g$ from one of the $O(1/\delta)$ values as an approximation of the true $p$-RPW distance. For this value, it executes only $O(p/\delta^2)$ iterations of the LMR algorithm [2]. It then returns the minimum estimate achieved across all the $O(1/\delta)$ guesses. In contrast, the algorithm in [1] runs the LMR algorithm [2] for $O(1/\delta^p)$ iterations.

[1] S. Raghvendra, P. Shirzadian, and K. Zhang. "A New Robust Partial $p$-Wasserstein-Based Metric for Comparing Distributions." ICML 2024.

[2] N. Lahn, D. Mulchandani, and S. Raghvendra. "A graph theoretic additive approximation of optimal transport." NeurIPS 2019.

We thank the reviewer for their thorough review and constructive feedback. We will add a short discussion and a reference about the convergence rate of the Wasserstein distance, as well as the formal definition of spread and the corrected inequality in Appendix A in our next version. We answer the main concern raised by the reviewer below.

My main reason for not providing a higher score is that the presentation seems a bit disjointed, with two related approximation problems solved with completely different techniques. Much of the intro describes noise sensitivity as an issue, but then the majority of text is spent on the first problem.

Response: Noise sensitivity of $p$-Wasserstein distance is the primary reason why both relative and additive approximations are challenging to design. We elaborate on this below and edit the paper to include this discussion.

Relative Approximations: Embedding a metric into a tree metric introduces noise, or distortion, to the edge costs. Due to the high sensitivity, the optimal $p$-Wasserstein distance with respect to these noisy edge costs fails to serve as a relative approximation to the $p$-Wasserstein distance with respect to the original costs, i.e., the approximation quality is unbounded. To address this, we reduce the distortion in edge costs by selecting the smallest distortion across $p$ different tree embeddings. We demonstrate that the optimal $p$-Wasserstein distance computed with respect to these edge costs is in fact a $O(\log n)$-approximation of the true $p$-Wasserstein distance.

Additive Approximations: A small noisy mass of $\delta^p$ can affect the $p$-Wasserstein distance by $\delta$. Consequently, any $\delta$-additive approximation algorithm must transport all but $\delta^p$ mass in an approximately minimum-cost way to ensure a bound of $\delta$ on the additive approximation, which, under reasonable assumptions takes $\Omega(n^2/\delta^p)$ time. In the extreme case, when $p =\infty$, any approximate solver must transport all the mass in an approximately minimum-cost way which requires $\Omega(n^{2.5})$ time. In essence, the sensitivity to noise is precisely why additive approximations for the $p$-Wasserstein distance require $\Omega(n^2/\delta^p)$ time.

In contrast, $p$-RPW is more robust, as a noise of $\delta$ impacts it by at most $\delta$. Our algorithm essentially transports all but $\delta$ mass in an approximate minimum-cost way and has an execution time of $O(n^2/\delta^3)$. Thus, it admits faster algorithms precisely because of its robustness to noise.

In line 207, lev(a,b) is never explicitly defined, I don't think. It is easy enough to infer, but please clarify.

Response: The notation is defined in line 186.

I don't think I understood the second inequality in the derivation 181-189 (second column, for inequality 1). It was my impression that the telescoping sum would eliminate $H_{k^a_0}$ and $H_{k^b_0}$ if $h \geq 2$. I believe the result still holds, but would have expected different terms. Please clarify.

Response: Thanks for raising this point. While the result of the equation in lines 181-189 is correct, the subscript of H in the first line of the equation has to be changed from $j-2$ to $j+2$. Note that in our notation, the root node (the largest cell) is at level 0. The slight typo resulted since the original paper presenting the HST construction [1] had the root at level $h$. We will change $j-2$ to $j+2$ in our next version.

[1] J. Fakcharoenphol, S. Rao, and K. Talwar. "A tight bound on approximating arbitrary metrics by tree metrics." STOC, 2003.

We thank the reviewer for their thorough review and constructive feedback. We answer the main concern raised by the reviewer below.

The importance of $O(\log n)$ approximation factor.

Response: A tree-embedding-based $O(\log n)$-approximation for the $1$-Wasserstein distance originally developed by [1, 2] has been effectively utilized in various applications. These applications include the design of the FlowTree algorithm [3], $k$-NN and LSH data structures [4, 5], $1$-Wasserstein barycenter computation [6, 7], streaming algorithms for $1$-Wasserstein [8], design of $(1+\varepsilon)$-relative approximation algorithms [9, 10] for the $1$-Wasserstein distance, improving the accuracy of additive approximations such as Sinkhorn for the $1$-Wasserstein distance [11], and defining related metrics such as the tree-sliced Wasserstein distance as a proxy for the $1$-Wasserstein distance [12]. For many of these applications, a bound of $O(\log n)$ on the approximation factor is critical.

Despite considerable efforts including some lower bounds [13], there are very few known approximation algorithms for the $p$-Wasserstein distance for $p > 1$ [14, 15]. In this paper, we introduce the first $O(\log n)$-approximation algorithm for computing the $p$-Wasserstein distance for any fixed $p \ge 1$. Extending our algorithm to applications such as $k$-NN under the $2$-Wasserstein distance or boosting the accuracy to $(1+\varepsilon)$-approximation remains an important future direction of research.

[1] M. S. Charikar. "Similarity estimation techniques from rounding algorithms". STOC, 2002.

[2] J. Kleinberg and E. Tardos. "Approximation algorithms for classification problems with pairwise relationships: Metric labeling and Markov random fields." JACM, 2002.

[3] A. Backurs, Y. Dong, P. Indyk, I. Razenshteyn, and T. Wagner. "Scalable nearest neighbor search for optimal transport." ICML, 2020.

[4] P. Indyk and N. Thaper. "Fast image retrieval via embeddings". International Workshop on Statistical and Computational Theories of Vision, 2003.

[5] T. Liu, A. Moore, K. Yang, and A. Gray. "An investigation of practical approximate nearest neighbor algorithms." NeurIPS 2004.

[6] T. Le, V. Huynh, N. Ho, D. Phung, and M. Yamada. "Tree-Wasserstein barycenter for large-scale multilevel clustering and scalable Bayes." ArXiv:1910.04483, 2019.

[7] P. K. Agarwal, S. Raghvendra, P. Shirzadian, and K. Yao. "Efficient Approximation Algorithm for Computing Wasserstein Barycenter under Euclidean Metric." SODA, 2025.

[8] X. Chen, R. Jayaram, A. Levi, and E. Waingarten. "New streaming algorithms for high dimensional EMD and MST." STOC, 2022.

[9] J. Sherman. "Generalized preconditioning and undirected minimum-cost flow." SODA, 2017.

[10] P. K. Agarwal, S. Raghvendra, P. Shirzadian, and K. Yao. "Fast and accurate approximations of the optimal transport in semi-discrete and discrete settings." SODA, 2024.

[11] P. K. Agarwal, S. Raghvendra, P. Shirzadian, and R. Sowle. "A higher precision algorithm for computing the 1-Wasserstein distance." ICLR, 2023.

[12] T. Le, M. Yamada, K. Fukumizu, and M. Cuturi. "Tree-sliced variants of Wasserstein distances." NeurIPS 2019.

[13] A. Andoni, A. Naor, and O. Neiman. "Impossibility of sketching of the 3d transportation metric with quadratic cost." ICALP, 2016.

[14] N. Lahn and S. Raghvendra. "An $O(n^{5/4})$ Time $\varepsilon$-Approximation Algorithm for RMS Matching in a Plane." SODA, 2021.

[15] P. K. Agarwal and J. M. Phillips. "On bipartite matching under the RMS distance." CCCG 2006.

We thank the reviewer for their thorough review and constructive feedback. We will update the paper to include your suggestions. We assume that distances can be queried in $O(1)$ time and will state it explicitly. We will also include additional dependencies of execution time on $p$ that arise due to taking the $p$th power of the costs.

Can the authors elaborate on the $n^{3/2}$ scaling they observe - in particular, do they think there is room for a tighter computational complexity bound?

Response: Yes, it is possible to achieve a bound of $\tilde{O}(n^{3/2})$ for the special case of the minimum-cost bipartite matching problem provided the HSTs are precomputed. Sharathkumar and Agarwal [1] showed that Gabow and Tarjan's algorithm [2] for min-cost bipartite matching can be implemented using $\tilde{O}(n^{3/2})$ queries to the bichromatic closest pair data structure. Combining this with the structure of Section 2, we can bound the execution time by $\tilde{O}(n^{3/2})$ time algorithm. We will include this discussion in the next version of our paper.

[1] R. Sharathkumar and P. K. Agarwal. "Algorithms for the transportation problem in geometric settings." SODA, 2012.

[2] H. Gabow and R. E. Tarjan. "Faster scaling algorithms for network problems." SIAM Journal on Computing 1989.

Since I am not as familiar with their robust variant, I am less sure of the significance of that result. E.g. it could be the case that it is easier to calculate for reasons that make it less useful in practice, though I am not claiming that to be the case.

Response: The $p$-Wasserstein distance's sensitivity to noise is a key factor limiting its practical utility and making the design of an additive approximation particularly challenging.

A small noisy mass of $\delta^p$ can affect the $p$-Wasserstein distance by $\delta$. Consequently, any $\delta$-additive approximation algorithm must transport all but $\delta^p$ mass in an approximately minimum-cost way to ensure a bound of $\delta$ on the additive approximation, which, under reasonable assumptions takes $\Omega(n^2/\delta^p)$ time. In the extreme case, when $p =\infty$, any approximate solver must transport all the mass in an approximately minimum-cost way which requires $\Omega(n^{2.5})$ time. In essence, the sensitivity to noise is precisely why additive approximations for the $p$-Wasserstein distance require $\Omega(n^2/\delta^p)$ time.

In contrast, $p$-RPW is more robust, as a noise of $\delta$ impacts it by at most $\delta$. Our algorithm essentially transports all but $\delta$ mass in an approximate minimum-cost way and has an execution time of $O(n^2/\delta^3)$. Therefore, the $p$-RPW distance is not only robust to noise but also allows for faster algorithms, a direct consequence of this robustness.

Can the authors please confirm that they will publish their code if the paper is accepted?

Response: Yes, we will make our GitHub repository public.

When $p=1$, how close is the robust W1 variant you consider to Dudley's bounded Lipschitz distance (IPM w.r.t Lipschitz and bounded functions)?

Response: For $p = 1$ and in a metric space with a unit diameter, the $1$-RPW can be written as $\sup_{Lip(f) \le 1-||f||_\infty} |\int f d\mu - \int f d\nu|$, which is upper-bounded by Dudley's bounded Lipschitz distance.