SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Nathaniel Lahn
Radford University
nlahn@radford.edu
Sharath Raghvendra
Virginia Tech sharathr@vt.edu
Kaiyi Zhang*
Virginia Tech kaiyiz@vt.edu
Optimal Transport is a popular distance metric for measuring similarity between distributions. Exact and approximate combinatorial algorithms for computing the optimal transport distance are hard to parallelize. This has motivated the development of numerical solvers (e.g. Sinkhorn method) that can exploit GPU parallelism and produce approximate solutions.
We introduce the first parallel combinatorial algorithm to find an additive  $\varepsilon$ -approximation of the OT distance. The parallel complexity of our algorithm is  $O(\log(n) / \varepsilon^2)$  where  $n$  is the total support size for the input distributions. In Massive Parallel Computation (MPC) frameworks such as Hadoop and MapReduce, our algorithm computes an  $\varepsilon$ -approximate transport plan in  $O(\log(\log(n / \varepsilon)) / \varepsilon^2)$  rounds with  $O(n / \varepsilon)$  space per machine; all prior algorithms in the MPC framework take  $\Omega(\log n)$  rounds. We also provide a GPU-friendly matrix-based interpretation of our algorithm where each step of the algorithm is row or column manipulation of the matrix. Experiments suggest that our combinatorial algorithm is faster than the state-of-the-art approximate solvers in the GPU, especially for higher values of  $n$ .
Optimal transport (OT) is a useful metric for measuring similarity between distributions and has numerous applications [2, 4, 5, 9, 12, 29], including image retrieval [28], GAN training [23], and interpolation between distributions [7]. Given two distributions  $\mu$  and  $\nu$ , this metric captures the minimum-cost plan for transporting mass from  $\mu$  to  $\nu$ .
More formally, in the optimal transport problem, we are given two discrete distributions  $\mu$  and  $\nu$  whose supports are the point sets  $A$  and  $B$ , respectively. For each point  $a \in A$  (resp.  $b \in B$ ), we associate a probability of  $\mu_{a}$  (resp.  $\nu_{b}$ ) with it such that  $\sum_{a \in A} \mu_{a} = \sum_{b \in B} \nu_{b} = 1$ . We refer to each point of  $A$  as a demand point and each point in  $B$  as a supply point. For any edge  $(a, b) \in A \times B$ , we are given a cost  $c(a, b)$ ; we assume that the costs are scaled so that the largest cost edge is 1. Let  $\beta c(a, b)$  be the cost of transporting a supply amount of  $\beta$  from  $b$  to  $a$ . A transport plan is a function  $\sigma: A \times B \to \mathbb{R}_{\geq 0}$  that assigns a non-negative value to each edge of  $G$ , indicating the amount of supply transported along the edge. The transport plan  $\sigma$  is such that the total supplies transported into (resp. from) any demand (resp. supply) node  $a \in A$  (resp.  $b \in B$ ) is bounded by the demand (resp. supply) at  $a$  (resp.  $b$ ). The cost of the transport plan, denoted by  $c(\sigma)$ , is given by  $\sum_{(a, b) \in A \times B} \sigma(a, b)c(a, b)$ . In this optimal transport problem, we are interested in finding a minimum-cost transport plan that transports all of the supply, denoted by  $\sigma^{*}$ . We also define an
*Following convention from Theoretical Computer Science, all authors are ordered in alphabetical order. The code used for the experiments reported in this paper is available at: https://github.com/kaiyiz/ Combinatorial-Parallel-OT
$\varepsilon$ -approximate transport plan to be any transport plan  $\sigma$  with a cost  $c(\sigma) \leq c(\sigma^{*}) + \varepsilon$  that transports all of the supply.
The special case where  $A$  and  $B$  each contain  $n$  points each and where every point in  $A$  (resp.  $B$ ) has a demand of  $1/n$  (resp. supply of  $1/n$ ) is called the assignment problem. In this special case, there is an optimal transport plan with a special structure; specifically, when there are  $n$  vertex-disjoint edges  $(a, b)$  with  $\sigma(a, b) = 1/n$ , they form a perfect matching. Let the cost of any matching  $M$ , denoted by  $c(M)$  be the total cost of all of its edges, i.e.,
$$
c (M) = \sum_ {(a, b) \in M} c (a, b).
$$
Given a perfect matching  $M$ , the cost of the corresponding transport plan is simply  $1 / n\sum_{(a,b)\in M}c(a,b) = (1 / n)c(M)$ . For simplicity in exposition, in the context of the assignment problem, we will uniformly scale all the demands and supplies from  $1 / n$  to 1. This does not change the optimal transport plan. It, however, increases the cost of the optimal transport plan to  $c(M)$ , an increase by a factor of  $n$ . Thus, for the assignment problem, finding the optimal transport plan is equivalent to finding a minimum-cost perfect matching  $M^{*}$ .
Similarly, for an  $\varepsilon > 0$ , after scaling the demands and supplies from  $1/n$  to  $1$ , an  $\varepsilon$ -approximate transport plan corresponds to a perfect matching  $M$  with cost  $c(M) \leq c(M^{*}) + \varepsilon n$ . We refer to such a perfect matching as  $\varepsilon$ -approximate matching. Thus, for the assignment problem, finding an  $\varepsilon$ -approximate transport plan corresponds to finding an  $\varepsilon$ -approximate matching.
Related Work: For discrete distributions, the optimal transport problem can be formulated as a minimum-cost flow problem and solved using any LP-solver. The best-known exact and approximate solver for optimal transport, as well as the assignment problem, are graph-based combinatorial algorithms [25, 19, 17, 27]. These solvers, however, are known to be difficult to parallelize. For instance, the best  $\varepsilon$ -approximate OT solver, in terms of sequential running time, was given by Lahn et al. [19]. This combinatorial algorithm runs in  $O(n^{2} / \varepsilon + n / \varepsilon^{2})$  time, and is a non-trivial adaptation of the classical combinatorial exact algorithm by Gabow and Tarjan (GT-algorithm) for the transportation problem [13]. The Lahn et al. algorithm runs no more than  $\lfloor 2 / \varepsilon \rfloor + 1$  iterations, where each iteration executes a Dijkstra's shortest path search to find and augment along a set of "augmenting paths". This algorithm is the state-of-the-art in the sequential setting; see [22] for a discussion on the various algorithms. Unfortunately, however, the flow augmentations have to be done in a sequential manner, making this algorithm hard to parallelize.
Motivated by the need for efficient scalable solutions, machine learning researchers have designed highly parallelizable approximate solvers that generate an  $\varepsilon$ -approximate transport plan in  $\tilde{O}(n^2/\varepsilon^{O(1)})$  sequential time and  $\tilde{O}(1/\varepsilon^{O(1)})$  parallel time. Perhaps the most successful among these is an entropy regularized version of the optimal transport, which can be solved using the Sinkhorn-Knopp method [11,8] and produces an  $\varepsilon$ -approximation to the optimal transport in  $\tilde{O}(n^2/\varepsilon^2)$  sequential time and  $\tilde{O}(1/\varepsilon^2)$  parallel time [10]. The simplicity of this algorithm has lead to an efficient GPU implementation. From a theoretical standpoint, Jambulapati et al. [16] designed a dual extrapolation algorithm using an area convex mapping [31] to achieve an improved parallel complexity of  $O((\log(n)\log\log n)/\varepsilon)$ . However, as noted by Lin et al. [22], despite its sound theoretical guarantees, the lack of simplicity and the difficulties of implementation make this algorithm by Jambulapati et al. less competitive, and the Sinkhorn algorithm remains the state-of-the-art for approximating the Optimal Transport on GPUs.
Despite being the state-of-the-art in sequential settings, combinatorial algorithms have remained difficult to parallelize and all known exact or approximate combinatorial algorithms run in only slightly sub-linear parallel time [14]. In this paper, we design the first parallel combinatorial algorithm that takes only  $O(\log (n) / \varepsilon^2)$  parallel time and finds an  $\varepsilon$ -approximation transport plan.
Our algorithm also improves upon existing algorithms in the massive parallel computation frameworks such as Hadoop, MapReduce, Dryad and Spark. In the Massively Parallel Computing (MPC) model, we are given a set of machines, where each machine has a bounded amount of memory. For this model, the base assumption is that communication between machines is the performance bottleneck, and the goal is to minimize the number of synchronized communication rounds, where each round consists of a period of local computation on each machine, followed by a period of message communication between machines. It is well known that any standard parallel algorithm that takes  $O(f(n))$  time can
be directly translated to an algorithm under the MPC model that runs in  $O(f(n))$  communication rounds. However, algorithms that are more specialized for the MPC model can achieve drastically faster computation times, often having a sub-logarithmic number of rounds. For example, it has long been known how to compute a maximal matching in  $O(\log n)$  parallel time [15], but only recently was a breakthrough made that shows how to compute a maximal matching in  $O(\log \log n)$  rounds under the MPC model [3]. Maximal matching is a substantially simpler problem than both the assignment and OT problems. For these OT problems, known parallel  $\varepsilon$ -approximation algorithms immediately yield an  $\tilde{O} (\log (n) / \varepsilon)$  round algorithm for the MPC model [16]. To our knowledge, no specialized MPC algorithms are known for the either problem. Thus, we provide the first sub-logarithmic round  $\varepsilon$ -approximation algorithm for both the assignment and OT problems. We obtain this bound by leveraging the recent breakthrough MPC algorithm for maximal matching by [3].
Our Results: In this paper, we present a very simple combinatorial algorithm to compute an  $\varepsilon$ -approximate transport plan in  $O(n^{2} / \varepsilon^{2})$  sequential time and  $O(\log (n) / \varepsilon^{2})$  parallel time. For the special case of the assignment problem, the sequential execution time of our algorithm improves to  $O(n^{2} / \varepsilon)$ . We also provide a GPU implementation of our algorithm that outperforms the implementation of the Sinkhorn algorithm provided by Python Optimal Transport library [11].
Our algorithm also extends to the well-known Massive Parallel Computation (MPC) frameworks such as MapReduce, Hadoop, Dryad and Spark. In the MPC model, our algorithm computes a  $\varepsilon$ -approximate transport plan in  $O(\log (\log n) / \varepsilon^2)$  rounds with  $O(n)$  memory per machine. Our algorithm is based on the popular push-relabel framework [14] for computing minimum-cost flow.
Theorem 1.1. Given an  $\varepsilon >0$ , there is an algorithm that computes an  $\varepsilon$ -approximate matching in  $O(n^{2} / \varepsilon)$  time. Furthermore, one can execute this algorithm in expected  $O(\log (n) / \varepsilon^2)$  parallel time or in expected  $O(\log (\log n) / \varepsilon^2)$  rounds, with  $O(n)$  memory per machine, in the MPC model.
Extension to the Optimal Transport problem: For any  $\varepsilon > 0$ , Lahn et al. [19] showed that computing an  $\varepsilon$ -approximation of the optimal transport between two discrete distributions containing  $n$  points in their support reduces to an instance of an unbalanced assignment problem with  $n / \varepsilon$  points. We apply this reduction and slightly adapt our algorithm from Theorem [1.1] to obtain our result for the optimal transport problem (Theorem [1.2]). In this paper, we present details of our algorithm for the assignment problem. Details of the adaptation of our algorithm to the optimal transport by using the reduction of [19] is presented in the Section B of the appendix.
Theorem 1.2. Given an  $\varepsilon >0$ , there is an algorithm that computes an  $\varepsilon$ -approximate transport plan in  $O(n^{2} / \varepsilon^{2})$  time. Furthermore, one can execute this algorithm in expected  $O(\log (n) / \varepsilon^2)$  parallel time or in  $O(\log (\log (n / \varepsilon)) / \varepsilon^2)$  rounds with  $O(n / \varepsilon)$  memory per machine in the MPC model.
From a theoretical stand-point, we provide the first parallel combinatorial algorithm for approximating the optimal transport with a expected parallel execution time of  $O(\log (n) / \varepsilon^2)$ . The sequential execution time of  $O(n^{2} / \varepsilon)$  for our algorithm for the assignment problem matches with the current state-of-the-art for the problem [16, 19]. We also provide the first sub-logarithmic round algorithm that approximates the optimal transport plan in the MPC model.
From a practical stand-point, for both the assignment problem and the OT problem, we provide an implementation that exploits GPU parallelism. Experiments suggest that both of our GPU implementations outperform the GPU implementation of the state-of-the-art Sinkhorn algorithm provided by the Python Optimal Transport library [11] in terms of running time, while achieving the same level of accuracy.
Our Approach: Our algorithmic approach is based on the popular push-relabel framework for computing network flows. For the assignment problem, our algorithm maintains a matching  $M$  and a set of dual weights  $y(\cdot)$  on vertices of  $A \cup B$ . The algorithm runs in  $O(1 / \varepsilon^2)$  iterations and in each iteration, it executes three steps. First, it greedily computes a maximal matching  $M'$ . In the second step, it uses  $M'$  to update the matching  $M$  (the push step). Finally, it updates the dual weights (relabel step). Our proof of correctness is based on the standard dual feasibility conditions used to compute minimum-cost maximum cardinality matchings, with some modifications made to better accommodate our additive-approximate setting. Our main technical difference from standard push-relabel techniques is the novel running time analysis for the additive approximate setting. In particular, we show that the number of iterations required by our algorithm is just  $O(1 / \varepsilon^2)$ . Within each iteration, the push and relabel steps take only  $O(n)$  sequential time and  $O(1)$  parallel time. The
only non-trivial step, therefore, is the computation of a maximal matching which can be done in  $O(n^{2})$  sequential time and  $O(\log n)$  parallel time [15]. Maximal matchings can also be computed in  $O(\log \log n)$  rounds in the massively parallel computation (MPC) model [3]. As a result, our algorithm can also be executed in  $O(\log (\log n) / \varepsilon^2)$  rounds in the MPC model, for the assignment problem. We extend our algorithm to also approximate the optimal transport plan by using the reduction of Lahn et al. [19] (see Section B of the appendix for details).
Organization: In Section 2.1 we present the definitions required to describe our algorithm. In Section 2.2 we present our algorithm for the assignment problem. For simplicity of exposition, we present an algorithm that computes a  $3\varepsilon$ -approximation of the optimal solution to the assignment problem. To obtain an  $\varepsilon$ -approximation, one can simply choose the error factor in the algorithm to be  $\varepsilon/3$ . In Section 3 we prove the sequential complexity of our algorithm for the assignment problem. In Section 4 we analyze the complexity of our algorithm in the parallel and MPC settings and also describe a GPU-friendly implementation of our matching algorithm. Finally, we present the experimental results in Section 5. In the appendix, Section B we extend our algorithm to the optimal transport problem. All missing proofs can also be found in the appendix, Section A.
2 Algorithm
In this section, given an input to the assignment problem and a value  $0 < \varepsilon < 1$ , we present an algorithm that computes a  $3\varepsilon$ -approximate matching.
We begin by introducing the terminologies required to understand our algorithm for the assignment problem. For any matching  $M$ , we say that any vertex  $v \in A \cup B$  is free if  $v$  is not matched in  $M$  and matched otherwise. Our algorithm critically uses the notion of a maximal matching which we introduce next. For any bipartite graph that is not necessarily complete, any matching  $M$  is maximal if and only if at least one end point of every edge in the graph is matched in  $M$ . Thus, if a matching is not maximal, there is at least one edge between two free vertices. One can, therefore, compute a maximal matching in a greedy fashion by iteratively picking such an edge and adding it to the matching.
For every edge  $(u,v)\in A\times B$ , we transform its cost so that it becomes an integer multiple of  $\varepsilon$  as follows:
$$
\bar {c} (u, v) = \varepsilon \left\lfloor c (u, v) / \varepsilon \right\rfloor \tag {1}
$$
The rounding of edge costs may introduce an error that is bounded by  $\varepsilon$  for each edge and by at most  $\varepsilon n$  for any matching. Our algorithm assigns a dual weight  $y(v)$  for every  $v\in A\cup B$  such that a set of relaxed dual feasibility conditions are satisfied. A matching  $M$  along with dual weights  $y(\cdot)$  is  $\varepsilon$ -feasible if, for every edge  $(a,b)\in A\times B$ ,
$$
y (a) + y (b) \leq \bar {c} (a, b) + \varepsilon \quad \text {i f} (a, b) \notin M \tag {2}
$$
$$
y (a) + y (b) = \bar {c} (a, b) \quad \text {i f} (a, b) \in M \tag {3}
$$
In Lemma 3.1, we show that any  $\varepsilon$ -feasible matching produced by our algorithm has a cost within an additive error of  $\varepsilon$  from the optimal solution with respect to the costs  $\overline{c}(\cdot, \cdot)$ . For any edge  $(u, v)$ , we define its slack  $s(u, v)$  to be 0 if  $(u, v) \in M$ . Otherwise, if  $(u, v) \notin M$ , we set its slack to be  $s(u, v) = \overline{c}(u, v) + \varepsilon - y(u) - y(v)$ . We say that  $(u, v)$  is admissible if the slack on the edge is 0.
We observe that any matching  $M$  whose cardinality is at least  $(1 - \varepsilon)n$  can be converted into a perfect matching simply by arbitrarily matching the remaining  $\varepsilon n$  free vertices. The cost of any edge is at most 1, and so, this increases the cost of the matching  $M$  by at most  $\varepsilon n$ . In addition to this, the rounding of costs from  $c(\cdot, \cdot)$  to  $\bar{c}(\cdot, \cdot)$  also introduces an increase of cost by  $\varepsilon n$ . Finally, the  $\varepsilon$ -feasibility conditions introduced an additional additive error of  $\varepsilon n$ , for a total error of  $3\varepsilon n$ , as desired. Thus, in the rest of this section, we present an algorithm that computes an  $\varepsilon$ -feasible matching of cardinality at least  $(1 - \varepsilon)n$ , which has a cost no more than  $\varepsilon n$  above the optimal matching's cost with respect to  $\bar{c}(\cdot, \cdot)$ .
2.2 Algorithm Details
Initially, we set the dual weight of every vertex  $b \in B$  to be  $\varepsilon$  and every vertex  $a \in A$  to be 0. We initialize  $M$  to  $\emptyset$ . Our initial choice of  $M$  and the dual weights satisfies (2) and (3). Our algorithm executes iterations, which we will call phases. Within each phase, the algorithm constructs the set  $B'$ , which consists of all free vertices of  $B$ . If  $|B'| \leq \varepsilon n$ , then  $M$  is an  $\varepsilon$ -feasible matching of cardinality at least  $(1 - \varepsilon)n$ , and the algorithm will arbitrarily match the remaining free vertices and return the resulting matching. Otherwise, the algorithm computes the subset  $E' \subseteq E$  of admissible edges with at least one end point in  $B'$ . Let  $A' = \{a \mid a \in A$  and  $(a,b) \in E'\}$ , i.e., the set of points of  $A$  that participate in at least one edge in  $E'$ . For each phase, the algorithm executes the following steps:
(I) Greedy step: Computing a maximal (i.e., greedy) matching  $M'$  in the graph  $G'(A' \cup B', E')$ .
(II) Matching Update: Let  $A''$  be the set of points of  $A'$  that are matched in both  $M$  and  $M'$  and let  $M''$  be the edges of  $M$  that are incident on some vertex of  $A''$ . The algorithm adds the edges of  $M'$  to  $M$  and deletes the edges of  $M''$  from  $M$ .
(III) Dual Update:
a. For every edge  $(a,b)\in M^{\prime}$  , the algorithm sets  $y(a)\gets y(a) - \varepsilon$  , and
b. For every vertex  $b \in B'$  that is free with respect to  $M'$ , the algorithm sets  $y(b) \gets y(b) + \varepsilon$ .
In each phase, the matching update step will add edges of  $M'$  to  $M$  and remove edges of  $M''$  from  $M$ . By construction, the updated set  $M$  is a matching. Furthermore, every vertex of  $A$  that was matched prior to the update continues to be matched after the update.
Lemma 2.1. The new set  $M$  of edges obtained after Step (II) is a matching. Furthermore, any vertex of  $A$  that was matched prior to Step (II) will continue to be matched after the execution of Step (II).
The dual update step increases or reduces dual weights by  $\varepsilon$ . Therefore, the dual weights always remain an integer multiple of  $\varepsilon$ .
The algorithm maintains the following invariants:
(I1) The dual weight of every vertex in  $B$  (resp.  $A$ ) is non-negative (resp. non-positive). Furthermore, every free vertex of  $A$  has a dual weight of 0.
(I2) The matching  $M$  and a set of dual weights  $y(\cdot)$  is  $\varepsilon$ -feasible.
Proofs of these invariants can be found in the appendix Section A.1
3 Analysis
Next, in Section 3.1, we use invariants (I1) and (I2) to show that the algorithm produces a matching with the desired accuracy. In Section 3.2, we use the invariants to bound the sequential and parallel execution times of our algorithm.
3.1 Accuracy
As stated in Section 2.1, the rounding of costs from  $c(\cdot, \cdot)$  to  $\overline{c}(\cdot, \cdot)$  introduces an error of  $\varepsilon n$ . Furthermore, after obtaining a matching of size at least  $(1 - \varepsilon)n$ , the cost of arbitrarily matching the last  $\varepsilon n$  vertices is no more than  $\varepsilon n$ . From the following lemma, we can conclude that the total error in the matching computed by our algorithm is no more than  $+3\varepsilon n$ . Proof of this Lemma can be found in the appendix Section A.2
Lemma 3.1. The  $\varepsilon$ -feasible matching of size at least  $(1 - \varepsilon)n$  that is produced by the main routine of our algorithm is within an additive error of  $\varepsilon n$  from the optimal matching with respect to the rounded costs  $\overline{c}(\cdot, \cdot)$
3.2 Efficiency
Suppose there  $t$  phases executed by the algorithm. We use  $n_i$  to denote the size of  $B'$  in phase  $i$ . By the termination condition, each phase is executed only if  $B'$  has more than  $\varepsilon n$  vertices, i.e.,
$n_i > \varepsilon n$ . First, in Lemma 3.2 (appendix A.3), we show that the magnitude of the dual weight of any vertex cannot exceed  $(1 + 2\varepsilon)$ . This means the total dual weight magnitude over all vertices is upper bounded by  $n(1 + 2\varepsilon)$ . Furthermore, in Lemma 3.3 (appendix A.4) we show that, during phase  $i$ , the total dual weight magnitude increases by at least  $\varepsilon n_i$ . From this, we can conclude that
$$
\sum_ {i = 1} ^ {t} n _ {i} \leq n (1 + 2 \varepsilon) / \varepsilon = O (n / \varepsilon). \tag {4}
$$
Note that, since each  $n_i \geq \varepsilon n$ , we immediately get  $t \varepsilon n \leq n(1 + 2\varepsilon) / \varepsilon$ , or  $t \leq (1 + 2\varepsilon) / \varepsilon^2 = O(1 / \varepsilon^2)$ . In order to get the total sequential execution time, we show, in Lemma 3.4 (appendix A.5) that each phase can be efficiently executed in  $O(n \times n_i)$  time. Combining this with equation (4) gives an overall sequential execution time of  $O(n(\sum_{i=1}^{t} n_i)) = O(n^2 / \varepsilon)$ .
Lemma 3.2. For any vertex  $v \in A \cup B$ , the magnitude of its dual weight cannot exceed  $1 + 2\varepsilon$ , i.e.,  $|y(v)| \leq (1 + 2\varepsilon)$ .
Lemma 3.3. The sum of the magnitude of the dual weights increases by at least  $\varepsilon n_{i}$  in each iteration.
Lemma 3.4. The execution time of each phase is  $O(n \times n_i)$  time.
3.3 Analysis for the Unbalanced Case