A Push-Relabel Based Additive Approximation for Optimal Transport
This provides a faster and more practical solution for machine learning practitioners needing efficient OT computations, though it is incremental as it builds on existing push-relabel methods.
The paper tackles the computational inefficiency of exact Optimal Transport algorithms by introducing a new combinatorial approach based on the push-relabel framework, achieving near-optimal execution times of O(n^2/ε^2) for OT distance and O(n^2/ε) for the assignment problem, with faster performance than the Sinkhorn algorithm in CPU and GPU implementations, especially for high accuracy.
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.