Communications in Mathematical Sciences

Volume 20 (2022)

Number 7

Fast Sinkhorn I: An $O(N)$ algorithm for the Wasserstein-1 metric

Pages: 2053 – 2067

(Fast Communication)

DOI: https://dx.doi.org/10.4310/CMS.2022.v20.n7.a11

Authors

Qichen Liao (Department of Mathematical Sciences, Tsinghua University, Beijing, China; and Theory Lab, Central Research Institute, 2012 Labs, Huawei Technologies Co. Ltd., Hong Kong)

Jing Chen (School of Physical and Mathematical Sciences, Nanyang Technological University, Singapore)

Zihao Wang (Department of Computer Science and Engineering, Hong Kong University of Science and Technology, Clear Water Bay, Hong Kong)

Bo Bai (Theory Lab, Central Research Institute, 2012 Labs, Huawei Technologies Co. Ltd., Hong Kong)

Shi Jin (School of Mathematical Sciences, Institute of Natural Sciences, and MOELSC, Shanghai Jiao Tong University, Shanghai, China)

Hao Wu (Department of Mathematical Sciences, Tsinghua University, Beijing, China)

Abstract

The Wasserstein metric is broadly used in optimal transport for comparing two probabilistic distributions, with successful applications in various fields such as machine learning, signal processing, seismic inversion, etc. Nevertheless, the high computational complexity is an obstacle for its practical applications. The Sinkhorn algorithm, one of the main methods in computing the Wasserstein metric, solves an entropy regularized minimizing problem, which allows arbitrary approximations to the Wasserstein metric with $O(N^2)$ computational cost. However, higher accuracy of its numerical approximation requires more Sinkhorn iterations with repeated matrix-vector multiplications, which is still unaffordable. In this work, we propose an efficient implementation of the Sinkhorn algorithm to calculate the Wasserstein‑1 metric with $O(N)$ computational cost, which achieves the optimal theoretical complexity. By utilizing the special structure of Sinkhorn’s kernel, the repeated matrix-vector multiplications can be implemented with $O(N)$ times multiplications and additions, using the Qin Jiushao or Horner’s method for efficient polynomial evaluation, leading to an efficient algorithm without losing accuracy. In addition, the log-domain stabilization technique, used to stabilize the iterative procedure, can also be applied in this algorithm. Our numerical experiments show that the newly developed algorithm is one to three orders of magnitude faster than the original Sinkhorn algorithm.

Keywords

optimal transport, Wasserstein-1 metric, Sinkhorn algorithm, FS-1 algorithm, fast algorithm

2010 Mathematics Subject Classification

49M25, 65K10

Full Text (PDF format)

This work was supported by National Natural Science Foundation of China grants 11871297 and 12031013, by the Tsinghua University Initiative Scientific Research Program, and by the Shanghai Municipal Science and Technology Major Project 2021SHZDZX0102.

Received 21 February 2022

Received revised 22 May 2022

Accepted 23 May 2022

Published 21 October 2022