A Push-Relabel Based Additive Approximation for Optimal Transport

TL;DR

This paper introduces a simple, parallelizable combinatorial algorithm based on push-relabel for efficiently approximating the Optimal Transport distance, outperforming existing methods like Sinkhorn in speed and providing additional transport plan insights.

Contribution

The authors develop a novel push-relabel based algorithm for approximate OT that is faster, parallelizable, and yields a transport plan, advancing computational efficiency in optimal transport problems.

Findings

Achieves near-optimal $O(n^2/\varepsilon^2)$ runtime for general OT

Improves to $O(n^2/\varepsilon)$ for assignment problems

Faster than Sinkhorn algorithm in CPU and GPU implementations

Abstract

Optimal Transport is a popular distance metric for measuring similarity between distributions. Exact algorithms for computing Optimal Transport can be slow, which has motivated the development of approximate numerical solvers (e.g. Sinkhorn method). We introduce a new and very simple combinatorial approach to find an $\varepsilon$-approximation of the OT distance. Our algorithm achieves a near-optimal execution time of $O(n^2/\varepsilon^2)$ for computing OT distance and, for the special case of the assignment problem, the execution time improves to $O(n^2/\varepsilon)$. Our algorithm is based on the push-relabel framework for min-cost flow problems. Unlike the other combinatorial approach (Lahn, Mulchandani and Raghvendra, NeurIPS 2019) which does not have a fast parallel implementation, our algorithm has a parallel execution time of $O(\log n/\varepsilon^2)$. Interestingly, unlike the Sinkhorn algorithm, our method also readily provides a compact transport plan as well as a solution to an approximate version of the dual formulation of the OT problem, both of which have numerous applications in Machine Learning. For the assignment problem, we provide both a CPU implementation as well as an implementation that exploits GPU parallelism. Experiments suggest that our algorithm is faster than the Sinkhorn algorithm, both in terms of CPU and GPU implementations, especially while computing matchings with a high accuracy.