Title: A Combinatorial Algorithm for the Semi-Discrete Optimal Transport Problem
Venue: NeurIPS 2024
Accepted: yes
Keywords: Semi-Discrete Optimal Transport

Approximation Algorithm

Cost-Scaling Framework
Abstract: Optimal Transport (OT, also known as the Wasserstein distance) is a popular metric for comparing probability distributions and has been successfully used in many machine-learning applications.
In the semi-discrete $2$-Wasserstein problem, we wish to compute the cheapest way to transport all the mass from a continuous distribution $\mu$ to a discrete distribution $\nu$ in $\mathbb{R}^d$ for $d\ge 1$, where the cost of transporting unit mass between points $a$ and $b$ is $d(a,b)=||a-b||^2$. When both distributions are discrete, a simple combinatorial framework has been used to find the exact solution (see e.g. [Orlin, STOC 1988]). 
In this paper, we propose a combinatorial framework for the semi-discrete OT, which can be viewed as an extension of the combinatorial framework for the discrete OT but requires several new ideas. We present a new algorithm that given $\mu$ and $\nu$ in $\mathbb{R}^2$ and a parameter $\varepsilon>0$, computes an $\varepsilon$-additive approximate semi-discrete transport plan in $O(n^{4}\log n\log \frac{1}{\varepsilon})$ time (in the worst case), where $n$ is the support-size of the discrete distribution $\nu$ and we assume that the mass of $\mu$ inside a triangle can be computed in $O(1)$ time. Our algorithm is significantly faster than the known algorithms, and unlike many numerical algorithms, it does not make any assumptions on the smoothness of $\mu$.
As an application of our algorithm, we describe a data structure to store a large discrete distribution $\mu$ (with support size $N$) using $O(N)$ space so that, given a query discrete distribution $\nu$ (with support size $k$), an $\varepsilon$-additive approximate transport plan can be computed in $O(k^{3}\sqrt{N}\log \frac{1}{\varepsilon})$ time in $2$ dimensions. Our algorithm and data structure extend to higher dimensions as well as to $p$-Wasserstein problem for any $p \ge 1$.

A Combinatorial Algorithm for the Semi-Discrete Optimal Transport Problem

Abstract

Optimal Transport (OT, also known as the Wasserstein distance) is a popular metric for comparing probability distributions and has been successfully used in many machine-learning applications. In the semi-discrete 2 -Wasserstein problem, we wish to compute the cheapest way to transport all the mass from a continuous distribution µ to a discrete distribution ν in R d for d ≥ 1 , where the cost of transporting unit mass between points a and b is d ( a, b ) = ∥ a -b ∥ 2 . When both distributions are discrete, a simple combinatorial framework has been used to find the exact solution (see e.g. [Orlin, STOC 1988]). In this paper, we propose a combinatorial framework for the semi-discrete OT, which can be viewed as an extension of the combinatorial framework for the discrete OT but requires several new ideas. We present a new algorithm that given µ and ν in R 2 and a parameter ε > 0 , computes an ε -additive approximate semi-discrete transport plan in O ( n 4 log n log 1 ε ) time (in the worst case), where n is the support-size of the discrete distribution ν and we assume that the mass of µ inside a triangle can be computed in O (1) time. Our algorithm is significantly faster than the known algorithms, and unlike many numerical algorithms, it does not make any assumptions on the smoothness of µ . As an application of our algorithm, we describe a data structure to store a large discrete distribution µ (with support size N ) using O ( N ) space so that, given a query discrete distribution ν (with support size k ), an ε -additive approximate transport plan can be computed in O ( k 3 √ N log 1 ε ) time in 2 dimensions. Our algorithm and data structure extend to higher dimensions as well as to p -Wasserstein problem for any p ≥ 1 .

1 Introduction

Optimal Transport (OT) is a powerful metric for comparing probability distributions and is used in many machine-learning applications. The semi-discrete optimal transport problem asks for the cheapest transport plan to transport mass from (a possibly continuous) distribution µ that is stored compactly (say, using a deep neural network) to a discrete distribution ν . In recent years, the semidiscrete OT has been used in data mining [31], image processing [28, 29, 33, 40], computational biology [53], variational inference [7], blue noise generation [21, 51], optics [46], solving PDEs [30, 34], and generative models [8, 9, 19].

More formally, a semi-discrete transport plan τ between a continuous probability distribution µ defined over a compact support A ⊂ R d and a discrete distribution ν with a support set B of n points in R d is a distribution over A × B whose marginals are dominated by µ and ν , i.e., τ : A × B → R ≥ 0 is a transport plan between µ and ν where ∑ b ∈ B τ ( a, b ) ≤ µ ( a ) for all a ⊆ A and ∫ A τ ( a, b ) da ≤ ν ( b ) for all b ∈ B . A transport plan τ is complete if τ transports all mass of µ to ν . For any fixed p ≥ 1 , in

the semi-discrete p -Wasserstein problem , the goal is to compute a complete transport plan τ between µ and ν minimizing the cost ¢ ( τ ) := ∫ A ∑ b ∈ B ∥ a -b ∥ p τ ( a, b ) da . An optimal plan is referred to as a W p -OT plan 1 . For any parameter ε > 0 , a transport plan τ is called an ε -close W p -OT plan if ¢ ( τ ) ≤ ¢ ( τ ∗ ) + ε , where τ ∗ is a W p -OT plan.

An optimal solution for the semi-discrete OT problem can be compactly represented as the weighted Voronoi diagram (also called the Laguerre diagram) with respect to a weight assignment y ( b ) on every point b ∈ B [12]. The choice of weights guarantees that mass at b is equal to the mass of µ inside the Voronoi cell (also called the Laguerre cell) of b and the optimal transport plan simply transports the mass at b to the mass inside its V oronoi cell. For 2 dimensions, the weighted Voronoi diagram under the squared Euclidean distance (also known as the power diagram) can be constructed in O ( n log n ) time [24, 57]. For higher dimensions d > 2 , the construction time would be O ( n ⌈ ( d +1) / 2 ⌉ ) [11].

The compact and connected nature of the semi-discrete optimal transport plans makes them attractive for many ML applications; for instance, they can help achieve stability in training GANs and avoid issues such as discontinuities and mode collapse [10, 19, 55], they improve the mapping between continuous latent spaces and discrete data in Variational Autoencoders (VAEs) [9], and also have been used in diffusion-based generative models [42]. However, there are no known exact algorithms for computing semi-discrete optimal transport. Additionally, computing an ε -close transport plan is known to be # P -Hard with respect to d and log 1 /ε , i.e., an algorithm with an execution time that is polynomial in d and log 1 /ε seems unlikely [58]. Due to the intractability of the semi-discrete optimal transport in high dimensions, researchers have considered taking n samples from the model, which is a continuous distribution µ , and computing an optimal discrete transport plan between the empirical distribution µ n defined on the n samples and ν [43]. It has been shown that this empirical p -Wasserstein distance converges to the true semi-discrete p -Wasserstein distance at a rate of n -1 / 2 p [22, 35, 50]; note that the rate of convergence does not depend on the dimension. However, the optimal transport plan from samples is not necessarily a good approximation of the semi-discrete transport plan and may cause biased gradients [14, 19].

For the discrete OT problem, there are several near-optimal scalable exact and approximation algorithms [1, 4, 5, 20, 25, 26, 36, 39, 52, 54], some of which extend to very high dimensions. For the semi-discrete OT in low dimensions, despite extensive work, scalable algorithms to find optimal semidiscrete transport plans remain elusive. There are algorithms to compute an ε -close semi-discrete transport plan using numerical solvers [12, 15, 17, 21, 37, 38, 41, 48], entropic regularization [6, 16, 32], and multiscale approaches [40, 44]. The execution time of all these algorithms is exponential in both d and log 1 /ε . Furthermore, their convergence relies on a smoothness parameter of µ . For instance, a notable numerical algorithm by Oliker and Prussner [48] assumes that for a point b ∈ B , a small change in y ( b ) will change the mass of µ inside the Voronoi cell of b by a proportionately small amount [45, Remark 22]. Under this assumption, their algorithm executes poly( n, 1 /ε ) iterations, where each iteration requires the computation of a weighted Voronoi diagram which takes n O ( d ) time. Their algorithm slows down when µ is non-smooth and does not even converge when µ is a discrete distribution. Furthermore, these methods approximate the transport cost but the transport plan that they compute may not be an approximation of the optimal weighted Voronoi diagram.

Recently, Agarwal et al. [5] described a cost-scaling paradigm to compute an ε -close semi-discrete transport plan. Their algorithm executes log ∆ ε scales, where ∆ is the diameter of A ∪ B . Within each scale, they create an instance of the discrete OT problem of size O ( n 5 ) in R 2 (and n O ( d ) in R d ). Using any existing strongly polynomial primal-dual discrete OT solver, such as the algorithm by Orlin [49], their algorithm computes an exact discrete OT plan for each instance and updates the weights for B . The OT plan computed in the final scale is an ε -close transport plan. They show that the weight assigned to any point b ∈ B is ε away from the optimal weight assignment. The overall execution time of this algorithm is O ( n 9 log ∆ ε ) in R 2 and n O ( d ) log ∆ ε in d -dimensions. Furthermore, in the limit, the transport plan of their algorithm converges to the optimal weighted V oronoi diagram. Note that their algorithm does not make any assumptions on µ .

Our Contributions. The following theorem states our main result.

Theorem 1.1. Let µ be a continuous probability distribution defined on a compact set A ⊂ R 2 , ν a discrete probability distribution with a support B ⊂ R 2 of size n , and ε > 0 a parameter. Suppose

there exists an oracle that, given a triangle φ , returns the mass of µ inside φ in Φ time. Then, an ε -close W 2 -OT plan between µ and ν can be computed in O ( n 3 (Φ + n log n ) log ∆ ε ) time in the worst-case, where ∆ is the diameter of A ∪ B .

Similar to the algorithm by Agarwal et al. [5], our algorithm is also based on a cost-scaling approach and executes O (log ∆ ε ) scales. However, the algorithm within each scale is different. Agarwal et al. create a discrete OT instance with n 4 vertices and n 5 edges and use an exact discrete OT algorithm to solve this instance. Instead, we extend the combinatorial primal-dual framework of discrete OT to the continuous space and present an algorithm to find the desired semi-discrete transport plan in O ( n 4 log n ) time in the worst case. There are several challenges in extending the combinatorial framework to semi-discrete settings, and overcoming these challenges is one of the main technical contributions of the paper.

In more detail, in each scale, our algorithm maintains a (partial) transport plan, iteratively computes a set of augmenting paths, and augments the transport plan along such paths until all of the mass is transported. In order to assist in finding augmenting paths, we define a residual graph of size O ( n 3 ) in the continuous space. Algorithms for discrete OT maintain dual weights for all points in A ∪ B . Unlike in the discrete setting where the vertex set is fixed, the vertex set of our residual graph includes 'continuous regions", which evolve over time, and the vertex set of the residual graph changes. Therefore, we are able to maintain weights only for points in B and not for the regions in A . Our primal-dual framework (especially the definition of admissibility in Section 2), as well as our algorithm (in particular Sections 3.2 and 3.4) contain a number of novel ideas, carefully designed to address various challenges that arise due to the dynamically changing continuous regions of the residual graph. Like all existing algorithms, our algorithm also requires access to an oracle that, given a query triangle, returns the mass of µ inside the triangle. We note that Dijkstra's shortest path algorithm has been extended to continuous space [47], but we are not aware of any previous work that extends a combinatorial discrete OT framework to continuous space.

Our algorithm extends to any dimension d ≥ 2 and any p ≥ 1 in a straightforward way. For d > 2 , the algorithm in Theorem 1.1 can be shown to have an execution time of O ( n d +1 (Φ+ n log n ) log ∆ ε ) . Note that the runtime of Oliker-Prussner's algorithm has a factor 1 /ε while ours has only log 1 /ε . Unlike their algorithm, ours does not make any assumptions on the smoothness of µ . Furthermore, similar to Agarwal et al. [5], our transport plan approximates an optimal weighted Voronoi diagram within a small additive factor.

One consequence of our algorithm is that we can use it to design a data structure that answers ε -close optimal transport queries efficiently. More precisely, consider a large distribution µ with support of size N . We can preprocess this distribution into a data structure that can return the total weight of points inside a query triangle in O ( N 1 -1 /d ) time [59]. By using Theorem 1.1, we can report an ε -additive approximate transport plan to any query distribution ν with support of size k in O ( k d +1 N 1 -1 /d log ∆ ε ) time. Note that this query time is sub-linear in N . There has been some work on computing 1 -Wasserstein distance approximately in sub-linear time and for answering nearest neighbor queries under 1 -Wasserstein distance [13], but these algorithms have larger error and do not extend to 2 -Wasserstein distance.

Theorem 1.2. Let µ be a discrete probability distribution with a support A ⊂ R d of size N . The distribution µ can be preprocessed, in O ( N log N ) time, into an O ( N ) size data structure so that for a discrete probability distribution ν with a support B ⊂ R d of size k , and a parameter ε > 0 , an ε -close W 2 -OT plan between µ and ν can be computed in O ( k d +1 ( N 1 -1 /d + k log k ) log ∆ ε ) time, where ∆ is the diameter of A ∪ B .

If the points in the supports of µ and ν have integer coordinates and the masses on them are rational numbers, we can adapt our data structure in Theorem 1.2 to compute an exact W p -OT plan between µ and ν in O ( k d +1 ( N 1 -1 /d + k log k ) log ∆ ε ) time for any fixed even value of p .

We also note that in 2 -dimensions, for k < N 1 / 4 , we can use the data structure to compute an exact discrete OT plan for any distribution ν in O ( k 3 N 1 / 2 log ∆) time that is faster than any existing OT algorithm that takes ( Nk ) 1+ o (1) time [3, 18].

For simplicity in presentation, we restrict our presentation to d = 2 and p = 2 .

b

b

2 Combinatorial Framework

In this section, we extend the combinatorial framework used by discrete OT algorithms to the semidiscrete settings. Let µ be a continuous probability distribution defined over a compact support A ⊂ R 2 and ν be a discrete distribution with a support set B of n points in R 2 .

