## HOT: An Efficient Halpern Accelerating Algorithm for Optimal Transport Problems

Guojun Zhang, Zhexuan Gu, Yancheng Yuan, Defeng Sun

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 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 ε -approximate solution to the optimal transport problem with M supports in O M ( 1 5 . /ε ) 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 map. 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.

Index Terms -Optimal transport, Kantorovich-Wasserstein distance, Halpern iteration, Acceleration, Computational complexity.

## I. INTRODUCTION

T HE Kantorovich-Wasserstein (KW) distance has become a primer choice for measuring the similarity between two probability distributions. It has demonstrated remarkable success in various applications, including color transfer [1][3], texture synthesis and mixing [4], registration and warping [5], transport-based morphometry [6], and hypothesis testing [7], among others. Despite its powerful geometric framework for comparing probabilities, the KW distance is computationally expensive in general [8], [9]. Specifically, it requires solving an optimal transport (OT) problem, which is a (largescale) linear programming (LP) in a discrete setting. Standard methods, such as the simplex method and the interior point method, suffer from high computational complexity relative to the problem size. Furthermore, these methods are difficult to parallelize, which can hardly benefit from the modern powerful graphics processing units (GPUs). Consequently, solving the LP problem of the discrete OT problem remains daunting in modern data-driven applications due to high computational and memory costs. This paper addresses these two challenges by proposing an efficient and easily parallelizable algorithm for solving an equivalent reduced model of the OT problem.

Guojun Zhang and Zhexuan Gu contribute equally to this manuscript. Corresponding author: Yancheng Yuan.

The authors are with The Hong Kong Polytechnic University. Email: guojun.zhang@connect.polyu.hk, zhexuan.gu@connect.polyu.hk, yancheng.yuan@polyu.edu.hk, defeng.sun@polyu.edu.hk.

A. Related work and existing challenges

When it comes to computing the KW distance between two discrete probability distributions with M supports, there are mainly two popular approaches: (i) Computing an approximated KW distance by solving the optimal transport problem with an additional entropy regularization [10]; (ii) Computing the KW distance via solving the corresponding LP problem [11]. The readers can refer to [9] and the references therein for a more detailed discussion of the algorithms for solving OT problems. Before introducing our new algorithm, we briefly discuss the challenges of the aforementioned approaches.

Challenges with the Entropy-regularized approach: Due to its scalability, the Sinkhorn algorithm, and its improved versions have been widely adopted to compute an approximation of the KW distance in applications [10], [12]-[14]. In particular, the Sinkhorn algorithm can efficiently solve the regularized OT problem when the regularization parameter is moderate (i.e., no less than 10 -2 ). However, a high-quality approximation of the KW distance is important for better performance in many applications, which requires solving the regularized OT problem with a small regularization parameter. Unfortunately, a small regularization parameter will usually cause numerical issues and a slower convergence for the Sinkhorn algorithm. Some stabilized and rescaling techniques [15] have been proposed to improve the robustness of solving the regularized OT problems, but the efficiency of the stabilized algorithms is unsatisfactory compared to the Sinkhorn algorithm.

Challenges with the LP approach: Along this line, the interior point method [16] and the network simplex method [17], [18] are popular choices for obtaining solutions with high accuracy to the moderate scale LP problem. Recently, a semismooth Newton based inexact proximal augmented Lagrangian method [19] has been proposed for solving linear programming problems which can solve the OT problem as a special case. The semismooth Newton based algorithm can exploit the sparsity of the solution by the generalized Jacobian and show superior numerical performance compared to Gurobi in some examples. However, these solvers are not applicable for solving the problem on a very large scale due to the high computational complexity. The urgent need to solve large-scale OT problems in applications inspires extensive research in designing first-order methods. Among them, the Douglas-Rachford splitting algorithm [20] and primal-dual hybrid gradient method (PDHG) [21]-[24]. Jambulapati et al. [25] proposed an algorithm based on a dual-extrapolation algorithm to achieve an ˜ O M /ε ( 2 ) computational complexity

Source

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00598v1-figure-0264ddd9-2e54-4aec-b1ae-2a43f3635a5f.jpg)

Target

Result

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00598v1-figure-0d394f64-ab59-400e-b798-cd9c7b88e588.jpg)

Fig. 1: Selected examples of color transfer based on the reduced optimal transport model with the optimal transport map recovered by Algorithm 1.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00598v1-figure-70ea3124-71e6-4981-af47-c9b682c588ff.jpg)

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00598v1-figure-03244a30-e1c2-48d3-8e6b-fc8a70dd0083.jpg)

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00598v1-figure-4b054d2e-c7e0-47ba-9c97-79919643fdae.jpg)

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00598v1-figure-8a6a4435-aa3f-483a-b43f-17405c2a63fd.jpg)

bound for obtaining an ε -approximate solution (in terms of objective function value) 1 . Recently, Zhang et al. [26] proposed an efficient Halpern-Peaceman-Rachford (HPR) algorithm for solving the OT model and the Wasserstein barycenter problem, which can obtain an ε -approximate solution (in terms of the Karush-Kuhn-Tucker (KKT) residual) of the OT model in O M /ε ( 2 ) flops. We summarize some known complexity results for solving the OT problem in Table I. Readers can refer to [26]-[28] and the references therein for a more detailed discussion.

TABLE I: Selected known complexity results for solving OT problem ( C represents the largest elements of the cost matrix, while R denotes the distance between the initial point and the solution set.)

| Algorithm                                                                                          | Complexity result                                             |
|----------------------------------------------------------------------------------------------------|---------------------------------------------------------------|
| Sinkhorn [13] APDAGD [13], [29] Greenkhorn [29] Sinkhorn AAM [14] Dual extrapolation [25] HPD [30] | ˜ O ( M 2 C 2 /ε 2 ) ˜ O ( M 2 . 5 C/ε ) ˜ O ( M 2 C 2 /ε 2 ) |
| Accelerated [29]                                                                                   | ˜ O ( M 7 / 3 C 4 / 3 /ε 4 / 3 )                              |
|                                                                                                    | ˜ O ( M 2 . 5 C/ε )                                           |
|                                                                                                    | ˜ O ( M 2 C/ε )                                               |
|                                                                                                    | ˜ O ( M 2 . 5 C/ε )                                           |
| HPR [26]                                                                                           | O ( M 2 R/ε )                                                 |
| HOT (This paper)                                                                                   | O ( M1 . 5 R /ε )                                             |

Beyond the challenges above in computational efficiency, all these algorithms for the original OT problem with M supports require at least a memory cost of O M ( 2 ) . This memory cost makes it forbidden to compute the KW distance of two distributions with massive supports (i.e., the OT problem for computing the KW distance of two 512 × 512 grey images has more than 6 8 . × 10 10 variables). This motivates researchers to explore possible reduced models for the OT problem. When the ground distances between supports in R d are calculated by L 1 -norm, Ling and Okada proposed an equivalent reduced model with O dM ( ) memory cost to calculate the earth mover distance (equivalent to the KW distance) [31]. Recently, Auricchio et al. [32] extended the idea to the case where the ground distances between supports are calculated by L 2 2 -norm and proposed an equivalent reduced model with O dM ( d +1 d ) memory cost. The authors in [32] adopted the Network Simplex method to solve the reduced model and demonstrated superior performance in terms of computational and memory efficiency compared to the Sinkhorn and the improved Sinkhon method [3], [33] on examples of moderate scale. Unfortunately, the efficiency becomes unsatisfactory for very large-scale problems (see Section IV for details). Moreover, the transport plan is not available if we solve the reduced model, which is critical for a wide class of applications, such as color transfer [1], [34] and texture synthesis [4].

1 Despite its better complexity bound, the empirical performance of this algorithm is not as efficient as other algorithms, such as accelerated gradient method [13], [14].

## B. Contributions

Motivated by the recent advancements in the accelerated algorithms based on Halpern iteration [26], [35]-[39], we propose an efficient H alpern accelerating method for solving the reduced OT problem, which is abbreviated as 'HOT' for convenience, to address the challenges in computing the KW distance with a finite number of supports. We particularly focus on an efficient implementation of the HOT algorithm for the case where the supports are in R 2 with ground distances calculated by L 2 2 -norm, which includes a wide class of applications as aforementioned. Specifically, HOT adopts a first-order algorithm with Halpern acceleration to solve the equivalent reduced OT model, which can obtain an ε -approximate solution in O (1 /ε ) iterations [26], [35]. More importantly, we design a fast procedure for solving the subproblems with linear time complexity by fully exploiting the problem structure. This also makes the popular alternating direction method of multiplier (ADMM) [40], [41] scalable for solving the reduced OT model. Overall, our proposed HOT algorithm can compute an ε -approximation of the KW distance between two histograms supported on M = m × n bins within O (( m n 2 + n m /ε 2 ) ) flops. This is the best-known computational complexity for computing an approximate KW distance to our knowledge. Moreover, we propose an efficient algorithm to recover a transport plan based on the obtained solution of the reduced model, which releases the power of the reduced model in applications. We implement HOT in PyTorch and extensive numerical results will be shown in Section IV to demonstrate the superior and robust performance of HOT for computing the KW distance, compared to state-of-the-art algorithms, including Sinkhorn [10], Improved Sinkhorn [3], [33], Network Simplex method [17], [18], ADMM [40], [41], interior point method (in Gurobi).

