A Scalable Constant-Factor Approximation Algorithm for $W_p$ Optimal Transport

Pankaj K Agarwal, Oliver Chubet, Sharath Raghvendra, Keegan Yao

Keywords: optimal transport, p-Wasserstein, min-cost matching

TL;DR: We give a randomized $(4+\varepsilon)$-approximate $W_p$-optimal transport plan in $O(n^2+\varepsilon^{-1}n^{3/2}\log\Delta)$, where $\Delta$ is the ratio of the diameter to smallest distance in the input.

Abstract:

Let $(X,d)$ be a metric space and let $\mu,\nu$ be discrete probability distributions supported on finite point sets $A,B \subseteq X$.
For any $p \in [1,\infty]$, the {\it $W_p$-distance} between $\mu$ and $\nu$, $W_p(\mu, \nu)$, is defined as the $p$-th root of the minimum cost of transporting all the probability mass from $\mu$ to $\nu$, where moving a probability mass of $\delta$ from $a \in A$ to $b \in B$ incurs a cost of $\delta d(a,b)^p$.
We give a (Las Vegas) randomized algorithm that computes a $(4+\varepsilon)$-approximate $W_p$ optimal-transport (OT) plan in $O(n^2 + (n^{3/2}\varepsilon^{-1}\log n\log\Delta)^{1+o(1)}\log U)$ time with probability at least $1-1/n$, for all $p \in [1,\infty]$, where $\varepsilon > 0$ is an arbitrarily small constant and $\Delta$ is the ratio between the largest and smallest interpoint distances in $A\cup B$.
The previous best result achieved an $O(\log n)$-approximation in $O(pn^2)$ time, for constant values of $p$.
Our algorithm significantly improves the approximation factor and, importantly, is the first quadratic-time method that extends to the $W_\infty$-distance.
In contrast, additive approximation methods such as Sinkhorn are efficient only for constant $p$ and fail to handle $p=\infty$. \changed{Our algorithm also extends to a query model where, for any integer $k > 1$, we give an algorithm that preprocesses $X$ into clusters in $O(n^2+kn^{1+1/k}\log n\log\Delta)$ time, after which a $O(k)$-approximate $W_p$ distance between any two distributions $\mu$ and $\nu$ with $X$ as support can be computed in $(n^{1+1/k}\log n\log\Delta)^{1+o(1)}$ time with probability at most $1-1/n$.}
Finally, for $p=\infty$, we show that obtaining a relative approximation factor better than $2$ in $O(n^2)$ time would resolve the long-standing open problem of computing a perfect matching in an arbitrary bipartite graph in quadratic time.

Primary Area: optimization

Submission Number: 21439