Weighted Voronoi diagram. For any pair of points a, b ∈ R 2 , let d( a, b ) = ∥ a -b ∥ 2 . Given a weight function y : B → R ≥ 0 on the points in B , the weighted distance d y ( a, b ) from a point b ∈ B to any point a ∈ R 2 is defined as d y ( a, b ) = d( a, b ) -y ( b ) . For a point b ∈ B , the Voronoi cell of b is the locus of all points with b as their weighted nearest neighbor; more formally,

The weighted Voronoi diagram VD y ( B ) of the points B with weights y ( · ) is the decomposition of R 2 induced by Voronoi cells. There exists a weight function y ( · ) for B such that µ (Vor y ( b )) = ν ( b ) for every point b ∈ B , and an optimal semi-discrete transport plan transports the mass of ν at each point b ∈ B to the mass of µ in Vor y ( b ) ∩ A [12].

δ -expanded Voronoi cell. Consider a weight function y ( · ) for the points in B . For any point b ∈ B and a parameter δ > 0 , consider the following weight function y δ b .

̸

The δ -expanded Voronoi cell of b , denoted by V δ b , is simply the Voronoi cell of b in the weighted Voronoi diagram VD y δ b ( B ) of the point set B with weights y δ b ( · ) . See Figure 1 (left).

δ -feasibility. Any (possibly partial) transport plan τ between µ and ν along with a weight function y ( · ) for the points in B is δ -feasible if

(F1) for any pair ( a, b ) ∈ A × B with τ ( a, b ) > 0 , the point a lies inside the 2 δ -expanded Voronoi cell V 2 δ b .

For any δ -feasible transport plan τ, y ( · ) , if τ is a complete transport plan between µ and ν , then τ, y ( · ) is called a δ -optimal transport plan . Recall that any optimal transport plan transports the mass at b to its weighted Voronoi cell. In a δ -optimal transport plan, however, the mass of each point b ∈ B is transported inside a 2 δ -expansion of the Voronoi cell of b . This introduces an additive increase of at most 2 δ in the cost of the transport plan.

Lemma 2.1. Any δ -optimal transport plan τ, y ( · ) between µ and ν is 2 δ -close.

See Appendix A for a proof. Next, given a δ -feasible transport plan τ, y ( · ) , we define a residual graph and an augmenting path. We also introduce the process of augmenting τ along an augmenting path, which allows us to increase the mass transported by τ .

Residual graph. Given a δ -feasible (possibly partial) transport plan τ, y ( · ) , we construct a residual graph G δ by first partitioning the support A of µ into regions to form the vertex set of G δ := G ( τ, y, δ ) and then defining a set of directed edges.

Vertex set. For each point b ∈ B , consider the Voronoi cell of b and its δ - and 2 δ -expansions V δ b and V 2 δ b . Let X δ denote the arrangement [2] of these 3 n cells across all n points of B . See Figure 1 (middle). For each region φ in this arrangement, pick an arbitrary representative point r φ inside φ and assign it a mass of µ r φ := µ ( φ ) , where µ ( φ ) = ∫ φ µ ( a ) da denotes the mass of µ inside φ . Let A δ denote the set of representative points of all regions in X δ . The vertex set of G δ is the point set A δ ∪ B along with a source vertex s and a sink vertex t . We refer to any point b ∈ B whose mass is not fully transported by τ as a free point and define its excess as the amount of mass of b that is not transported by τ , i.e., ex( b ) = ν ( b ) -∫ a ∈ A τ ( a, b ) da . Similarly, any point r φ ∈ A δ is free if τ does not fully transport the mass into the region φ , and its excess is defined as ex( r φ ) = µ r φ -∑ b ′ ∈ B τ ( φ, b ′ ) .

Edge set. For each pair ( r φ , b ) ∈ A δ × B , if τ ( φ, b ) > 0 , we add a backward edge directed from r φ to b in the residual graph. Furthermore, if r φ ∈ V 2 δ b , we add a forward edge directed from b to r φ in G δ . Additionally, we add a forward edge from the source s to every free point b ∈ B and a backward edge directed from every free vertex r φ to t . This completes the description of the residual graph.

Lemma 2.2. For any δ > 0 and a δ -feasible transport plan ˆ τ, y ( · ) , the residual graph G δ has O ( n 2 ) vertices and O ( n 3 ) edges.

While describing our algorithm, it is useful to have the definition of weighted distance for all the backward edges, including those incident on t . Therefore, we extend the definition of weighted distance to any edge ( r φ , t ) as follows. Let b φ denote the weighted nearest neighbor of r φ in B , i.e., b φ := min b ∈ B d y ( r φ , b ) . Define d y ( r φ , t ) := d y ( r φ , b φ ) + δ .

Compressing a semi-discrete transport plan. Given a semi-discrete transport plan τ that transports mass from B to A , we construct a transport plan ˆ τ from B to A δ as follows. For each pair ( r φ , b ) ∈ A δ × B , let ˆ τ ( r φ , b ) := τ ( φ, b ) , i.e., we assign the entire mass transported from b to φ to the pair ( r φ , b ) (Figure 1 (right)). We refer to the transport plan ˆ τ as the compressed transport plan .

Lemma 2.3. For any δ -feasible semi-discrete transport plan τ, y ( · ) , the compressed transport plan ˆ τ along with weights y ( · ) is also δ -feasible.

Conversely, consider a transport plan ˆ τ, y ( · ) from B to A δ . One can compute a semi-discrete transport plan τ that, given a pair ( r φ , b ) ∈ A δ × B , arbitrarily transports a mass of ˆ τ ( r φ , b ) from b to φ .

Lemma 2.4. Any δ -feasible transport plan ˆ τ, y ( · ) from B to A δ can be converted into a δ -feasible semi-discrete transport plan τ from B to A .

We say that any compressed transport plan ˆ τ is a forest if the edges transporting a positive mass in ˆ τ do not create an undirected cycle.

Augmentation. Given the residual graph G δ for a δ -feasible transport plan ˆ τ, y ( · ) , an alternating path (resp. alternating cycle ) is a directed path (resp. directed cycle) in G δ . Note that, in any directed path (resp. cycle) in the residual graph, the edges alternate between forward and backward edges. An augmenting path P = ⟨ s, b 1 , r 1 , . . . , r k , t = b k +1 ⟩ is a directed path from the source vertex s to the sink vertex t in the residual graph. By construction, the vertex b 1 and the vertex r k are free vertices in the residual graph. One can augment ˆ τ along P as follows. Define the bottleneck capacity of the augmenting path P as bc( P ) := min { ex( b 1 ) , ex( r k ) , min i ∈ [1 ,k -1] ˆ τ ( r i , b i +1 ) } . To augment ˆ τ along P , set ˆ τ ( r i , b i ) ← ˆ τ ( r i , b i ) + bc( P ) for each forward edge ( b i , r i ) ∈ P and ˆ τ ( r i , b i +1 ) ← ˆ τ ( r i , b i +1 ) -bc( P ) for each backward edge ( r i , b i +1 ) ∈ P .

Lemma 2.5. The transport plan obtained after augmenting a δ -feasible transport plan ˆ τ, y ( · ) along any augmenting path P in the residual graph is δ -feasible.

Consider the following straightforward way to augment a semi-discrete transport plan. Given a δ -feasible semi-discrete transport plan τ , y ( · ) , we can compute an augmenting path P in the residual graph and augment the compressed transport plan ˆ τ along P . From Lemmas 2.3 and 2.5, ˆ τ remains δ -feasible and from Lemma 2.4, the updated δ -feasible transport plan ˆ τ can be converted to a δ -feasible semi-discrete transport plan as desired. To obtain a complete transport plan, one can iteratively apply this procedure until there are no augmenting paths in the residual graph; however, this may result in an unbounded number of iterations. To obtain an efficient algorithm, we iteratively compute a set of

2

S

by

a

3

special augmenting paths called the 'admissible" augmenting paths and augment the transport plan along these paths. By doing so, we can reduce the number of iterations to O ( n ) .

Admissibility. Suppose ˆ τ, y ( · ) is a δ -feasible transport plan between µ and ν . For any pair of points b 1 , b 2 ∈ B and any region φ ∈ X δ such that ˆ τ ( r φ , b 2 ) > 0 , the triple ( b 1 , r φ , b 2 ) is admissible if d y ( r φ , b 1 ) < d y ( r φ , b 2 ) 2 . See Figure 2 (left). Note that an admissible triple ( b 1 , r φ , b 2 ) forms by a forward edge followed by a backward edge in the residual graph satisfying d y ( r φ , b 1 ) < d y ( r φ , b 2 ) . Intuitively, for any admissible triple ⟨ b 1 , r φ , b 2 ⟩ , the mass of µ inside r φ is transported from b 2 but b 1 is nearer to r φ than b 2 (with respect to the weighted distances).

We extend the definition of admissibility to augmenting paths and alternating cycles as follows. Any augmenting path (resp. alternating cycle) P = ⟨ b 1 , r 1 , b 2 , . . . , r k , b k +1 ⟩ is admissible if all triples ( b i , r i , b i +1 ) , i ∈ [1 , k ] , are admissible. In Figure 2(right), b 2 transports mass to r 1 , and b 1 is its weighted nearest neighbor. By augmenting along this admissible path, we increase ˆ τ ( r 1 , b 1 ) and reduce ˆ τ ( r 1 , b 2 ) , thereby transporting more mass to r 1 from its weighted nearest neighbor.

3 Algorithm

In this section, we present our cost-scaling algorithm that uses the combinatorial framework from Section 2 to compute an ε -close semi-discrete OT plan. Classical discrete OT algorithms assign weights to points in A ∪ B and use them to identify a large set of augmenting paths. Their efficiency critically relies on the acyclicity of the 'search' graph. In contrast, during the execution of our algorithm, a change in the weights of B creates a new weighted Voronoi diagram, which in turn changes A δ , the discrete representation of A , and thus the vertex set of the residual graph. Since A δ may change significantly in each iteration during the execution of our algorithm, we cannot maintain weights for them. This creates significant challenges as the algorithm searches for augmenting paths (See Section 3.2 for details). Furthermore, the updated residual graph may have cycles. We introduce additional steps in our algorithm to eliminate these cycles (See Section 3.4 for details). First, we present an overview of our algorithm and then present the details.

3.1 Overview

The algorithm runs for O (log ∆ ε ) scales, where ∆ is the diameter of A ∪ B . In each scale δ , our algorithm maintains a transport plan ˆ τ δ and a weight y ( b ) for every point b ∈ B . Initially, set δ = ∆ 2 and define y ( b ) = 0 for all b ∈ B . In each scale δ , execute the following steps.

Initialization: Set τ δ to be an empty transport plan. Compute the residual graph G δ and the compressed transport plan ˆ τ δ with respect to τ δ , y ( · ) .

Iterations: While ˆ τ δ is not a complete transport plan:

(i) Compute a set of admissible augmenting paths in the residual graph G δ and augment ˆ τ δ along these paths using the SEARCHANDAUGMENT procedure described in Section 3.2. At the end of this step, there are no admissible augmenting paths in the residual graph.

(ii) Adjust the weights of all points of B that are reachable from the source by admissible paths by δ and recompute the set A δ , the residual graph G δ , and the compressed transport plan ˆ τ δ using the INCREASEWEIGHTS procedure described in Section 3.3.

(iii) Update the compressed transport plan ˆ τ δ and the residual graph using the ACYCLIFY procedure described in Section 3.4, so that the transport plan ˆ τ δ is a forest and the residual graph does not have any admissible cycles.

Scale Update: Set δ ← δ/ 2 .

After the execution of a scale δ ≤ ε/ 2 , our algorithm terminates by returning a complete semi-discrete transport plan τ δ obtained from the compressed complete transport plan ˆ τ δ (using Lemma 2.4).

Invariants. As shown in Section 4 below, our algorithm iteratively updates the weights y ( · ) and the transport plan ˆ τ δ while maintaining the following invariants:

(I1) The transport plan ˆ τ δ , y ( · ) is δ -feasible, and

(I2) At the start of each iteration, the transport plan ˆ τ δ is a forest and there are no admissible cycles in the residual graphs.

Remark. The algorithm by Agarwal et al. [5] creates a discrete instance in each scale of the algorithm by computing the arrangement of the δ -, 2 δ -, . . . , and (4 n +1) δ -expanded Voronoi cells of each point b ∈ B . Instead of using such a fine partition to create a discrete instance with O ( n 5 ) edges, we work directly with the continuous space, maintain a much smaller residual graph with O ( n 3 ) edges, and use our semi-discrete combinatorial framework to find a transport plan.

3.2 SEARCHANDAUGMENT Procedure

The SEARCHANDAUGMENT procedure executes a partial DFS-style search to identify a set of admissible augmenting paths and augments the transport plan along these paths. The SEARCHANDAUGMENT procedure is somewhat similar in style to the blocking flow procedure in Dinic's max-flow algorithm [23] or the partial-DFS procedure in Gabow-Tarjan's algorithm [27], both of which rely on the property that there are no admissible cycles in the residual graph. Unlike these algorithms, in our case, if we augment along an arbitrary admissible augmenting path, we may create an admissible cycle. See Figure 3.

We overcome this challenge by carefully calibrating the search algorithm in two ways. First, we begin our search from the free regions instead of the free points of B . Thus, we reverse the direction of all the edges of the residual graph and begin our search from the sink t . Second, we explore all forward edges incident on a region in the increasing order of their weighted distance. This order of processing edges ensures that no admissible cycles are created and that there are no more admissible augmenting paths in the residual graph after the SEARCHANDAUGMENT procedure terminates. We provide the details below.

Let ← -G δ be the graph formed by reversing the direction of all the edges of G δ . We conduct our search starting from the sink t in the graph ← -G δ . Initially, mark all points of B and all backward edges as unvisited, define U := B as the set of unvisited points of B , and Q = ⟨ t = b 0 ⟩ as the search path that the procedure grows. Execute the following steps until the search path Q becomes empty.

If Q = ⟨ t = b 0 , r 1 , b 1 , . . . , r i , b i ⟩ for some i ≥ 0 ,

(a) If there is an edge from b i to s in ← -G δ (i.e., b i is a free point), then P = ⟨ s, b i , r i , . . . , r 1 , t ⟩ is an augmenting path in G δ . Augment ˆ τ δ along P and set Q = ⟨ t = b 0 ⟩ .

(b) Assume there is not edge from b i to s . If there is an unvisited edge ( b i , r ) in ← -G δ , then add r = r i +1 to Q . Otherwise, mark b i as visited and remove b i from U and Q .

t

bi

P1

P2

b2

b4

b3

S

