ICML2025
Scalable Approximation Algorithms for p-Wasserstein Distance and Its Variants
Nathaniel Lahn, Sharath Raghvendra, Emma Saarinen, Pouyan Shirzadian
摘要
The p-Wasserstein distance measures the cost of optimally transporting one distribution to another, where the cost of moving a unit mass from a to b is the p th power of the ground distance d(a, b) between them. Despite its strong theoretical properties, its use in practice -especially for p ≥ 2is limited due to two key challenges: sensitivity to noise and a lack of scalable algorithms. We identify noise sensitivity as a key reason why some existing approximation algorithms for p = 1 fail to generalize to p ≥ 2 and then present new algorithms for approximating the p-Wasserstein distance and its variant. First, when d(•, •) is a metric, for any constant p ≥ 2, we present a novel relative O(log n)approximation algorithm to compute the p-Wasserstein distance between any two discrete distributions of size n. The algorithm runs in O(n 2 log U log ∆ log n) time, where log U is the bit-length of the input probabilities and ∆ is the ratio of the largest to the smallest pairwise distance. We use p hierarchically well-separated trees to define a distance that approximates the p-Wasserstein cost within a factor of O(log n) and then present a simple primal-dual algorithm to compute the p-Wasserstein cost with respect to this distance. Second, due to the noise sensitivity of the p-Wasserstein distance, we show that existing combinatorial approaches require Ω(n 2 /δ p ) time to approximate the p-Wasserstein distance within an additive error of δ. In contrast, we show that, for any arbitrary distance d(•, •), a recent noise-resistant variant of the p-Wasserstein distance, called the p-RPW distance, can be approximated in O(n 2 /δ 3 ) time.