Efficient Algorithms for Robust and Partial Semi-Discrete Optimal Transport

We appreciate your thoughtful review and comments. We answer your main questions and concerns below.

How long does it take to compute the $\Phi$ factor, both for known distributions (e.g. normal, exponential) and unknown distributions?

The factor of $\Phi$ depends on how the distribution is represented. If the distribution is a parametric distribution with a constant number of parameters, e.g. normal or exponential (mixture), then a constant degree polygonal query region can be integrated over in constant time using numerical methods. On the other hand if the continuous distribution is represented as a histogram or piecewise-linear function, then the time depends on the complexity of the distribution inside the query region. In low dimensions, one could use some geometric data structures to expedite this computation, but in high dimensions, a simple method would be more practical. Finally, if the distribution is represented as a large representative sample, then the $\Phi$ factor simply corresponds to assigning each sample point to the region of the partition it is contained in. In turn, our algorithm successfully computes a transport plan between a discrete distribution and a large sample of an unknown continuous distribution.

Can your methods handle distributions for which only sample data is available?

Our algorithm can handle sample data from a distribution, as long as one global sample is used to approximate the continuous mass inside each region of the arrangement (rather than sampling from the distribution for each mass query).

Are the obtained convergence rates optimal with respect to their dependencies on parameters such as $n, d$?

A: As noted in Section 2, as well as references [9] and Bourne et al., the optimal semi-discrete transport plan requires computing a (restricted) weighted Voronoi diagram. There are known lower bounds of $n^{\Omega(d)}$ for the complexity of Voronoi diagrams in the worst case in $d$ dimensions (see e.g. Har-Peled "Geometric Approximation Algorithms"), and therefore our high-precision algorithm is near-optimal. While we do not have a rigorous lower bound for the optimal runtime of a low-precision algorithm in constant dimension, we suspect that optimal transport is at least as hard as nearest neighbor search (a naive plan would assign each point to its nearest neighbor). The runtime of our low-precision algorithm is near linear in $n$ and incurs comparable dependence in $\varepsilon$ and $d$ to the best known algorithms for approximate nearest neighbor search.

The paper lacks practical validation of the proposed methods, particularly comparisons with deterministic counterparts in scenarios where robustness is essential. Even basic 2D experiments on synthetically corrupted data could have provided a proof of concept while demonstrating the methods' robustness and quality.

While the original submission emphasized our theoretical contributions, we appreciate your interest in practical implementation and experimental evaluation. In response, we include below preliminary experimental results. A more extensive evaluation will be added in the next version.

Implementation Details. We provide a Python implementation of our algorithm. However, it relies on two core procedures that lack efficient implementations in Python:

1. Estimating continuous mass within regions defined by the arrangement of the restricted Voronoi cells and their expansions.
2. Maintaining a dynamic tree data structure as per Sleator and Tarjan (reference [44] in the appendix).

To address these challenges in our prototype:

• For (1), we approximate the mass within each region using a fine grid, summing the masses of grid points lying inside the region.
• For (2), we simulate the dynamic tree by naïvely iterating over all edges.

An efficient implementation will require optimized versions of these components, which remains a technical challenge in Python.

Preliminary Results on Algorithm Performance. We evaluated our implementation using a continuous distribution sampled from a Gaussian distribution inside the unit square. The discrete support consists of $n$ samples drawn from the same distribution, each assigned a mass of $1/n$. We set $\lambda = 0.2$ and $\varepsilon = 0.02$. Results are averaged over 10 runs for each $n$.

Key Findings:

• The additive error stays within the target $\varepsilon = 0.02$.
• The number of iterations in the two steps of our algorithm (Step 1 and 2) remains well below our upper bounds proved in the paper ($6n$ and $12n$ as described in Lemma 4.3 and D.9).
• The total path and cycle lengths computed in the two steps grow subquadratically in $n$, aligning with our $O(n^2)$ complexity analysis (Section 4.3).

The details of our experiments are presented in the following table.

Robustness of $\alpha$-OPT.

We examined the effect of 10% noise in the continuous distribution (added in the top-right corner) and measured how much noise was transported under various $\alpha$ values. Results confirm the robustness of $\alpha$-OPT: for $\alpha$ < 0.9, almost no noise mass is transported.

Robustness of $\lambda$-ROT. We conducted similar experiments for $\lambda$-ROT. As $\lambda$ increases, total mass transported grows, but noise transport remains limited when the transported mass is under 0.9.

