HOT : An Efficient Halpern Accelerating Algorithm for Optimal Transport Problems

This is the official implementation of the HOT solver. Our paper has been published online at IEEE Transactions on Pattern Analysis and Machine Intelligence.

Overview

HOT is designed to solve reduced 2D optimal transport model which can be formulated as a standard linear programming problem:

$$ \begin{array}{ll} \min\limits_{x \in \mathbb{R}^N} & \langle{c}, {x}\rangle + \delta_{\mathbb{R}_{+}^{N}}(x)\\ \text { s.t. } &{A} {x}={b}. \end{array} $$

Getting Started

Prerequisites

First, install the necessary packages listed in requirements.txt. For the PyTorch package, we encourage you to visit the official PyTorch website to download the version compatible with your CUDA Toolkit.

pip install -r requirements.txt

Example: Solving an OT Problem Between Two Grayscale Images

python -m demo.py

The pipeline of this demo is like following:

1. Specify the paths to the two images used for the OT problem.
2. Specify the desired crop size $(m, n)$ of two images.
3. Read the images, convert them to grayscale images, resize and perform $\ell_{1}$ normalization.
4. Generate $b$ and $c$.
5. Exucute the HOT Solver.

For computing Kantorovich-Wasserstein distance between two images, you only need to $${\color{blue}\textbf{update the image paths and desired crop size in demo.py accordingly}}$$. You can also adapt this pipeline to suit your project's needs.

Results

The HOT solver runs $${\color{red}\textbf{fast}}$$! 😄

Apart from the high computational efficiency, HOT also enjoys an $${\color{red}\textbf{advantage on memory}}$$ compared to most other popular algorithms. 😁

$${\color{red}\textbf{The experiments were conducted on RTX 4090.}}$$

Color Transfer

The ColorTransfer folder contains a Jupyter notebook, color_transfer.ipynb, which provides a detailed walkthrough of performing color transfer using the HOT Solver.

Citations

@article{zhang2025hot,
  title={HOT: An efficient Halpern accelerating algorithm for optimal transport problems},
  author={Zhang, Guojun and Gu, Zhexuan and Yuan, Yancheng and Sun, Defeng},
  journal={IEEE Transactions on Pattern Analysis and Machine Intelligence},
  year={2025},
  publisher={IEEE}
}

Star History