A Combinatorial Algorithm for Approximating the Optimal Transport in the Parallel and MPC Settings
Authors: Nathaniel Lahn
Sharath Raghvendra
Kaiyi Zhang
Venue: NeurIPS 2023 poster (2023)
Keywords: Optimal Transport
Combinatorial Optimization

Abstract: 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$.

A Combinatorial Algorithm for Approximating the Optimal Transport in the Parallel and MPC Settings
Nathaniel Lahn Radford University nl nlahn@radf ord. edu
Sharath Raghvendra Virginia Tech sharathrOvt.edu
Kaiyi Zhang" Virginia Tech kaiyiz@vt. edu
Abstract
Optimal Transport is 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 e- approximation of the OT distance. The parallel complexity of algorithm is (log(n))EE) where n is the total support size for the input distributions. Massive Parallel Computation (MPC) frameworks such Hadoop and MapReduce, our algorithm computes an E-approximate transport plan in ((log(ogg(n/))) (n)) rounds with Onn/e) space per machine; all prior algorithms in the MPC framework take 2(log n) rounds. We also provide a GPU-friendly matrix-based interpretation of our algorithm where each step of the algorithm is 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.
1 Introduction
Optimal transport (OT) is a useful metric for measuring similarity between distributions and has numerous applications 2 4 2
including image retrieval [88] GAN training [2]] and Z Given two distributions μ and v, this metric captures the
20
interpolation between distributions
minimum-cost plan for transporting mass from
More formally, the optimal transport problem, we are given two discrete distributions μ and whose supports are the point sets A and B, respectively. For each point E A (resp. b E B), we associate probability of P. (resp. 1b)) with such that Enen Hr, = >be 7,, We refer each point of A as a demand point and each point B supply point. For any edge (a. 6) E A 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 Ba(a, b) be the cost of transporting a supply amount of from function A x B + Rgo that assigns non-negative value to each edge of G, indicating the amount of supply transported along the edge. The transport plan o is such that the total supplies transported into (resp. from) any demand (resp. supply) node a A (resp. b E B) is bounded by the demand (resp. supply) at (resp. b). The cost of transport plan, denoted by c(o). is given by (a.M) (a, ()c b). In this optimal transport problem, we are interested in finding minimum-cost transport plan that transports all of the supply, denoted by o². We also define an
a. A transport plan is
FFllowiig 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: Paa.iatttrral-PPPallllll
com/kaiyizz
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

e-approximate transport plan to be any transport plan or with a cost c(o) S c(o²) ) + that transports all of the supply.
The special case where and B each contain n points each and where every point in A (resp. B) has demand of 1 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-disjon edges (a. b) with (a, 6) 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)
Σ c(a,b).
(a.d)EM
Given perfect matching M. the cost of the corresponding transport plan is simply 1/n E(a,)jeM c(a. = (1/ For simplicity in exposition, in the context of the assign- ment 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 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 e > 0, after scaling the demands and supplies from 1/n to 1, an e-approximate transport plan corresponds to perfect matching M with cost c(M) c(M)) + en. We refer to such a perfect matching as E-approximate matching. Thus, for the assignment problem, finding an e-approximate transport plan corresponds to finding e-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 p2) These solvers, however, are known to be difficult to parallelize. For instance, the best e-approximate OT solver, in terms of sequential running time, was given by Lahn et al. [19] This combinatorial algorithm runs in / +n/"") time, and is non-trivial adaptation of the classical combinatorial exact algorithm by Gabow and Tarjan (GT-algorithm) for the transportation problem [3.. The Lahn et al. algorithm runs more than [2/=] + 1 iterations, where each iteration exccutes a Dijkstra's s shortest path search to find and augment along a set of "augmenting paths". This algorithm is the state-of-the-ar 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.
algorithms 25 19
17
Motivated by the need for efficient scalable solutions, machine learning researchers have de- signed highly parallelizable approximate solvers that generate an e-approximate transport plan in Ó(n² ((1) sequential time and (I))) 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 [I] [ and produces an e-approximation to the optimal transport in Ö /s²)
sequential time and O(1/=²) parallel time 10. The simplicity of this algorithm has lead to an ef- ficient GPU implementation. From a theoretical standpoint, Jambulapati et al. [66 designed dual extrapolation algorithm using an area convex mapping [31 to achieve an improved parallel complexity of ((log(n) log log nE) However, noted by Lin et al. [22) despite its sound theo- retical guarantees, the lack of simplicity and the difficulties of implementation make this algorithm by Jambulapati al. less competitive, and the Sinkhorn algorithm remains the state-of-the -frtt--rttttttt for approximating the Optimal Transport on GPUs.
Despite being the state-of-the 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) e²2) parallel time and finds an s-approximation transport plan.
in sequential settings, combinatorial algorithms have remained
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 set of machines, where each machine has 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 rownds, where each round consists of 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 (f)))) time can
2

---