Please let us know if any further clarification or experiments would strengthen the paper. We are happy to incorporate additional feedback into the next version of our paper.

We appreciate the reviewer’s comment regarding the density of technical details in the introduction. In view of this comment, we will fix the introduction to present the motivation, high-level ideas, and contributions in a clearer and more accessible manner, deferring technical details to later sections. We will also do our best to make the later sections more accessible.

Could you conduct some experiments justifying your theoretical conclusions?

Summary of Contributions. Our paper presents the following key contributions:

• A characterization of partial and robust variants of the optimal transport (OT) problem in the semi-discrete setting,
• Theoretical guarantees on the approximation error when using an approximate $\lambda$-ROT solver to compute an approximate $\alpha$-OPT plan, and,
• Two algorithmic solutions for the $\lambda$-ROT problem.

While the original submission emphasized our theoretical contributions, we appreciate your interest in practical implementation and experimental evaluation. In response, we include below preliminary experimental results. A more extensive evaluation will be added in the next version.

Implementation Details. We provide a Python implementation of our algorithm. However, it relies on two core procedures that lack efficient implementations in Python:

1. Estimating continuous mass within regions defined by the arrangement of the restricted Voronoi cells and their expansions.
2. Maintaining a dynamic tree data structure as per Sleator and Tarjan (reference [44] in the appendix).

To address these challenges in our prototype:

• For (1), we approximate the mass within each region using a fine grid, summing the masses of grid points lying inside the region.
• For (2), we simulate the dynamic tree by naïvely iterating over all edges.

An efficient implementation will require optimized versions of these components, which remains a technical challenge in Python.

Preliminary Results on Algorithm Performance. We evaluated our implementation using a continuous distribution sampled from a Gaussian distribution inside the unit square. The discrete support consists of $n$ samples drawn from the same distribution, each assigned a mass of $1/n$. We set $\lambda = 0.2$ and $\varepsilon = 0.02$. Results are averaged over 10 runs for each $n$.

Key Findings:

• The additive error stays within the target $\varepsilon = 0.02$.
• The number of iterations in the two steps of our algorithm (Step 1 and 2) remains well below our upper bounds proved in the paper ($6n$ and $12n$ as described in Lemma 4.3 and D.9).
• The total path and cycle lengths computed in the two steps grow subquadratically in $n$, aligning with our $O(n^2)$ complexity analysis (Section 4.3).

The details of our experiments are presented in the following table.

Reduction from $\alpha$-OPT to $\lambda$-ROT. We evaluated our reduction from $\alpha$-OPT to $\lambda$-ROT by using the Sinkhorn algorithm to approximate $\lambda$-ROT plans, followed by binary search over $\lambda$ to compute an $\alpha$-OPT plan. The continuous distribution in our experiments is an exponential distribution over the unit square, centered at the bottom-left corner, with 10% noise added to the top-right corner. The discrete set $B$ consists of $n=400$ samples from the same exponential distribution, each assigned mass $1/n$. We target a transported mass of $\alpha$ = 0.9.

Additive Error: For any fixed regularization parameter, we evaluate the additive error (with respect to the optimal $\lambda$-ROT solution) incurred by Sinkhorn, and compare it to the additive error of our computed $\alpha$-OPT solution (with respect to the optimal $\alpha$-OPT). We observe that the error in $\alpha$-OPT closely mirrors that of $\lambda$-ROT across different choices of regularization parameter. Moreover, as the regularization parameter decreases, the additive errors in $\alpha$-OPT also diminish. This behavior is consistent with the theoretical guarantee established in Theorem 3.5.

Approximation factor: Similarly, we evaluate the approximation factor for each method using the same setup. Once again, we find that the approximation factor in $\alpha$-OPT closely mirrors that of $\lambda$-ROT over a range of regularization parameters. As the regularization parameter decreases, the approximation factor in $\alpha$-OPT also approaches one, further corroborating the theoretical prediction in Theorem 3.5.

Please let us know if any further clarification or experiments would strengthen the paper. We are happy to incorporate additional feedback into the next version of our paper.

Could you please provide a comment regarding the originality of your approach which follows the ideas of the one suggested in Agarwal et al., 2024?

