HOT: An Efficient Halpern Accelerating Algorithm for Optimal Transport Problems

TL;DR

This paper introduces HOT, an accelerated algorithm for optimal transport problems with finite supports, achieving faster computation and efficient plan recovery, especially in 2D with $L_2^2$ ground distances.

Contribution

The paper develops a Halpern accelerated method with linear-time linear system solutions, significantly improving OT problem-solving efficiency and enabling plan recovery.

Findings

HOT algorithm achieves $O(M^{1.5}/\varepsilon)$ complexity.

The linear-time linear system solver enhances computational speed.

Numerical results demonstrate HOT's superior performance over existing methods.

Abstract

This paper proposes an efficient HOT algorithm for solving the optimal transport (OT) problems with finite supports. We particularly focus on an efficient implementation of the HOT algorithm for the case where the supports are in $\mathbb{R}^2$ with ground distances calculated by $L_2^2$-norm. Specifically, we design a Halpern accelerating algorithm to solve the equivalent reduced model of the discrete OT problem. Moreover, we derive a novel procedure to solve the involved linear systems in the HOT algorithm in linear time complexity. Consequently, we can obtain an $\varepsilon$-approximate solution to the optimal transport problem with $M$ supports in $O(M^{1.5}/\varepsilon)$ flops, which significantly improves the best-known computational complexity. We further propose an efficient procedure to recover an optimal transport plan for the original OT problem based on a solution to the reduced model, thereby overcoming the limitations of the reduced OT model in applications that require the transport plan. We implement the HOT algorithm in PyTorch and extensive numerical results show the superior performance of the HOT algorithm compared to existing state-of-the-art algorithms for solving the OT problems.