P3

xt

b1

P1

P2

b2

×

S

b4

b3

P3

If Q = ⟨ t = b 0 , r 1 , b 1 , . . . , b i , r i +1 ⟩ for some i ≥ 0 , let b := arg min b ′ ∈ U d y ( r i +1 , b ′ ) denote the unvisited point with the minimum weighted distance to r i +1 3 .

(a) If ( b, r i +1 , b i ) is admissible, i.e., d y ( r i +1 , b ) < d y ( r i +1 , b i ) , then add b as b i +1 to Q .

(b) Otherwise, remove r i +1 from Q and mark the edge ( b i , r i +1 ) as visited.

The algorithm terminates when the search path Q becomes empty, i.e., the procedure marked t as visited and removed it from Q . The following lemma shows the properties of the SEARCHANDAUGMENT procedure.

Lemma 3.1. Suppose invariants (I1) and (I2) hold at the start of the SEARCHANDAUGMENT procedure. Then, during the execution of the SEARCHANDAUGMENT procedure,

(S1) the transport plan ˆ τ δ , y ( · ) remains δ -feasible,

(S2) any point b ∈ B (resp. backward edge ( r, b ) ) marked as visited will not form an admissible augmenting path during the execution of the procedure, and

(S3) there are no admissible cycles in the residual graph.

3.3 INCREASEWEIGHTS Procedure

The INCREASEWEIGHTS procedure adjusts the weights of B leading to the formation of admissible augmenting paths. The INCREASEWEIGHTS procedure performs a DFS on the residual graph G δ to compute a set K of all vertices of B that are reachable from the source vertex s by admissible paths. It increases the weights of all points in K by δ (expand their Voronoi cells). Recall that the set A δ , and thus G δ as well as the compressed transport plan ˆ τ δ depend on the weights of B . The procedure then recomputes A δ , G δ , and ˆ τ δ from τ δ . See Appendix B.2 for complete details of the INCREASEWEIGHTS procedure. The following lemma states the properties of the INCREASEWEIGHTS procedure.

Lemma 3.2. Suppose invariant (I1) holds at the start of the INCREASEWEIGHTS procedure. Then, during the execution of the INCREASEWEIGHTS procedure,

(W1) the transport plan ˆ τ δ , y ( · ) remain δ -feasible,

(W2) the weight of each free point b ∈ B increases by δ , and

(W3) the weight of each point b ∈ B with free regions inside V δ b remains unchanged.

3.4 ACYCLIFY Procedure

The change in the weights of B in the INCREASEWEIGHTS procedure requires us to recompute the residual graph and the compressed transport plan. This recomputation may potentially create a cycle in the transport plan or an admissible cycle in G δ . The ACYCLIFY procedure eliminates such cycles and ensures that the invariant (I2) holds at the start of the next iteration. Converting the transport plan into a forest is critical for the efficiency of the SEARCHANDAUGMENT procedure while eliminating admissible cycles is essential for the correctness of the SEARCHANDAUGMENT procedure. The procedure executes the following steps. First, use the dynamic tree structure by Sleator and Tarjan [56] to make the transport plan ˆ τ δ a forest. Then, execute a partial DFS search from

each unvisited point b ∈ B similar to the one described in the SEARCHANDAUGMENT procedure to detect admissible cycles. Upon finding an admissible cycle, cancel the cycle right away, remove the vertices of the cycle from the search path, and continue the search. When all vertices are visited, no admissible cycles are remaining in the residual graph. Since canceling admissible cycles could have introduced new cycles in the transport plan, repeat the first step to update ˆ τ δ and make it a forest. See Appendix B.3 for details. The following lemma shows the properties of the ACYCLIFY procedure.

Lemma 3.3. Suppose invariant (I1) holds at the start of the ACYCLIFY procedure. Then, during the execution of the ACYCLIFY procedure,

(A1) the transport plan ˆ τ δ , y ( · ) remains δ -feasible, and

(A2) the transport plan ˆ τ δ is a forest and there are no admissible cycles in the residual graph.

This completes the description of our algorithm.

4 Algorithm Analysis

In this section, we first prove the correctness of our algorithm and then analyze its runtime.

Proof of invariants (I1) and (I2). For any scale δ , the initial transport plan ˆ τ δ is empty. Therefore, ˆ τ δ along with the weights y ( · ) is δ -feasible. By properties (S1), (W1), and (A1), the transport plan ˆ τ δ , y ( · ) remains δ -feasible in each iteration of our algorithm, and therefore, invariant (I1) holds. The invariant (I2) is a direct consequence of property (A2) in Lemma 3.3.

Proof of Correctness. From Invariant (I1), in each scale δ , our algorithm maintains a δ -feasible transport plan ˆ τ δ , y ( · ) during its execution. The while loop in Step 2 breaks when ˆ τ δ is a complete transport plan. Therefore, ˆ τ δ along with weights y ( · ) is δ -optimal. From Lemma 2.4, one can convert ˆ τ δ into a δ -optimal semi-discrete transport plan τ δ . Given that our algorithm terminates when δ ≤ ε/ 2 , from Lemma 2.1, the transport plan returned by our algorithm is ε -close, as desired.

Efficiency of the algorithm. The SEARCHANDAUGMENT procedure runs a partial DFS on the residual graph with O ( n 3 ) edges. Upon finding an augmenting path P , the procedure augments the transport plan along P in O ( | P | ) time. In Lemma C.1, we use invariant (I2) to show that the total length of all augmenting paths found by the procedure is O ( n 3 ) . Hence, the running time of the SEARCHANDAUGMENT procedure is O ( n 3 ) . The INCREASEWEIGHTS procedure stores a sorted list of neighbors for each region and executes a DFS procedure in the residual graph, which takes O ( n 3 log n ) time. After increasing the weight of a subset of points of B , the procedure recomputes the residual graph G δ and the transport plan ˆ τ δ , which can be done in O ( n 2 (Φ + n log n )) time. See Lemma C.2 in the appendix for details. In the ACYCLIFY procedure, converting a transport plan into a forest can be done using a dynamic tree data structure [56] in O ( n 3 log n ) time. To eliminate the admissible cycles, the ACYCLIFY procedure relies on a partial DFS that runs in O ( n 3 ) time. Therefore, the total time by the ACYCLIFY procedure is O ( n 3 log n ) (see Lemma C.3 for details). Combining the running times of all three procedures, each iteration of step 2 of our algorithm takes O ( n 2 (Φ + n log n )) time. In the following lemma, we show that in each scale, the total number of iterations of step 2 of our algorithm is at most O ( n ) .

Lemma 4.1. For each scale δ , the total number of iterations of step 2 of our algorithm is O ( n ) .

Proof Sketch. Let τ 2 δ , y 2 δ ( · ) denote the 2 δ -feasible semi-discrete transport plan computed by our algorithm for scale 2 δ , and let τ δ , y δ ( · ) denote a partial semi-discrete transport plan maintained during the execution of step 2 of our algorithm. Let X 2 δ (resp. X δ ) denote the partitioning of the set A with respect to weights y 2 δ ( · ) (resp. y δ ( · ) ). Let Y be the arrangement of all 3 n cells used to construct X 2 δ with all 3 n cells used to construct X δ . Let ˆ τ 2 δ (resp. ˆ τ δ ) denote the compressed transport plan for τ 2 δ (resp. τ δ ) using the partitioning Y . It is well-known that one can transform ˆ τ δ to ˆ τ 2 δ by augmenting ˆ τ δ along a set of augmenting paths P on Y × B and rearrange the transported mass along a set of cycles C on Y × B . Consider an augmenting path P = ⟨ r 1 , b 1 , . . . , r k , b k ⟩ in P . Since P is a simple path, it contains each point of B at most once and therefore, it has a length at most 2 n -1 . Additionally, for all i ∈ [1 , k ] , ˆ τ 2 δ ( r i , b i ) > 0 and for each i ∈ [2 , k ] , ˆ τ δ ( r i , b i -1 ) > 0 . For each edge ( r i , b i ) , since ˆ τ 2 δ ( r i , b i ) > 0 , by 2 δ -feasibility of τ 2 δ , y 2 δ ( · ) (condition (F1)), the point b i is a 4 δ -weighted nearest neighbor of r i with respect to weights y 2 δ ( · ) . Similarly, for each edge ( b i -1 , r i ) , since ˆ τ δ ( r i , b i -1 ) > 0 , by δ -feasibility of τ δ , y δ ( · ) (condition (F1)), the point b i -1 is a 2 δ -weighted nearest neighbor of r i with respect to weights y δ ( · ) . Since the length of P is at most 2 n -1 and the

weight of b 1 does not change (by property (W3) in Lemma 3.3), we show that for the free point b k in P , y δ ( b k ) -y 2 δ ( b k ) ≤ 6 nδ . Our algorithm increases the weight of b k by δ in each iteration (property (W2) in Lemma 3.3), and therefore after 6 n iterations, the point b k cannot be free, i.e., after O ( n ) iterations, all points of B are fully transported in ˆ τ δ . We provide the full proof in Section C.4.

Using Lemma 4.1, the total time spent in step 2 in each scale of our algorithm is O ( n 3 (Φ + n log n )) . Since there are O (log ∆ ε ) scales, the overall runtime is O ( n 3 (Φ + n log n ) log ∆ ε ) , thereby proving Theorem 1.1.

5 Applications to the Discrete Optimal Transport Problem

In this section, we extend our combinatorial semi-discrete OT algorithm to the discrete OT problem and design a data structure that preprocesses and stores a large discrete distribution µ and efficiently computes an ε -close OT plan between µ and any query distribution ν in sub-linear time relative to the support size of µ . More precisely, given a discrete distribution µ with a (possibly large) support A of N points in R 2 , we design a data structure that, given a query discrete distribution ν with a support B of k points, computes an ε -close transport plan between µ and ν in O ( k 3 ( √ N + k log k ) log ∆ ε ) time. Additionally, we show that if the support points have bounded integer coordinates and the masses are rational numbers, our data structure can efficiently compute an exact discrete OT plan.

At a high level, our data structure interprets the large discrete distribution µ as a continuous distribution and uses a simplex range-searching data structure as an oracle to compute the mass of µ inside a query triangle. In this way, for any query distribution ν , one can execute the steps of our combinatorial semi-discrete algorithm to compute an ε -close transport plan between µ and ν . More formally, our data structure preprocesses the distribution µ into a simplex range-searching data structure RS-DS, which takes O ( N ) space, can be built in O ( N log N ) time, and returns the mass of µ inside a query triangle in Φ = O ( √ N ) time [59]. Given a query discrete distribution ν , one can use our algorithm from Section 3 in conjunction with the RS-DS to compute an ε -close transport plan between µ and ν in O ( k 3 ( √ N + k log k ) log ∆ ε ) time leading to Theorem 1.2.

Consider the special case where the points in A ∪ B have positive integer coordinates bounded by λ , the mass of µ (resp. ν ) on each point a ∈ A (resp. b ∈ B ) is a rational number of form x a T (resp. x b T ) for positive integers T and x a (resp. x b ), and p is an even number. In this case, the p -Wasserstein cost of any transport plan between µ and ν is an integer multiple of 1 T , and therefore, any 1 2 T -close transport plan between µ and ν would have a minimum cost. Thus, one can compute an exact discrete OT plan between µ and ν by setting ε = 1 2 T in our data structure, which would have a query time of O ( k 3 ( √ N + k log k ) log( λT )) , leading to the following corollary.

Corollary 5.1. When the points in the supports of the distributions µ and ν have integer coordinates bounded by λ and the mass on each point is a rational number of form x T , our data structure computes, for any even number p , an optimal solution for the p -Wasserstein problem between µ and ν in O ( k 3 ( √ N + k log k ) log( λT )) time.

6 Conclusion

In this paper, we designed a novel combinatorial framework for the semi-discrete optimal transport problem and used it to compute an ε -close semi-discrete transport plan. We also used this framework to design a data structure that stores a discrete distribution µ over a large support of size N and can compute ε -close OT cost between µ and a query discrete distribution ν in a time that is sub-linear in N . We conclude with the following question: Can we use our combinatorial framework to compute an ε -close semi-discrete transport plan between high dimensional distributions in poly( n, d, 1 /ε ) time.

References

Pankaj K. Agarwal and R. Sharathkumar. Approximation algorithms for bipartite matching with metric and geometric costs. In Proc. 46th Annual ACM Sympos. on Theory of Computing , STOC '14, pages 555-564, 2014.

Pankaj K Agarwal and Micha Sharir. Arrangements and their applications. In Handbook of Computational Geometry , pages 49-119. Elsevier, 2000.

Pankaj K Agarwal, Kyle Fox, Debmalya Panigrahi, Kasturi R Varadarajan, and Allen Xiao. Faster algorithms for the geometric transportation problem. In 33rd Internat. Symposium on Computational Geometry , 2017.

Pankaj K Agarwal, Sharath Raghvendra, Pouyan Shirzadian, and Rachita Sowle. A higher precision algorithm for computing the 1 -Wasserstein distance. In Proc. 11th Internat. Conference on Learning Representations , 2023.

Pankaj K Agarwal, Sharath Raghvendra, Pouyan Shirzadian, and Keegan Yao. Fast and accurate approximations of the optimal transport in semi-discrete and discrete settings. In Proc. 2024 Annual ACM-SIAM Sympos. on Discrete Algorithms , pages 4514-4529. SIAM, 2024.

Jason M Altschuler, Jonathan Niles-Weed, and Austin J Stromme. Asymptotics for semidiscrete entropic optimal transport. SIAM Journal of Mathematical Analysis , 54(2):1718-1741, 2022.

Luca Ambrogioni, Umut Guclu, and Marcel van Gerven. Wasserstein variational gradient descent: From semi-discrete optimal transport to ensemble variational inference. arXiv preprint arXiv:1811.02827 , 2018.

Luca Ambrogioni, Umut Güçlü, and Marcel van Gerven. k-gans: Ensemble of generative models with semi-discrete optimal transport. arXiv preprint arXiv:1907.04050 , 2019.

Dongsheng An, Yang Guo, Na Lei, Zhongxuan Luo, Shing-Tung Yau, and Xianfeng Gu. AE-OT: A new generative model based on extended semi-discrete optimal transport. In Proc. 32nd Internat. Conference on Learning Representations , 2019.