Challenges in Extending Cost-Scaling to Partial Transport: Cost-scaling algorithms have proven effective for solving the complete optimal transport (OT) problem. However, extending these techniques to partial optimal transport introduces substantial challenges—even in fully discrete settings (see, e.g., Ramshaw and Tarjan, FOCS 2012 [37])—and these difficulties are further amplified in the semi-discrete setting. A central obstacle lies in the behavior of the dual weights: in complete OT, the optimal solution is invariant under translations of the dual weight vector for the discrete set $B$. In contrast, in partial OT, the absolute magnitudes of the dual weights are crucial, as they control the truncation of Voronoi cells and thus directly determine which regions of the continuous distribution are assigned to each discrete point.

Limitations of Existing Scaling Algorithms. Recent cost-scaling algorithms (e.g., Agarwal et al., SODA 2024; NeurIPS 2024) increase the magnitudes of the dual weights for the discrete set $B$, and implicitly for the continuous domain $A$, over successive scales. However, these methods lack a mechanism to reduce or correct the dual weights once they are overestimated. This renders existing scaling algorithms unsuitable for partial OT, where accurate dual weight magnitudes are critical.

Our Solution. To overcome this limitation, we introduce a novel step that uses consolidating paths and augmenting paths to reroute mass in the residual graph and correct overestimated dual weight magnitudes. The addition of this step makes it non-trivial to establish both the correctness of the modified algorithm and its efficiency matching that of Agarwal et al., NeurIPS 2024. These challenges necessitate several new ideas, as detailed in Lemmas 4.2 and 4.3 of the main text and Lemma D.9 in the Appendix.

Geometric Interpretation of this limitation and consolidating paths. At the start of a scale $\delta$, some mass from the continuous distribution may lie well inside the restricted Voronoi cells of the discrete points $B$, yet remain untransported—these are what we call violating deficit regions (see Figure 3). In the complete OT setting, such violations are naturally resolved as the plan evolves to transport all mass. In the partial OT setting, however, this issue is non-trivial and must be explicitly resolved to ensure correctness. Our consolidating paths are specifically designed to eliminate all violating deficit regions by restructuring the transport plan to ensure that the total transported mass lies strictly within the restricted Voronoi regions, thereby guaranteeing correctness of the partial OT plan.

We appreciate your thoughtful review and comments. Below, we answer your main questions and concerns.

Could you provide numerical experiments (even in simple cases) that could validate experimentally your algorithm and show its practical utility?

Summary of Contributions. Our paper presents the following key contributions:

• A characterization of partial and robust variants of the optimal transport (OT) problem in the semi-discrete setting,
• Theoretical guarantees on the approximation error when using an approximate $\lambda$-ROT solver to compute an approximate $\alpha$-OPT plan, and,
• Two algorithmic solutions for the $\lambda$-ROT problem.

While the original submission emphasized our theoretical contributions, we appreciate your interest in practical implementation and experimental evaluation. In response, we include below preliminary experimental results. A more extensive evaluation will be added in the next version.

Implementation Details. We provide a Python implementation of our algorithm. However, it relies on two core procedures that lack efficient implementations in Python:

1. Estimating continuous mass within regions defined by the arrangement of the restricted Voronoi cells and their expansions.
2. Maintaining a dynamic tree data structure as per Sleator and Tarjan (reference [44] in the appendix).

To address these challenges in our prototype:

• For (1), we approximate the mass within each region using a fine grid, summing the masses of grid points lying inside the region.
• For (2), we simulate the dynamic tree by naïvely iterating over all edges.

An efficient implementation will require optimized versions of these components, which remains a technical challenge in Python.

Preliminary Results on Algorithm Performance. We evaluated our implementation using a continuous distribution sampled from a Gaussian distribution inside the unit square. The discrete support consists of $n$ samples drawn from the same distribution, each assigned a mass of $1/n$. We set $\lambda = 0.2$ and $\varepsilon = 0.02$. Results are averaged over 10 runs for each $n$.

Key Findings:

• The additive error stays within the target $\varepsilon = 0.02$.
• The number of iterations in the two steps of our algorithm (Step 1 and 2) remains well below our upper bounds proved in the paper ($6n$ and $12n$ as described in Lemma 4.3 and D.9).
• The total path and cycle lengths computed in the two steps grow subquadratically in $n$, aligning with our $O(n^2)$ complexity analysis (Section 4.3).

The details of our experiments are presented in the following table.

Please let us know if any further clarification or experiments would strengthen the paper. We are happy to incorporate additional feedback into the next version of our paper.

How does this work relate to Bourne et al., 2018? Could you clarify how your results relate to the duality findings of Caffarelli and McCann [9]?

