SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In this section, we describe how the analysis of our matching algorithm can be extended to work for the unbalanced case, where  $|A| \neq |B|$ . This analysis is critical for proving the correctness of our optimal transport version of the algorithm. Without loss of generality, assume  $|B| \leq |A| = n$ . The overall description of the algorithm remains the same, except for the main routine of our algorithm produces an  $\varepsilon$ -feasible matching of size at least  $(1 - \varepsilon)|B|$ . The asymptotic running time of both the parallel and sequential algorithms remains unchanged. In the following lemma, we bound the additive error of our algorithm for the unbalanced case; the argument is very similar to Lemma 3.1.
Lemma 3.5. Given an unbalanced input to the assignment problem with  $|B| \leq |A|$ , the  $\varepsilon$ -feasible matching of cardinality at least  $(1 - \varepsilon)|B|$  that is returned by our algorithm is within an additive error of  $\varepsilon |B|$  from the optimal matching with respect to the cost function  $\bar{c}(\cdot, \cdot)$
4 Parallel Algorithm
In this section, we describe how to parallelize the matching algorithm of Section 2.2 leading to the result of Theorem 1.1. Recall that each phase of this algorithm has three steps: (I) Greedy step, (II) Matching update, and (III) Dual update. Steps (II) and (III) are easily parallelizable using  $O(1)$  time. However, step (I) is nontrivial to parallelize. Fortunately, Isreali and Itai gave a  $O(\log n)$  randomized parallel algorithm for computing a maximal matching on an arbitrary graph [15]. Therefore, we 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, we use the Israeli Itai algorithm to bound the parallel complexity of our algorithm. In Section 4.2, we further parallelize the phases of our algorithm leading to a simplified variation of the Israeli Itai algorithm that is more suited for a practical implementation. Finally, in Section 4.3, we provide a matrix-based interpretation of our simplified algorithm, allowing the algorithm to be easily implemented for GPUs.
4.1 Analysis of our Algorithm for Parallel and MPC Models
The Israeli Itai algorithm is designed to work on an arbitrary graph  $G(V, E)$ , which may not be bipartite. It computes a maximal matching on  $G$  in  $O(\log n)$  iterations. By directly using their  $O(\log n)$  algorithm for step (I) of our algorithm, we obtain an  $O(\log n)$  parallel running time for each phase. Since our algorithm executes  $O(1 / \varepsilon^2)$  phases, we obtain a worst-case theoretical bound of  $O(\log n / \varepsilon^2)$  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 (\Delta)))$  rounds and  $O(n)$  space per machine, where  $\Delta$  is the maximum degree of the graph. [3]. As a result, we are able to achieve an algorithm in the MPC model that requires only  $O(\log (\log n) / \varepsilon^2)$  rounds with  $O(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 a black-box, we use a simpler adaptation of their algorithm for bipartite graphs. The Israeli Itai algorithm executes  $O(\log n)$  iterations, where each iteration executes the following steps, using  $O(1)$  parallel time:
(i) Each vertex  $u \in V$  selects an incident edge  $(u, v)$  at random, and directs it from  $u$  to  $v$ , yielding a directed subgraph  $R \subseteq E$ .
(ii) Each vertex  $u \in 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 a maximum vertex degree of 2.
(iii) 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 to 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; all vertices must be handled using a symmetric process. This results in a subgraph  $S$ , where each vertex could have a degree of 2, and an additional step is required to eliminate some of these edges in order to form a matching (see step (iii)).
However, in our situation, we are solely working with bipartite graphs. In this situation, the vertices are divided into two sets  $A$  and  $B$ , and, as a result, we do not need to process each vertex in a symmetric fashion. Instead, we can allow one side to make proposals and the other side to receive proposals. As a 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 a 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 at most one incident edge selected. This alleviates the need for step (iii), since steps (a) and (b) alone immediately result in a matching.
In addition to our simplified approach to each iteration of the Israeli Itai algorithm, we make a second optimization: Within our algorithm, instead of waiting for the Israeli Itai algorithm to complete, which could require  $O(\log 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  $O(1)$  time, resulting in practical improvements. We believe that this additional source of parallelization could lead to asymptotic improvements to 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, in practice, it is important for our algorithm to have an efficient GPU-based implementation. In this section, we provide a matrix-based implementation of our algorithm, using the simplified maximal matching approach described in Section 4.2
In Algorithm 1, we provide a pseudocode 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  $\varepsilon$ . The matching returned by the algorithm has cardinality at least  $(1 - \varepsilon)n$  and a cost at most  $\varepsilon n$  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.
Algorithm 1 Approximate Bipartite Matching
1: Input:  $W \in \mathbb{R}_{\geq 0}^{n \times n}, \delta \in \mathbb{R}_{>0}$   
2:  $M \gets \mathbf{0}_{n \times n}, S \gets W$   
3: while  $M$  has more than  $\varepsilon n$  columns with all 0's do  
4:  $P \gets \mathbf{0}_{n \times n}$   
5: for all columns  $b$  with all zero entries in  $M$  do  
6: Randomly select a row  $a$  such that  $S_{a,b} = 0$   
7:  $P_{a,b} \gets 1$   
8: for all  $a, b \in [1..n]$  do  
9:  $M_{a,b} \gets 1$  if  $M_{a,b} = 1$  and row  $a$  of  $P$  has all 0's  
10: otherwise  $M_{a,b} \gets P_{a,b}$   
11: for all rows  $a$  in  $P$  with at least one 1 do  
12: Add  $\varepsilon$  to every entry in row  $a$  of  $S$   
13: for all columns  $b$  in  $M$  with all 0 entries do  
14:  $\rho \gets$  the minimum entry in column  $b$  of  $S$   
15: Decrease every entry in column  $b$  of  $S$  by  $\rho$   
16: return  $M$
The algorithm takes as input an  $n \times n$  cost matrix  $W$  and a value for the error parameter  $\delta$ , and returns an  $n \times n$  bit matrix that describes the matching, where an edge  $(i,j)$  is in the matching if and only if the value at row  $i$  and column  $j$  is a 1. We use  $\mathbf{1}_n$  to represent a row vector containing all 1's, and  $\mathbf{1}_{m \times n}$  to represent an  $m$  by  $n$  matrix of all 1's. We use similar notation for vectors and matrices of all 0's. When considering an  $n \times n$  matrix, we follow a convention that each row corresponds to a 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_{i,j} = W_{i,j} + \varepsilon - 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  $\varepsilon n$  free vertices of  $B$  remain. Each iteration of this main loop is a phase. Lines 5-7 compute a matrix  $P$ , which represents a 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 a slight change to Step (3). Instead of increasing the dual weight of free vertices of  $B$  by exactly  $\varepsilon$ , we increase it as much as possible. For some free vertices of  $B$ , the increase will be 0 (since  $P$  is not maximal), but it is also possible, in practice, for the increase to be larger than  $\varepsilon$ .
In this section, we present our 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 PyTorch library, which supports GPU operations. We compare these implementations of our algorithm to the Sinkhorn algorithm implementation in the Python Optimal Transport (POT) library [11]. This Sinkhorn implementation also uses PyTorch. Additionally, we compare our algorithm to a CUDA based GPU implementation and present the results of this comparison in appendix section C.3
Our experiments are run with an Intel Xeon E5-2680v4 2.4GHz chip and an Nvidia Tesla V100 GPU. We ran our algorithms using both real and synthetic data, including four different settings:
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  $\varepsilon$ . 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  $\varepsilon$  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 on the CPU. We only present for values of  $\varepsilon$  that are 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.1. 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  $n = 10,000$  vertices,  $A$  and  $B$ , in a 2-dimensional unit square. For the assignment problem, the demand or supply of every vertex is  $1/n$ . For the OT problem, the capacity of each vertex is initially chosen by selecting, uniformly at random from the interval [0, 1]. Then, the capacities are normalized such that the total supply and demand are both equal to 1. For any pair of points  $(a,b) \in A \times B$ , the cost  $c(a,b)$  was set as the squared Euclidean distance between them. For each value of  $\varepsilon$ , we executed 10 runs. For each combination of  $\varepsilon$ , 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 1(a)(b) and Figure 2(a)(b) respectively.
Figure 1: Plots of running times on GPU for the synthetic inputs (a)(b) and the real data inputs (c)(d)(e)(f).
MNIST: 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 a  $28 \times 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 \in A$  and  $b \in B$  is computed as follows: Let  $a(i,j)$  (resp.  $b(i,j)$ ) be the value of the pixel in row  $i$  and column  $j$  of image  $a$  (resp.  $b$ ). First, the two images are normalized so that the sum of all pixel values is equal to 1 for each image, i.e.,  $\sum_{i,j \in [1,28]} a(i,j) = 1$  and  $\sum_{i,j \in [1,28]} b(i,j) = 1$ . Then, the cost  $c(a,b)$  is given by the  $L_1$  distance between the resulting normalized images:  $c(a,b) = \sum_{i,j \in [1,28]} |a(i,j) - b(i,j)|$ . Note that an upper
bound on the largest cost is 2. For each algorithm and for each value of  $\varepsilon$ , 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 1(c) and the plot in Figure 2(c).
Figure 2: Plots of parallel rounds on GPU for the synthetic inputs (a)(b) and the real data inputs (c)(d)(e)(f).
NLP: In our final experiment, we considered OT with natural language data. We calculate the document distances based on OT following the procedure of previous work [18]. 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 [6], 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 GloVe word embedding [26]. The cost of any edge  $(a, b) \in A \times B$  is then given by the Euclidean distance between the corresponding points in this embedding.
In the last three plots in Figure 1 and Figure 2 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 are created, using different sections of the text, and the results are averaged over all 5 runs. These experiments can be found in Figure 1(d)(e)(f) and Figure 2(d)(e)(f).
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.
In this work, we provided a fast, highly parallelizable combinatorial algorithm for computing an  $\varepsilon$ -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 running 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  $O(\log (n) / \varepsilon^2)$  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?
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.
[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. arXiv:1701.07875v3 [stat.ML], 2017.
[3] Soheil Behnezhad, Mohammad Taghi Hajiaghayi, and David G Harris. Exponentially faster massively parallel maximal matching. In 2019 IEEE 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 Peyre. Iterative bregman projections for regularized transportation problems. SIAM Journal on Scientific Computing, 37(2):A1111-A1138, 2015.
[5] Jérémie Bigot, Raul Gouet, Thierry Klein, Alfredo López, et al. Geodesic PCA in 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. "O'Reilly 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 2011 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 is better than by sinkhorn'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élie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, Ievgen 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, 2021.
[12] Rémi Flamary, Marco Cuturi, Nicolas Courty, and Alain Rakotomamonjy. Wasserstein discriminant analysis. Machine Learning, 107(12):1923-1945, 2018.
[13] Harold N Gabow and Robert E Tarjan. Faster scaling algorithms for network problems. SIAM Journal on Computing, 18(5):1013-1036, 1989.
[14] Andrew V Goldberg, Serge A Plotkin, and Pravin M Vaidya. Sublinear-time parallel algorithms for matching and related problems. Journal of Algorithms, 14(2):180-213, 1993.
[15] Amos Israeli and Alon Itai. A fast and simple randomized parallel algorithm for maximal matching. Information Processing Letters, 22(2):77-80, 1986.