Dongsheng An, Yang Guo, Min Zhang, Xin Qi, Na Lei, and Xianfang Gu. AE-OT-GAN: Training gans from data specific latent distribution. In Proc. 16th European Conference on Computer Vision , pages 548-564. Springer, 2020.

Franz Aurenhammer. Power diagrams: properties, algorithms and applications. SIAM journal on computing , 16(1):78-96, 1987.

Franz Aurenhammer, Friedrich Hoffmann, and Boris Aronov. Minkowski-type theorems and least-squares clustering. Algorithmica , 20(1):61-76, 1998.

Arturs Backurs, Yihe Dong, Piotr Indyk, Ilya Razenshteyn, and Tal Wagner. Scalable nearest neighbor search for optimal transport. In Internat. Conference on Machine Learning , pages 497-506. PMLR, 2020.

Marc G Bellemare, Ivo Danihelka, Will Dabney, Shakir Mohamed, Balaji Lakshminarayanan, Stephan Hoyer, and Rémi Munos. The Cramer distance as a solution to biased Wasserstein gradients. arXiv preprint arXiv:1705.10743 , 2017.

Jean-David Benamou and Yann Brenier. A computational fluid mechanics solution to the Monge-Kantorovich mass transfer problem. Numerische Mathematik , 84(3):375-393, 2000.

Bernard Bercu and Jérémie Bigot. Asymptotic distribution and convergence rates of stochastic algorithms for entropic optimal transportation between probability measures. Annals of Statistics , 49(2):968-987, 2021.

Rick Chartrand, Brendt Wohlberg, Kevin Vixie, and Erik Bollt. A gradient descent solution to the Monge-Kantorovich problem. Applied Math. Sci. , 3(22):1071-1080, 2009.

Li Chen, Rasmus Kyng, Yang P Liu, Richard Peng, Maximilian Probst Gutenberg, and Sushant Sachdeva. Maximum flow and minimum-cost flow in almost-linear time. In Proc. 63rd IEEE Annual Sympos. Foundations of Computer Science , pages 612-623. IEEE, 2022.

Yucheng Chen, Matus Telgarsky, Chao Zhang, Bolton Bailey, Daniel Hsu, and Jian Peng. A gradual, semi-discrete approach to generative network training via explicit Wasserstein minimization. In Internat. Conference on Machine Learning , pages 1071-1080. PMLR, 2019.

Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems , pages 2292-2300, 2013.

Fernando De Goes, Katherine Breeden, Victor Ostromoukhov, and Mathieu Desbrun. Blue noise through optimal transport. ACM Transactions of Graphics , 31(6):1-11, 2012.

Eustasio Del Barrio, Alberto González Sanz, and Jean-Michel Loubes. Central limit theorems for semi-discrete Wasserstein distances. Bernoulli , 30(1):554-580, 2024.

Efim A Dinic. Algorithm for solution of a problem of maximum flow in networks with power estimation. In Soviet Math. Doklady , volume 11, pages 1277-1280, 1970.

Steven Fortune. Voronoi diagrams and Delaunay triangulations. In Handbook of Discrete and Computational Geometry , pages 225-265, 1995.

Emily Fox and Jiashuai Lu. A deterministic near-linear time approximation scheme for geometric transportation. In Proc. 64th IEEE Annual Sympos. on Foundations of Computer Science , pages 1301-1315. IEEE, 2023.

Kyle Fox and Jiashuai Lu. A near-linear time approximation scheme for geometric transportation with arbitrary supplies and spread. In Proc. 36th Annual Sympos. on Computational Geometry , pages 45:1-45:18, 2020.

H. N. Gabow and R. E. Tarjan. Faster scaling algorithms for network problems. SIAM Journal of Computing , 18(5):1013-1036, October 1989. ISSN 0097-5397.

Bruno Galerne, Arthur Leclaire, and Julien Rabin. Semi-discrete optimal transport in patch space for enriching Gaussian textures. In Proc. 3rd Internat. Conference on Geometric Science of Information , pages 100-108, 2017.

Bruno Galerne, Arthur Leclaire, and Julien Rabin. A texture synthesis model based on semidiscrete optimal transport in patch space. SIAM Journal of Imaging Sciences , 11(4):2456-2493, 2018.

Thomas O Gallouët and Quentin Mérigot. A Lagrangian scheme à la Brenier for the incompressible Euler equations. Foundations of Computational Mathematics , 18(4):835-865, 2018.

Jing Gao, Peng Li, Asif Ali Laghari, Gautam Srivastava, Thippa Reddy Gadekallu, Sidra Abbas, and Jianing Zhang. Incomplete multiview clustering via semidiscrete optimal transport for multimedia data mining in IoT. ACM Transactions of Multimedia Computing, Communications and Applications , 20(6):1-20, 2024.

Aude Genevay, Marco Cuturi, Gabriel Peyré, and Francis Bach. Stochastic optimization for large-scale optimal transport. Advances in Neural Information Processing Systems , 29, 2016.

Antoine Houdard, Arthur Leclaire, Nicolas Papadakis, and Julien Rabin. A generative model for texture synthesis based on optimal transport between feature distributions. Journal of Mathematical Imaging and Vision , 65(1):4-28, 2023.

Martin Huesmann and Karl-Theodor Sturm. Optimal transport from Lebesgue to poisson. The Annals of Probability , 41(4):2426-2478, 2013.

Shayan Hundrieser, Thomas Staudt, and Axel Munk. Empirical optimal transport between different measures adapts to lower complexity. In Annales de l'Institut Henri Poincare (B) Probabilites et statistiques , volume 60, pages 824-846, 2024.

Andrey Boris Khesin, Aleksandar Nikolov, and Dmitry Paramonov. Preconditioning for the geometric transportation problem. arXiv preprint arXiv:1902.08384 , 2019.

Jun Kitagawa. An iterative scheme for solving the optimal transportation problem. Calculus of Variations and Partial Differential Equations , 51(1):243-263, 2014.

Jun Kitagawa, Quentin Mérigot, and Boris Thibert. Convergence of a Newton algorithm for semi-discrete optimal transport. Journal of European Math. Society , 21(9):2603-2651, 2019.

Nathaniel Lahn, Deepika Mulchandani, and Sharath Raghvendra. A graph theoretic additive approximation of optimal transport. Advances in Neural Information Processing Systems , 32, 2019.

Arthur Leclaire and Julien Rabin. A stochastic multi-layer algorithm for semi-discrete optimal transport with applications to texture synthesis and style transfer. Journal of Mathematical Imaging and Vision , 63(2):282-308, 2021.

Bruno Lévy and Erica L Schwindt. Notions of optimal transport theory and how to implement them on a computer. Computers & Graphics , 72:135-148, 2018.

Zezeng Li, Shenghao Li, Zhanpeng Wang, Na Lei, Zhongxuan Luo, and David Xianfeng Gu. DPM-OT: a new diffusion probabilistic model based on optimal transport. In Proc. IEEE/CVF Internat. Conference on Computer Vision , pages 22624-22633, 2023.

Huidong Liu, GU Xianfeng, and Dimitris Samaras. A two-step computation of the exact GAN Wasserstein distance. In Internat. Conference on Machine Learning , pages 3159-3168, 2018.

Quentin Mérigot. A multiscale approach to optimal transport. Computer Graphics Forum , 30 (5):1583-1592, 2011.

Quentin Merigot and Boris Thibert. Optimal transport: discretization and algorithms. In Handbook of Numerical Analysis , volume 22, pages 133-212. Elsevier, 2021.

Jocelyn Meyron. Initialization procedures for discrete and semi-discrete optimal transport. Computer-Aided Design , 115:13-22, 2019.

Joseph SB Mitchell, David M Mount, and Christos H Papadimitriou. The discrete geodesic problem. SIAM Journal of Computing , 16(4):647-668, 1987.

Vladimir I Oliker and Laird D Prussner. On the numerical solution of the equation ∂ 2 z ∂x 2 ∂ 2 z ∂y 2 -( ∂ 2 z ∂x∂y ) = f and its discretizations, i. Numerische Mathematik , 54(3):271-293, 1989.

James Orlin. A faster strongly polynomial minimum cost flow algorithm. In Proc. 20th Annual ACM Sympos. Theory of Computing , pages 377-387, 1988.

Aram-Alexandre Pooladian, Vincent Divol, and Jonathan Niles-Weed. Minimax estimation of discontinuous optimal transport maps: The semi-discrete case. In Internat. Conference on Machine Learning , pages 28128-28150. PMLR, 2023.

Hongxing Qin, Yi Chen, Jinlong He, and Baoquan Chen. Wasserstein blue noise sampling. ACM Transactions of Graphics , 36(5):1-13, 2017.

Sharath Raghvendra and Pankaj K Agarwal. A near-linear time ε -approximation algorithm for geometric bipartite matching. Journal of the ACM , 67(3):1-19, 2020.

Nathan A Seifert, Kirill Prozument, and Michael J Davis. Computational optimal transport for molecular spectra: The semi-discrete case. The Journal of Chemical Physics , 156(13), 2022.

R Sharathkumar and Pankaj K Agarwal. Algorithms for the transportation problem in geometric settings. In Proc. 23rd Annual ACM-SIAM Sympos. on Discrete Algorithms , pages 306-317, 2012.

Ilia Shumailov, Zakhar Shumaylov, Yiren Zhao, Nicolas Papernot, Ross Anderson, and Yarin Gal. Ai models collapse when trained on recursively generated data. Nature , 631(8022): 755-759, 2024.

Daniel D Sleator and Robert Endre Tarjan. A data structure for dynamic trees. In Proc. 13th Annual ACM Sympos. on Theory of Computing , pages 114-122, 1981.

Kokichi Sugihara. Laguerre Voronoi diagram on the sphere. Journal for Geometry and Graphics , 6(1):69-81, 2002.

Bahar Ta¸ skesen, Soroosh Shafieezadeh-Abadeh, and Daniel Kuhn. Semi-discrete optimal transport: Hardness, regularization and numerical solution. Mathematical Programming , 199 (1):1033-1106, 2023.

Emo Welzl. Partition trees for triangle counting and other range searching problems. In Proc. 4th Annual Sympos. Computational Geometry , pages 23-33, 1988.

A Missing Proofs of Section 2

In this section, we present the missing proofs and details of our combinatorial framework described in Section 2.

Lemma 2.1. Any δ -optimal transport plan τ, y ( · ) between µ and ν is 2 δ -close.

Proof. Define the weighted cost of a transport plan τ ′ as ¢ y ( τ ′ ) := ∑ b ∈ B ∫ A d y ( a, b ) τ ′ ( a, b ) da . For any complete transport plan τ ′ ,

For any point a ∈ A , let b a denote the weighted nearest neighbor of a , i.e., b a := arg min b ∈ B d y ( a, b ) . By property (F1) in the definition of δ -feasibility,

Let τ ∗ denote an optimal transport plan from µ to ν . Then,

Combining Equations (3) and (4) and plugging τ and τ ∗ in Equation (2),

Therefore, ¢ ( τ ) ≤ ¢ ( τ ∗ ) + 2 δ and τ is 2 δ -close.

Residual graph. Next, we show that for any δ > 0 and any δ -feasible transport plan ˆ τ, y ( · ) , the residual graph G δ has O ( n 2 ) vertices and O ( n 3 ) edges, leading to Lemma 2.2. To do so, we show below that the partitioning X δ consists of O ( n 2 ) regions. We then conclude that the number of vertices of G δ is O ( n 2 ) . Furthermore, since the residual graph is a bipartite graph between set B of size n and set A δ of size O ( n 2 ) , the number of edges of G δ would be at most O ( n 3 ) .

Recall that the partitioning X δ is constructed as the arrangement of all Voronoi cells, δ -expansions, and 2 δ -expansions of the Voronoi cells of all points in B . Let V denote the set of vertices of this arrangement. Since the arrangement is a planar graph, the number of faces (i.e., regions) in X δ is O ( |V| ) . Therefore, to show that X δ has O ( n 2 ) regions, we show that the number of vertices of this arrangement is O ( n 2 ) .

For each point b ∈ B , let V 0 b (resp. V δ b , V 2 δ b ) denote the Voronoi cell (resp. δ -expanded Voronoi cell, 2 δ -expanded Voronoi cell) of b , and let y 0 b (resp. y δ b , y 2 δ b ) denote the weights of B used to compute the cell. Note that each Voronoi cell V i b for each b ∈ B and i ∈ { 0 , δ, 2 δ } has at most n vertices. Hence, the total number of Voronoi vertices in the arrangement X δ is O ( n 2 ) . Next, we count the number of intersection points of these Voronoi cells. For each pair of points b 1 , b 2 ∈ B , consider the pair of cells V i b 1 and V j b 2 , for some i, j ∈ { 0 , δ, 2 δ } . We show that V i b 1 and V j b 2 intersect each other in at most two points.

Define the weighted bisector of two points b and b ′ with respect to weights y ( · ) as the locus of points that have the same weighted distance to b and b ′ , i.e., all points x ∈ R 2 such that d y ( x, b ) = d y ( x, b ′ ) . Note that under the squared Euclidean distances, the weighted bisector of two points is a straight line. Let v denote an intersection point of V i b 1 and V j b 2 , and suppose v lies on the segment of V i b 1

na

corresponding to the weighted bisector of b 1 and a point b 3 ∈ B (See Figure 4). Note that if b 3 = b 2 , then V i b 1 and V j b 2 share a segment containing v , which means that

and also

̸

Therefore, in this case, i = j = 0 , and the two endpoints of the shared segment is counted toward the number of Voronoi vertices of the arrangement. Hence, we assume b 3 = b 2 . In this case,

By the construction of the weight function y i b 1 ( · ) in Equation (1), for any point b ∈ B \ { b 1 , b 2 } , y i b 1 ( b ) = y j b 2 ( b ) = y ( b ) ; thus, using Equation (5),

Using Equation (6), the point v also lies in the segment of V j b 2 that corresponds to the weighted bisector of b 2 and b 3 with respect to weights y j b 2 . Thus,

In other words, the point v satisfies

i.e., the point v lies on the weighted bisector of b 1 and b 2 (the blue dashed line in Figure 4). Since the weighted bisector is a straight line, it intersects the convex polygon V i b 1 in at most 2 points; hence, the two Voronoi cells V i b 1 and V j b 2 intersect each other in at most 2 points. Since the are O ( n 2 ) pairs of such Voronoi cells, each intersecting each other in at most 2 points, the total number of intersection points is at most O ( n 2 ) , as claimed.

