A Fast and Accurate Splitting Method for Optimal Transport: Analysis and Implementation

TL;DR

This paper introduces a novel splitting method for large-scale optimal transport problems that is faster, more accurate, and more robust than existing techniques, with proven linear convergence and efficient GPU implementation.

Contribution

The authors develop a direct, non-regularized splitting algorithm for OT that achieves linear convergence and is computationally efficient, outperforming Sinkhorn in speed and accuracy.

Findings

Achieves $O(1/\epsilon)$ iteration complexity, better than Sinkhorn's $O(1/\epsilon^2)$

Provides sparse transport plans and avoids entropic regularization issues

Demonstrates superior performance and robustness in experiments

Abstract

We develop a fast and reliable method for solving large-scale optimal transport (OT) problems at an unprecedented combination of speed and accuracy. Built on the celebrated Douglas-Rachford splitting technique, our method tackles the original OT problem directly instead of solving an approximate regularized problem, as many state-of-the-art techniques do. This allows us to provide sparse transport plans and avoid numerical issues of methods that use entropic regularization. The algorithm has the same cost per iteration as the popular Sinkhorn method, and each iteration can be executed efficiently, in parallel. The proposed method enjoys an iteration complexity $O(1/\epsilon)$ compared to the best-known $O(1/\epsilon^2)$ of the Sinkhorn method. In addition, we establish a linear convergence rate for our formulation of the OT problem. We detail an efficient GPU implementation of the proposed method that maintains a primal-dual stopping criterion at no extra cost. Substantial experiments demonstrate the effectiveness of our method, both in terms of computation times and robustness.