We summarize the main contributions of this paper as follows:

- 1) We propose an efficient HOT algorithm for solving the reduced model of the OT problem with an attractive O (1 /ε ) iteration complexity guarantee with respect to the KKT residual.
- 2) We designed a highly efficient algorithm for solving the subproblems of the HOT algorithm with linear time complexity.
- 3) We propose an efficient algorithm to recover a transport plan based on the obtained solution of the reduced model, which removes the restriction of the reduced model in applications requiring a transport plan.
- 4) We implement the HOT algorithm in PyTorch, which supports both CPU and GPU computation and is userfriendly for researchers in the machine learning community.
- 5) Extensive numerical testings are conducted and presented to demonstrate the efficiency of the HOT algorithm.

The rest of the paper is organized as follows. We introduce the equivalent reduced OT model in Section II. This section also includes an efficient procedure for recovering the transport plan from a solution to the reduced OT problem. The HOT algorithm and its computational complexity guarantees will be presented in Section III. We present extensive numerical results in Section IV and conclude the paper in Section V.

Notation. We denote the n -dimensional real Euclidean space as R n and the nonnegative orthant of R n as R n + . For any x ∈ R n and y ∈ R n , we define ⟨ x, y ⟩ := ∑ n i =1 x y i i and ∥ x ∥ := √∑ n i =1 x 2 i , respectively. Additionally, let 1 m (resp. 0 m ) denote the m -dimensional vector with all entries being 1 (resp. 0). For a given matrix A ∈ R m n × , we denote A ⊤ ∈ R n × m as its transpose. For a collection of matrices { A , . . . , A 1 m } , we denote the block diagonal matrix with diagonal blocks A i as diag( A , . . . , A 1 m ) . A 1 ⊗ A 2 stands for the Kronecker product of matrices A 1 and A 2 . Moreover, for a closed convex set C , we denote the indicator function of C and the Euclidean projector over C as δ C and Π ( ) := arg min C x z ∈ C ∥ x - ∥ z , respectively.

## II. KANTOROVICH-WASSERSTEIN DISTANCES

In this section, we first introduce an equivalent reduced model of the OT problem for computing the KW distance between two-dimensional histograms. Subsequently, we present a fast and easily implementable algorithm to reconstruct the optimal transport plan of the original OT model from a solution of this reduced model.

## A. An equivalent reduced model of the OT problem

In the following discussion, we assume two-dimensional histograms for simplicity. As previously mentioned, these histograms are widely used in applications as shape and image descriptors. Without loss of generality, we adopt the following assumptions and notations:

- 1) Histograms have supports in M = m × n bins with m rows and n columns;
- 2) The index set for bins is defined as I = ( { i, j ) | 1 ≤ i ≤ m, 1 ≤ j ≤ n } . We use ( i, j ) to denote a bin or a node corresponding to it;
- 3) µ 1 and µ 2 are the two histograms to be compared, where each histogram µ k is defined as { µ k i,j | µ k i,j ≥ 0 ( , i, j ) ∈ I , ∑ ( i,j ) ∈I µ k i,j = 1 } for k = 1 2 , .

With these notations and assumptions, the discrete OT problem for computing KW distance between histograms µ 1 and µ 2 can be defined as follows:

$$& \text{and} \, \mu \prec \text{ can be dense as follows:} \\ & \min _ { \pi } \quad \sum _ { ( i, j ) \in \mathcal { I } } \sum _ { ( k, l ) \in \mathcal { I } } c _ { i, j ; k, l } \pi _ { i, j ; k, l } \\ & \quad \sum _ { ( k, l ) \in \mathcal { I } } \pi _ { i, j ; k, l } = \mu _ { i, j } ^ { 1 }, \quad \forall ( i, j ) \in \mathcal { I }, \\ & \text{s.t.} \quad \sum _ { ( i, j ) \in \mathcal { I } } \pi _ { i, j ; k, l } = \mu _ { k, l } ^ { 2 }, \quad \forall ( k, l ) \in \mathcal { I }, \\ & \pi _ { i, j ; k, l } \geq 0, \quad \forall ( i, j ) \in \mathcal { I }, \ \forall ( k, l ) \in \mathcal { I }, \\ & \text{where} \, \pi \text{ is the transport plan between histograms } \mu ^ { 1 } \text{ and } \mu ^ { 2 }.$$

where π is the transport plan between histograms µ 1 and µ 2 . The ground distance c i,j k,l ; is commonly defined by the L p p distance:

$$\text{tion} _ { \text{sort} } \quad c _ { i, j ; k, l } = \| ( i, j ) ^ { \top } - ( k, l ) ^ { \top } \| _ { p } ^ { p } = ( | i - k | ^ { p } + | j - l | ^ { p } ). \ \ ( 2 )$$

In this paper, we focus on the case where p = 2 . By exploiting the separable structure of the ground distance, Auricchio et al.

[32] proposed the following equivalent model in terms of the optimal objective function value:

$$\int _ { f ^ { ( 1 ) }, f ^ { ( 2 ) } } \sum _ { \substack { ( i, j ) \in \mathcal { I } \\ k = 1 } } \left [ \sum _ { \substack { ( i, j ) \in \mathcal { I } \\ k = 1 } } ^ { m } ( k - i ) ^ { 2 } f ^ { ( 1 ) } _ { i, k, j } + \sum _ { l = 1 } ^ { n } ( j - l ) ^ { 2 } f ^ { ( 2 ) } _ { k, j, l } \right ] \\ \sum _ { i = 1 } ^ { m } f ^ { ( 1 ) } _ { i, k, j } = \sum _ { l = 1 } ^ { n } f ^ { ( 2 ) } _ { k, j, l }, \quad \forall ( k, j ) \in \mathcal { I }, \\ \sum _ { k = 1 } ^ { m } f ^ { ( 1 ) } _ { i, k, j } = \mu ^ { 1 } _ { i, j }, \quad \forall ( i, j ) \in \mathcal { I }, \\ \sum _ { j = 1 } ^ { n } f ^ { ( 2 ) } _ { k, j, l } = \mu ^ { 2 } _ { k, l }, \quad \forall ( k, l ) \in \mathcal { I }, \\ f ^ { ( 1 ) } _ { i, k, j } \geq 0, f ^ { ( 2 ) } _ { k, j, l } \geq 0, \quad \forall ( i, j ), ( k, l ) \in \mathcal { I }, \\ \intertext { w e n t e } f ^ { ( 1 ) } _ { i, k, j } \text{ denotes the input flow from bin } ( i, j ) \text{ to } ( k, j ) \text{ and } \text{ as } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{\ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text { } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text{ } \text`}$$

(3)

where f (1) i,k,j denotes the input flow from bin ( i, j ) to ( k, j ) and f (2) k,j,l denotes the output flow from bin ( k, j ) to ( k, l ) . Compared to formulation (1), the formulation (3) offers substantial computational benefits. Specifically, the reduced problem (3) only has mn 2 + m n 2 variables, whereas the original model has m n 2 2 variables. Moreover, the reduced model remains an LP problem. Consequently, popular algorithms for LP problems can be applied to solve this reduced model, such as the network-simplex method and the interior point method. Although the computation and memory costs of these mentioned algorithms are lower for the reduced model, it remains a challenge for solving large-scale problems. In this paper, we focus on addressing the challenges by designing a fast algorithm to solve the reduced model (3).

To facilitate the design of the algorithm, we reformulate the model (3) into the following standard form of linear programming:

$$\min _ { \begin{smallmatrix} x \in \mathbb { R } ^ { N } \\ s. t. \end{smallmatrix} } & \langle c, x \rangle + \delta _ { \mathbb { R } _ { + } ^ { N } } ( x ) & \text{OT} \, \left | \\ \hat { A } x = \hat { b }, & \text{map} \\ \text{issue}$$

where

$$\begin{array} { c } 1 ) \ M _ { 3 } = 3 M - 1, N = m ^ { 2 } n + m n ^ { 2 } ; \\ 2 ) \ x = [ f ^ { ( 1 ) } ; f ^ { ( 2 ) } ] \in \mathbb { R } ^ { m ^ { 2 } n } \times \mathbb { R } ^ { m n ^ { 2 } } \text{ with } \\ & \begin{cases} f ^ { ( 1 ) } = \{ f ^ { ( 1 ) } _ { i, k, j }, & \forall ( i, j ) \in \mathcal { I }, & k = 1, \dots, m \}, \\ \ e \omega \end{cases}. \end{array}$$

$$\begin{cases} f ^ { ( 1 ) } = \{ f ^ { ( 1 ) }$$

3) c = [ c 1 ; c 2 ] ∈ R m n 2 × R mn 2 with

$$\begin{cases} c ^ { 1 } = \{ c ^ { ( 1 ) } _ {$$

$$\begin{array} {$$

$$\begin{cases} A _ { 1 } = I _ { M } \otimes \mathbf$$

For notational convenience, let ¯ I m = [ I m -1 , 0 m -1 ] ∈ R ( m -1) × m . We define

$$\stackrel { \mathbb { m } } { k, j, l } } \int _$$

with

$$& \mathcal { L }, & A _ { 4 } = \text{diag} \left ( \mathbf 1 _ { n } ^ { \top } \otimes I _ { m }, \dots, \mathbf 1 _ { n } ^ { \top } \otimes I _ { m }, \mathbf 1 _ { n } ^ { \top } \otimes \bar { I } _ { m } \right ) \in \mathbb { R } ^ { ( M - 1 ) \times m n ^ { 2 } }. \\ & \quad \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot$$

Similar to [42, Lemma 7.1], we can obtain that A defined in (6) has full row rank, and

$$\{ x \in \mathbb { R } ^ { N } \ | \ A x = b \} = \{ x \in \mathbb { R } ^ { N } \ | \ \hat { A } x = \hat { b } \}.$$

As a result, the linear programming problem (4) is equivalent to

$$\min _ { x \in \mathbb { R } ^ { N } } \quad & \langle c, x \rangle + \delta _ { \mathbb { R } ^ { N } _ { + } } ( x ) \\ \text{s.t.} \quad & A x = b.$$

Furthermore, the dual problem of (8) takes the form:

$$\min _ { y \in \mathbb { R } ^ { M _ { 3 } }, z \in \mathbb { R } ^ { N } } \left \{ - \langle b, y \rangle + \delta _ { \mathbb { R } ^ { N } _ { + } } ( z ) \ | \ A ^ { \top } y + z = c \right \}. \quad ( 9 )$$

The KKT conditions associated with (8) and (9) can be given by

$$\mathfrak { z }, \quad \ A ^ { * } y + z = c, \ \ A x = b, \ \ \mathbb { R } _ { + } ^ { N } \ni x \perp z \in \mathbb { R } _ { + } ^ { N }, \quad ( 1 0 )$$

where x ⊥ z means x is perpendicular to z , i.e., ⟨ x, z ⟩ = 0 .

## B. Reconstruct the transport plan from the reduced model

The absence of the transport plan π makes the reduced OT model less favorable in applications where the transport map is necessary (i.e., color transfer [1]-[3]). We address this issue by proposing a fast algorithm (shown in Algorithm 1) to reconstruct an optimal transport plan of the original model from an optimal solution of the reduced model (3).

Algorithm 1 A fast algorithm for reconstructing transport plan π from the network flows f (1) and f (2) .

$$\begin{array} { l l$$

The following proposition shows that the output of Algorithm 1 is an optimal transport plan for of the original OT model.

Proposition 1. Given an optimal solution ( f (1) , f (2) ) to problem (3) , the output π of Algorithm 1 is an optimal solution to the optimal transport problem (1) .

Proof. Since ( f (1) , f (2) ) is an optimal solution to problem (3), we have

$$f _ { i, k, j } ^ { ( 1 ) } \geq 0, \ \ f$$

which implies

$$\pi _ { i, j ; k, l } \geq 0, \ \forall ( i$$

Additionally, since

$$f _ { i, k, j } ^ { ( 1 ) } \leq \sum _ { i$$

by fixing i, j, k , we obtain from Algorithm 1 that

$$\sum _ { l = 1 } ^ { n } \pi _ { i, j ; k, l$$

Similarly, by fixing j, k, l , we also obtain

$$\sum _ { i = 1 } ^ { m } \pi _ { i, j ; k, l } = f _ { k, j, l } ^ { ( 2 ) }. \quad \quad ( 1 3 ) \begin{array} { c } \text{itera} \\ \text{fall} \\ \text{HO} \end{array}$$

Hence, according to the constraints in problem (3), we have

$$\dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \dots \delta.$$

which, together with (11), shows that π is a feasible solution to problem (1). Furthermore, from (12), (13), and the definition of c in (2), we can obtain

$$& \sum _ { ( ( i, j ), ( k, l ) ) } c _ { i, j ; k, l } \pi _ { i, j ; k, l } \\ = & \sum _ { ( i, j ) \in \mathcal { I } } [ \sum _ { k = 1 } ^ { m } ( k - i ) ^ { 2 } f _ { i, k, j } ^ { ( 1 ) } + \sum _ { l = 1 } ^ { n } ( j - l ) ^ { 2 } f _ { k, j, l } ^ { ( 2 ) } ].$$

According to [32, Theorem 1], we know that the optimal objective function values of problems (1) and (3) are equivalent. Therefore, π is an optimal solution to problem (1).

## III. A HALPERN ACCELERATING ALGORITHM FOR SOLVING OT PROBLEM

In this section, we first introduce an efficient Halpern accelerating method for solving problem (9), which includes the equivalent reduced OT problem (8) as a special case. Subsequently, we present an efficient implementation of the proposed algorithm by designing a novel procedure to solve the involved linear system in linear time complexity.

A. HOT: A Halpern accelerating method for solving OT problem

Given σ &gt; 0 , the augmented Lagrange function corresponding to the dual problem (9) is defined by, for any ( y, z, x ) ∈ R M 3 × R N × R N ,

$$L _ { \sigma } ( y, z ; x ) \coloneqq - \langle b, y \rangle + \delta _ { \mathbb { R } ^ { N } _ { + } } ( z ) + \frac { \sigma } { 2 } \| A ^ { \top } y + z - c + \frac { 1 } { \sigma } x \| ^ { 2 } - \frac { 1 } { \sigma } \| x \| ^ { 2 }. \text{Remark 1} \\ a n \ O ( 1 / \sqrt { }$$

For ease of notation, denote w := ( y, z, x ) . A fast Halpern accelerating method [35], [43] for solving OT problems is

Algorithm 2 HOT: A Halpern accelerating method for solving the OT problem (9).

- 1: Input: Choose an initial point w 0 = ( y , z 0 0 , x 0 ) ∈ R M 3 × R N × R N . Set parameters σ &gt; 0 . For k = 0 1 , , . . . , perform the following steps in each iteration.
- 2: Step 1. ¯ y k = arg min { L σ ( y, z k ; x k )} .

y

∈

Y

- 3: Step 2. ¯ x k = x k + σ A ( ⊤ ¯ y k + z k -c ) .
- 4: Step 3. ¯ z k = arg min { L σ ( ¯ y k , z ; ¯ x k )} .
- z ∈ Z 5: Step 4. w k +1 = 1 k +2 w 0 + k +1 k +2 (2 ¯ w k -w k ) .

presented in Algorithm 2. A detailed derivation of the algorithm in its current form and more discussions can be found in [35] and the references therein.

Note that Step 4 in Algorithm 2 is from the Halpern iteration with a stepsize of 1 k +2 . Without Step 4, the HOT algorithm reduces to the ADMM with a unit step size. The following proposition shows the global convergence of the HOT algorithm. The proof of this proposition follows from [35, Corollary 3.5].

Proposition 2. The sequence { ¯ w k } = { (¯ y k , z ¯ k , x ¯ k ) } generated by the HOT algorithm in Algorithm 2 converges to the point w ∗ = ( y ∗ , z ∗ , x ∗ ) , where ( y ∗ , z ∗ ) is a solution to problem (9) and x ∗ is a solution to problem (8) .

Next, we analyze the iteration complexity of the HOT algorithm for obtaining an ε -approximate solution, where an appropriate measure for the quality of the solution is crucial. In this paper, we consider the residual mapping associated with the KKT system (10):

$$\mathcal { R } ( w ) = \left ( \begin{array} { c } b - A x \\ z - \Pi _ { \mathbb { R } _ { + } ^ { N } } ( z - x ) \\ c - A ^ { \top } y - z \end{array} \right )$$

for any w = ( y, z, x ) ∈ R M 3 × R N × R N . Note that R ( w ∗ ) = 0 is equivalent to the facts that x ∗ ∈ R N and ( y ∗ , z ∗ ) ∈ R M 3 × R N are the solution to problems (8) and (9), respectively. The KKT residual ∥R · ∥ ( ) is a commonly used and practical measure for the quality of the approximation solution to (8). It follows from [35, Theorem 3.7] that the HOT algorithm enjoys an appealing O (1 /k ) nonergodic convergence rate in terms of the KKT residual for solving (8), which is summarized in the following proposition.

Proposition 3. Let { (¯ y k , z ¯ k , x ¯ k ) } be the sequence generated by Algorithm 2, and let w ∗ = ( y ∗ , z ∗ , x ∗ ) be the limit point of the sequence { (¯ y k , z ¯ k , x ¯ k ) } and R 0 = ∥ x 0 -x ∗ + ( σ z 0 -z ∗ ) ∥ . For all k ≥ 0 , we have the following bounds:

$$\| \mathcal { R } ( \bar { w } ^ { k } ) \| \leq \left ( \frac { \sigma + 1 } { \sigma } \right ) \frac { R _ { 0 } } { ( k + 1 ) }. \quad \quad ( 1 4 )$$

. Remark 1. Note that, without acceleration, the ADMM has an O (1 / √ k ) non-ergodic rate in terms of both the objective function value gap and feasibility violations [44], [45]. In contrast, the HOT algorithm in Algorithm 2 achieves an O (1 /k )

non-ergodic convergence rate, offering significant advantages for solving large-scale OT problems.

## B. A fast implementation of the HOT algorithm

In this section, we present a fast implementation of the HOT algorithm. Through direct calculations, we obtain the following updates of ¯ z k and ¯ y k for any k ≥ 0 :

- 1) Update of ¯ z k :

$$\bar { z } ^ { k } = \Pi _ { \mathbb { R } _ { + } ^ { N } } \left ( c - A ^ { \top } \bar { y } ^ { k } - \bar { x } ^ { k } / \sigma \right ) ;$$

2) Update of ¯ y k :

$$A A ^ { \top } \bar { y } ^ { k } = \frac { b } { \sigma } - A \left ( \frac { x ^ { k } } { \sigma } + z ^ { k } - c \right ). \quad ( 1 5 ) ^ { \quad w h e r } _ { 1 ) \ \ 1 }$$

Therefore, the main computational bottleneck of the HOT algorithm for solving (8) is solving the linear system (15). Note that the dimension of the matrix AA ⊤ is M 3 × M 3 with M 3 defined in (4). In applications, M 3 is usually a huge number. As an illustrative example, for an image with 256 × 256 pixels, M 3 = 196 607 , . As a result, it is not computationally affordable for computing a (sparse) Cholesky decomposition for the matrix AA ⊤ or solving the linear system (15) with standard direct solver. Indeed, it is computationally expensive even for computing the matrix AA ⊤ . Instead, in the remaining part of this subsection, we will derive a linear time complexity procedure for solving the linear equation AA y ⊤ = R with a given vector R ∈ R M 3 . It is worthwhile mentioning that our procedure does not require calculating nor storing the matrix AA ⊤ . Note that AA ⊤ can be written in the following form:

$$\Lambda \\ A A ^ { \top } = \left [ \begin{array} { c c c } E _ { 1 } & E _ { 2 } & E _ { 3 } \\ E _ { 2 } ^ { \top } & E _ { 4 } & 0 \\ E _ { 3 } ^ { \top } & 0 & E _ { 5 } \end{array} \right ], \quad \quad ( 1 6 ) \quad \text{where} \\ \quad \text{is to} \\ \text{direct}$$

where

$$\text{where} \\ 1 ) & \, E _ { 1 } = ( m + n ) I _ { M } \in \mathbb { R } ^ { M \times M } ; \\ 2 ) & \, E _ { 2 } = \text{diag} \left ( 1 _ { m } 1 _ { m } ^ { \top }, \dots, 1 _ { m } 1 _ { m } ^ { \top }, 1 _ { m } 1 _ { m } ^ { \top } \right ) \in \mathbb { R } ^ { M \times M } ; \\ 3 ) & \, E _ { 3 } = - 1 _ { n } \otimes ( I _ { m }, \dots, I _ { m }, \bar { I } _ { m } ^ { \top } ) \in \mathbb { R } ^ { M \times ( M - 1 ) } ; \\ 4 ) & \, E _ { 4 } = m I _ { M } \in \mathbb { R } ^ { M \times M } ; \\ 5 ) & \, E _ { 5 } = A _ { 4 } A _ { 4 } ^ { \top } = n I _ { M - 1 } \in \mathbb { R } ^ { ( M - 1 ) \times ( M - 1 ) }. \\ \text{To better explore the structure of the linear system } A A ^ { \top } \mu =$$

To better explore the structure of the linear system AA y ⊤ = R , we rewrite it equivalently as

$$A A ^ { \top } y = \left [ \begin{array} { c c c } \cdot & \cdot & \cdot \\ E _ { 1 } & E _ { 2 } & E _ { 3 } \\ E _ { 2 } ^ { \top } & E _ { 4 } & 0 \\ E _ { 3 } ^ { \top } & 0 & E _ { 5 } \end{array} \right ] \left [ \begin{array} { c } y _ { 1 } \\ y _ { 2 } \\ y _ { 3 } \end{array} \right ] = \left [ \begin{array} { c } R _ { 1 } \\ R _ { 2 } \\ R _ { 3 } \end{array} \right ], \quad ( 1 7 ) \\ \quad \text{Denote} \\ \cdot \end{array}$$

where y := ( y 1 ; y 2 ; y 3 ) ∈ R M × R M × R M -1 and R := ( R 1 ; R 2 ; R 3 ) ∈ R M × R M × R M -1 . To further explore the block structure of the linear system, we can denote y i := ( y 1 i ; . . . ; y n i ) ∈ R m × · · · × R m for i = 1 2 , , and y 3 = ( y 1 3 ; . . . ; y n 3 ) ∈ R m ×···× R m × R m -1 . Correspondingly, we write R i = ( R 1 i ; . . . ; R n i ) for i = 1 2 3 , , . The next proposition gives an explicit formula of the solution to the linear equation in (17).

Proposition 4. Consider A ∈ R M 3 × N defined in (6) . Given R ∈ R M 3 , the solution y to AA y ⊤ = R in the form (17) is given by:

$$y _ { 2 } ^ { j } = \frac { 1 } { m } ( R _ { 2 } ^ { j } - \underset { n } { \mathbf i } ^ { \top } y _ { 1 } ^ { j } ), \ \ j = 1, \dots, n, \quad \ ( 1 8 )$$

$$y _ { 3 } ^ { j } = \frac { 1 } { n } ( R _ { 3 } ^ { j } + \sum _ { j = 1 } ^ { n } y _ { 1 } ^ { j } ), \ \ j = 1, \dots, n - 1, \ \ ( 1 9 )$$

$$y _ { 3 } ^ { n } = \frac { 1 } { n } ( R _ { 3 } ^ { n } + \bar { I } _ { m } \sum _ { j = 1 } ^ { n } y _ { 1 } ^ { j } ),$$

$$y _ { 1 } ^ { j } = \hat { y } _ { 1 } ^ { j } - \hat { y } _ { 1 } ^ { a }, \ \ j = 1, \dots, n,$$

where

$$( 1 5 ) \quad & \text{where} \\ \vdots \text{HOT} \\ \times M _ { 3 } \\ \quad & \text{all} \, \text{$a_{1}^{j}=\frac$$

Proof. By some direct calculations, we can solve (17) equivalently as:

$$y _ { 2 } = \frac { R _ { 2 } - E _ { 2 } ^ { \top } y _ { 1 } } { m _ { \infty ^ { \top } } },$$

$$y _ { 3 } = \frac { R _ { 3 } - \stackrel { \dots \nu } { E _ { 3 } ^ { \top } } y _ { 1 } } { n }, \quad \quad \quad$$

$$\tilde { E } _ { 1 } y _ { 1 } = R _ { 1 } - \frac { 1 } { m } E _ { 2 } R _ { 2 } - \frac { 1 } { n } E _ { 3 } R$$

where ˜ E 1 = ( E 1 -1 m E E 2 ⊤ 2 -1 n E E 3 ⊤ 3 ) . As a result, the key is to obtain y 1 by solving (24). Let ˆ E 1 := E 1 -1 m E E 2 ⊤ 2 . By direct calculations, we have:

$$\hat { E } _ { 1 } = \text{diag} \left ( \hat { E } _ { 1 } ^ { 1 }, \dots, \hat { E } _ { 1 } ^ { n } \right ) \text$$

with ˆ E j 1 = ( m + n I ) m -1 m m 1 ⊤ , j = 1 , . . . , n. By the Sherman-Morrison-Woodbury formula, we directly get:

$$\hat { E } _ { 1 } ^ { - 1 } = \text{diag} \left ( ( \hat { E } _ { 1 } ^ { 1 } ) ^ { - 1 }, \dots, ( \hat { E } _ { 1 } ^ { n } ) ^ { - 1 } \right ) \quad \text{ (2.6)}$$

with ( ˆ E j 1 ) -1 = 1 m n + ( I m + 1 n 1 m m 1 ⊤ ) , j = 1 , . . . , n. On the other hand, let:

$$Q = \frac { n - 1 } { n } I _ { m } + \frac { 1 } { n } \bar { I } _ { m } ^ { \top } \bar { I } _ { m } \in \mathbb { R } ^ { m \times m }.$$

Denote ˆ := Q Q 1 / 2 such that ˆ ˆ = QQ Q , and ¯ := Q 1 n ⊗ ˆ Q . We can obtain:

$$\tilde { E } _ { 1 } = \left ( \hat { E } _ { 1 } - \frac { 1 } { n } E _ { 3 } E _ { 3 } ^ { \top } \right ) = \hat { E } _ { 1 } - \bar { Q } \bar { Q } ^ { \top }.$$

Hence, by the Sherman-Morrison-Woodbury formula, we can derive:

$$\tilde { E } _ { 1 } ^ { - 1 } \ = \ \left ( \hat { E } _ { 1 } - \frac { 1 } { n } E _ { 3 } E _ { 3 } ^ { \top } \right ) ^ { - 1 } \\ = \ \hat { E } _ { 1 } ^ { - 1 } - \hat { E } _ { 1 } ^ { - 1 } \bar { Q } W ^ { - 1 } \bar { Q } ^ { \top } \hat { E } _ { 1 } ^ { - 1 }$$

with W = -I m + ¯ Q E ⊤ ˆ -1 1 ¯ Q . Note that:

$$W \ & = - I _ { m } + \sum _ { j = 1 } ^ { n } \hat { Q } ( E _ { 1 } ^ { j } ) ^ { - 1 } \hat { Q } \\ & = \frac { 1 } { m + n } \left ( \text{diag} \left ( - m I _ { m - 1 }, - ( m + 1 ) \right ) + d _ { 0 } d _ { 0 } ^ { \top } \right ), \quad \text{i} \\ \ t h$$

where d 0 = [ 1 m -1 ; √ 1 -1 n ] ∈ R m . Applying the Sherman-Morrison-Woodbury formula to W , we can obtain:

$$W ^ { - 1 } \ = ( m + n ) \left ( \text{diag} ( - \frac { 1 } { m } I _ { m - 1 }, - \frac { 1 } { m + 1 } ) - \frac { 1 } { w } d _ { 1 } d _ { 1 } ^ { \top } \right ) \quad i t e r a$$

with d 1 = [ 1 m 1 m -1 ; 1 m +1 √ 1 -1 n ] . It follows that

$$\bar { Q } W ^ { - 1 } \bar { Q } ^ { \top } \ & = \mathbf 1 _ { n } \mathbf 1 _ { n } ^ { \top } \otimes ( \hat { Q } W ^ { - 1 } \hat { Q } ) \\ & = \mathbf 1 _ { n } \mathbf 1 _ { n } ^ { \top } \otimes ( m + n ) \hat { W }. \quad \begin{array} { c } o v e r a \end{array} \\ \quad ( 2 8 ) \ \ A l g o r. \quad \\ /$$

To explore the block structure of R , we denote:

$$\hat { R } _ { 1 } = R _ { 1 } { - \frac { 1 } { m } } E _ { 2 } R _ { 2 } { - \frac { 1 } { n } } E _ { 3 } R _ { 3 } = ( \hat { R } _ { 1 } ^ { 1 } ; \dots ; \hat { R } _ { 1 } ^ { n } ) \in \mathbb { R } ^ { m } { \times } \cdots { \times } \mathbb { R } ^ { m }, \text{ Remark } \colon$$

which implies

$$\hat { R } _ { 1 } ^ { j } = R _ { 1 } ^ { j } - \frac { 1 } { m } \mathbf 1 _ { m } ^ { \top } R _ { 2 } ^ { j } + \frac { 1 } { n } \left ( \sum _ { j = 1 } ^ { n - 1 } R _ { 3 } ^ { j } + \bar { I } _ { m } ^ { \top } R _ { 3 } ^ { n } \right ), \ \ j = 1, \dots, n. \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \$$

Then, for j = 1 , . . . , n ,

$$\hat { y } _ { 1 } ^ { j } \ & = ( \hat { E } _ { 1 } ^ { - 1 } \hat { R } _ { 1 } ) ^ { j } \\ & = \frac { 1 } { m + n } \left ( \tilde { R } _ { 1 } ^ { j } + \tilde { R } _ { 2 } ^ { j } + \tilde { R } _ { 3 } \right ).$$

Define

$$\hat { y } _ { 1 } ^ { a } \coloneqq \left ( I _ { m } + \frac { 1 } { n } \mathbf 1 _ { m } \mathbf 1 _ { m } ^ { \top } \right ) \hat { W } \sum _ { j = 1 } ^ { n } \hat { y } _ { 1 } ^ { j }.$$

From (24), (27), and (28), we have:

$$y _ { 1 } ^ { j } = \hat { y } _ { 1 } ^ { j } - \hat { y } _ { 1 } ^ { a }, \ \ j = 1, \dots, n.$$

Substituting y 1 into (22) and (23), we can obtain the results for y 2 and y 3 . This completes the proof.

According to the explicit formula in Proposition 4, we can immediately derive the complexity result for solving the linear equation in (17).

Corollary 1. Consider A ∈ R M 3 × N defined in (6) . The linear system AA y ⊤ = R in the form (17) can be solved in O M ( 3 ) flops.

Based on this corollary, we can determine the per-iteration computational cost of the HOT algorithm in Algorithm 2 for each iteration.

Corollary 2. The per-iteration computational complexity of the HOT algorithm in Algorithm 2 in terms of flops is O N ( ) .

Proof. Since A has at most 2 N nonzero elements, the computational cost for calculating Ax and A y ⊤ is only O N ( ) . Except for solving linear systems, which can be done in O M ( 3 ) from Corollary 1, Algorithm 2 primarily involves matrix-vector multiplications and vector additions. Hence, the computational cost of Algorithm 2 for each iteration is O N ( ) .

Combining the iteration complexity in Proposition 3 and the computational cost for each iteration in Corollary 2, we can

,

derive the following overall computational complexity result of the HOT algorithm for solving problem (3).

Theorem 1. Let { ¯ y k , z ¯ k , x ¯ k } be the sequence generated by the HOT algorithm in Algorithm 2. For any given tolerance ε &gt; 0 , the HOT algorithm needs at most

$$\frac { 1 } { \varepsilon } \left ( \frac { 1 + \sigma } { \sigma } \left ( \left \| x ^ { 0 } - x ^ { * } \right \| + \sigma \left \| z ^ { 0 } - z ^ { * } \right \| \right ) \right ) - 1$$

iterations to return a solution to the equivalent OT problem (8) such that the KKT residual ∥ ∥ R ( ¯ w k )∥ ∥ ≤ ε , where ( x , z ∗ ∗ ) is the limit point of the sequence { ¯ x , z k ¯ k } . In particular, the overall computational complexity of the HOT algorithm in Algorithm 2 to achieve this accuracy in terms of flops is

$$O \left ( \left ( \frac { 1 + \sigma } { \sigma } \left ( \left \| x ^ { 0 } - x ^ { * } \right \| + \sigma \left \| s ^ { 0 } - s ^ { * } \right \| \right ) \right ) \frac { m ^ { 2 } n + m n ^ { 2 } } { \varepsilon } \right ).$$

Remark 2. Note that the complexity of the Sinkhorn method for solving the original OT problem (1) is O ( m n 2 2 ε 2 ) to achieve an ε -accuracy solution in terms of objective function value [13]. Even by leveraging the separable structure of the cost function c defined in (2) , the improved Sinkhorn method, as mentioned in [3] and [32], still requires O ( m n 2 + mn 2 ε 2 + m n 2 2 ) to compute the KW distance. In contrast, the HOT algorithm exhibits lower overall computational complexity, offering a significant advantage in calculating the KW distance between large-scale histograms.

## IV. EXPERIMENTS

In this section, we comprehensively compare the HOT algorithm with five other state-of-the-art methods on the popular DOTmark dataset [46]. Additionally, employing the transport plan derived from Algorithm 1, we test the performance of the reduced OT model on the color transfer task. The findings reveal that HOT provides significant advantages over state-ofthe-art methods in both memory and computational efficiency, particularly for large-scale problems.

## A. Numerical comparison on the DOTmark dataset

The DOTmark dataset [46] is a comprehensive collection of benchmark instances for evaluating and comparing algorithms in the field of optimal transport. It consists of a variety of instances categorized into different classes, such as Classic Images, Shapes, and Gaussian Distributions. In this experiment, we selected eight images each from the Classic Images and Shapes categories, as illustrated in Fig.2. These selected images were resized to five different resolutions: 64 × 64 , 128 × 128 , 256 × 256 , and 512 × 512 . Finally, we randomly selected 10 pairs from each category and computed the KW distance for each pair.

To exhibit the superiority of HOT, we compare it with the following five state-of-the-art methods:

- · Sinkhorn [10] is a widely used algorithm for computing an approximate KW distance by solving an entropyregularized OT problem. As the Introduction Section discusses, its performance is sensitive to the regularization parameter, denoted as λ . To achieve varying levels of

Fig. 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.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00598v1-figure-019655be-d7e7-4763-89ee-12e11f7520f6.jpg)

solution accuracy, we selected the λ to be 1% , 0 1% . , and 0 01% . of the median transport cost, following the setup from [3]. Due to the potential numerical instability caused by small λ values, we employed the Log-domain Sinkhorn algorithm implemented by POT [47] as the baseline method.

of the Halpern iteration.

- · Improved Sinkhorn [3], [33] is an improved version of the Sinkhorn algorithm, optimized for computing distances over regular two-dimensional grids. Specifically, since the kernel matrix used by the Sinkhorn algorithm can be constructed using a Kronecker product of two smaller matrices, a matrix-vector product using a matrix of dimension M × M can be replaced by two matrixmatrix products over matrices of dimension √ M × √ M , resulting in a significant improvement in computational efficiency. In our experiment, based on the MATLAB implementation of the improved Sinkhorn in [32], we further developed the Log-domain improved Sinkhorn using PyTorch. The regularization parameter λ is kept the same as that of the original Sinkhorn method.
- · Gurobi (11.0.1) is a popular optimization solver designed to address a wide range of mathematical programming problems, including linear and quadratic programming. In our experiment, we employ the interior point method (IPM) implemented in Gurobi to solve the reduced model (8). Since it is unnecessary to obtain a basic solution, we disable the cross-over strategy.
- · Network Simplex implemented in the Lemon C++ graph library 2 , is a highly efficient algorithm for solving uncapacitated minimum cost flow problems [48]. It has demonstrated the computational advantages over the Sinkhorn type methods in [32] for small size images by solving the reduced model (8).
- · ADMM [40], [41] is a popular first-order primal-dual method for solving large-scale optimization problems. It has shown great potential in solving large-scale optimal transport problems (1) in a GPU setting [20]. In this numerical experiment, we use a generalized ADMM [49], [50] to solve the equivalent model (8) (replacing Step 4 in Algorithm 2 with w k +1 = (1 -ρ w ) k + ρw ¯ k and setting ρ = 1 7 . ) as the baseline to evaluate the acceleration effect

All experiments are conducted on an Ubuntu 22.04 server equipped with an Intel(R) Xeon(R) Platinum 8480C processor and an Nvidia GeForce RTX 4090 GPU with 24 GB of RAM. Due to hardware specifications, we limit the maximum memory usage of each algorithm to 24 GB. Additionally, we set the maximum one-call running time of each algorithm to 3600 seconds. For the HOT and ADMM methods, we adopt a stopping criterion based on the relative KKT residual:

$$\underset { t \text{ of two } } { \text{gorithm} } \underset { a \text{ matrix} } { \text{gorithm} } \ \overset { \quad \quad \quad } { \text{kKT} } _ { \text{res} } = \max \left \{ \frac { \| A ^ { \top } y + z - c \| } { 1 + \| c \| }, \frac { \| \min ( x, z ) \| } { 1 + \| x \| + \| z \| }, \frac { \| A x - b \| } { 1 + \| b \| } \right \}.$$

Other methods use their default stopping criteria. We terminate all tested algorithms, except for Network Simplex which is an exact algorithm with a stopping tolerance of 1E-6. Finally, since different methods have varying stopping criteria, we use the following metrics to fairly evaluate the quality of the solutions: the 'relative objective gap' (gap) and the 'relative primal feasibility error' (feaserr). These metrics are defined as follows:

$$& \text{gap} = \frac { | \langle c, x \rangle - \langle c, x _ { b } \rangle | } { | \langle c, x _ { b } \rangle | + 1 }, \\ & \text{feaser} = \max \left \{ \frac { \| \min ( x, 0 ) \| } { 1 + \| x \| }, \frac { \| A x - b \| } { 1 + \| b \| } \right \},$$

where x b is the solution obtained using Gurobi with the tolerance set to 1E-8.

We present the average results of 10 pairs within each category for all tested algorithms in Table II. Since Gurobi runs out of memory for images sized from 256 × 256 to 512 × 512 , we only report the feaserr for the HOT, Network Simplex, and ADMM. Additionally, Table II only shows the results of Sinkhorn-type methods with λ = 0 01% . of the median transport cost, as this parameter returns a solution of comparable quality to other methods. For more results of the improved Sinkhorn with different regularization parameters, refer to Table III. Due to space constraints, we omit the results of the Sinkhorn with varying parameters, which exhibit similar performance to the improved Sinkhorn in terms of the gap.

We summarized some key findings in Table II from the perspective of computational efficiency and memory cost:

- 1) Computational efficiency: HOT can return a comparable solution in terms of the gap and feaserr in the shortest time. Although the Network Simplex and Gurobi have

computational advantages for computing the KW distance for small-scale images, these methods cannot handle large-scale problems effectively. IPM implemented in Gurobi requires solving a linear equation that relies heavily on Cholesky decomposition, causing the computational cost of each iteration to increase rapidly with image size. Additionally, the inherent sequential nature of the Network Simplex algorithm makes it challenging to parallelize effectively. In contrast, HOT and improved Sinkhorn benefit from low per-iteration costs and are easily parallelizable. However, the improved Sinkhorn needs to recover the solution to the original problem (1) to compute the KW distance, making it unsuitable for largescale problems. As a result, for images sized 128 × 128 , HOT achieves a 17.44x speedup over Network Simplex, a 15.83x speedup over Gurobi, and a 19.54x speedup over improved Sinkhorn. Additionally, compared to ADMM, HOT benefits from a superior O (1 /k ) iteration complexity and saves 40% of iterations for images sized 512 × 512 .

- 2) Memory cost: HOT demonstrates a significant advantage in memory efficiency by exploiting the sparse structure of A in (5) and utilizing the explicit solution of the linear equation presented in Proposition 4 to avoid sparse Cholesky decomposition. In contrast, Sinkhorntype methods need to recover the solution to the original OT problem (1) and maintain the transport cost to calculate the KW distance, which involves M 2 variables. Given that these variables are stored using 64-bit floatingpoint representation, the Sinkhorn-based methods require at least 32GB of memory for images sized 256 × 256 , which far exceeds the available 24GB. While the IPM in Gurobi solves the reduced model (8), each iteration requires performing a sparse Cholesky decomposition, making it unsuitable for images sized 256 × 256 or larger.

To further illustrate the benefit of the explicit solution of the linear system presented in Proposition 4, we conduct a comparison between solving the linear equation using Proposition 4 and sparse Cholesky decomposition. Given that the average iteration number of HOT is around 1500, we solve the linear system (17) 1500 times for different image sizes. The results are shown in Fig.3. It is clear that the computational costs of Cholesky decomposition increase rapidly as image size increases. Additionally, the forward-backward substitution method used to solve the linear system is also time-consuming because it cannot be efficiently parallelized. In contrast, solving linear systems using Proposition 4 can be easily parallelized. It consumes significantly less time and remains constant as the size varies because it only requires O M ( 3 ) flops, as shown in Corollary 1.

## B. An Application in Color Transfer

Color transfer [1]-[3] based on the OT model has found important applications in various fields, including digital image processing, computer graphics, and visual arts. It involves transferring the color characteristics from a target image to a source image to achieve a desired visual effect. Inspired by its convincing performance in color grading and color histogram

Fig. 3: Comparison of solving the linear system (17) using Proposition 4 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 × 256 and 512 × 512 cases, Cholesky decomposition is out-of-memory in the test.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00598v1-figure-c028ff4e-558e-4045-ab33-7e9fad2ed095.jpg)

manipulation [2], [3], we conduct the color transfer over the CIE-Lab domain by applying the optimal transport model to the 1D luminance channel and the 2D chrominance channel independently. Note that the 1D optimal transport plan can be found efficiently [8]. To efficiently conduct the optimal transport over the 2D chrominance channel, we first solve the equivalent reduced OT model using the HOT Algorithm 2 and then recover an optimal transport map using Algorithm 1. The performance of color transfer for selected image pairs can be found in Fig.1 and Fig.4.

## V. CONCLUSION

In this paper, we proposed an efficient and scalable HOT algorithm for computing the KW distance with finite supports. In particular, the HOT algorithm solves an equivalent reduced OT model where the involved linear systems are solved by a novel procedure in linear time complexity. Consequently, we can obtain an ε -approximate solution to the OT problem with M supports in R 2 in O M ( 1 5 . /ε ) flops, significantly enhancing the best-known computational complexity. Additionally, we have designed an efficient algorithm to recover an optimal transport map from a solution to the reduced OT model, thereby overcoming the limitations of the reduced OT model in applications that require the transport map. Extensive numerical results presented in this paper demonstrated the superior performance of the HOT algorithm. For future research directions, we aim to design an efficient implementation of the HOT algorithms for solving the OT problems with discrete supports in the more general R d space. We also consider designing an efficient algorithm for solving the Wasserstein barycenter problem with discrete supports.

## ACKNOWLEDGMENT

The research of Yancheng Yuan was supported by The Hong Kong Polytechnic University under the grant P0045485.

TABLE II: The numerical results of different algorithms on the DOTmark dataset.

| Category   | Resolution   |                      | HOT                | Network Simplex           | Gurobi          | ADMM                 | Improved Sinkhorn   | Sinkhorn        |
|------------|--------------|----------------------|--------------------|---------------------------|-----------------|----------------------|---------------------|-----------------|
| Classic    | 64 × 64      | time(s)              | 0.67               | 2.73                      | 2.16            | 1.77                 | 16.18               | 174.82          |
|            | 64 × 64      | gap                  | 8.26E-04           | 3.46E-10                  | 1.20E-04        | 2.67E-04             | 1.69E-04            | 2.12E-04        |
|            | 64 × 64      | feaserr              | 4.58E-07           | 4.88E-32                  | 2.55E-11        | 3.09E-07             | 7.90E-07            | 9.75E-07        |
|            | 64 × 64      | iter                 | 1700               | -                         | 13              | 3420                 | 64126               | 62474           |
|            | 64 × 64      | time(s)              | 1.58               | 36.18                     | 29.15           | 3.53                 | 39.40               | 2632.17         |
|            | 128 × 128    | gap                  | 6.24E-03           | 8.74E-10                  | 1.07E-04        | 1.72E-03             | 6.98E-04            | 6.35E-04        |
|            | 128 × 128    | feaserr              | 7.27E-07           | 9.67E-32                  | 7.24E-12        | 3.73E-07             | 8.34E-07            | 9.87E-07        |
|            | 128 × 128    | iter                 | 1170               | -                         | 14              | 3240                 | 58446               | 57010           |
|            | 256 × 256    | time(s)              | 12.98              | 2562.92                   | Memory Overflow | 20.80 6.04E-07       | Memory Overflow     | Memory Overflow |
|            | 256 × 256    | feaserr iter         | 8.05E-07 1140      | 1.35E-31 -                | Memory          | 2250 116.92          | Memory              | Memory          |
|            | 512 × 512    | time(s) feaserr iter | 81.02 3.28E-07 900 | Over Maximum Running Time | Overflow        | 4.32E-07 1610        | Overflow            | Overflow        |
|            | 64 × 64      | time(s)              | 0.64               | 1.48                      | 1.33            | 3.92                 | 9.60                | 103.74          |
|            | 64 × 64      | gap                  | 3.78E-04           | 1.81E-10                  | 2.28E-05        | 5.85E-05             | 4.86E-05            | 6.07E-05        |
|            | 64 × 64      | feaserr              | 5.77E-07           | 7.24E-32                  | 1.88E-10        | 2.58E-07             | 7.95E-07            | 9.68E-07        |
|            | 64 × 64      | iter                 | 1610               | -                         | 15              | 10430                | 37986               | 37077           |
|            | 128 × 128    | time(s) gap          | 1.68               | 20.70 2.46E-09            | 22.46 2.19E-05  | 2.32                 | 24.32               | 1616.34         |
|            | 128 × 128    |                      | 2.51E-03           |                           |                 | 4.11E-04             | 3.28E-04            | 3.09E-04        |
| Shapes     | 128 × 128    | feaserr              | 1.01E-06           | 1.16E-31                  | 2.01E-10        | 7.74E-07             | 8.01E-07            | 9.83E-07        |
|            | 128 × 128    | iter                 | 1240               | -                         | 18              | 2130                 | 36080               | 35009           |
|            | 256 × 256    | time(s) feaserr      | 14.87              | 959.77 1.59E-31           | Memory Overflow | 23.30 7.17E-07       | Memory Overflow     | Memory Overflow |
|            | 256 × 256    | iter                 | 6.68E-07 1310      | -                         |                 | 2530                 |                     |                 |
|            | 512 × 512    | time(s) feaserr iter | 87.12 3.54E-07 970 | Over Maximum Running Time | Memory Overflow | 118.10 5.71E-07 1630 | Memory Overflow     | Memory Overflow |

TABLE III: The numerical results of the improved Sinkhorn method with different λ .

| Solver            | Category   | Resolution   |         |   λ = 0 . 01% |   λ = 0 . 1% |     λ = 1% |
|-------------------|------------|--------------|---------|---------------|--------------|------------|
| Improved Sinkhorn | Classic    | 64 × 64      | time(s) |     16.18     |     1.6      |   0.19     |
| Improved Sinkhorn | Classic    | 64 × 64      | gap     |      0.000169 |     0.0243   |   0.329    |
| Improved Sinkhorn | Classic    | 64 × 64      | feaserr |      7.9e-07  |     8.15e-07 |   7.1e-07  |
| Improved Sinkhorn | Classic    | 64 × 64      | iter    |  64126        |  6400        | 650        |
| Improved Sinkhorn | Classic    | 64 × 64      | time(s) |     39.4      |     3.94     |   0.39     |
| Improved Sinkhorn | Classic    | 128 × 128    | gap     |      0.000698 |     0.0334   |   0.351    |
| Improved Sinkhorn | Classic    | 128 × 128    | feaserr |      8.34e-07 |     8.29e-07 |   7.29e-07 |
| Improved Sinkhorn | Classic    | 128 × 128    | iter    |  58446        |  5850        | 594        |

The research of Defeng Sun was supported by grants from the Research Grants Council of the Hong Kong Special Administrative Region, China (GRF Project No. 15303720).

- [4] A. Dominitz and A. Tannenbaum, 'Texture mapping via optimal mass transport,' IEEE Transactions on Visualization and Computer Graphics , vol. 16, no. 3, pp. 419-433, 2009.

## REFERENCES

[1] F. Piti´ e and A. Kokaram, 'The linear Monge-Kantorovitch linear colour mapping for example-based colour transfer,' 4th European Conference on Visual Media Production , 2007.

[2] N. Bonneel, K. Sunkavalli, S. Paris, and H. Pfister, 'Example-based video color grading,' ACM Transactions on Graphics (ToG) , vol. 32, no. 4, pp. 1-12, 2013.

[3] J. Solomon, F. De Goes, G. Peyr´ e, M. Cuturi, A. Butscher, A. Nguyen, T. Du, and L. Guibas, 'Convolutional Wasserstein distances: Efficient optimal transportation on geometric domains,' ACM Transactions on Graphics (ToG) , vol. 34, no. 4, pp. 1-11, 2015.

- [5] S. Haker, L. Zhu, A. Tannenbaum, and S. Angenent, 'Optimal mass transport for registration and warping,' International Journal of Computer Vision , vol. 60, pp. 225-240, 2004.
- [6] S. Basu, S. Kolouri, and G. K. Rohde, 'Detecting and visualizing cell phenotype differences from microscopy images using transport-based morphometry,' Proceedings of the National Academy of Sciences , vol. 111, no. 9, pp. 3448-3453, 2014.
- [7] E. Del Barrio, J. A. Cuesta-Albertos, C. Matr´ an, and J. M. Rodr´ ıguezRodr´ ıguez, 'Tests of goodness of fit based on the L2-Wasserstein distance,' Annals of Statistics , pp. 1230-1239, 1999.
- [8] C. Villani, 'Topics in optimal transportation,' Graduate Studies in Mathematics , 2003.
- [9] G. Peyr´ e, M. Cuturi et al. , 'Computational Optimal Transport: With applications to data science,' Foundations and Trends® in Machine

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00598v1-figure-9f54ce7f-b284-41bf-8486-dad1f7b012bc.jpg)

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00598v1-figure-ac0ee025-528f-4afc-90a8-20db128033be.jpg)

Target

Result

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00598v1-figure-ed1120aa-e5cf-4086-a164-26499947d00a.jpg)

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00598v1-figure-a5222a12-63de-4117-a252-050c662f7eae.jpg)

Fig. 4: More examples of color transfer.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00598v1-figure-eda2525a-aefd-487b-bdc3-6f275904849d.jpg)

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00598v1-figure-f4c54ac4-b3a2-460a-af21-4b0e7b05dba6.jpg)

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00598v1-figure-538c755d-57ef-4c08-9770-621aa2994521.jpg)

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00598v1-figure-7a44f040-6d9e-4482-b4c7-7808245944ef.jpg)

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00598v1-figure-fd4c3576-61bf-46d3-bbc4-fff712fb662f.jpg)

Learning , vol. 11, no. 5-6, pp. 355-607, 2019.

- [10] M. Cuturi, 'Sinkhorn distances: Lightspeed computation of optimal transport,' Advances in neural information processing systems , vol. 26, 2013.
- [18] H. N. Gabow and R. E. Tarjan, 'Faster scaling algorithms for general graph matching problems,' Journal of the ACM (JACM) , vol. 38, no. 4, pp. 815-853, 1991.
- [11] J. Orlin, 'A faster strongly polynomial minimum cost flow algorithm,' in Proceedings of the Twentieth annual ACM symposium on Theory of Computing , 1988, pp. 377-387.
- [12] T. Lin, N. Ho, and M. Jordan, 'On efficient optimal transport: An analysis of greedy and accelerated mirror descent algorithms,' in International Conference on Machine Learning . PMLR, 2019, pp. 3982-3991.
- [13] P. Dvurechensky, A. Gasnikov, and A. Kroshnin, 'Computational optimal transport: Complexity by accelerated gradient descent is better than by Sinkhorn's algorithm,' in International conference on machine learning . PMLR, 2018, pp. 1367-1376.
- [14] S. Guminov, P. Dvurechensky, N. Tupitsa, and A. Gasnikov, 'On a combination of alternating minimization and Nesterov's momentum,' in International conference on machine learning . PMLR, 2021, pp. 3886-3898.
- [15] B. Schmitzer, 'Stabilized sparse scaling algorithms for entropy regularized transport problems,' SIAM Journal on Scientific Computing , vol. 41, no. 3, pp. A1443-A1481, 2019.
- [16] O. Pele and M. Werman, 'Fast and robust earth mover's distances,' in 2009 IEEE 12th International Conference on Computer Vision . IEEE, 2009, pp. 460-467.
- [17] A. V. Goldberg, E. Tardos, and R. Tarjan, 'Network flow algorithm,' ´ Cornell University Operations Research and Industrial Engineering, Tech. Rep., 1989.
- [19] X. Li, D. F. Sun, and K.-C. Toh, 'An asymptotically superlinearly convergent semismooth Newton augmented Lagrangian method for linear programming,' SIAM Journal on Optimization , vol. 30, no. 3, pp. 2410-2440, 2020.
- [20] V. V. Mai, J. Lindb¨ ack, and M. Johansson, 'A fast and accurate splitting method for optimal transport: Analysis and implementation,' International Conference on Learning Representations , 2022.
- [21] M. Zhu and T. Chan, 'An efficient primal-dual hybrid gradient algorithm for total variation image restoration,' UCLA Cam Report , vol. 34, no. 2, 2008.
- [22] E. Esser, X. Zhang, and T. F. Chan, 'A general framework for a class of first order primal-dual algorithms for convex optimization in imaging science,' SIAM Journal on Imaging Sciences , vol. 3, no. 4, pp. 10151046, 2010.
- [23] A. Chambolle and T. Pock, 'A first-order primal-dual algorithm for convex problems with applications to imaging,' Journal of mathematical imaging and vision , vol. 40, pp. 120-145, 2011.
- [24] D. Applegate, M. D´ ıaz, O. Hinder, H. Lu, M. Lubin, B. O'Donoghue, and W. Schudy, 'Practical large-scale linear programming using primaldual hybrid gradient,' Advances in Neural Information Processing Systems , vol. 34, pp. 20 243-20 257, 2021.
- [25] A. Jambulapati, A. Sidford, and K. Tian, 'A direct ˜ (1 O /ϵ ) iteration parallel algorithm for optimal transport,' Advances in Neural Information Processing Systems , vol. 32, 2019.

- [26] G. Zhang, Y. Yuan, and D. F. Sun, 'An efficient HPR algorithm for the Wasserstein barycenter problem with O (Dim(P)/ ε ) computational complexity,' arXiv preprint arXiv:2211.14881 , 2022.
- [27] N. Tupitsa, P. Dvurechensky, D. Dvinskikh, and A. Gasnikov, 'Numerical methods for large-scale optimal transport,' arXiv preprint arXiv:2210.11368 , 2022.
- [28] A. Khamis, R. Tsuchida, M. Tarek, V. Rolland, and L. Petersson, 'Scalable optimal transport methods in machine learning: A contemporary survey,' IEEE Transactions on Pattern Analysis and Machine Intelligence , 2024.
- [29] T. Lin, N. Ho, and M. I. Jordan, 'On the efficiency of entropic regularized algorithms for optimal transport,' Journal of Machine Learning Research , vol. 23, no. 137, pp. 1-42, 2022.
- [30] A. Chambolle and J. P. Contreras, 'Accelerated Bregman primal-dual methods applied to optimal transport and Wasserstein barycenter problems,' SIAM Journal on Mathematics of Data Science , vol. 4, no. 4, pp. 1369-1395, 2022.
- [31] H. Ling and K. Okada, 'An efficient earth mover's distance algorithm for robust histogram comparison,' IEEE Transactions on Pattern Analysis and Machine Intelligence , vol. 29, no. 5, pp. 840-853, 2007.
- [32] G. Auricchio, F. Bassetti, S. Gualandi, and M. Veneroni, 'Computing Kantorovich-Wasserstein distances on d -dimensional histograms using ( d + 1) -partite graphs,' Advances in Neural Information Processing Systems , vol. 31, 2018.
- [33] J. Solomon, 'Optimal transport on discrete domains,' AMS Short Course on Discrete Differential Geometry , 2018.
- [34] S. Ferradans, N. Papadakis, G. Peyr´ e, and J.-F. Aujol, 'Regularized discrete optimal transport,' SIAM Journal on Imaging Sciences , vol. 7, no. 3, pp. 1853-1882, 2014.
- [35] D. F. Sun, Y. Yuan, G. Zhang, and X. Zhao, 'Accelerating preconditioned ADMM via degenerate proximal point mappings,' arXiv preprint arXiv:2403.18618 , 2024.
- [36] B. Yang, X. Zhao, X. Li, and D. F. Sun, 'An accelerated proximal alternating direction method of multipliers for optimal decentralized control of uncertain systems,' arXiv preprint arXiv:2304.11037 , 2023.
- [37] D. Kim, 'Accelerated proximal point method for maximally monotone operators,' Mathematical Programming , vol. 190, no. 1, pp. 57-87, 2021.
- [38] Q. Tran-Dinh and Y. Luo, 'Halpern-type accelerated and splitting algorithms for monotone inclusions,' arXiv preprint arXiv:2110.08150 , 2021.
- [39] F. Lieder, 'On the convergence rate of the Halpern-iteration,' Optimization Letters , vol. 15, no. 2, pp. 405-418, 2021.
- [40] R. Glowinski and A. Marroco, 'Sur l'approximation, par ´ el´ ements finis d'ordre un, et la r´solution, e par p´ enalisation-dualit´ e d'une classe de probl` emes de dirichlet non lin´aires,' e Revue franc ¸aise d'automatique, informatique, recherche op´ erationnelle. Analyse num´ erique , vol. 9, no. R2, pp. 41-76, 1975.
- [41] D. Gabay and B. Mercier, 'A dual algorithm for the solution of nonlinear variational problems via finite element approximation,' Computers &amp; mathematics with applications , vol. 2, no. 1, pp. 17-40, 1976.
- [42] G. B. Dantzig and M. N. Thapa, Linear Programming 2: Theory and Extensions . Springer, 2003.
- [43] B. Halpern, 'Fixed points of nonexpanding maps,' Bulletin of the American Mathematical Society , vol. 73, no. 6, pp. 957-961, 1967.
- [44] D. Davis and W. Yin, 'Convergence rate analysis of several splitting schemes,' in Splitting Methods in Communication, Imaging, Science, and Engineering . Springer, 2016, pp. 115-163.
- [45] Y. Cui, X. Li, D. F. Sun, and K.-C. Toh, 'On the convergence properties of a majorized alternating direction method of multipliers for linearly constrained convex optimization problems with coupled objective functions,' Journal of Optimization Theory and Applications , vol. 169, pp. 1013-1041, 2016.
- [46] J. Schrieber, D. Schuhmacher, and C. Gottschlich, 'Dotmark-a benchmark for discrete optimal transport,' IEEE Access , vol. 5, pp. 271-282, 2016.
- [47] R. Flamary, N. Courty, A. Gramfort, M. Z. Alaya, A. Boisbunon, S. Chambon, L. Chapel, A. Corenflos, K. Fatras, N. Fournier, L. Gautheron, N. T. Gayraud, H. Janati, A. Rakotomamonjy, I. Redko, A. Rolet, A. Schutz, V. Seguy, D. J. Sutherland, R. Tavenard, A. Tong, and T. Vayer, 'Pot: Python optimal transport,' Journal of Machine Learning Research , vol. 22, no. 78, pp. 1-8, 2021.
- [48] P. Kov´ acs, 'Minimum-cost flow algorithms: An experimental evaluation,' Optimization Methods and Software , vol. 30, no. 1, pp. 94-127, 2015.
- [49] J. Eckstein and D. P. Bertsekas, 'On the Douglas-Rachford splitting method and the proximal point algorithm for maximal monotone operators,' Mathematical programming , vol. 55, no. 1, pp. 293-318, 1992.
- [50] Y. Xiao, L. Chen, and D. Li, 'A generalized alternating direction method of multipliers with semi-proximal terms for convex composite conic programming,' Mathematical Programming Computation , vol. 10, pp. 533-555, 2018.
Guojun Zhang Zhexuan Gu Yancheng Yuan Defeng Sun
HOT: An Efficient Halpern Accelerating Algorithm for Optimal Transport Problems
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.