Fast Algorithms for Computational Optimal Transport and Wasserstein Barycenter

[edit]

Wenshuo Guo, Nhat Ho, Michael Jordan ;
Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics, PMLR 108:2088-2097, 2020.

Abstract

We provide theoretical complexity analysis for new algorithms to compute the optimal transport (OT) distance between two discrete probability distributions, and demonstrate their favorable practical performance compared to state-of-art primal-dual algorithms. First, we introduce the \emph{accelerated primal-dual randomized coordinate descent} (APDRCD) algorithm for computing the OT distance. We show that its complexity is $\bigOtil(\frac{n^{5/2}}{\varepsilon})$, where $n$ stands for the number of atoms of these probability measures and $\varepsilon > 0$ is the desired accuracy. This complexity bound matches the best known complexities of primal-dual algorithms for the OT problems, including the adaptive primal-dual accelerated gradient descent (APDAGD) and the adaptive primal-dual accelerated mirror descent (APDAMD) algorithms. Then, we demonstrate the improved practical efficiency of the APDRCD algorithm through extensive comparative experimental studies. We also propose a greedy version of APDRCD, which we refer to as \emph{accelerated primal-dual greedy coordinate descent} (APDGCD), to further enhance practical performance. Finally, we generalize the APDRCD and APDGCD algorithms to distributed algorithms for computing the Wasserstein barycenter for multiple probability distributions.

Related Material