PDOT: a Practical Primal-Dual Algorithm and a GPU-Based Solver for Optimal Transport

TL;DR

This paper introduces PDOT, a practical primal-dual algorithm with theoretical guarantees and a GPU-based solver for large-scale optimal transport problems, offering high accuracy and efficiency compared to existing methods.

Contribution

The paper presents a novel primal-dual algorithm with proven complexity bounds and an open-source GPU solver for optimal transport, outperforming traditional algorithms in speed and accuracy.

Findings

PDOT achieves high-accuracy solutions efficiently on GPUs.

Theoretical complexity bounds are established for PDOT.

Numerical experiments show PDOT outperforms Gurobi and Sinkhorn algorithms.

Abstract

In this paper, we propose a practical primal-dual algorithm with theoretical guarantees and develop a GPU-based solver, which we dub PDOT, for solving large-scale optimal transport problems. Compared to Sinkhorn algorithm or classic LP algorithms, PDOT can achieve high-accuracy solution while efficiently taking advantage of modern computing architecture, i.e., GPUs. On the theoretical side, we show that PDOT has a data-independent $\widetilde O(mn(m+n)^{1.5}\log(\frac{1}{\epsilon}))$ local flop complexity where $\epsilon$ is the desired accuracy, and $m$ and $n$ are the dimension of the original and target distribution, respectively. We further present a data-dependent $\widetilde O(mn(m+n)^{3.5}\Delta + mn(m+n)^{1.5}\log(\frac{1}{\epsilon}))$ global flop complexity of PDOT, where $\Delta$ is the precision of the data. On the numerical side, we present PDOT, an open-source GPU solver based on the proposed algorithm. Our extensive numerical experiments consistently demonstrate the well balance of PDOT in computing efficiency and accuracy of the solution, compared to Gurobi and Sinkhorn algorithms.