Lemma 2.2. For any δ > 0 and a δ -feasible transport plan ˆ τ, y ( · ) , the residual graph G δ has O ( n 2 ) vertices and O ( n 3 ) edges.

Next, we extend our bounds for the number of vertices and edges of the residual graph in Lemma 2.2 to any dimension d ≥ 3 and show that for any pair of d -dimensional distributions µ and ν and any δ -feasible transport plan ˆ τ, y ( · ) , the residual graph G δ has O ( n d ) vertices and O ( n d +1 ) edges. We prove our bounds by showing that the number of vertices of the arrangement X δ is O ( n d ) . Using this

bound, we get that the number of regions in the arrangement is upper-bounded by O ( n d ) ; therefore, | A δ | = O ( n d ) and the residual graph has O ( n d ) vertices. Furthermore, since G δ is a bipartite graph between set B with n points and set A δ with O ( n d ) points, the number of edges would be bounded by O ( n d +1 ) .

For each point b ∈ B and each i ∈ { 0 , δ, 2 δ } , the V oronoi cell V i b has O ( n ⌈ d/ 2 ⌉ ) vertices; since 3 n such cells are used in the construction of the arrangement, the total number of Voronoi vertices in the arrangement is O ( n ⌈ d/ 2 ⌉ +1 ) . Next, we bound the number of intersection points of these Voronoi cells. For each d -tuple ( V i 1 b 1 , V i 2 b 2 , . . . , V i d b d ) of Voronoi cells, for d distinct points b 1 , . . . , b d and d values i 1 , . . . , i d ∈ { 0 , δ, 2 δ } , let v denote a point in the intersection of these cells. Similar to our analysis for 2 -dimensional distributions, we show that v lies on the weighted bisector of b 1 , . . . , b d , where the weight of each point b j is y ( b j ) + i j . Since this weighted bisector is a straight line, and since a straight line intersects a convex polytope in at most 2 points, each d -tuple of Voronoi cells intersect in at most 2 points, where the number of such d -tuples are O ( n d ) 4 . Therefore, we conclude that the number of intersection points is O ( n d ) , as desired.

Let v be an intersection point of the d Voronoi cells V i 1 b 1 , V i 2 b 2 , . . . , and V i d b d . Suppose v lies on the ( d -1) -dimensional hyperplane of V i 1 b 1 that is shared between b 1 and a point b d +1 ∈ B . In this case,

We claim that for any k ∈ { 1 , . . . , d } , y i k b k ( b d +1 ) = y ( b d +1 ) . Consider the following two cases:

If b d +1 = b j for some j ∈ { 2 , . . . , d } , then

and also

Therefore, i 1 = i j = 0 and for any k ∈ { 1 , . . . , d } , y i k b k ( b d +1 ) = y ( b d +1 ) .

Otherwise, b d +1 / ∈ { b 1 , . . . , b d } and by the construction of y i k b k ( · ) , for any k ∈ { 1 , . . . , d } , y i k b k ( b d +1 ) = y ( b d +1 ) .

Therefore, from Equation (7),

and for any k ∈ { 1 , . . . , d } , since v lies on a ( d -1) -dimensional hyperplane of V i k b k , it has to lie on the hyperplane of V i k b k that is shared between b k and b d +1 ; therefore, d( v, b d +1 ) -y ( b d +1 ) = d( v, b k ) -( y ( b k ) + i k ) . Thus,

From Equation (9), the point v lies on the weighted bisector of b 1 , . . . , b d with weights y ( b 1 ) + i 1 , . . . , y ( b d ) + i d , and the d -tuple ( V i 1 b 1 , V i 2 b 2 , . . . , V i d b d ) intersect each other in at most 2 points.

Lemma A.1. For any dimension d ≥ 2 , a parameter δ > 0 , and a δ -feasible transport plan ˆ τ, y ( · ) , the residual graph G δ has O ( n d ) vertices and O ( n d +1 ) edges.

Compressing the transport plan. Recall that for any semi-discrete transport plan τ from B to A , a set of weights y ( · ) , and a parameter δ > 0 , the compressed transport plan ˆ τ is a discrete transport plan from B to A δ . Note that the definition of δ -feasibility (and more precisely, condition (F1)) is applicable for the discrete OT as well; the discrete transport plan ˆ τ along with weights y ( · ) for B is δ -feasible if ˆ τ transports the mass of each point b ∈ B to points of A δ within its 2 δ -expanded Voronoi cell. Any complete transport plan ˆ τ from B to A δ that is δ -feasible along with weights y ( · ) is called δ -optimal.

Lemma A.2. Any δ -optimal compressed transport plan ˆ τ, y ( · ) from B to A δ is 2 δ -close.

Proof. For any transport plan τ ′ from B to A δ define ¢ y ( τ ′ ) := ∑ ( r,b ) ∈ A δ × B d y ( r, b ) τ ′ ( r, b ) . Then,

For any point r ∈ A δ , let b r denote the weighted nearest neighbor of r , i.e., b r := arg min b ∈ B d y ( r, b ) . By the definition of δ -feasibility, for any pair ( r, b ) with ˆ τ ( r, b ) > 0 , we have d y ( r, b ) ≤ d y ( r, b r )+2 δ . Therefore,

Let τ ∗ denote an optimal transport plan from B to A δ . Then,

Combining Equations (11) and (12) and plugging ˆ τ and τ ∗ in Equation (10),

Therefore, ¢ (ˆ τ ) ≤ ¢ ( τ ∗ ) + 2 δ and ˆ τ is 2 δ -close.

The following observation, which is straightforward from the construction of the partitioning X δ , helps in proving Lemmas 2.3 and 2.4.

Observation A.3. For any region φ ∈ X δ and any point b ∈ B , the region φ either completely lies inside V 2 δ b or it is completely outside V 2 δ b .

Lemma 2.3. For any δ -feasible semi-discrete transport plan τ, y ( · ) , the compressed transport plan ˆ τ along with weights y ( · ) is also δ -feasible.

Proof. Consider any point b ∈ B and any region φ ∈ X δ such that τ ( φ, b ) > 0 , i.e., in the compressed graph, ˆ τ ( r φ , b ) > 0 . To prove this lemma, we show that r φ ∈ V 2 δ b . Since τ ( φ, b ) > 0 , there exists a point a ∈ φ such that τ ( a, b ) > 0 , and by the δ -feasibility of τ, y ( · ) , we have a ∈ V 2 δ b ; therefore, using Observation A.3, the region φ has to completely lie inside V 2 δ b , and the representative point r φ , which is a point inside φ also lies inside V 2 δ b . Hence, ˆ τ, y ( · ) is δ -feasible.

Lemma 2.4. Any δ -feasible transport plan ˆ τ, y ( · ) from B to A δ can be converted into a δ -feasible semi-discrete transport plan τ from B to A .

Proof. Consider any transport plan τ from B to A such that for any point b ∈ B and any region φ ∈ X δ , transports a mass of ˆ τ ( r φ , b ) from b to the mass of µ inside the region φ . One such construction is to assign τ ( a, b ) = ˆ τ ( r φ ,b ) µ ( φ ) µ ( a ) for each point a ∈ φ . We next show that the transport plan τ is δ -feasible.

Consider any point b ∈ B and any point a ∈ A such that τ ( a, b ) > 0 . We prove this lemma by showing that a ∈ V 2 δ b . Suppose φ ∈ X δ is the region containing the point a . Since τ ( a, b ) > 0 , we should have ˆ τ ( r φ , b ) > 0 , and by the δ -feasibility of ˆ τ, y ( · ) , we have r φ ∈ V 2 δ b . By Observation A.3, the whole region φ has to lies inside V 2 δ b and therefore, the point a also lies inside V 2 δ b . Hence, τ, y ( · ) is δ -feasible.

Augmentation. In this section, we first show that upon augmenting a transport plan ˆ τ along an augmenting path P , the transport plan remains valid. We then show that upon augmentation, either a backward edge gets removed from the transport plan or a free point will become fully transported. We finally prove Lemma 2.5.

For an augmenting path P = ⟨ s, b 1 , r 1 , . . . , r k , t ⟩ from a free point b 1 to a free point r k , recall that the bottleneck capacity of P is defined as

where ex( b 1 ) (resp. ex( r k ) ) denotes the excess mass of b 1 (resp. r k ). To show that the transport plan after augmentation is a valid one, we show that no point v in G δ transports more mass than the mass at v and that each edge transports a non-negative amount of mass.

In the augmentation process, for any i ∈ { 1 , . . . , k -1 } , we decrease ˆ τ ( r i , b i +1 ) by bc( P ) , where by definition, bc( P ) ≤ ˆ τ ( r i , b i +1 ) ; hence, ˆ τ ( r i , b i +1 ) ≥ 0 after augmentation. Furthermore, for any forward edge ( b i , r i ) , i ∈ { 1 , . . . , k } , we increases ˆ τ ( r i , b i ) by bc( P ) , and ˆ τ ( r i , b i ) remains non-negative.

For any i ∈ { 2 , . . . , k } , we increase ˆ τ ( b i , r i ) (resp. decrease ˆ τ ( r i -1 , b i ) ) by bc( P ) ; hence, the total amount of mass transported from b i remains unchanged. Similarly, for each i ∈ { 1 , . . . , k -1 } , we increase ˆ τ ( b i , r i ) (resp. decrease ˆ τ ( r i , b i +1 ) ) by bc( P ) and the total amount of mass transported into r i remains unchanged. For the endpoint b 1 (resp. r k ), we only increase the amount of mass transported from b 1 (resp. into r k ) by bc( P ) , where by definition, bc( P ) ≤ ex( b 1 ) (resp. bc( P ) ≤ ex( r k ) ). Therefore, the total mass transport from b 1 (resp. into r k ) after augmentation is at most ν ( b ) (resp. µ r k ), as desired.

Note that if a backward edge ( r, b ) ∈ P determines the bottleneck capacity of the augmenting path P , then ˆ τ ( r, b ) = 0 after augmentation and the backward edge is removed from the residual graph. Otherwise, if the endpoint b 1 (reps. r k ) determines the bottleneck capacity of P , then b 1 (resp. r k ) will be fully transported after augmentation.

Lemma 2.5. The transport plan obtained after augmenting a δ -feasible transport plan ˆ τ, y ( · ) along any augmenting path P in the residual graph is δ -feasible.

Proof. Let P = ⟨ s, b 1 , r 1 , . . . , b k , r k , t ⟩ denote an augmenting path in the residual graph. When augmenting the transport plan along P , we increase the mass transportation on forward edges ( b i , r i ) for each i ∈ [1 , k ] and decrease the mass transportation on the backward edges ( r i , b i +1 ) for each i ∈ [1 , k -1] . Therefore, any pair ( r, b ) that transports mass after augmentation but was not transporting mass before augmentation has to be a forward edge of P . Since we only add forward edges from the point b to the representative points in V 2 δ b ∩ A δ , the edge ( r, b ) satisfies the δ -feasibility condition (F1), and ˆ τ, y ( · ) remains feasible after augmentation.

B Missing Details of Section 3

In this section, we provide the missing details of the implementation of the INCREASEWEIGHTS and ACYCLIFY procedures and also prove the properties of the three procedures.

B.1 Missing Proofs of the SEARCHANDAUGMENT Procedure

Lemma 3.1. Suppose invariants (I1) and (I2) hold at the start of the SEARCHANDAUGMENT procedure. Then, during the execution of the SEARCHANDAUGMENT procedure,

(S1) the transport plan ˆ τ δ , y ( · ) remains δ -feasible,

(S2) any point b ∈ B (resp. backward edge ( r, b ) ) marked as visited will not form an admissible augmenting path during the execution of the procedure, and

(S3) there are no admissible cycles in the residual graph.

Proof. We prove the properties separately in the following.

Property (S1). By the construction of the search path, any augmenting path computed by the SEARCHANDAUGMENT procedure is an admissible augmenting path. From Lemma 2.5, augmenting

b4

b2

b6

b1

b4

t

b2

S

b6

b1

b3

ˆ τ δ along an admissible augmenting path does not violate the δ -feasibility condition; hence, the transport plan ˆ τ δ , y ( · ) is a δ -feasible transport plan during the execution of the SEARCHANDAUGMENT procedure and (S1) holds.

Next, we present an overview of a new property of the SEARCHANDAUGMENT procedure, which we formally state and prove in Lemma B.1 below. Using that, we first prove (S3) and then prove (S2).

In Figure 5 (left), suppose the blue edges show an admissible augmenting path found by the SEARCHANDAUGMENT procedure, and suppose the green (resp. red) points show the unvisited (resp. visited) points of B . For the region φ , suppose ⟨ b 1 , b 2 , . . . , b 6 ⟩ denote the set of neighbors of r φ in G δ , sorted in increasing order of their weighted distance to r φ . When the partial DFS procedure processes r φ , the two weighted nearest neighbors of φ (i.e., b 1 and b 2 ) are already marked as visited, leading the procedure to add b 3 to the search path. After augmentation (Figure 5 (right)), for the newly created backward edge ( r φ , b 3 ) , the only admissible triples containing ( r φ , b 3 ) are ( b 1 , r φ , b 3 ) and ( b 2 , r φ , b 3 ) (the triples highlighted in pink), where b 3 is unvisited and both b 1 and b 2 are visited. More formally, as shown in Lemma B.1, for an augmenting path P found by the procedure, assuming that (S3) holds before augmentation along P , for any newly created admissible triples ( b, r φ , b ′ ) after augmenting the transport plan along P , the point b (resp. b ′ ) is marked as visited (resp. unvisited).

Property (S3). We begin by presenting an overview of our proof. Consider any augmenting path P found by the SEARCHANDAUGMENT procedure. Assuming that (S3) holds before augmentation along P , all vertices that are reachable from a visited point b by an admissible path in the procedure are also visited, since those points were also added to the search path, did not lead to an admissible augmenting path, marked as visited and removed from the search path. Hence, all points having an admissible path to the visited points (i.e., all points that are reachable from the visited points in our backward DFS) are also visited. Therefore, there are no admissible paths from an unvisited point to a visited point in the residual graph. After augmenting along P , by Lemma B.1, for any newly formed admissible triple ( b, r φ , b ′ ) , the point b (resp. b ′ ) is visited (resp. unvisited), and since there are no admissible paths from any unvisited point to any visited point, the newly formed admissible triple ( b, r φ , b ′ ) does not form a cycle of admissible triples. Hence, (S3) holds after augmentation as well. We provide the details of the proof below.

