SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Akshaykumar G. Gattani<sup>1</sup>, Sharath Raghvendra<sup>1</sup>, and Pouyan Shirzadian<sup>1</sup>
<sup>1</sup>Department of Computer Science, Virginia Tech
Algorithms for the minimum-cost bipartite matching can be used to estimate Wasserstein distance between two distributions. Given two sets  $A$  and  $B$  of  $n$  points in a 2-dimensional Euclidean space, one can use a fast implementation of the Hungarian method to compute a minimum-cost bipartite matching of  $A$  and  $B$  in  $\tilde{O}(n^2)$  time. Let  $\Delta$  be the spread, i.e., the ratio of the distance of the farthest to the closest pair of points in  $A \cup B$ . In this paper, we present a new algorithm to compute a minimum-cost bipartite matching of  $A$  and  $B$  with a similar worst-case execution time of  $\tilde{O}(n^2 \log \Delta)$ . However, when  $A$  and  $B$  are drawn independently and identically from a fixed distribution that is not known to the algorithm, the execution time of our algorithm is, in expectation,  $\tilde{O}(n^{7/4} \log \Delta)$ .
To the best of our knowledge, our algorithm is the first one to achieve a subquadratic execution time even for stochastic point sets with real-valued coordinates. Our algorithm extends to any dimension  $d$ , where it runs in  $\tilde{O}(n^{2 - \frac{1}{2d}}\Phi(n))$  time for stochastic point sets  $A$  and  $B$ ; here  $\Phi(n)$  is the query/update time of a dynamic weighted nearest neighbor data structure. Our algorithm can be seen as a careful adaptation of the Hungarian method in the geometric divide-and-conquer framework.
Given two distributions  $\mu$  and  $\nu$  defined on the sets  $S_{\mu}, S_{\nu} \subseteq \mathbb{R}^{d}$ , respectively, and an integer  $p \geq 1$ , the  $p$ -Wasserstein distance between  $\mu$  and  $\nu$  is the minimum cost required to transport mass from one to the other. More formally, let  $\Gamma(\mu, \nu)$  be the set of all probability measures on  $S_{\mu} \times S_{\nu}$  with marginal distributions  $\mu$  and  $\nu$ . The  $p$ -Wasserstein distance between  $\mu$  and  $\nu$  is defined as
$$
W _ {p} (\mu , \nu) = \left(\inf  _ {\gamma \in \Gamma (\mu , \nu)} \int_ {S _ {\mu} \times S _ {\nu}} \| s _ {1} - s _ {2} \| ^ {p} d \gamma (s _ {1}, s _ {2})\right) ^ {1 / p},
$$
where  $\| s_1 - s_2\|$  denotes the Euclidean distance of  $s_1$  and  $s_2$ .
In many applications, the distribution  $\mu$  (resp.  $\nu$ ) may not be known, but one may have access to  $n$  i.i.d samples  $A$  (resp.  $B$ ) from  $\mu$  (resp.  $\nu$ ). The empirical distribution  $\mu_{n}$  (resp.  $\nu_{n}$ ) is a discrete distribution with  $A$  (resp.  $B$ ) as support, where each sample has a mass of  $1/n$ . The empirical  $p$ -Wasserstein distance is simply the  $p$ -Wasserstein distance between the empirical distributions  $\mu_{n}$  and  $\nu_{n}$ . As  $n \to \infty$ , the empirical  $p$ -Wasserstein distance between  $\mu$  and  $\nu$  converges to the real  $p$ -Wasserstein distance between them [7, 15, 16, 46] with several results showing sharp upper and lower bounds on the empirical  $p$ -Wasserstein distances [10, 17, 19, 45, 49]. See also [36]
*Following convention from Theoretical Computer Science, all authors are ordered alphabetically.
for a survey on such results. Due to these properties, empirical  $p$ -Wasserstein distance has found applications in training generative adversarial networks [8, 9, 24, 33-35], image retrieval [39, 42], graph predictions [27, 37], clustering stability validation [18, 31], and two-sample tests [13, 14, 23, 25, 41]. The convergence of the empirical distribution to the real distribution, however, exhibits the "curse of dimensionality"; that is, the convergence rate decreases as the dimension increases [16, 49]. The empirical  $p$ -Wasserstein distance therefore is most useful in low-dimensional settings.
Computing the empirical  $p$ -Wasserstein distance can be done by solving an instance of the minimum-cost bipartite matching problem as follows. Consider the complete bipartite graph on  $A$  and  $B$  where each edge from  $a \in A$  to  $b \in B$  has a cost of  $\| a - b \|^p$ . A bipartite matching (or simply a matching)  $M$  on  $A \cup B$  is a set of vertex-disjoint edges of  $\mathcal{G}(A, B)$ . The matching  $M$  is said to be a perfect matching if  $|M| = n$ . We define the cost of a matching  $M$ , denoted by  $w_p(M)$ , as  $w_p(M) = \sum_{(a, b) \in M} \| a - b \|^p$ . The  $p$ th power Euclidean bipartite matching is a perfect matching  $M_p^*$  with the minimum-cost under  $w_p(\cdot)$ . The empirical  $p$ -Wasserstein distance between  $\mu$  and  $\nu$  is  $W_p(\mu_n, \nu_n) = \left( \frac{1}{n} w_p(M_p^*) \right)^{1/p}$ .
The minimum-cost matching  $M_p^*$  and therefore, the empirical  $p$ -Wasserstein distance can be computed in  $O(n^3)$  time using the Hungarian algorithm [28]. In this paper, we adapt the Hungarian algorithm in a geometric divide-and-conquer framework. The execution time of our algorithm is similar to that of the Hungarian algorithm (Remark 3.2). However, when the input point sets  $A$  and  $B$  are samples drawn from the same distribution, the algorithm is asymptotically faster than the Hungarian algorithm.
Related Work. Computing an optimal minimum-cost bipartite matching can be done in  $O(n^{3})$  time using the classical Hungarian algorithm [28]. In geometric settings, the efficiency of the Hungarian algorithm can be improved to  $\tilde{O}(n^{2}\Phi(n))$ , where  $\Phi(n)$  is the query and update time of a dynamic weighted nearest neighbor data structure [3, 44, 48]. For two dimensions,  $\Phi(n) = \log^{O(1)} n$  and therefore, we get a  $\tilde{O}(n^{2})$  time exact algorithm. For higher dimensions, however, this leads to only slightly sub-cubic execution time.
For graphs with  $n$  vertices and  $m$  edges, there are faster algorithms that compute an optimal minimum-cost matching [11, 21, 22] provided the edge costs are integers that are bounded by  $C^1$ . However, in our setting, the coordinates of the input points are real-valued and the Euclidean distances, due to the presence of a square-root, can be irrational. As a result, these algorithms only provide an approximate solution and the Hungarian algorithm and its fast implementations remain the only known exact algorithm for the problem.
For 2-dimensional points with integer coordinates bounded by  $\Gamma$ , Sharathkumar [43] presented a weakly polynomial  $\tilde{O}(n^{3/2}\log \Gamma)$  time algorithm to compute the exact 1-Wasserstein distance. The main idea was to use the approximation algorithm by Sharathkumar and Agarwal [44] to initially find a planar sub-graph that traps all the edges of the optimal matching. Then, they use the planar separator based algorithm by Lipton and Tarjan [32] to find a minimum-cost matching in this subgraph. This algorithm, however, is restricted only to two dimensions and for  $p = 1$ . The design of a faster algorithm to compute the  $p$ th power Euclidean bipartite matching for  $d$ -dimensional point sets remains a major open problem. We would also note the extensive work on the design of approximation algorithms for the 1-Wasserstein distance [2, 4-6, 12, 20, 26, 30, 38, 40] as well as the 2-Wasserstein distance [1, 29].
Our Results. For any point set  $P$  in the Euclidean space, the spread of  $P$  is the ratio of the distance of the farthest pair to the closest pair in  $P$ . The main result of this paper is presented in Theorem 1.1.
Theorem 1.1. There exists a randomized algorithm that, given any two point sets  $A$  and  $B$  sampled independently and identically from a distribution  $\mu$  inside the unit  $d$ -dimensional hypercube, where  $|A| = |B| = n$  and  $\mu$  is not known to the algorithm, computes an exact Euclidean bipartite matching between  $A$  and  $B$  and has an expected running time of  $\tilde{O}(n^{2 - \frac{1}{2d}}\Phi(n)\log \Delta)$ ; here,  $\Delta$  is the spread of the points in  $A \cup B$ .
When  $A$  and  $B$  are i.i.d samples from a distribution  $\mu$  in the unit square  $(d = 2)$  and  $\mu$  is not known to the algorithm, our algorithm achieves a weakly polynomial sub-quadratic execution time of  $\tilde{O}(n^{7/4} \log \Delta)$ . Our algorithm easily extends to any constant  $p > 1$  and any fixed  $d$ -dimensional
The execution time of these algorithms have a dependence on  $\log C$ , making them weakly polynomial.
(b)
Figure 1: (a) An infeasible edge, (b) a feasible matching edge satisfying constraint (2), (c) a feasible non-matching edge satisfying constraint (1).
space (we present our result for any arbitrary  $p$  as well as the proof of all our claims in the Appendix). For instance, for  $p = 2$  and  $d = 2$  as well as for  $p = 1$  and  $d = 3$ , the execution time of our algorithm can be bounded by  $\tilde{O}(n^{11/6} \Phi(n) \log \Delta)$ . However, for simplicity in exposition of the ideas, we restrict the presentation of our algorithm to two-dimensional ( $d = 2$ ) and Euclidean settings ( $p = 1$ ). To the best of our knowledge, this is the first sub-quadratic weakly polynomial exact algorithm for computing the  $p$ th power Euclidean bipartite matching for 2-dimensional stochastic point sets. Furthermore, for many distributions, such as the uniform distribution on the unit square, one can show that, with high probability, the spread of the point set is bounded by a polynomial in  $n$ . For all such distributions, we achieve a strongly polynomial sub-quadratic exact algorithm.
Our algorithm can be seen as an elegant adaptation of the classical Hungarian algorithm in a quadtree based divide-and-conquer framework. The previous speed-ups of the Hungarian algorithm in geometric settings are based on sophisticated data structures that are hard to implement. In contrast, our geometric improvements are relatively simple, which allows us to implement and compare our algorithm with the standard implementation of the Hungarian algorithm. Experiments suggest that our algorithm outperforms the Hungarian algorithm for synthetic data samples drawn from two (possibly different) distributions  $\mu$  and  $\nu$  as long as the 1-Wasserstein distance between  $\mu$  and  $\nu$  is small. Our algorithm also outperforms the Hungarian algorithm on data samples drawn from the New York Taxi dataset, where one sample is drawn from the set of request pick-up locations  $(\mu)$  and the other sample is drawn from the request drop-off locations  $(\nu)$ . Experiments also suggest that our analysis may not be tight, at least for the case where  $\mu = \nu$  is the uniform distribution inside the unit square. Obtaining a tighter analysis of our algorithm remains an important open question. Next, we present an overview of our divide-and-conquer algorithm.
Our Approach. The classical Hungarian algorithm is based on the primal-dual framework where, along with a matching  $M$ , for every point  $v \in A \cup B$ , the algorithm maintains a non-negative dual weight  $y(v)$ . The matching  $M$  and dual weights  $y(\cdot)$  are feasible if, for every edge  $(a, b) \in A \times B$ ,
$$
y (b) - y (a) \leq \| a - b \|, \quad (a, b) \notin M, \tag {1}
$$
$$
y (b) - y (a) = \| a - b \|, \quad (a, b) \in M. \tag {2}
$$
The Hungarian algorithm initializes  $M$  to be an empty matching and the dual weights of every vertex  $v \in A \cup B$  to be 0. It then incrementally builds a perfect matching  $M$  while maintaining its feasibility and returns the constructed feasible perfect matching. It is well-known that a perfect matching that is feasible is also a minimum-cost matching.
In the Euclidean setting, we interpret the dual weight  $y(v)$  of any point  $v$  as the radius of a disc centered around  $v$ . We refer to these discs as the dual discs. For any edge  $(a, b)$ , if the dual disc of  $a$  is in the interior of the dual disc of  $b$  (Figure 1(a)), then  $y(b) - y(a) > \| a - b \|$ , making the dual assignment infeasible. Therefore, for any feasible dual weight assignment, the dual disc of  $a$  cannot be completely inside the dual disc of  $b$ . The condition (2) corresponds to Figure 1(b) where the boundary of the disc of  $a$  touches the boundary of the disc of  $b$  from inside. Figure 1(c) is an example of an edge that only satisfies (1).
In order to implement the Hungarian algorithm in the geometric divide-and-conquer framework, we can use a quadtree  $\mathcal{Q}$  with the unit square as the root node. Each node of this tree represents a square  $\square$  and if  $\square$  has more than one input point, then it has four children that are obtained by splitting  $\square$  into four equal squares. Squares with no more than one point become a leaf node in this tree.
(b)
Figure 2: (a) Feasible matchings for children that is infeasible when combined at the parent (the edge  $(b, a')$  is infeasible), (b) Restricted feasible matchings inside the children, which is also restricted feasible when combined at parent.
Consider the following naïve implementation of the Hungarian algorithm in a quadtree based geometric divide-and-conquer framework. Given a square  $\square$  of the quadtree  $\mathcal{Q}$ , we recursively find a feasible matching in each of the four children. Doing so, however, may lead to a dual weight assignment that violates the feasibility condition for bridge edges, i.e., edges that cross the boundary of the children squares. See Figure 2(a) that illustrates a dual feasible assignment at two children squares, which when combined causes the edge  $(a', b)$  to violate the feasibility condition (1). Thus, implementing the conquer step becomes challenging since we cannot simply combine the feasible dual assignments at the children to achieve a feasible dual assignment at the parent.
To overcome this difficulty, we add an additional condition that restricts the dual disc of a point  $v$  inside any child square  $\square'$  to always stay inside  $\square'$ , i.e.,  $y(v) \leq d(v, \square')$ , where  $d(v, \square')$  is the Euclidean distance of  $v$  to the boundary of  $\square'$ . Figure 2(b) shows a restricted dual assignment at the two children squares, which remains feasible when combined at the parent. Such a restricted feasible matching inside each of the four children squares of  $\square$  trivially combines to form a feasible matching at  $\square$ . Thus, one can apply the divide-and-conquer framework to compute such a restricted feasible matching. Unfortunately, however, the restricted feasible matchings computed at the children may be of an extremely low cardinality. This is particularly true when most of the input points are close to the boundaries of the children squares. See, for instance, in Figure 3(a), the restricted feasible matching is an empty one since most of the blue points are close to the boundary. As a result, the divide step may not make any progress in computing a perfect matching and the conquer step may have to do significant amount of work to find a perfect matching, causing the execution time of the conquer step to be the same as that of the standard implementation of the Hungarian algorithm.
We observe that if the quadtree  $\mathcal{Q}$  is constructed using a random-shift, then, in expectation, only sub-linearly many points are close to the boundary of any square of  $\mathcal{Q}$ . By combining this with the fact that the average length of a matching edge for  $n$  points drawn from an arbitrary distribution  $\mu$  inside unit square is  $\approx 1 / \sqrt{n}$ , we are able to bound the expected number of unmatched points in the restricted matchings by  $\tilde{O}(n^{3/4})$ , leading to an execution time of  $\tilde{O}(n^{7/4})$  for the conquer steps across all cells at each level of the quadtree and an overall execution time of  $\tilde{O}(n^{7/4} \log \Delta)$ .
We begin by introducing the notations necessary to describe our algorithm. Given  $A \cup B$  and any square  $\square$ , let  $A_{\square} = A \cap \square$ ,  $B_{\square} = B \cap \square$ , and  $n_{\square} = |A_{\square} \cup B_{\square}|$ . Let  $\ell_{\square}$  denote the side-length of  $\square$ . For any point  $v \in A_{\square} \cup B_{\square}$  inside  $\square$ , let  $\mathrm{d}(v, \square)$  denote the Euclidean distance of  $v$  to its closest point on the boundary of  $\square$ . We say that  $v$  is  $\delta$ -close to  $\square$  if  $\mathrm{d}(v, \square) \leq \delta \ell_{\square}$ . Let  $n_{\square}^{\delta}$  denote the number of points of  $A_{\square} \cup B_{\square}$  that are  $\delta$ -close to  $\square$ .
Randomly Shifted Quadtree. Given the input points  $A \cup B$  inside a unit square, a randomly shifted quadtree on the input can be constructed as follows. Let  $\xi$  be a point chosen uniformly at random from the unit square  $[0,1]^2$ . Define the square  $\square^* := [-4,4]^2 + \xi$  to be the root of the quadtree  $\mathcal{Q}$ . Recursively construct  $\mathcal{Q}$  by decomposing any square  $\square$  with  $n_\square > 1$  into four equal squares, each of which become the children of  $\square$ . Any square with exactly one point becomes a leaf square of the quadtree. Given that the spread of the points in  $A \cup B$  is  $\Delta$ , the height of  $\mathcal{Q}$  is  $O(\log \Delta)$ .
(c)
Figure 3: (a) A  $\square$ -MCM where all points of  $B$  (blue points) are unmatched, (b) the boundaries of the cells of the randomly shifted quadtree (solid line) and the  $\delta$ -close points to the cells of the quadtree (points in the shaded area), (c) the matchings constrained to the cells of the quadtree.
For any square  $\square$  of the quadtree  $\mathcal{Q}$ , in the next lemma, we show that due to the random shift, there are not many points that are very close to the boundary of  $\square$ .
Lemma 2.1. For any square  $\square$  of a randomly shifted quadtree and any  $\delta \in (0,1/2)$ ,  $\mathbb{E}\left[n_{\square}^{\delta}\right] = O(\delta \mathbb{E}[n_{\square}])$ .
Constrained Matching. For any square  $\square$ , we introduce a variant of the minimum-cost matching problem, which we refer to by the minimum-cost  $\square$ -constrained matching problem. Given a square  $\square$ , consider a matching  $M_{\square}$  of the points in  $A_{\square} \cup B_{\square}$ . Any point  $v \in A_{\square} \cup B_{\square}$  is a free point with respect to  $M_{\square}$  if  $v$  is not matched in  $M_{\square}$ . Let  $A_{\square}^{F}$  (resp.  $B_{\square}^{F}$ ) denote the set of free points of  $A_{\square}$  (resp.  $B_{\square}$ ) with respect to  $M_{\square}$ . We define the  $\square$ -constrained cost of the matching  $M_{\square}$ , denoted by  $w_{\square}(M_{\square})$ , as
$$
w _ {\square} (M _ {\square}) := \sum_ {(a, b) \in M _ {\square}} \| a - b \| + \sum_ {b \in B _ {\square} ^ {F}} \mathrm {d} (b, \square). \tag {3}
$$
For any square  $\square$ , the minimum-cost  $\square$ -constrained matching ( $\square$ -MCM), denoted by  $M_{\square}^{*}$ , is simply a matching with the minimum  $\square$ -constrained cost. In Lemma 2.2, we show that for the root square  $\square^{*}$  of  $\mathcal{Q}$ , the minimum-cost  $\square^{*}$ -constrained matching is also a minimum-cost matching on  $A \cup B$ .
For any arbitrary square  $\square$  of  $\mathcal{Q}$ ,  $M_{\square}^{*}$  might be of a very small cardinality. See, for instance in Figure 3(a) where  $B$  is the set of blue points, Equation (3) is minimized when all blue points are free points; i.e.,  $w_{\square}(M_{\square}^{*}) = \sum_{b \in B} \mathrm{d}(b, \square)$  and  $M_{\square}^{*}$  is an empty matching. Nonetheless, by using the bound on the number of  $\delta$ -close points to any square  $\square$  of a randomly-shifted quadtree, we can bound the expected number of free points with respect to any  $\square$ -MCM  $M_{\square}^{*}$ . In the example of Figure 3(b), the solid lines show the boundary of the squares of a randomly-shifted quadtree, and as shown in Figure 3(c), only a few points will be free with respect to the optimal matchings constrained to the squares of the quadtree.
Lemma 2.2. For any square  $\square$  of a randomly shifted quadtree and any minimum-cost  $\square$ -constrained matching  $M_{\square}^{*}$ , (i) the expected number of free points of  $B_{\square}$  with respect to  $M_{\square}^{*}$  is  $\tilde{O}\left(n^{3/4}\right)$  in 2-dimensions (and  $\tilde{O}(n^{1-\frac{1}{2d}})$  in  $d$  dimensions), and (ii) if  $\square$  is the root square, then  $M_{\square}^{*}$  is a minimum-cost perfect matching on  $A \cup B$ .
For the rest of the paper, we design a primal-dual method to compute a  $\square^{*}$ -MCM, which by Lemma 2.2 is also a minimum cost matching on  $A \cup B$ .
Constrained Feasibility. Similar to the Hungarian algorithm, we devise a primal-dual method to compute a  $\square$ -MCM. For any square  $\square$  of  $\mathcal{Q}$ , we say that a matching  $M_{\square}$  on  $A_{\square} \cup B_{\square}$  along with a set of non-negative dual weights  $y(\cdot)$  for the points in  $A_{\square} \cup B_{\square}$  is  $\square$ -feasible if,
$$
y (b) - y (a) \leq \| a - b \|, \quad \forall (a, b) \in A \times B, \tag {4}
$$
$$
y (b) - y (a) = \| a - b \|, \quad \forall (a, b) \in M _ {\square}. \tag {5}
$$
$$
y (b) \leq \mathrm {d} (b, \square), \quad \forall b \in B _ {\square}, \tag {6}
$$
$$
y (a) = 0, \quad \forall a \in A _ {\square} ^ {F}. \tag {7}
$$
(c)
Figure 4: (a) A  $\square$ -feasible matching  $M_{\square}$  where  $b_{1}$  is a  $\square$ -free point and  $b_{2}$  is a free point that is not  $\square$ -free, (b) an admissible path  $P$  (dashed line) from  $b_{1}$  to a point  $b_{3}$  whose vertex slack is zero, and (c) the  $\square$ -optimal matching obtained from augmenting  $M_{\square}$  along  $P$ .
Note that (4) and (5) are conditions identical to the ones maintained by the Hungarian algorithm. Condition (6) ensures that the dual discs stay inside the square and condition (7) ensures that the  $\square$ -MCM is of minimum cost.
Given a  $\square$ -feasible matching  $M_{\square}, y(\cdot)$ , we say that any point  $b \in B_{\square}$  is a  $\square$ -free point with respect to  $M_{\square}$  if  $b$  is not matched in  $M_{\square}$  and  $y(b) < \mathrm{d}(b, \square)$ . Note that a free point may not be  $\square$ -free. For instance, in Figure 4(a), the blue point  $b_{2}$  is free but not a  $\square$ -free point since its dual weight is equal to its distance to  $\square$ . Let  $\mathcal{B}_{\square}^{F}$  denote the set of all  $\square$ -free points in  $B_{\square}$ . We say that a  $\square$ -feasible matching  $M_{\square}, y(\cdot)$  is  $\square$ -optimal if it does not have any  $\square$ -free points, i.e.,  $\mathcal{B}_{\square}^{F} = \emptyset$ . The following lemma shows that any  $\square$ -optimal matching  $M_{\square}, y(\cdot)$  is a  $\square$ -MCM. Thus, we focus on computing a  $\square$ -optimal matching.
Lemma 2.3. Let  $M_{\square}, y(\cdot)$  be a  $\square$ -optimal matching on  $A_{\square} \cup B_{\square}$ . Then,  $M_{\square}$  is a minimum-cost  $\square$ -constrained matching.
Given a square  $\square$  of  $\mathcal{Q}$  and its four children  $\square_1, \square_2, \square_3$ , and  $\square_4$ , we can obtain a  $\square$ -feasible matching at  $\square$  by simply combining the  $\square_i$ -optimal matchings from each child  $\square_i$ . We use this property to design our divide-and-conquer algorithm for computing  $\square$ -optimal matching.
Lemma 2.4. For any square  $\square$ , let  $\square_i$ ,  $i \in [1,4]$  be the set of all children of  $\square$  and let  $M_i, y(\cdot)$  denote a  $\square_i$ -optimal matching. Then, the matching  $\bigcup_{i=1}^{4} M_i, y(\cdot)$  is a  $\square$ -feasible matching.
Next, we define slack and an augmenting path with respect to a  $\square$ -feasible matching.