Faster Sinkhorn’s Algorithm with Small Treewidth

Boyu Liu, Zhao Song, Jiale Zhao, Tianyi Zhou

Keywords: Optimal Transport, Sinkhorn Algorithm, Machine Learning

Abstract: While approximating optimal transport (OT) distances such as the earth mover's distance is a fundamental problem in statistics and machine learning, it is computationally expensive. Given the cost matrix $C=AA^\top$ where $A \in \mathbb{R}^{n \times d}$, the state-of-the-art results [Dvurechensky, Gasnikov, and Kroshnin ICML 2018] cost $\widetilde{O}(\epsilon^{-2} n^2)$ time to approximate OT distance, where $n$ is the size of given two discrete distributions and $\epsilon$ is the error. In this paper, we proposed a faster Sinkhorn's Algorithm to approximate the OT distance when matrix $A$ has treewidth $\tau$, which is usually very small. Our algorithm achieves a running time of $\widetilde{O}(\epsilon^{-2} n \tau)$, improving upon the previous $\widetilde{O}(\epsilon^{-2} n^2)$ time complexity. To the best of our knowledge, our paper is the first work to improve the OT distance approximating problem running time to $\widetilde{O}(\epsilon^{-2} n \tau)$.

Primary Area: learning theory

Submission Number: 9219