Let P 1 , . . . , P k denote the sequence of augmenting paths computed by the SEARCHANDAUGMENT procedure, and let ˆ τ 0 δ , ˆ τ 1 δ , . . . , ˆ τ k δ denote the sequence of transport plans computed by the procedure, i.e., ˆ τ 0 δ is the transport plan computed in the previous iteration, and for each i ∈ [1 , k ] , ˆ τ i δ is obtained by augmenting ˆ τ i -1 δ along P i . Let G i denote the residual graph corresponding to ˆ τ i δ for each i ∈ [0 , k ] . Let V i (resp. U i ) denote the set of visited (resp. unvisited) points in B when the procedure augments ˆ τ i -1 δ along P i .

Initially, from invariant (I2), there are no cycles of admissible triples in the residual graph, and (S3) holds for G 0 . For any i ∈ [1 , k ] , assuming G 0 , . . . , G i -1 satisfies (S3), we show that G i also satisfies (S3). Suppose ( b, r, b ′ ) is any admissible triple in G i formed after augmenting ˆ τ i -1 δ along P i , i.e., the triple ( b, r, b ′ ) is admissible in G i but not in G i -1 . We show that the triple ( b, r, b ′ ) does not participate in any admissible cycles; hence, using property (S3) on G i -1 , there are no admissible cycles in G i and property (S3) holds for G i as well.

Since (S3) holds for G 0 , . . . , G i -1 , for any visited point b ∈ V i -1 , the point b has been added to the search path Q by the SEARCHANDAUGMENT procedure, did not lead to the computation of an augmenting path, marked as visited and removed from the path. In this case, any point b ′′ that is reachable from b by an admissible path in ←- -G i -1 (and therefore, is reachable by our partial DFS procedure from b ) would have also been added to the path, marked as visited and removed from Q , i.e., any point b ′′ ∈ B that has an admissible path to the visited point b in the residual graph G i -1 is also visited. By this observation, there are no admissible paths from any unvisited point in B to any visited point in G i -1 . From Lemma B.1, for any newly formed admissible triple ( b, r, b ′ ) , we have b ∈ V i -1 and b ′ ∈ U i -1 . Thus, all admissible triples formed after augmenting ˆ τ i -1 δ along P i are from a visited point to an unvisited point, while there are no admissible paths from any unvisited point to any visited point; therefore, the newly formed admissible triples do not form any admissible cycles and (S3) holds for G i as well.

Property (S2). We use the property (S3) to show that (S2) holds. For any point b ∈ B that is marked as visited, as discussed above, if (S3) holds, all vertices that are reachable from b by an admissible path in our backward DFS (i.e., all points having an admissible path to b in the residual graph) are also visited. Since any free point b f ∈ B is unvisited, b f does not have an admissible path to any visited point b ∈ B and therefore, the visited points do not participate in an admissible augmenting path. Furthermore, the procedure marks a backward edge ( r, b ) as visited if, for each admissible triple ( b ′ , r, b ) , the point b ′ is visited. Since the point b ′ cannot be included in an admissible augmenting path, the visited backward edge ( r, b ) also does not form an admissible augmenting path.

Lemma B.1. During the execution of the SEARCHANDAUGMENT procedure, suppose P is an admissible augmenting path found by the procedure, and let G (resp. G ′ ) denote the residual graph before (resp. after) augmenting the transport plan along P . Let ( b, r, b ′ ) denote an admissible triple in G ′ that is not admissible in G . Assuming that (S3) holds prior to augmentation along P , the point b is marked as visited, and b ′ is marked as unvisited.

Proof. Consider any triple ( b, r, b ′ ) that is admissible in G ′ but not in G . Recall that by the definition of the admissible triples, ( r, b ′ ) is a backward edge in G ′ and d y ( r, b ) > d y ( r, b ′ ) . Since the SEARCHANDAUGMENT procedure does not change the weights y ( · ) , the only case where ( b, r, b ′ ) is not admissible in G is when ( r, b ′ ) is not a backward edge in G , i.e., the pair ( b ′ , r ) is in P as a forward edge, and augmenting ˆ τ δ along P results in transporting mass from b ′ to r . On the other hand, by step 2(a) of the SEARCHANDAUGMENT procedure, a forward edge ( b ′ , r ) will be added to the search path only if b ′ is the weighted nearest unvisited neighbor of r ; in other words, since d y ( r, b ) > d y ( r, b ′ ) and the procedure added b ′ to the search path (instead of b ), the point b was marked as visited by the procedure. Therefore, for any newly formed admissible triple ( b, r, b ′ ) , point b (resp. b ′ ) is marked as visited (resp. unvisited).

B.2 Missing Details and Proofs of the INCREASEWEIGHTS Procedure

After the execution of the SEARCHANDAUGMENT procedure, no admissible augmenting paths remain in the residual graph, i.e., there are no admissible paths from the source vertex s to the sink vertex t . The INCREASEWEIGHTS procedure increases the weights of the subset of points in B that are reachable from s by admissible paths to expand their Voronoi cells and to create new admissible triples in the residual graph. For instance, in Figure 6 (left), the path from s to t is not admissible (the triple ( b 2 , r 2 , b 3 ) is not admissible, as b 3 has a lower weighted distance to r 2 than b 2 ). As shown in Figure 6 (right), the INCREASEWEIGHTS procedure then increases the weight of the points b 1 and b 2 (which are reachable from s ), leading to the formation of an admissible augmenting path (note that upon updating the weights, the regions corresponding to r 1 and r 2 have slightly changed). The details of the INCREASEWEIGHTS procedure are described below.

For each point r ∈ A δ , let N ( r ) ⊆ B denote the set of points b ∈ B with ˆ τ δ ( r, b ) > 0 , sorted in decreasing order of their weighted distance to r , i.e., N ( r ) = { b 1 , . . . , b k } where d y ( r, b 1 ) ≥ d y ( r, b 2 ) ≥ . . . d y ( r, b k ) . Mark all points b ∈ B and all forward edges ( b, r ) as unvisited and set K = ∅ , let U = B denote the set of unvisited points of B , and define Q := ⟨ s ⟩ as the search path that the algorithm grows. Execute the following steps until Q becomes empty.

b4

b2f

_b4

b

If Q = ⟨ s ⟩ , then if there exists an unvisited point b ∈ U such that ( s, b ) ∈ G δ , then add b to Q as b 1 . Otherwise, remove s from Q .

If Q = ⟨ s, b 1 , r 1 , . . . , b i ⟩ for some i ≥ 1 ,

(a) If there exists an unvisited forward edge ( b i , r ) in G δ , add r to Q as r i +1 .

(b) Otherwise, mark b i as visited, remove b i from U , add b i to K , and remove b i from Q .

If Q = ⟨ s, b 1 , r 1 , . . . , b i , r i ⟩ for some i ≥ 1 , let b := arg min b ′ ∈ U ∩N ( r ) d y ( r i , b ′ ) denote the unvisited point of b with the minimum weighted distance to r among all points of B that transport mass to r .

(a) If ( b i , r i , b ) is admissible, i.e., d y ( r i , b ) > d y ( r i , b i ) , then add b as b i +1 to Q .

(b) Otherwise, remove r i from Q and mark ( b i , r i ) as visited.

After the DFS procedure terminates, for each point b ∈ K , set y ( b ) ← y ( b ) + δ . This completes the description of the DFS step. We next describe how to recompute the residual graph and the compressed transport plan with respect to the updated weights.

Let y ( · ) (resp. y ′ ( · ) ) denote the weights of the points in B after (resp. before) the weight updates, let X δ (resp. X ′ δ ) denote the partitioning of the set A with respect to weights y ( · ) (resp. y ′ ( · ) ), and let A δ (resp. A ′ δ ) denote the set of representative points of the regions in X δ (resp. X ′ δ ). Furthermore, let ˆ τ ′ δ denote the transport plan maintained by the algorithm for partitioning X ′ δ . To compute the new transport plan ˆ τ δ for the point set A δ , the INCREASEWEIGHTS procedure first computes the arrangement Y of all 3 n cells used to construct X ′ δ with all 3 n cells used to construct X δ , i.e., Y is the arrangement of Voronoi cell, δ -expanded Voronoi cell, and 2 δ -expanded Voronoi cell of each point b ∈ B both before and after weight updates. See Figure 7. For each region φ ∈ X ′ δ ∪ X δ , let C ( φ ) ⊆ Y denote the set of regions of Y that lie inside φ . For each region ϱ ∈ Y , pick an arbitrary representative point r ϱ inside ϱ . We denote the set of all representative points of the regions in Y by Y .

The INCREASEWEIGHTS procedure first converts ˆ τ ′ δ to a transport plan ˆ τ over the finer partitioning Y × B by simply splitting each region φ ′ in X ′ δ (and its mass transportation) to fine regions of Y inside φ ′ . The procedure then uses ˆ τ to construct a transport plan ˆ τ δ over A δ × B by merging the regions of Y (and their mass transportation) to regions of X δ . The details are provided next.

For each point b ∈ B , each region φ ′ ∈ X ′ δ with ˆ τ ′ δ ( r φ ′ , b ) > 0 , and each ϱ ′ ∈ C ( φ ′ ) , set ˆ τ ( ϱ ′ , b ) = µ ( ϱ ′ ) µ ( φ ′ ) ˆ τ ′ δ ( r φ ′ , b ) . This completes the description of the split step and the construction of ˆ τ . Next, the procedure constructs ˆ τ δ by setting, for each point b ∈ B and each region φ ∈ X δ , ˆ τ δ ( r φ , b ) := ∑ ϱ ∈C ( φ ) ˆ τ ( ϱ, b ) . The transport plan ˆ τ δ is defined over A δ × B .

V

X

Finally, for each region φ ∈ X δ , our algorithm stores a list N ( r φ ) of all points b ∈ B with r φ ∈ V 2 δ b , sorted in increasing order of their weights distance to r φ , i.e., N ( r φ ) = ⟨ b 1 , . . . , b k ⟩ , r φ ∈ V 2 b i δ for each i ∈ [1 , k ] , and d y ( r φ , b i ) ≤ d y ( r φ , b j ) for each 1 ≤ i < j ≤ k . This completes the description of the INCREASEWEIGHTS procedure.

Lemma 3.2. Suppose invariant (I1) holds at the start of the INCREASEWEIGHTS procedure. Then, during the execution of the INCREASEWEIGHTS procedure,

(W1) the transport plan ˆ τ δ , y ( · ) remain δ -feasible,

(W2) the weight of each free point b ∈ B increases by δ , and

(W3) the weight of each point b ∈ B with free regions inside V δ b remains unchanged.

Proof. Let y ( · ) (resp. y ′ ( · ) ) denote the weights of the points after (resp. before) the execution of the INCREASEWEIGHTS procedure, and let X δ (resp. X ′ δ ) be the partitioning with respect to weights y ( · ) (resp. y ′ ( · ) ). For any point a ∈ A , let φ a (resp. φ ′ a ) be the region in X δ (resp. X ′ δ ) that contains a , and let b a (resp. b ′ a ) denote the weighted nearest neighbor of a in B with respect to weights y ( · ) (resp. y ′ ( · ) ). For any pair of points ( a, b ) ∈ A × B with τ ′ δ ( a, b ) > 0 , by the δ -feasibility of τ ′ δ , y ′ ( · ) , we have d y ′ ( a, b ) -2 δ ≤ d y ′ ( a, b ′ a ) . To prove property (W1), we show that d y ( a, b ) -2 δ ≤ d y ( a, b a ) .

Recall that the INCREASEWEIGHTS procedure finds the subset K ⊂ B of points that have admissible paths from the source vertex s of the residual graph and increases the weights of all points in K by δ . Consider the following cases:

If b ∈ K is among the points whose weights are increased by the procedure, then

where the second inequality holds from δ -feasibility of τ ′ δ , y ′ ( · ) , the third inequality holds since b a is the weighted nearest neighbor of a with respect to y ′ ( · ) , and the last inequality holds since y ( b a ) ≤ y ′ ( b a ) + δ . Consequently, d y ( a, b ) -2 δ ≤ d y ( a, b a ) .

Otherwise, b / ∈ K and d y ( a, b ) = d y ′ ( a, b ) .

-If b a ∈ K , then d y ′ ( r φ a , b ) ≤ d y ′ ( r φ a , b a ) (since otherwise, the triple ( b a , r φ a , b ) would have been an admissible triple and b a ∈ K would have resulted in b ∈ K ). In this case, since the weighted nearest neighbor of a and r φ a are the same,

where the last inequality holds since increasing the weight of b ′ a by δ made it the weighted nearest neighbor of a . Therefore, the region φ a and consequently, the point a lie inside V δ b and V δ b a . Thus,

-Otherwise, if b a is also not in K , then d y ( a, b a ) = d y ′ ( a, b a ) and,

b3

b3

bi

b3

b

5

As a result, the transport plan τ δ along with the updated weights y ( · ) is δ -feasible and property (W1) holds.

Note that each free point b f ∈ B has an edge from the source vertex s ; therefore, b f ∈ K and the INCREASEWEIGHTS increases the weight of b f by δ , proving (W2). Furthermore, for each point b ∈ B with a free region φ ∈ V δ b , since we defined d y ( r φ , t ) = min b ′ ∈ B d y ( r φ , b ′ ) , the triple ( b, r φ , t ) is admissible. Therefore, the SEARCHANDAUGMENT procedure should have added r φ and b to the search path, not found an admissible augmenting path, and marked b (resp. ( r, t ) ) as visited. Therefore, there are no admissible paths from the source to b , or equivalently b / ∈ K , and the weight of b remains unchanged, leading to (W3).

B.3 Missing Details of the ACYCLIFY Procedure

The goal of the ACYCLIFY procedure is to ensure that the invariant (I2) holds, i.e., that at the beginning of each iteration of our algorithm, the transport plan is a forest and there are no admissible cycles in the residual graph. The procedure runs in three steps: (1) make ˆ τ δ a forest, as described in Section B.3.1, (2) cancel any admissible cycles from the residual graph, as described in Section B.3.2, and (3) acyclify the transport plan again, as described in Section B.3.1. Note that our procedure acyclifies the transport plan (to make it a forest) twice, in steps (1) and (3). Making the transport plan a forest in the first step is essential for the efficiency of the second step, and making it a forest in the third step is essential for invariant (I2), as canceling admissible cycles might introduce cycles in the transport plan.