A: First, we appreciate the reviewer for bringing the work of Bourne et al. to our attention. We will cite their work in the final version. These papers also form connections among $\alpha$-OPT, $\lambda$-ROT, and weighted Voronoi diagrams. But we respectfully submit that we give a richer characterization that provides deeper insights into these problems. For example, [9] works with the weighted Voronoi diagram under the capped distance function and argues that the mass transported to $b$ in a $\lambda$-ROT plan is contained within the weighted Voronoi cell centered at $b$. We note that the Voronoi diagram under the capped distance is counter-intuitive, e.g., the boundary of a Voronoi cell may not be connected, and thus hard to interpret. Similarly, Bourne et al. prove an optimal $\alpha$-OPT plan routes the mass of $b$ to a subset of the mass within its weighted Voronoi cell, and only find experimentally that the mass transported to each discrete point is contained within a ball (they do not prove the second claim). In contrast, we work with the weighted Voronoi diagram under the Euclidean (or squared Euclidean) distance, where the Voronoi diagram is more standard, and prove that the $\alpha$-OPT plan routes $b$ to all mass of its restricted Voronoi cell, i.e. its weighted Voronoi cell intersected with the ball of radius equal to the dual weight. With our method, we additionally prove that every point whose mass is not fully routed in the optimal $\alpha$-partial transport plan has a maximal dual weight, which is an even stronger guarantee. Thus we provide a more detailed characterization. For the sake of space, we did not go into as much detail on the differences between Section 2 and prior work. We appreciate the question and will make this difference more clear in a final version.

Furthermore, [9] only studies the duality between exact solutions of $\lambda$-ROT and $\alpha$-OPT. In contrast, our main contribution (see Section 3) is to show that an approximate solver for one could be used to compute an approximate solution for the other. This result is nontrivial, as shown by the pathological example in Section 3.

We finally note that our work provides two new combinatorial algorithms to approximate the $\lambda$-ROT problem to varying degrees of precision.

We appreciate your thoughtful review and comments. We address your concerns below.

Summary of Contributions. Our paper presents the following key contributions:

• A characterization of partial and robust variants of the optimal transport (OT) problem in the semi-discrete setting,
• Theoretical guarantees on the approximation error when using an approximate $\lambda$-ROT solver to compute an approximate $\alpha$-OPT plan, and,
• Two algorithmic solutions for the $\lambda$-ROT problem.

While the original submission emphasized our theoretical contributions, we appreciate your interest in practical implementation and experimental evaluation. In response, we include below preliminary experimental results. A more extensive evaluation will be added in the next version.

Implementation Details. We provide a Python implementation of our algorithm. However, it relies on two core procedures that lack efficient implementations in Python:

1. Estimating continuous mass within regions defined by the arrangement of the restricted Voronoi cells and their expansions.
2. Maintaining a dynamic tree data structure as per Sleator and Tarjan (reference [44] in the appendix).

To address these challenges in our prototype:

• For (1), we approximate the mass within each region using a fine grid, summing the masses of grid points lying inside the region.
• For (2), we simulate the dynamic tree by naïvely iterating over all edges.

An efficient implementation will require optimized versions of these components, which remains a technical challenge in Python.

Preliminary Results on Algorithm Performance. We evaluated our implementation using a continuous distribution sampled from a Gaussian distribution inside the unit square. The discrete support consists of $n$ samples drawn from the same distribution, each assigned a mass of $1/n$. We set $\lambda = 0.2$ and $\varepsilon = 0.02$. Results are averaged over 10 runs for each $n$.

Key Findings:

• The additive error stays within the target $\varepsilon = 0.02$.
• The number of iterations in the two steps of our algorithm (Step 1 and 2) remains well below our upper bounds proved in the paper ($6n$ and $12n$ as described in Lemma 4.3 and D.9).
• The total path and cycle lengths computed in the two steps grow subquadratically in $n$, aligning with our $O(n^2)$ complexity analysis (Section 4.3).

The details of our experiments are presented in the following table.

Robustness of $\alpha$-OPT.

We examined the effect of 10% noise in the continuous distribution (added in the top-right corner) and measured how much noise was transported under various $\alpha$ values. Results confirm the robustness of $\alpha$-OPT: for $\alpha$ < 0.9, almost no noise mass is transported.

Robustness of $\lambda$-ROT. We conducted similar experiments for $\lambda$-ROT. As $\lambda$ increases, total mass transported grows, but noise transport remains limited when the transported mass is under 0.9.

Please let us know if any further clarification or experiments would strengthen the paper. We are happy to incorporate additional feedback into the next version of our paper.