be directly translated to an algorithm under the MPC model that runs in (n)( communication rounds. However, algorithms that are more specialized for the MPC model can acheive 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 rounds under the MPC model [ Maximal matching is a substantially simpler problem than both the assignment and OT problems. For these OT problem, known parallel s-approximation algorithms immediately specialized MPC algorithms are known for the either problem. Thus, we provide the first sub-logarithmic round e-approximation algorithm for both the assignment and OT problems. We obtain this bound by leveraging the recent breakthrough MPC algorithm for maximal matching by B
yield an (1log(n) /e) round algorithm for the MPC model [1 To our knowledge,
Our Results: In this paper, we present a very simple combinatorial algorithm compute an e-approximate transport plan in O( /s9) sequential time and (llog(n)/s* parallel time. For the special case of the assignment problem, the sequential execution time of our algorithm im- proves to On /e). We also provide GPU implementation of our algorithm that outperforms the implementation of the Sinkhorn algorithm provided by Python Optimal Transport library [-]-
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 e-approximate transport plan in Ollg n)/=?) rounds with (nn) memory per machine. Our algorithm is based on the popular mush-relabel framework 414 for computing minimum- flow.
Theorem 1.1. Given an E O, there an algorithm that computes an E e-approximate matching in O time. Furthermore, one can execute this algorithm in expected ([log(n)/"* parallel time or in expected O(log(logn)/ //??) rounds, with (n) memory per machine, in the MPC model.
Extension to the Optimal Transport problem: For any E >> O, Lahn et al. [19 showed that computing an e-approximation of the optimal transport between two discrete distributions containing n points in their support reduces an instance of an unbalanced assignment problem with n/e points. We apply this reduction and slightly adapt our algorithm from Theorem [.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 of the appendix.
Theorem 1.2. Given an s > O, there an algorithm that computes an e-approximate transport plan in O(1² fs²) time. Furthermore, one can execute this algorithm in expected O (log(n)/e² parallel time or in (log({og(n//s))/² rounds with ((n/s) 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 (llog())/&²) The sequential exccution time of O(n² /s) for our algorithm for the assignment problem matches with the current state-of-the-art for the problem [16] [1]. 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 tfate-o-tth-aar art Sinkhorn algorithm provided by the Python Optimal Transport library [II] in terms of running time, while achieving the same level of accuracy.
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--) on vertices of A U B. The algorithm runs in OLL/²²) 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-com maximum cardinality matchings, with some modifications made to better accommodate our additive-approxiimat
Our Approach:
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/E²). Within each iteration, the push and relabel steps take only (n) sequential time and ((1) parallel time. The
3

---

only non-trivial step, therefore, is the computation of a maximal matching which can be done in Onn)) sequential time and ( n) parallel time [TS]. Maximal matchings can also be computed in Ollog log u) rounds in the massively parallel computation (MPC) model 3 As a result, our algorithm can also be executed in (logglog n)/=²) 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. 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 algorithm that computes 3e-approximation of the optimal solution to the assignment problem. To obtain an e-approximation, one can simply choose the error factor in the algorithm to be e/3. In Section we prove the sequential complexity of our algorithm for the assignment problem. In Section we analyze the complexity of our algorithm in the parallel and MPC settings and also describe 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
2 Algorithm
In this section, given an input to the assignment problem and a value 0 < < 1, we present an algorithm that computes 3c-approximate matching.
2.1 Preliminaries
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 E A U B is free if v is not matched in M and matched otherwise. Our algorithm eritically uses the notion of a maximal matching which 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 maximal matching in a greedy fashion by iteratively picking such an edge and adding it to the matching.
For every edge (u, E AA B, transform its cost so that it becomes an integer multiple of e as follows:
ε(u, u) = E[c(u, w)/e]
The rounding of edge costs may introduce an error that is bounded by for each edge and by at most en for any matching. Our algorithm assigns dual weight y(v) for every v E A U B such that set of relaxed dual feasibility conditions are satisfied. A matching M along with dual weights v(-) is e-feasible if, for every edge (a,b) E A x B,
v(a) + v(b) S c(a,b)+ y(a) + ybb = c(a,b)
if (a M if (a,
(2)
(3)
we show that any e-feasible matching produced by our algorithm has cost within an additive error of from the optimal solution with respect to the costs ((. ). For any edge (u, v). we define its slack s(u, v) to be 0 if (u, v) E M. Otherwise, if (u, v) & M, we set its slack to be s(u, v) a(u, v) + E y(u) v(v). We say that (u, v) is admissible if the slack on the edge is 0.
In Lemma 3.1
We observe that any matching A whose cardinality is at least (1 =)n can be converted into perfect matching simply by arbitrarily matching the remaining en free vertices. The cost of any edge at most 1, and so, this increases the cost of the matching M by at most en. In addition to this, the rounding of costs from -) to z((--) also introduces an increase of cost by En. Finally, the e-feasibility conditions introduced an additional additive error of en, for a total error of 3en, as desired. Thus, in the rest of this section, we present an algorithm that computes e-feasible matching of cardinality least (1 - e)n, which has cost no more than en above the optimal matching's cost with respect to -..

---

2.2 Algorithm Details
Initially, we set the dual weight of every vertex b B be e and every vertex a E A be O. We initialize . to 0. Our initial choice of M and the dual weights satisfies 2 and B.. 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 B3] S en, then M is e-feasible matching of cardinality least (1 E) n)n, and the algorithm will arbitrarily match the remaining free vertices and return the resulting matching. Otherwise, the algorithm computes the subset F' C E of admissible edges with least one end point B'. Let A' {a a E A and (a, b) E i.e., the set of points of that participate in at least one edge in E'. For each phase, the algorithm executes the following steps:
()) Greedy step: Computes maximal (i.e., greedy) matching M' the graph G' U B', E"). 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) E M', the algorithm sets y(a) + y(a) - E, and
b. For every vertex b E B' that is free with respect to M', the algorithm sets y(b) <-
y(b) +.
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 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 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 reduces dual weights by Therefore, the dual weights always remain an integer multiple e.
The algorithm maintains the following invariants:
(II) The dual weight of every vertex in B (resp. A) is non-negative (resp. non-positive). Furthermore, every free vertex of A has dual weight of O.
(12) The matching M and a set of dual weights y(-) is e-feasible.
Proofs of these invariants can be found the appendix Section A.1
Analysis
invariants (I1) and (I2) show that the algorithm produces matching with the desired accuracy. In Section 3.2 we use the invariants to bound the sequential and parallel execution times of our algorithm.
Next, in Section 3.1
3.1 Accuracy
As stated Section 2.1 the rounding of costs from el...) to c(.:)) introduces an error of en. Furthermore, after obtaining matching of size at least (1 E)n, the cost of arbitrarily matching the last en vertices is no more than en. From the following lemma, we can conclude that the total error in the matching computed by our algorithm is more than +3en. Proof of this Lemma can be found in the appendix Section A.2
Lemma 3.1. The e-feasible matching of size at least (1 e)n that is produced by the main routine of our algorithm is within an additive error of en fromrom the optimal matching with respect to the rounded costs 2
3.2 Efficiency
Suppose there phases executed by the algorithm. We use n; to denote the size of B' phase i. By the termination condition, each phase is executed only if B' has more than vertices, i.c.
5

---

n; En. First, in Lemma
3.2 (appendix A.3 we show that the magnitude of the dual weight of any vertex cannot exceed (1 1-+22 2E). This means the total dual weight magnitude over all vertices is upper bounded by r21 + 2c). Furthermore, in Lemma 3.3 (appendix A.4) we show that, during phase i, the total dual weight magnitude increases by at least Eng. From this, we can conclude that
T4 MaS m((1+2))/E O(n/s).
Note that, since each n, 2 en, we immediately get len S nn1++e)/e. or S (1+2))/e²= (142))/===0(1''
In order to get the total sequential execution time, we show, in Lemma 3.44 (appendix [A.5) that each
phase can be efficiently executed in On n4) time. Combining this with equation (1 gives overall sequential execution time of O (n n.)) = O
Lemma 3.2. For any vertex E AU B, the magnitude of its dual weight cannot exceed 1 2e, i.e. [v(t)] & (1 + 28).
Lemma 3.3. The sum of the magnitude of the dual weights increases by least en; each iteration.
Lemma 3.4. The execution time of each phase is O(n nu.) time.
3.3 Analysis for the Unbalanced Case
In this section, we describe how the analysis of our matching algorithm can extended to work for the unbalanced case, where |A| * [B). This analysis is critical for proving the correctness of our optimal transport version of the algorithm. Without loss of generality, assume [B] S |A| = n. The overall description of the algorithm remains the same, except for the main routine of our algorithm produces an e-fcasible matching of size at least (1 B). The asymptotic running time of both the parallel and sequential algorithms remains unchanged. In the following lemma, bound the additive error of our algorithm for the unbalanced case; the argument is very similar to Lemma 3.T
Lemma 3.5. Given an unbalanced input to the assignment problem with |B| S IA], the e-feasible matching of cardinality at least (1 E) B| that is returned by our algorithm is within an additive error of B| from the optimal matching with respect to the cost function c(
4 Parallel Algorithm
In this section, we describe how to parallelize the matching algorithm of Section 2.2 leading to the result of Theorem [.I] Recall that each phase of this algorithm has three steps (1) Greedy step, (II) Matching update, and (III) Dual update. Steps (II) and (III) are easily parallelizable using ((1) time. However, step (I) is nontrivial to parallelize. Fortunately, Isreali and Itai gave Ollog n) randomized parallel algorithm for computing maximal matching
an arbitrary graph (I5] Therefore, can complete step (I) by applying their algorithm as a black box. However their algorithm is very generic, applying even for non-bipartite graphs. In Section 4.1
the Isracli Itai algorithm to bound the parallel complexity of our algorithm. In Scction 4.2 we further parallelize the phases of algorithm leading to simplified variation of the Israeli Itai algorithm that is more suited for practical implementation. Finally, in Section
3 we provide a matrix- interpretation of our simplified algorithm, allowing the algorithm be easily implemented for GPUs.
4.1 Analysis of our Algorithm for Parallel and MPC Models
The Isracli Itai algorithm is designed to work on an arbitrary graph G(V, which may not be bipartite. It computes maximal matching on G in O(log n) iterations. By directly using their Olog n.) algorithm for step (1) of our algorithm, we obtain an Ologn) parallel running time for each phase. Since our algorithm executes O(1/s²) phases, we obtain a worst-case theoretical bound of O(log n/ for our algorithm.
We would also like to note that specialized algorithms for maximal matching exist for the MPC model as well. For example, it is possible to compute a maximal matching under the MPC model using just O(log(log( A))) rounds and ((n) space per machine, where A is the maximum degree of the graph. 3. As result, we are able to achieve an algorithm in the MPC model that requires only (log(log =²)= rounds with (n) space per machine.

---

4.2 Simplifying the Parallel Implementation of our Algorithm
Instead of using Israeli Itai algorithm (which works for any arbitrary graph) as black-box, we use simpler adaptation of their algorithm for bipartite graphs. The Israeli Itai algorithm executes (log iterations, where each iteration executes the following steps, using (11) parallel time:
(i) Each vertex u E V selects an incident edge (u, 2) at random, and directs it from u to v, yielding a directed subgraph R C E.
(ii) Each vertex u V selects, at random, one incoming edge from R. Let S be the set of edges chosen by this step, with all directions removed. The graph S has maximum vertex degree of 2.
(iil) Each vertex selects an edge of S at random. Any edge that is selected by both of its endpoints is added to M, and any vertex matched by this step is removed for the next phase.
The Israeli Itai algorithm is designed work for any graph that is not necessarily bipartite. In this situation, there is no way to partition the vertices into sets A and B, and so it is necessary to consider edges as having two directions; a vertex u could 'propose" an outgoing edge (u, v) (see step (i)) and also 'receive' multiple proposals as incoming edges (see step (ii). In the non-bipartite case, it is necessary for every vertex to both send and receive proposals; symmetric process. This results in subgraph S, where each vertex could have degree of 2, and an additional step is required to eliminate some of these edges in order to form a matching (see step (ii)).
vertices must be handled using a
However, in our situation, we are solely working with bipartite graphs. this situation, the vertices are divided into two sets A and B, and, as result, we do not need to process each vertex in a symmetric fashion. Instead, we can allow one side make proposals and the other side to receive proposals. As result, for each iteration of the maximal matching algorithm, we can execute the following steps, which correspond to steps (i) and (ii) in the Israeli Itai algorithm.
(a) Each vertex of B selects, at random, an incident edge, yielding subgraph S. (b) Each vertex of A, with degree at least one in S, arbitrarily selects an edge from S and adds it to the matching.
Note that, after step (b) in our approach, each vertex has most one incident edge selected. This alleviates the need for step (iii), since steps (a) and (b) alone immediately result in matching.
In addition to our simplified approach to each iteration of the Israeli Itai algorithm, we make second optimization: Within our algorithm, instead of waiting for the Israeli Itai algorithm to complete, which could require Olog n) iterations, our implementation immediately updates the matching and dual weights after each iteration before moving to the next iteration of the Israeli Itai algorithm. While this increases the number of phases, each phase becomes very simple, taking only (1)) time, resulting in practical improvements. We believe that this additional source of parallelization could lead to asymptotic improvements the parallel complexity of our algorithm. However, the proofs of Israeli Itai do not readily extend to this modified algorithm. Obtaining a tight bound on the parallel complexity of our modified algorithm is an important open question.
4.3 A GPU-Friendly Implementation
Thus far, we have described our algorithm using graph theoretic notations. However, practice, is important for our algorithm to have an efficient GPU-based implementation. In this section, we provide a matrix-base based implementation of our algorithm, using the simplified maximal matching approach described in Section 4.2
In Algorithm I we provide a pscudocode of our simplified algorithm. The algorithm resembles the one described in Section 2.2 except for the differences described in Section 4.2 It assumes that the costs were already rounded to each be an even multiple of e. The matching returned by the algorithm has cardinality at least (1 E)n and cost at most en above the optimal cost. The algorithm, as written, does not maintain dual weights explicitly, but if dual weights are required as part of the output, then, as discussed below, the algorithm can be modified to keep track of them. Note that all operations in this psuedocode are based on relatively simple matrix operations, and can be implemented easily on a GPU.
7

---

Algorithm Approximate Bipartite Matching
1: Input: W E RR 8 E Rys
2: M + Onxns S W
3: while M has more than En columns with all O's do
4: P + 0nxn
5: for all columns b with all zero entries in M do
6: Randomly select row a such that Sa,b = 0
7: Pa,b -
8:
for all a, b E [1...r] do
9: Ma. 1 if Ma,b = 1 and row of P has all O's
10: otherwise Ma. 5- Pa.s
11: for all rows a in P with at least one 1 do
12: Add to every entry in row of S
for all columns b in A with all 0 entries do
14:
f the minimum entry in column b of 5 Decrease every entry in column b of S by p
16:
return M
The algorithm takes as input an X cost matrix W and a value for the error parameter 5, and returns an n x bit matrix that describes the matching, where an edge (i. is in the matching if and only if the value at row and column j is 1. We use 1, to represent a row vector containing all I's, and xm to represent an m by n matrix of all I's. We use similar notation for vectors and matrices of all 0's. When considering an n x n matrix, we follow convention that each row corresponds to vertex of A, and each column corresponds to a vertex of B.
Next, we explain further each part of the algorithm. Line 3 initializes the slack matrix S to reflect the initial slacks of all edges, which are initially non-matching. Throughout the algorithm, the matrix S will reflect the slack with respect to the edge 'as if it were a non-matching edge', i.e., S - W., + y(a)- - y(b). regardless of the matching status of the edge. This makes the slacks easier to track without the need to explicitly maintain dual weights.
The main loop of the algorithm, beginning at line 4, specifies the stopping condition of the algorithm. The algorithm terminates once at most en free vertices of B remain. Each iteration of this main loop is phase. Lines 5-7 compute matrix P. which represents set of edges that will be added to M during the current phase. This edge set will be a matching on the admissible edges that are incident on free vertices of B. This corresponds to Step (I) of the algorithm from Section 2.2 except for P is not necessarily maximal. Lines 8-10 update the matching M by adding edges of P to M, and removing any preexisting edges of M that are matched in P. This corresponds to Step (2) of the algorithm from Section 2.2 Finally, lines 11-15 update the slacks to reflect dual weight adjustments, corresponding to Step (3) of the algorithm from Section 2.2 However, instead of tracking the dual weights explicitly, we simply update the slacks directly. Note that, when updating the slacks on edges incident on free vertices of B, we include slight change to Step (3). Instead of increasing the dual weight of free vertices of B by exactly we increase it as much as possible. For some free vertices of B, the increase will be 0 (since P is not maximal), but is also possible, in practice, for the increase to be larger than E.
5 Experiments
In this section, we present experimental results. We implemented the parallel version of our algorithm for both the assignment problem as well as the OT problem. Both implementations are written in Python using the Py PyTocc library, which supports GPU operations. We compare these implementations of our algorithm to the Sinkhorn algorithm implementation in the Python Optimal Transport (POT) library (I]. This Sinkhom implementation also PyTorch. Additionally, we compare our algorithm to a CUDA based GPU implementation and present the results of this comparison in appendix section
Our experiments are run with an Intel Xeon E5-268004 2.4GHz chip and an Nvidia Tesla V100 GPU. We ran our algorithms using both real and synthetic data, including four different settings:
8

---

an assignment problem between randomly generated synthetic point sets, an OT problem between randomly generated point sets, each having randomly assigned demands and supplies, an assignment problem formed from two sets of MNIST images, and an OT problem between two text word embeddings. For each setup, we generated input data and computed the assignment or OT cost using our algorithm with different values of e. Then, we determined the appropriate regularization parameter of the Sinkhorn algorithm, ensuring the Sinkhorn distance is close but no lower than the cost of the solution generated by our algorithm. We recorded the running time and the number of parallel rounds for both Sinkhorn and our algorithm. We also repeated each experiment using a reversed process by fixing the regularization parameter of Sinkhorn and searching for the value for our algorithm, which guarantees the our cost is similar to, but no more than, the Sinkhorn distance; the results for these reversed experimental setups can be found in the technical appendix Note that, we also recorded the execution time of solving for the exact solution using POT"s EMD function, which runs the CPU. We only present for values of & that large enough such that either our algorithm or Sinkhorn algorithms run faster than the exact algorithm. We also record the additive error, relative to the optimal solution, of both Sinkhorn and our algorithm in appendix section C.. Additionally, we conducted experiments comparing the sequential performances of our algorithm and Sinkhorn on the CPU, which can be found in the appendix section C.4
Synthetic Data:
For synthetic data generation, for both the assignment and OT experiments, we randomly sampled the location of two groups of 10, 000 vertices, A and B, in 2-dimensional unit square. For the assignment problem, the demand or supply of every vertex is 1 1/n. For the OT problem, the capacity of each vertex is initially chosen by selecting, uniformly random from the interval [0, 1]. Then, the capacities are normalized such that the total supply and demand are both equal 1. For any pair of points ba, E A B, the cost c(a, b) was set as the squared Euclidean distance between them. For each value of e, we executed 10 runs. For each combination of e, and algorithm choice, we averaged both the running times as well as the number of parallel rounds over all 10 runs and recorded the results. The results of running times and parallel rounds can be seen in Figure I (b) and Figure 2 (b)) respectively.
Syntheti Resulss
Syrthaic Reselts Avsigrment
MNIST Rouls Aniggment
100
100
100
s-ihhom
-arkkhern
sinkheen -oer sigonihm
algocthun
goritbn
10
10
10
Time's
Timels
0.1
0.1
0.0007
0,002
0.006
0.01
0.0007
0.005
0.01
0.08
0.05
0.1
0.15
0.2
(6)
INP Results OT tome
NLP Rewlls OT 20sews
NLP Resuls OT- IMDB
10
sinihom
sinkhora
stakhom
our algociti
our algorittm
-ur algootiin
Timnss
Tino's
0.1
0.1
0.1
0.01
0.01
0.01
02
04
0.6
0.12
0.2
0.6
0.8
0.2
08
(5)
66)
(b
Figure Plots of running times on GPU for the synthetic inputs (b) and the real data inputs (t).)
Next, we ran a similar experiment using real-world image data. We generated our inputs using the MNIST dataset of hand-written digit images 21.. Each image consists of 28 x 28 pixel gray-scale image. The sets A and B each consist of n = 10, 000 images from the MNIST dataset, selected at random. The cost c(a, b) between two images a E A and b E B is computed as follows: Let a(i.j) (resp. b(i,j) be the value of the pixel in row and column j of image a (resp. b). First, the two images are normalized so that the sum of all pixel values is equal I for each image, i.e.. (e a(i, j) = and Exje(1. b(i,j) = 1. Then, the cost c(a, b) is given by the L, distance between the resulting normalized images: c(a, b) Ex.e((1.3) ja(i,j) b(i.j)). Note that an upper
MNIST:
9

---

bound on the largest cost is 2. For each algorithm and for each value of E, we averaged both the running times as well as the number of parallel rounds over 10 runs. The results for these experiments can be found in the plot in Figure (c) and the plot in Figure c).
Systhetic OOr
Systheic Romilts Aovignment
MNIST Aniggment
10000
100000
sinkbeen Aigochana
sinkkern
siibbo
aigottitoo
our algecitt
10000
Puulid Rounks
1000
Paulle PanledR
Preuitd
1000
1000
100
100
100
10
0.002
0.0007
0.000
0.006
0.01
0.2
(4)
(b)
NLP Roulls -Ime
NIP Roulls- OT 20news
NLP Beaths
DMDD
10000
10000
10000
sinkhocz
sitktor
OLE dlgorttm
algonihm
our alginim
1000
Porillt Bounds
1000
Pualee Rouds
1000
Pasild
100
100
10
10
0.2
04
0.6
088
0.12
0.2
0.5
0.6
0.8
0.3
0.2
0.4
0.6
08
Figure 2: Plots of parallel rounds on GPU for the synthetic inputs (a)(b) and the real data inputs (f).
In our final experiment, considered OT with natural language data. We calculate the document distances based on OT following the procedure of previous work [8]. Each OT instance was generated by selecting two discrete sections of text with fixed lengths. Each unique word in the first (resp. second) section of text corresponds to a vertex of A (resp. B) in the OT problem. To acquire the supply and demand for each vertex, we tokenized each text section with NLTK (. and counted the number of appearances of each unique token. Then the counts were normalized so that the total supply and demand were each equal to 1. Next, to generate the costs between vertices, we represent each unique token using a 100-dimensional Glo word embedding [o The cost of any edge (a, b) e A B is then given by the Euclidean distance between the corresponding points in this embedding.
NLP:
In the last three plots in Figure T and Figure we show the results of applying this experimental setup, using sections of the text from The Count of Monte Cristo, 20News, IMDB. For each dataset, five different OT instances created, using different sections of the text, and the results averaged over all 5 runs. These experiments can be found in Figure (I) and Figure 2 d)()(().
In all our experiments, our new parallel combinatorial algorithm almost always runs significantly faster than the Sinkhorn algorithm for the OT problem often with significantly fewer parallel rounds. Unlike the POT's highly optimized implementation of the Sinkhorn method, the implementation of our algorithm is new and may benefit significantly from further optimizations.
6 Conclusion
In this work, we provided fast, highly parallelizable combinatorial algorithm for computing an e-approximate solution to the assignment problem and the OT problem. We also provided a practical implementation of a slight variation of our algorithm, which outperforms the Sinkhorn algorithm, in terms of rumning times, in our experimental comparison. In light of this work, we would like to propose the following open question: Is it possible to improve the (llg(n))EE) parallel running time of our algorithm, possibly by introducing an appropriate modification? In particular, our simplified practical algorithm presented in Section 4.2 seems to execute fewer parallel rounds than the our worst-case theoretical analysis might suggest. Can the simplifications used in our practical implementation be used to improve our worst-case theoretical running times?
10

---

Acknowledgement
We would like to acknowledge Advanced Research Computing (ARC) at Virginia Tech, which provided us with the computational resources used to run the experiments. Research presented in this paper was funded by NSF CCF-1909171 and NSF CCF-2223871, We would like to thank the anonymous reviewers for their useful feedback.
References
[1] Jason Altschuler, Jonathan Weed, and Philippe Rigollet. Near-linear time approximation algorithms for optimal transport via sinkhorn iteration. In NIPS, pages 1961-1971, 2017.
[2] Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein GAN. 1711::170. [stat.ML], 2017.
[3] Soheil Behnezhad, Mohammad Taghi Hajiaghayi, and David G Harris. Exponentially faster massively parallel maximal matching. In 2019 JEEE 60th Annual Symposium on Foundations of Computer Science (FOCS), pages 1637-1649. IEEE, 2019.
[4] Jean-David Benamou, Guillaume Carlier, Marco Cuturi, Luca Nenna, and Gabriel Peyré. Itera- tive bregman projections for regularized transportation problems. SIAM Journal on Scientific Computing, 37(2):A1111AA
2015.
[5] Jérémie Bigot, Raúl Gouet, Thierry Klein, Alfredo López, et Geodesic PCA the wasserstein space by convex PCA. In Annales de l'Institut Henri Poincaré, Probabilités et Statistiques, volume 53, pages 1-26. Institut Henri Poincaré, 2017.
[6] Steven Bird, Ewan Klein, and Edward Loper. Natural language processing with Python: analyzing text with the natural language toolkit. "'Reeilly Media, Inc.", ", 2009.
[7] Nicolas Bonneel, Michiel Van De Panne, Sylvain Paris, and Wolfgang Heidrich. Displacement interpolation using lagrangian mass transport. In Proceedings of the 201 SIGGRAPH Asia Conference, pages 1-12, 2011.
[8] Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In NIPS, pages 2292-2300, 2013.
[9] Marco Cuturi and Arnaud Doucet. Fast computation of wasserstein barycenters. In International Conference on Machine Learning, pages 685-693, 2014.
[10] Pavel Dvurechensky, Alexander Gasnikov, and Alexey Kroshnin. Computational optimal transport: Complexity by accelerated gradient descent better than by sinkhorn's 's algorithm. In International Conference on Machine Learning, pages 1367-1376. PMLR, 2018.
[11] Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélic Boisbunon, Stanislas Chambon, Lactitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1-8, 8, 2021.
[12] Rémi Flamary, Marco Cuturi, Nicolas Courty, and Alain Rakotomamonjy. Wasserstein discrim- inant analysis. Machine Learning, 107(12)1123-1945, 2018.
[13] Harold N Gabow and Robert E Tarjan. Faster scaling algorithms for network problems. SIAM Journal on Computing, 1813):0013 1989.
[14] Andrew Goldherg, Serge A Plotkin, and Pravin M Vaidya. Sublinear-time parallel algorithms for matching and related problems. Journal of Algorithms, 14(2):80-213,, 1993.
[15] Amos Israeli and Alon Itai. fast and simple randomized parallel algorithm for maximal matching. Information Processing Letters, 22(2):77-80, 1986.

---

[16] Arun Jambulapati, Aaron Sidford. and Kevin Tian. A direct tilde O(1/E) iteration parallel algorithm for optimal transport. Advances in Neural Information Processing Systems, 32, 2019.
[17] Harold Kuhn. Variants of the hungarian method for assignment problems. Naval Research Logistics, (4))2533-588 1956.
[18] Matt Kusner, Yu Sun, Nicholas Kolkin, and Kilian Weinberger. From word embeddings to document distances. In International conference on machine learning, pages 957-966. PMLR, 2015.
[19] Nathaniel Lahn, Deepika Mulchandani, and Sharath Raghvendra. graph theoretic additive approximation of optimal transport. In Advances in Neural Information Processing Systems 32, pages 13813-13823, 2019.
[20] Nathaniel Lahn and Sharath Raghvendra. A weighted approach the maximum cardinality bipartite matching problem with applications in geometric settings. Journal of Computational Geometry, (1),), 2021. Special Issue of Selected Papers from SoCG 2019.
[21] Yann LeCun. The MNIST database handwritten digits. http://yann. lecun. com/exdb/nnistl, 1998. License: Creative Commons Attribution-Shurt Alike 3.0.
[22] Tianyi Lin, Nhat Ho, and Michael Jordan. On the efficiency of entropic regularized algorithms for optimal transport. Journal of Machine Learning Research, 23(137): 2022.
[23] Huidong Liu, GU Xianfeng, and Dimitris Samaras. A two-step computation of the exact gan wasserstein distance. In International Conference on Machine Learning, pages 3159-3168. PMLR, 2018.
[24] Vien V Mai, Jacob Lindbäck, and Mikael Johansson. A fast and accurate splitting method for optimal transport: Analysis and implementation. arXiv preprint arXiv:2110.1173, 2021.
[25] James B Orlin. A polynomial time primal network simplex algorithm for minimum cost flows. Mathematical Programming, 78(2):109-129, 1997.
[26] Jeffrey Pennington, Richard Socher, and Christopher D. Manning. Glove: Global vectors for word representation. In Empirical Methods in Natural Language Processing (EMNLP), pages 1532-1543, 2014.
[27] Abhijeet Phatak, Sharath Raghvendra, Chittaranjan Tripathy, and Kaiyi Zhang. Computing all optimal partial transports. In The Eleventh International Conference on Learning Representa- tions, 2022.
[28] Yossi Rubner, Carlo Tomasi, and Leonidas J Guibas. The earth mover's distance as a metric for image retrieval. International journal of computer vision, 40(2):99-121, 2000.
[29] Roman Sandler and Michael Lindenbaum. Nonnegative matrix factorization with earth mover's distance metric for image analysis. IEEE Transactions on Pattern Analysis and Machine Intelligence,
2011.
[30] R. Sharathkumar and P. K. Agarwal. Algorithms for transportation problem in geometric settings. In Proceedings of the 23rd Annual ACM-SIAM Symposium on Discrete Algorithmss pages 306-317, 2012.
[31] Jonah Sherman. Area-convexity, I-infinitty regularization, and undirected multicommodity flow. In Proceedings of the 49th Annual ACM SIGACT Symposium on Theory of Computing, pages 452-460,4 2017.
12


Reviews:
[{"review_id": "etVIV6gUAM", "reviewer": "Reviewer_mqPb", "summary": "The paper considers the Optimal Transport problem, which is a metric for the distance between distributions, used by some ML/DL methods. It presents a combinatorial algorithm for that problem, which can be parallelized and implemented efficiently both in map-reduce systems (MPC) and on a GPU. It also addresses a special case of that problem, the assignment problem. \nThe paper proves that the worst-case run-time complexity of the presented algorithm is O(n^2/ \\epsilon^2) (for assignment O(n^2/ \\epsilon)) in sequential time. The parallel time is O(\\log n / \\epsilon ^2). This improves upon the Sinkhorn-Knopp method, which is the state-pf-the-art method for this problem used in the Python optimal transport package. Although this run-time does not improve upon the theoretic results of Jambulapati et al., the authors explain that their algorithm is difficult to implement and not used. For MPC the paper shows that O(\\loglog n / \\epsilon ^2) rounds are required with O(n) memory per machine, the first sub-logarithmic combinatorial algorithm for that problem.\nFinally, the paper presents experimental results on a GPU, showing significant run-time improvements (sometimes above 10x) over the Sinkhorn-Knopp method for six datasets in most of the relevant ranges of epsilon (two for assignment, four for optimal transport).", "questions": "Questions to the authors:\n- Can you share the code, to enable reproducibility and benefit the community?\n- Can you present an experimental comparison with the algorithm of Jambulapati et al.?\n- Can you present an experimental comparison in one of the real scenarios in which optimal transport was used in ML/DL? (e.g., from one of the papers cited)", "limitations": "Yes, I don't see an issue regarding this.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper considers an interesting problem and clearly presents a new algorithm for solving it, showing both theoretic and experimental results. The new results improve upon the previous methods published in NeurIPS several years ago. The algorithm creatively combines and modifies existing ideas. The presented results are comprehensive, addressing different systems on which the algorithm may run. The special case of the assignment problem is interesting on its own. ", "weaknesses": "The paper has several weaknesses:\n- The code used for the experiments was not provided, so the results cannot be reproduced and might not benefit the community.\n- The algorithm of Jambulapati, which is the best asymptotically based on theoretical analysis (for parallel time), was not a part of the experimental comparison.\n- The experimental results use \"toy examples\" of datasets with some ranges of epsilon. It is unclear if the presented algorithm is the best choice for practical scenarios, such as those in the cited papers that use optimal transport (or whether linear programming or the exact algorithm could be the best choice there).\n- Although this paper improves upon papers published in NeurIPS several years ago, it may have a bigger relevant audience in conferences focused on combinatorial algorithms.\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\n*The code used for the experiments was not provided...* Actually, our code was provided as part of the supplement. You can find it in the provided zip file, alongside the appendix PDF. The code also includes a detailed README that explains how to run the experiments, specifically meant to facilitate reproducibility.\n\n*The algorithm of Jambulapati ... was not a part of the experimental comparison.* Although the work by Jambulapati et al. gives the asymptotically fastest bound in theory, there is good reason to believe that the overhead of the algorithm causes it to be less efficient than Sinkhorn in practice. Lin et al. [21] agree with this assessment, stating: \"Despite the theoretically sound complexity bound, the lack of simplicity and ease-of implementation make this algorithm less competitive with Sinkhorn and Greenkhorn which remain the baseline solution methods in practice.'' Even the Jambulapati et al. paper itself does not claim to outperform Sinkhorn.  Overall, since Sinkhorn remains a better baseline in practice, we did not see a need to compare to the Jambulapati et al. Furthermore, we are not aware of any publicly available GPU implementation of their algorithm.\n\n*Can you present an experimental comparison in one of the real scenarios in which optimal transport was used in ML/DL?* This was the intent of our current NLP experiments, which are similar in principle to the so-called `Word Movers Distance', a very well known method of comparing documents. See the paper ``From Word Embeddings To Document Distances'', which appeared in the 32nd International Conference on Machine Learning (ICML) and has over 2000 citations.\n"}, {"review_id": "5W3yFUWmu0", "reviewer": "Reviewer_pKrQ", "summary": "The paper presents a parallel combinatorial algorithm for the optimal transport (OT) and minimum weight perfect matching. The algorithm computes an additive epsilon approximation.\n\nThe algorithm runs in O(n^2 / epsilon^O(1)) time and (in a straightforward way) gives an algorithm running in O(log log n / epsilon^O(1)) MPC rounds, using linear space per machine.\n\nThe paper also compares a GPU implementation of the algorithm with the Sinkhorn method showing improved running time.", "questions": "1. What is the work and depth (aka span) of the algorithm described in 4.2? (in the shared-memory work-depth model)\n1. What is the reason for not including DROT results in section 5? How does DROT compare to Sinkhorn?", "limitations": "Yes", "rating": 6, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "* The paper gives a simple and clean algorithm for a practically relevant problem.\n* The algorithm delivers good speedups over the baseline", "weaknesses": "The main weakness of the paper is a very limited empirical evaluation. The 1.5-page long section on empirical evaluation contains mostly the dataset descriptions (which in my opinion could easily go to the supplementary material if space is a concern) + a short running time comparison to a single baseline. I'd like to see a more in-depth comparison, for example (I'm trying to suggest ideas here, not define a list of requirements):\n* a comparison between CPU implementations\n* analysis of the number of rounds used\n* evaluation on non-dense graphs\n* quality analysis: what is the actual absolute / relative error as a function of epsilon\n\nMinor suggestions:\n * I don't think describing the Israeli-Itai algorithm is a good use of the space (I'd suggest removing it), given that the algorithm ends up using a different idea (which I consider folklore)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\n*The paper presents a parallel combinatorial algorithm for the optimal transport (OT) and maximum weight matching.*\nWe believe that the reviewer has misunderstood the problem being solved in this paper. \nWe present an algorithm for the *minimum* weight matching (i.e., the assignment problem) which is a significantly harder to approximate than the *maximum* weight matching. For instance, the greedy algorithm is known to provide a $1/2$ approximation for the maximum weight matching problem. However, for minimum weight matching, the greedy approach has an unbounded approximation ratio and, even for points on a line, the approximation ratio improves only to $\\Omega(n^{0.59})$ (see citation below). Given the reviewer's misunderstanding of the problem, we believe they might have given an inaccurate evaluation of the paper.\n\n*It is possible to black-box reduce maximal matching to $(1+\\varepsilon)$ approximate matchings in MPC (see Corollary 1.3 of ...* Again, this result is for approximating the maximum weighted matching, which, as we discussed previously is a different problem than minimum-weight matching.\n\n*I'd like to see a more in-depth comparison, for example ... CPU implementations.*\nWe have already provided a CPU comparison results for our sequential algorithm; see appendix section C.3.\n\n*Analysis of the number of rounds used ...* \nThis is a good suggestion, one that we should have included, because it gives an implementation-independent gauge of algorithm convergence. We reran all of our experiments. For OT, we found that our algorithm consistently executes something like 10x fewer iterations than Sinkhorn. For assignment problem, it is more mixed. This corresponds to what we observe in our actual running times, and demonstrates that our speedup is intrinsic, rather than simply a matter of implementation. For the final version of the paper, we will add iteration counts for all experiments. \n\n*Evaluation on non-dense graphs ...* \nThe OT and assignment problems have been formulated on complete graphs. So, testing the algorithm on sparse graphs is not applicable for the problem.\n\n*Quality analysis ...* In each of our experiments, we already do compare the actual errors produced by the algorithms involved, not just the input values of the error parameters. In the Sinkhorn comparisons, we select the error parameter input value so that the actual error produced by Sinkhorn is within 1E-6 of our algorithm’s actual error. Thus, the current comparison is fair in terms of the actual produced solution quality. However, we understand that including concrete values for the actual errors produced in comparison to the optimal may still be valuable information, and we will update the figures to include this information. \n\n*What is the work and depth (aka span) of the algorithm described in 4.2?*\n\nAs mentioned in Section 4.1 L262--266, plugging in the parallel algorithm by Israeli-Itai to our algorithm from Section 3 leads to a parallel implementation with a depth of $O(\\log n/\\varepsilon^2)$. The work is dependent on the total work done by all invocations of the Itai Israeli algorithm. While they do not specifically give a bound on work, they do assume at most $|E|$ processors are used, so the work is $|E| \\log |E|$, where $E$ is the set of edges in the admissible graph. Using equation (4), the total work can therefore be bounded by $O(n^2 \\log(n) / \\varepsilon)$. \n\nThe modifications presented in Section 4.2 are motivated by a need for practical implementation on the GPU. In particular, we adapt and simplify the Israeli-Itai's algorithm for general graphs to the bipartite setting. Although, we haven't formally analyzed the span of this practical implementation, we believe it will be the same as that of the standard implementation. After a thorough check, we will include a discussion on the span of this implementation in the next version of the paper.\n\n*What is the reason for not including DROT results in section 5? How does\nDROT compare to Sinkhorn?*\nWe have included a comparison of DROT with our algorithm in appendix section C.2 and our algorithm generally outperforms DROT.\n\nThe implementation of DROT is in CUDA whereas our algorithm as well as Sinkhorn are implemented in PyTorch. CUDA implementations are faster than those on PyTorch, and despite this, our algorithm generally outperformed DROT. Nonetheless, due to a difference in implementation platform, we decided not include these results in Section 5 but to move it to the appendix. For the same reason, we decided not to compare DROT against Sinkhorn.  \n\nCitation : \nReingold, E. M., Tarjan, R. E. (1981). On a greedy heuristic for complete matching. SIAM Journal on Computing, 10(4), 676-681."}, {"review_id": "veeqYe5ysO", "reviewer": "Reviewer_dyx3", "summary": "This paper presents a novel combinatorial algorithm that finds an $\\epsilon$-optimal transport plan for the optimal transport problem. Additionally, it introduces a variant of this algorithm in the MPC model, exhibiting an expected $O(\\log\\log n)$ communication rounds, with each machine having a memory of $O(n)$ order. This algorithm offers speedups over previous approaches, which required at least $\\Omega(\\log n)$ rounds, in the large-$n$ scenarios. Furthermore, the authors provided a GPU-friendly interpretation of their algorithm utilizing matrix operations, and supported their results using numerical experiments.", "questions": "Correspondingly, I have two questions that, if addressed, could make the results even stronger:\n\n1. Using similar combinatoric techniques, is it possible to obtain an algorithm that has better performance than previous works in parallel time?\n2. Can further improvements be made to reduce the dependence on $\\epsilon$, or are there practical reasons why the focus remains primarily on the dependence on $n$?\n\nMinor comments:\n1. It would be good to have a discussion on open questions.\n2. Page 2, line 40: \"one\" -> \"when\"\n3. Please check the format of reference [16,22,23,24]. In particular, the title of [16] and the author lists of [22-24] are not presented correctly.", "limitations": "Not relevant in my opinion.", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The primary contribution of this paper is a combinatorial algorithm for approximate optimal transport, which exhibits efficient parallelization and significantly outperforms previous algorithms in the Massive Parallel Computation (MPC) model with respect to $n$, the total support size of the input distributions. The algorithm is simple and elegant, allowing for a GPU-friendly interpretation based on matrix operations, thereby enhancing its competitiveness in real-world scenarios.\n\nThe analysis of the algorithm is relatively straightforward and easy to follow, while the conducted numerical experiments provide compelling evidence for its effectiveness.", "weaknesses": "1. This paper only improves upon previous works in the MPC model rather than the broader parallel settings. In the context of the parallel time, both this work and previous algorithms based on Sinkhorn-Knopp methods are of order $O(\\log n)$. Consequently, this limitation might restrict the applicability of this algorithm in real-world scenarios with different parallelization settings.\n\n2. Compared to previous works, this algorithm has a larger dependence on $\\epsilon$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\n*Using similar combinatoric techniques, is it possible to obtain an algorithm that has better performance than previous works in parallel time?*\n\nWe believe the improving our parallel OT epsilon dependency from $O(1/\\varepsilon^2)$ to $O(1/\\varepsilon)$ is plausible, but difficult – we have spent considerable time trying to do so. Designing such an algorithm remains an interesting open question. \n\n*Can further improvements be made to reduce the dependence on $\\varepsilon$, or are there practical reasons why the focus remains primarily on the dependence on $n$?*\n\nTypically, in many applications, one may be happy with obtaining OT-cost with a precision of two decimal points, i.e., $\\varepsilon = 0.01$. In contrast, the input size may vary by large numbers. For instance, an image with $n$ pixels can be treated as a distribution with $n$ points in its support. The number of pixels in different images may vary by a lot. Therefore, typically, it is desirable to have algorithms with better dependence on the input size as opposed to a better dependence on the error parameter. \n\n*It would be good to have a discussion on open questions.* Yes, we can add one using the extra page of space in the final version."}, {"review_id": "cpSQAuQ5ri", "reviewer": "Reviewer_i1bb", "summary": "This paper gives the first parallel, combinatorial algorithm to approximate the optimal transport distance. Moreover, the algorithm nearly matches the run-time of the best known parallel algorithms for the problem, with expected parallel runtime of O(log(n)/eps^2) (for eps*n the additive error of the OT cost). The motivation for this work comes from the fact that the previous best theoretical guarantee for this problem—an O(log n/eps) runtime parallel algorithm by Jambulapti et al (Neurips 2019)—is not really implementable. Thus practitioners normally use the much easier to implement Sinkhorn algorithm, which has parallel complexity \\tilde{O}(log n/eps^O(1)). Additionally, while there is another combinatorial for the sequential problem, it was not able to be parallelized.\n\nFurther, in the MPC setting, the authors show that one can  execute their algorithm in the MPC setting with O(log log (n/eps)/ eps^2) rounds with O(n/eps) memory per machine.  Lastly the authors implement their algorithm to exploit GPU parallelism and show that in nearly all settings reported, they beat the Sinkhorn algorithm.", "questions": "Comments to the authors:\n- Remove sentence in line 21-23, you already say this at a more relevant part later\n- R_{>0} -> \\mathb{R}_{>0} in line 30\n- Line 72 “achieve a improved” -> “achieve an improved”\n- Is the Sinkhorn runtime tight around O(log n/eps^O(1))? Or empirically is it better than that guarantee? Also can you explicitly state whatever constant the O(1) term is hiding at least once?\n- Line 179 “which will has a ” -> “which has a ” \n- This algorithm is so reminiscent of push-relabel that I kept looking for more explicit direct comparisons to push-relabel in section 2. In fact step 2 of the algorithm is just an augmenting path and you keep an eps-feasible dual, similar to the heights in the push-relabel algorithm. Then even your proof of correctness is the same as push-relabel. Can you say more about the inspiration from push-relabel?\n- You should add a figure for the algorithm, maybe 3 sub-figures, one for each of the 3 steps\n- Line 308 “we have describe our” ->  “we have described our” \n- I think the lines you refer to in algorithm 1 are all messed up. Double check me on this but I think where you say “Lines 5-9” should be “Lines 5-7”, where you say “Lines 10-13” should be “Lines 8-10”, and where you say “Lines 14-20” should be “Lines 11-15”.\n- Lime 694 in appendix needs to be rewritten", "limitations": "na", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- I find the motivation for this paper to be very interesting. While it is in theory important that we push worst-case guarantees (such as the Jambulapti et al work), it is important to design algorithms that are implementable and still have good results. This combinatorial, parallelizeable algorithm certainly does this.\n- As someone more combinatorially minded, I liked the algorithm. It’s basically a push-relabel type algorithm.\n- Empirically, the proposed algorithm does well enough compared to the Sinkhorn algorithm that it might even be that the analysis in the paper is a bit lossy. Empirical evidence was convincing enough for me.\n- Well Written", "weaknesses": "- In terms of raw numbers, there’s no big theoretical improvement. (But again, that’s not the main point of the paper.)\n- The presentation in the paper could be better at times. To be clear, the writing is very very good, but the algorithm is so simple in a nice way, that I think they could’ve made points clearer. I will elaborate in comments below. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\n*Is the Sinkhorn runtime tight around $O(\\log n/\\varepsilon^{O(1)})$? ... Also can you explicitly state whatever constant the $O(1)$ term is hiding at least once?* Over time, the analysis of the Sinkhorn algorithm has seen multiple improved bounds to its $\\varepsilon$ dependency. To our knowledge, the best bound is given by Dvurechensky et al. (see below for citation), which gives a sequential $\\tilde{O}(n^2/\\varepsilon^2)$ bound, and is easily parallelizable to $\\tilde{O}(1/\\varepsilon^2)$, basically matching our algorithm's provable bounds. (We use $\\tilde{O}$ to hide logarithms here). We will edit the paper to clarify this.\n\n*Or empirically is it (Sinkhorn) better than that guarantee?* This seems to be the case, based on our results, and prior work with Sinkhorn that we have seen. Our algorithm seems to also outperform its own theoretical expectations.\n\n*This algorithm is so reminiscent of push-relabel that ...*\nThank you for pointing this out. Indeed, our algorithm is using the popular push-relabel framework and our proof of correctness is also based on the standard feasibility conditions for matchings. Our main technical novelty is in the efficiency analysis of our algorithm. Using the fact that edge costs are small integer multiples of $\\varepsilon$, we bound the number of push-relabel iterations by $1/\\varepsilon^2$. This allows us to achieve a parallel complexity of $O(\\log n/\\varepsilon^2)$. We will include this in the discussion of \"Our Approach\" (L131-141).  \n\nDvurechensky, P., Gasnikov, A., Kroshnin, A. (2018, July). Computational optimal transport: Complexity by accelerated gradient descent is better than by Sinkhorn’s algorithm. In International conference on machine learning (pp. 1367-1376). PMLR."}]

Author rebuttal:
We would like to thank all the reviewers for their helpful suggestions. The reviewers pointed out some minor grammatical and presentation changes; we appreciate these suggestions and will address them in the final paper version. For questions and comments that warrant a more significant discussion, we have addressed our responses directly to the relevant reviewer. When addressing such questions or comments, we will quote directly from the reviewer, using italics, in order to concisely reference the relevant portion of the review.

Meta-review: The paper provides an elegant parallel algorithm that is easy to implement and it outperforms the popular Sinkhorn algorithm in the experiments. Compared to the state of the art theoretical algorithm, the proposed algorithm is more practical and easier to implement. The reviewers appreciated the algorithm's simplicity and practical merit. Overall, the paper makes a valuable contribution to this line of work that is beneficial to the community.
Decision: Accept (poster)