B.3.1 Acyclifying the Transport Plan

Similar to the Acyclify procedure introduced in [54, Section 3.3], we use a dynamic tree structure to make ˆ τ δ a forest as follows. Let E = ⟨ e 1 , e 2 , . . . , e u ⟩ denote the set of all edges e = ( r, b ) ∈ A δ × B with ˆ τ δ ( r, b ) > 0 . For any k ≤ u , let E k := ⟨ e 1 , e 2 , . . . , e k ⟩ . Define F 0 := ∅ as an empty forest and ˆ τ ′ 0 ( r, b ) = 0 for all pairs ( r, b ) ∈ A δ × B . Starting from k = 1 , for any k ≤ u , the algorithm computes a forest F k and a transport plan ˆ τ ′ k defined over F k using F k -1 and ˆ τ ′ k -1 as follows. If adding the edge e k to F k -1 does not create a cycle, then the algorithm simply sets F k ← F k -1 ∪{ e k } , ˆ τ ′ k ( e k ) ← ˆ τ δ ( e k ) and ˆ τ ′ k ( e ) ← τ ′ k -1 ( e ) for all edges e ∈ F k -1 . Otherwise, adding e k to F k -1 results in the creation of an even-length cycle C . Let c denote the minimum capacity of the edges in C , and let e ∗ denote the edge with the minimum capacity. Consider an ordering of the edges of the cycle C that starts with e ∗ , i.e., C = ⟨ e ∗ = e ′ 1 , e ′ 2 , . . . , e ′ 2 j ⟩ . The algorithm increases (resp. reduces) the mass transported along the edge e ′ 2 i (resp. e ′ 2 i -1 ) by c for each i ∈ [1 , j ] . Finally, the algorithm sets F k ← F k -1 ∪{ e k }\{ e ∗ } . This completes the description of step 1. Using the dynamic tree structure by Sleator and Tarjan [56], each operation takes O (log n ) amortized time, and since |E| = O ( n 3 ) , this process takes a total of O ( n 3 log n ) time.

S

by

S

S

b

b3

b2

b2

b

B.3.2 Acyclifying the Admissible Triples

To remove all admissible cycles from the residual graph, we use a partial DFS similar to the one described in the SEARCHANDAUGMENT procedure. Let ← -G δ be the graph formed by reversing the direction of all the edges of G δ . The procedure first marks all points of B and all backward edges as unvisited and defines U := B as the set of unvisited points. While there exists an unvisited point b ∈ B , the procedure initializes a partial DFS by setting Q = ⟨ b = b 1 ⟩ and searches as follows until Q becomes empty.

If Q = ⟨ b 1 , r 1 , . . . , b i ⟩ for some i ≥ 1 ,

(a) If there are no unvisited backward edges ( b i , r ) in ← -G δ , then mark b i as visited and remove b i from Q and U .

(b) Otherwise, there exists an unvisited backward edge ( b i , r φ ) . Add r φ to Q as r i +1 .

If Q = ⟨ b 1 , r 1 , . . . , b i , r i ⟩ for some i ≥ 1 , let b := arg min b ′ ∈ U d y ( r i , b ′ ) be the unvisited point with the minimum weighted distance to r i .

(a) If ( b, r i , b i ) is admissible, i.e., d y ( r i , b ) < d y ( r i , b i ) ,

-If b already exists in the path Q as b j , then C = ⟨ b j , r j , . . . , b i , r i , b i +1 = b = b j ⟩ is an admissible cycle. Cancel the cycle C (as described below) and set Q = ⟨ b 1 , r 1 , . . . , b j ⟩ .

-Otherwise, add b as b i +1 to Q .

(b) Otherwise, remove r i from Q and mark the backward edge ( r i , b i ) as visited.

Given an admissible cycle C = ⟨ b j , r j , . . . , b i , r i , b i +1 = b j ⟩ , the procedure cancels the cycle as follows. Let bc ( C ) := min t ∈ [ j,i ] ˆ τ δ ( r t , b t ) denote the bottleneck capacity of the cycle C . For any forward edge ( b t +1 , r t ) (resp. backward edge ( r t , b t ) ) in C , set ˆ τ δ ( r t , b t +1 ) ← ˆ τ δ ( r t , b t +1 ) + bc ( C ) (resp. ˆ τ δ ( r t , b t ) ← ˆ τ δ ( r t , b t ) -bc ( C ) ). In this way, at least one of the backward edges of the cycle C is removed from the residual graph and the cycle has vanished.

Lemma 3.3. Suppose invariant (I1) holds at the start of the ACYCLIFY procedure. Then, during the execution of the ACYCLIFY procedure,

(A1) the transport plan ˆ τ δ , y ( · ) remains δ -feasible, and

(A2) the transport plan ˆ τ δ is a forest and there are no admissible cycles in the residual graph.

Proof. In the first step of the ACYCLIFY procedure, the algorithm makes the transportation network to be a DAG. Note that the resulting transportation network is a subset of the transportation network before the ACYCLIFY procedure, and therefore, for any transporting edge ( r, b ) ∈ A δ × B , the point b is a 2 δ -weighted nearest neighbor of r , i.e., the transport plan obtained after the first step of the ACYCLIFY procedure is δ -feasible.

Next, we show property (A1) and (A2) for the second step of the procedure. By construction, all triples on the search path maintained by the procedure are admissible, and therefore, any cycle

computed by the second step of ACYCLIFY procedure is admissible. Since all forward edges on the computed cycles are from points b ∈ B to regions in V 2 δ b , canceling a cycle C does not violate the δ -feasibility condition (F1); hence, the transport plan ˆ τ δ , y ( · ) is a δ -feasible transport plan during the execution of the second step of the ACYCLIFY procedure and (A1) holds.

To prove property (A2), we show that

(A3) any point b ∈ B (resp. backward edge ( r, b ) ) marked as visited will not form an admissible cycle during the execution of the procedure, and

(A4) during the execution of the second step of the ACYCLIFY procedure, the subgraph of the residual graph induced by visited vertices and their neighboring regions does not have any cycles of admissible triples. Furthermore, there are no admissible paths from an unvisited point to a visited point.

Assuming property (A3) holds, any point b ∈ B that is marked as visited does not participate in an admissible cycle. Since the ACYCLIFY procedure stops when all points in B are visited, there are no admissible cycles in the residual graph. Furthermore, since the transport plan is maintained by a dynamic tree structure the transport plan is a forest. Therefore, to prove property (A2), we show that (A3) holds.

Furthermore, the property (A3) is a direct corollary of property (A4), as explained next: For each visited point b ∈ B , all points b ′ ∈ B having admissible paths to b are also visited (note that the ACYCLIFY procedure searches on the residual graph in the reverse direction of the edges), and if (A4) holds, there are no admissible cycles solely formed by visited vertices; hence, b does not form admissible cycles. Furthermore, the procedure marks a backward edge ( r, b ) as visited if, for each admissible triple ( b ′ , r, b ) , the point b ′ is visited. Since b ′ is not a part of any admissible cycles (assuming (A4) holds), the triple ( b ′ , r, b ) also cannot be a part of an admissible cycle. Therefore, to prove property (A3), we prove that (A4) holds in the following.

We use an inductive argument to prove (A4). Let C 1 , . . . , C k denote the sequence of admissible cycles found by the procedure, and let ˆ τ 0 δ , . . . , ˆ τ k δ denote the sequence of transport plans, where ˆ τ 0 δ is the transport plan maintained by the algorithm at the beginning of the second step of the procedure and ˆ τ i δ is obtained by canceling ˆ τ i -1 δ along C i . The property (A4) trivially holds at the beginning of the execution of the second step of the ACYCLIFY procedure.

Note that the ACYCLIFY procedure marks a point b ∈ B as visited only if the search from b did not lead to finding an admissible cycle, i.e., for all pairs ( b ′ , r ) ∈ B × A δ such that ( b ′ , r, b ) is admissible, the point b ′ is marked as visited and the backward edge ( r, b ) is marked as visited. Hence, all points having an admissible path to the visited points (i.e., all points that are reachable from the visited points in our backward DFS) are also visited. Therefore, Assuming that property (A4) holds before marking b as visited, property (A4) holds after marking b as visited as well. Next, we show the same when we cancel a cycle C i .

First, note that all points of B that are on the search path are unvisited. Therefore, for any admissible cycle found by the procedure, the points of B on the cycle are unvisited. Using an identical proof as in Lemma B.1, one can show that after canceling a cycle, for any newly formed admissible triple ( b, r, b ′ ) , the point b is visited and the point b ′ is unvisited. In other words, canceling a cycle C i only creates additional admissible paths from visited points to unvisited points. Assuming that (A4) holds before canceling C i , since no new admissible triples are created from a visited point to another visited point, the subgraph induced by visited points and their neighboring regions remain free of admissible cycles after canceling C i . Furthermore, since no new admissible triples are created from an unvisited point to a visited point upon canceling C i , there will be no admissible paths from an unvisited point to a visited point after cancellation; hence, (A4) remains satisfied.

C Missing Details of Section 4

In this section, we analyze the efficiency of the three procedures SEARCHANDAUGMENT, INCREASEWEIGHTS, and ACYCLIFY.

C.1 Efficiency of the SEARCHANDAUGMENT Procedure

The SEARCHANDAUGMENT procedure runs a partial DFS on the residual graph to find a set of admissible augmenting paths. The partial DFS procedure, upon backtracking from a point b ∈ B (resp. r ∈ A δ ), marks the point b (resp. the backward edge ( b ′ , r ) used to reach r ) as visited and does not visit the point b (resp. the backward edge ( b ′ , r ) ) again in the same execution. Upon finding an augmenting path P , the procedure augments the transport plan along P in O ( | P | ) time. Let ⟨ P 1 , . . . , P k ⟩ denote the set of all augmenting paths found by the SEARCHANDAUGMENT procedure. In 2 dimensions (resp. d dimensions), since the residual graph has O ( n 3 ) edges (resp. O ( n d +1 ) ), the running time of the procedure would be O ( n 3 + ∑ k i =1 | P i | ) (resp. O ( n d +1 + ∑ k i =1 | P i | ) ). In Lemma C.1, we show that ∑ k i =1 | P i | = O ( n 3 ) (resp. ∑ k i =1 | P i | = O ( n d +1 ) ). Hence, each execution of the SEARCHANDAUGMENT procedure takes O ( n 3 ) (resp. O ( n d +1 ) ) time.

Lemma C.1. The total length of augmenting paths computed during the execution of the SEARCHANDAUGMENT procedure is O ( n 3 ) in 2 dimensions and O ( n d +1 ) in d dimensions.

Proof. Let ˆ τ 0 δ denote the transport plan maintained by the algorithm at the beginning of execution of the SEARCHANDAUGMENT procedure. To prove this lemma, we categorize the augmenting paths found by the procedure based on the source of their bottleneck capacity, namely (1) set P v consisting of augmenting paths whose bottleneck capacity is determined based on the residual capacity of its endpoints, and (2) set P e consisting of augmenting paths whose bottleneck capacity is determined based on mass transportation over its backward edges. We first show that |P v | = O ( n 2 ) and then show the same bound for P e . Since each augmenting path has a length of at most 2 n , we then conclude that the total length of all augmenting paths is O ( n 3 ) .

Let P be an augmenting path in P v . If the bottleneck capacity of P is determined by a free point b ∈ B (resp. free region r ∈ A δ ), then the mass of b (resp. r ) will be fully transported after augmentation; therefore, since | A δ ∪ B | = O ( n 2 ) , we have |P v | = O ( n 2 ) . Next, let P be an augmenting path in P e ; in this case, the backward edge ( r, b ) determining the bottleneck capacity of P will be removed from the transport plan after augmentation. Note that by Lemma B.1, for any newly formed admissible triples ( b ′ , r, b ) , the point b ′ is already marked as visited. By property (S2) in Lemma 3.1, the point b ′ cannot form an admissible augmenting path during the same execution of the SEARCHANDAUGMENT procedure. Hence, the edge ( r, b ) determining the bottleneck capacity of P was a backward edge of the initial transport plan ˆ τ 0 δ and augmentation along each path P ∈ P e removes one of the transporting edges of the transport plan ˆ τ 0 δ . Using invariant (I2), |P e | = O ( n 2 ) , as claimed. Hence, the total number of augmenting paths found by the procedure is O ( n 2 ) , and since each augmenting path has a length of at most 2 n , their total length is O ( n 3 ) .

Next, we extend our analysis to d dimensional space, for any d ≥ 2 . Note that in d dimensions, the residual graph has O ( n d ) vertices and O ( n d +1 ) edges. Hence, |P v | = O ( n d ) . Since the transport plan ˆ τ 0 δ is a forest over the point set A δ ∪ B , the total number of edges transporting a positive mass in ˆ τ 0 δ would be O ( n d ) ; since augmenting the transport plan along each augmenting path in P e eliminates one of the edges transporting positive mass in ˆ τ 0 δ , |P e | = O ( n d ) . Finally, since each augmenting path has a length of at most 2 n , the total length of all augmenting paths found by the SEARCHANDAUGMENT procedure would be O ( n d +1 ) .

C.2 Efficiency of the INCREASEWEIGHTS Procedure

In this section, we show that for 2 -dimensional (resp. d -dimensional) distributions, the INCREASEWEIGHTS procedure runs in O ( n 2 (Φ + n log n )) (resp. O ( n d (Φ + n log n )) ) time. The INCREASEWEIGHTS procedure runs a DFS that visits each edge of the residual graph at most once and has a total running time of O ( n 3 ) (resp. O ( n d +1 ) ). Furthermore, in the arrangement used to construct partitioning Y , each point b ∈ B has at most 6 Voronoi cells (three cells that are used in the construction of X δ and three that are used in the construction of X ′ δ ). Using a similar discussion as Section A, one can show that the total number of vertices in the arrangement used to construct Y is O ( n 2 ) (resp. O ( n d ) ), and the number of regions in Y is at most O ( n 2 ) (resp. O ( n d ) ). The construction of the transport plan ˆ τ can be done in O ( n 2 (Φ + n )) (resp. O ( n d (Φ + n )) ) time since (1) the mass of all regions in Y can be determined in O ( n 2 Φ) (resp. O ( n d Φ) ) time (partitioning the regions in Y into simplices remains an arrangement with O ( n 2 ) (resp. O ( n d ) ) vertices and therefore,

have O ( n 2 ) (resp. O ( n d ) ) regions), and (2) the mass transported on each pair ( ϱ, b ) ∈ Y × B can be determined in O (1) time. Converting ˆ τ to ˆ τ δ , as is done in the merge step, also takes O ( n 3 ) (resp. O ( n d +1 ) ) time, given that the total complexity of ˆ τ is O ( n 3 ) (resp. O ( n d +1 ) ). Finally, storing a sorted list of neighbors for each region r ∈ A δ takes O ( n 3 log n ) (resp. O ( n d +1 ) log n ) time in total. Hence, the execution of the INCREASEWEIGHTS procedure takes O ( n 2 (Φ + n log n )) (resp. O ( n d (Φ + n log n )) ) time.

