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

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.