HOT: An Efficient Halpern Accelerating Algorithm for Optimal Transport Problems

TL;DR

This work tackles the computational bottleneck of the Kantorovich–Wasserstein distance for discrete 2D histograms by solving an equivalent reduced OT model with a Halpern-accelerated method (HOT). HOT achieves an $O(1/\varepsilon)$ iteration complexity for the reduced problem, while a novel linear-time procedure solves the involved linear systems without forming $AA^{\top}$, yielding per-iteration cost $O(N)$ and memory $O(N)$, and overall complexity $O(M^{1.5}/\varepsilon)$ for $M$ supports. A transport-plan recovery method extends the reduced-model solution to the original OT problem, enabling practical use in color transfer and related tasks. Implemented in PyTorch, HOT demonstrates strong empirical performance on large-scale benchmarks (e.g., DOTmark), outperforming Sinkhorn-type, Network Simplex, and ADMM approaches in both speed and memory usage, with substantial gains at higher resolutions and in plan-based applications. The results suggest HOT as a scalable tool for OT-based similarity measures and transport-based applications in computer vision and data analysis, with potential extensions to higher dimensions and Wasserstein barycenters.

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.

Figures (8)

• Figure 1: Selected examples of color transfer based on the reduced optimal transport model with the optimal transport plan recovered by Algorithm \ref{['alg:transportplan']}.
• Figure 1: A visualization of the sample dataset.
• Figure 2: A visualization of the selected images from the DOTmark Dataset is presented. The upper row features images from the Classic Images category, while the bottom row contains images from the Shapes category.
• Figure 2: W2NeuralDual on CheckerBoard Using ICNN Potential with default parameter settings.
• Figure 3: Comparison of solving the linear system \ref{['equ:normal']} using Proposition \ref{['prop:linsol_OT']} and the sparse Cholesky decomposition. The time for solving the linear system using Cholesky decomposition is divided into two parts: the time for Cholesky decomposition (orange part) and the time for forward-backward substitution (blue part). For the $256 \times 256$ and $512 \times 512$ cases, Cholesky decomposition is out-of-memory in the test.

Theorems & Definitions (21)

• Proposition 1
• proof
• Remark 1
• Proposition 2
• Proposition 3
• Remark 2
• Proposition 4
• proof
• Corollary 1
• Corollary 2