Lemma C.2. Each execution of the INCREASEWEIGHTS procedure takes O ( n 2 (Φ + n log n )) time in 2 dimensions and O ( n d (Φ + n log n )) time in d dimensions.

C.3 Efficiency of the ACYCLIFY Procedure

The first step of this procedure uses a dynamic tree structure to acyclify the transport plan ˆ τ δ . Using the dynamic tree structure by Sleator and Tarjan [56], since the total number of edges of the graph is O ( n 3 ) in 2 dimensions (resp. O ( n d +1 ) in d dimensions), the running time of this step would be O ( n 3 log n ) (resp. O ( n d +1 log n ) ). In the second step, the procedure runs a partial DFS procedure on the residual graph and cancels the admissible cycles. The partial DFS procedure, upon backtracking from a point b ∈ B (resp. r ∈ A δ ), marks the point b (resp. the backward edge ( b ′ , r ) used to reach r ) as visited and does not visit the point b (resp. the edge ( b ′ , r ) ) again in the same execution. Furthermore, upon finding an admissible cycle C , it cancels the cycle in O ( | C | ) time. Let ⟨ C 1 , . . . , C k ⟩ denote the set of all cycles found in the execution of the ACYCLIFY procedure. In Lemma C.4, we show that ∑ k i =1 | C i | = O ( n 3 ) in 2 dimensions (resp. ∑ k i =1 | C i | = O ( n d +1 ) in d dimensions). Given that the size of the residual graph is at most O ( n 3 ) (resp. O ( n d +1 ) ), the second step of the ACYCLIFY procedure takes a total of O ( n 3 + ∑ k i =1 | C i | ) = O ( n 3 ) (resp. O ( n d +1 ) ) time, leading to the following lemma.

Lemma C.3. Each execution of the ACYCLIFY procedure takes O ( n 3 log n ) time in 2 dimensions and O ( n d +1 log n ) time in d dimensions.

Lemma C.4. The total length of admissible cycles computed during the execution of the second step of the ACYCLIFY procedure is O ( n 3 ) in 2 dimensions and O ( n d +1 ) time in d dimensions.

Proof. Let ˆ τ 0 δ denote the transport plan maintained by the algorithm at the beginning of execution of the ACYCLIFY procedure. To prove this lemma, we show that the ACYCLIFY procedure finds O ( n 2 ) admissible cycles, where each cycle has a length of at most 2 n ; hence, the total length of all cycles found by the procedure would be O ( n 3 ) .

Let C be an admissible cycle found by the procedure; in this case, the backward edge determining the bottleneck capacity of C will be removed from the transport plan after cancellation. For any admissible triple ( b, r, b ′ ) formed after canceling C , using an identical argument as in Lemma B.1, one can show that the edge ( r, b ′ ) is a backward edge that was on the cycle C as a forward edge, and the point b is marked as visited; hence, by Lemma 3.3, the point b does not form an admissible cycle in the same execution of the ACYCLIFY procedure and therefore, the newly formed backward edge ( r, b ′ ) cannot be included in any admissible cycles. Therefore, each cycle cancellation removes one of the backward edges of ˆ τ 0 δ , which is the transport plan obtained after the first step of the ACYCLIFY procedure, i.e., the transportation network of ˆ τ 0 δ is a forest and the number of its transporting edges is O ( n 2 ) . Therefore, the total number of cycles found by the ACYCLIFY procedure is O ( n 2 ) , and their total length is O ( n 3 ) , as claimed.

We next show that the total length of admissible cycles in d dimensions is O ( n d +1 ) in d dimensions, for any d ≥ 2 . Note that in d dimensions, the residual graph has O ( n d +1 ) edges. Since the transport plan ˆ τ 0 δ is a forest over the point set A δ ∪ B , the total number of edges transporting a positive mass in ˆ τ 0 δ would be O ( n d ) ; since canceling each admissible cycle eliminates one of the edges transporting positive mass in ˆ τ 0 δ , |C| = O ( n d ) . Since each admissible cycle has a length of at most 2 n , the total length of all admissible cycles found by the INCREASEWEIGHTS procedure would be O ( n d +1 ) .

C.4 Number of Iterations

Lemma 4.1. For each scale δ , the total number of iterations of step 2 of our algorithm is O ( n ) .

Proof. Let τ 2 δ , y 2 δ ( · ) denote the 2 δ -feasible transport plan computed by our algorithm for scale 2 δ , and let τ δ , y δ ( · ) denote a partial transport plan maintained during the execution of step 2 of our algorithm. Let X 2 δ (resp. X δ ) denote the partitioning of the set A with respect to weights y 2 δ ( · ) (resp. y δ ( · ) ). Let Y be the arrangement of all 3 n cells used to construct X 2 δ with all 3 n cells used to construct X δ , i.e., Y is the arrangement of Voronoi cell, δ -expanded Voronoi cell, and 2 δ -expanded Voronoi cell of each point b ∈ B with respect to weights y δ ( · ) along with the Voronoi cell, 2 δ -expanded Voronoi cell, and 4 δ -expanded Voronoi cell of each point b ∈ B with respect to weights y 2 δ ( · ) . For each region ϱ ∈ Y , pick an arbitrary representative point r ϱ inside ϱ . We denote the set of all representative points of the regions in Y by Y . Let ˆ τ 2 δ (resp. ˆ τ δ ) denote the compressed transport plan for τ 2 δ (resp. τ δ ) using the partitioning Y . Note that the partitioning Y is a refinement of both partitionings X δ and X 2 δ . Define τ ′ := ˆ τ 2 δ -ˆ τ δ .

We construct a bipartite graph G ′ over Y × B , where for any pair ( r, b ) ∈ Y × B , there exists an edge directed from r to b if τ ′ ( r, b ) < 0 and directed from b to r if τ ′ ( r, b ) > 0 . Consider any directed path P = ⟨ r 1 , b 1 , . . . , r k , b k ⟩ from a free region r ∈ Y to a free point b ∈ B (with respect to ˆ τ δ ). The path P is an augmenting path in the residual graph corresponding to ˆ τ δ , y δ ( · ) .

Similar to the standard graph algorithms, we define the net-cost of the path P as ϕ ( P ) := ∑ k i =1 d ( r k , b k ) -∑ k -1 i =1 d ( r k , b k +1 ) . Let b 0 := b r 1 be the weighted nearest neighbor of r 1 . Then, we can rewrite the net-cost of P as

Due to δ -feasibility of the transport plan ˆ τ δ , y δ ( · ) , for all edges ( r i , b i -1 ) , i ∈ [1 , k ] , the point b i -1 is a 2 δ -weighted nearest neighbor of r i ; hence, d y δ ( r i , b i ) -d y δ ( r i , b i -1 ) ≥ -2 δ . Since the length of P is at most 2 n -1 ,

Similarly, we can rewrite the net-cost of P using weights y 2 δ as follows.

According to the 2 δ -feasibility of ˆ τ 2 δ , y 2 δ ( · ) , for all edges ( r i , b i ) , i ∈ [1 , k ] , the point b i is a 4 δ -weighted nearest neighbor of r i (with respect to weights y 2 δ ( · ) ); hence, d y 2 δ ( r i , b i ) -d y 2 δ ( r i , b i -1 ) ≤ 4 δ . Since the length of P is at most 2 n -1 ,

By property (W3) in Lemma 3.3, during the execution of step 2 of our algorithm, the weights of the points in B with free regions in their δ -expansion remain unchanged; therefore, since b 0 contains free regions inside its δ -expanded Voronoi cell, y 2 δ ( b 0 ) = y δ ( b 0 ) and d y 2 δ ( r 1 , b 0 ) = d y δ ( r 1 , b 0 ) . Combining with Equations (13) and (14),

Equivalently,

By property (W2) in Lemma 3.3, our algorithm increases the weight of the point b k by δ in each iteration while its mass is not fully transported by ˆ τ δ ; therefore, the point b k cannot remain free after 6 n iterations, i.e., after O ( n ) iterations, there cannot be any remaining free points in B and the step 2 of our algorithm terminates after O ( n ) iterations.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: This paper presents a new combinatorial framework for the semi-discrete OT problem. It is used to design an algorithm for computing a ε -close semi-discrete transport plan, as well as a data structure for the discrete OT problem. Our main results are stated in Theorems 1.1 and 1.2 in the introduction and are summarized in the abstract.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: Our semi-discrete OT algorithm relies on the existence of an oracle that computes the mass of µ inside a triangle, which is common in existing algorithms for the semi-discrete optimal transport problem. We clearly state this assumption in the abstract and Theorem 1.1.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: All theorems and lemmas contain the full set of assumptions, and their proofs are either provided in the main text or the appendix.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [NA]

Justification: This paper does not include experiments.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [NA]

Justification: This paper does not include experiments.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [NA]

Justification: This paper does not include experiments.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [NA]

Justification: This paper does not include experiments.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [NA]

Justification: This paper does not include experiments.

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: There are no data or human subjects involved in the research process of this paper and there are no known societal impacts and potentially harmful consequences associated with the research conducted in this paper.

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: None of the categories described by the NeurIPS Code of Ethics, namely, safety, security, discrimination, surveillance, deception and harassment, environment, human rights, and bias and fairness is impacted by the research conducted in this paper.

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper poses no such risks.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [NA]

Justification: The paper does not use existing assets.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: The paper does not release new assets.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: This paper does not involve crowdsourcing nor research with human subjects.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: This paper does not involve crowdsourcing nor research with human subjects.

Peer reviews:
comments: SUMMARY:
The paper presents a primal-dual algorithm for approximately solving the semi-discrete optimal transport problem. The algorithm runs $\log \Delta/\epsilon $ scales, starting with the scale $\delta = \Delta^2$ and halving it in each round. The idea is to maintain a $\delta$-feasible weight function during the course of the algorithm, which in each round constructs a residual graph from the three levels of expansion of the Voronoi diagram, and augments the transport plan along augmenting paths. To ensure the invariants and that the algorithm terminates, the algorithm considers only admissible paths and performs a procedure to eliminate cycles after each round. The proposed algorithm improves the running time from $n^9$ in Agarwal et al. to $n^4$ and the size of the graph from $n^5$ to $n^3$. The algorithm also extends to any dimension $d > 2$ and any $p\ge 1$-Wasserstein distance.

STRENGTHS:
- The paper presents a set of strong results that significantly improve known results on semi-discrete optimal transport problem. 
- Although I didn't verify all the proofs, I think the algorithm and procedures and the idea of using admissibility for DFS to avoid and remove cycles are all reasonable.

WEAKNESSES:
- There isn't any big weakness I can see.

QUESTIONS:
- Could the authors provide some more background on the construction of Voronoi diagram, ie, the complexity of the construction?
- Is there any lower bound on the runtime of algorithms for the semi-discrete OT problem?
- How is Theorem 1.2 compared with other algorithms for the discrete OT problem?

LIMITATIONS:
The authors adequately addressed the limitations of the paper.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-06-25

comments: SUMMARY:
The paper studies the semi-discrete optimal transport problem, i.e. a generalization of bipartite matching where one side of the graph is not finite, but instead a probability distribution on some subset of the R^d. 
It gives an algorithm that computes an optimal transport plan (up to some additive $\varepsilon) to such instances subject to an oracle that can integrate the density function of the continuous distribution in constant time over some given triangle. 
The algorithm proposed is an adaption of the solution proposed by Agarwal et al. [SODA 24] which makes some improvements to depress the degree of the polynomial dependence of the runtime on the size of the discrete distribution in the input.
The authors further note that their methods yield a kind of sub-linear time online algorithm for discrete optimal transport where one (large) side of the bipartite graph is fixed and the smaller side is subject to updates. This is achieved by representing the larger side as a continuous distribution which can be efficiently sampled.

Note that I was not able to check the appendix where essentially all proofs live.

STRENGTHS:
The paper is clear and easy to follow, and improves the state of the art on a very relevant topic (optimum transport).

WEAKNESSES:
- The work is largely a reworking of other recent results in the field, in particular the paper of Agarwal et al. at SODA24. It's not clear that improving the polynomial dependence here is highly relevant.

QUESTIONS:
No questions

LIMITATIONS:
Limitiations and impact have been adressed fully, this is a theoretical result with no negative consequences to be expected.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-09

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a novel combinatorial algorithm for the problem known as semi-discrete optimal transport. The proposed method constructs a residual graph by considering the cells of a $\\delta$-expanded Voronoi diagram, a relaxed concept of a weighted Voronoi diagram, as vertices. The algorithm performs augmentation on the residual graph while scaling $\\delta$. This approach significantly reduces the theoretical computational complexity compared to existing methods. Furthermore, the proposed method can be applied to discrete optimal transport with large supports, enabling sublinear time responses to OT queries with respect to support size through preprocessing.\n\nSTRENGTHS:\n* The paper makes a significant contribution to the important problem of semi-discrete optimal transport in the field of machine learning.\n* The proposed method drastically reduces computational complexity compared to existing methods. Additionally, its application to query responses through preprocessing also demonstrates excellent computational efficiency.\n* The paper is exceptionally well-written, making it easy to follow the ideas despite the challenging content.\n\nWEAKNESSES:\n* Understanding that this is a theoretical paper, it is important to note that the lack of numerical experiments makes it difficult to assess the practical applicability of the proposed method.\n* While this is not a weakness of the paper, I am not able to fully verify the validity or the details of the proofs and methodology, and therefore cannot guarantee their correctness.\n\nQUESTIONS:\n* Is this method an implementable and practically applicable algorithm, or is it challenging to apply in real-world scenarios due to large constant factors, thus possessing only theoretical value?\n\nLIMITATIONS:\nThe authors properly address the limitations of the method within the main text.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}