One for all and all for one: Efficient computation of partial Wasserstein distances on the line

Could you print the transportation plan for the new algorithm in the following partial OT problem?

It seems my question wasn’t fully addressed.

Could you present the transportation plan in every iteration of your new algorithm for this simple 1-dimensional partial OT problem?

• $\mu$ is an empirical distribution supported on points $0, 1, 2$.
• $\nu$ is an empirical distribution supported on points $-2, -1, 0$.
• $k = s = 3$. Manhattan distance

I would like to know the transportation plan at every iteration. Could you verify whether the transportation plan in your algorithm (from $\mu$ to $\nu$) in each iteration matches the following?

• Iteration 1: $0 \to 0$
• Iteration 2: $0 \to -1, 1 \to 0$
• Iteration 3: $0 \to -2, 1 \to -1, 2 \to 0$

If this is correct, I’m confused by the statement (Proposition 7) that the complexity is $O(n \ln(n))$. Even if all other parts of the computational cost were $0$, updating the transportation plan in iteration $k$ requires $k$ computations. This suggests that the computational cost would be $1 + 2 + \dots + n = O(n^2)$.

Could you clarify this?

Response to the following comments:

• The sliced partial Wasserstein is defined based on SWGG, however, it is still possible to define many other variants of sliced partial Wasserstein by using more projections
• There is no experiment for testing the robustness of the proposed methods and other baselines.

Thank you for pointing out the possibility of defining alternative formulations for sliced partial Wasserstein distances using different projection schemes. While we acknowledge that multiple definitions are possible, we argue that SWGG is particularly well-suited to the partial setting for the following reasons:

• Sparsity in Partial Optimal Transport Maps: SWGG ensures sparsity in the resulting transport maps, which is a desirable feature of partial optimal transport. This sparsity aligns with the partial nature of the problem, where only a subset of the mass is transported.

• Straightforward Estimation of Transport Maps: The SWGG framework directly provides an estimated transport map rather than a general transport plan. This property is particularly beneficial in applications such as label propagation, where having an explicit map simplifies downstream tasks.

• Straightforward Maps for all masses $s$: by taking the minimum over different sampled lines, solutions for all masses $s$ can be obtained: it suffices to choose, for each amount of mass $s$, the line that provides the minimum OT cost.

• Experimental Validation of SWGG’s Advantages: To further demonstrate the utility of SWGG, we have included a new experiment in the revised manuscript that complements the gradient flow experiment. In this experiment, we compare SWGG to unbalanced optimal transport (UOT) with $L_2$ regularization in the context of outlier detection. Specifically, Figure 7 in the updated paper shows that UOT-based approaches often require setting a threshold on the transport plan to decide if a sample has been effectively transported. This thresholding can lead to inaccuracies, as evidenced by the fact that accuracies for UOT $L_2$​ do not converge to 0.5 (the baseline for a balanced classification). In contrast, the SWGG-based approach inherently provides a transport map, enabling exact label propagation without the need for thresholding, which leads to superior performance in this task.

• Robustness of our Sliced-PW Scheme: projection-based OT variants, such as Sliced-Wasserstein, is known to be more robust to data contamination than Vanilla Wasserstein, especially when considering $L_1$ norm (a thorough discussion is provided in [7]). Our Sliced-PW places itself in this context; nevertheless, we leave the theoretical analysis of this property for future work.

We believe these points highlight the advantages of SWGG for the partial optimal transport problem and justify its use in our proposed framework. We thank the reviewer for their insightful comment, which has allowed us to strengthen the discussion in the revised manuscript.

Response to the following comment and question:

• The application seems to be limited to low-dimension i.e., <= 3.
• Could elbow variant provide the true contamination percentage when the data contains additive noise?

It is at the moment unclear to us how we could theoretically study the ability of the elbow method to capture the contamination ratio in the presence of additive noise. Note that Partial or Unbalanced OT do not perform outlier detection: they remove samples that have large distance to the samples of the other distribution. Hence, the effectiveness of the elbow method probably depends on "how far" the outliers are from the original distribution. To our knowledge, this is still an open question in the optimal transport community.

We have included a new experiment that focuses on Domain Adaptation in the presence of outliers that further empirically showcases the benefits of computing solutions to all partial problems at once in practice in a high dimensional context.

[6] Luis A Caffarelli and Robert J McCann. Free boundaries in optimal transport and Monge-Ampere obstacle problems. Annals of mathematics, pp. 673–730, 2010.

[7] Nietert, S., Goldfeld, Z., Sadhu, R., & Kato, K. (2022). Statistical, robustness, and computational guarantees for sliced wasserstein distances. Advances in Neural Information Processing Systems, 35, 28179-28193.

Response to the following comment and question:

• Might not be friendly to backpropagation algorithm limiting use in context of deep learning / neural nets
• Is the algorithm usable in training neural net architectures? What is the notion of gradient in this case? We thank the reviewer for raising this important concern regarding the compatibility of PAWL with gradient-based optimization and its use in neural network training.

As noted in the literature (e.g., [3]), back-propagating through sorting-based algorithms like PAWL poses challenges due to the non-smooth nature of the operations involved. Specifically, if PAWL is treated as a function of some parameter $\theta$, its gradient can be expressed as:

$\nabla_\theta \text{PAWL}= \langle \pi^\star , \frac{\partial C}{\partial \theta} \rangle$

where $\pi^\star$ is the optimal transport plan and $C$ is the cost matrix. However, this gradient is often unstable, making it unsuitable for direct use in standard gradient-based optimization methods.
To address this, we can for instance leverage perturbed optimizers (e.g., [4]), which introduce noise to smooth the optimization landscape.

For illustration, consider a 2D setup where we compute the PAWL distance between two distributions $\mu$ and $\nu$, projected onto 1D as a function of the orientation $\theta$ of the projection direction. Using a perturbed optimizer, the smoothed objective $\text{PAWL}_\epsilon (\theta)$ is defined as:

$\text{PAWL}_\epsilon (\theta) = \mathbb{E} _{z \sim N _{0, 1}} [\text{PAWL} (\theta + \epsilon z)]$

where $\epsilon$ controls the magnitude of the perturbation. One can then use Stein's lemma to get the following expression for its gradient (cf. [5], Section 14.4.5):

$\nabla _{\theta} \text{PAWL} _\epsilon = \mathbb{E} _{z \sim N _{0, 1}} [(\text{PAWL} (\theta + \epsilon z) - \text{PAWL} (\theta)) \cdot z / \epsilon ]$

The figure 9 in Appendix illustrates this approach in a simple 2D example, showing the smoothed PAWL in function of the orientation $\theta$ (top) and its gradient (bottom) with Monte Carlo estimation using 1k samples and $\epsilon=0.1$. We can see that this perturbation is sufficient to smooth the gradients, making them more stable and suitable for gradient-based optimization in neural network contexts.

While this approach demonstrates feasibility, further exploration is warranted to optimize stability and scalability in high-dimensional neural network settings. We appreciate the reviewer's feedback and believe this opens interesting avenues for future work.

This discussion is included in Appendix A.3 of the updated paper.

[1] Laetitia Chapel, Rémi Flamary, Haoran Wu, Cédric Févotte, and Gilles Gasso. Unbalanced optimal transport through non-negative penalized linear regression. Advances in Neural Information Processing Systems, 34:23270–23282, 2021.

[2] Abhijeet Phatak, Sharath Raghvendra, Chittaranjan Tripathy, and Kaiyi Zhang. Computing all optimal partial transports. In International Conference on Learning Representations, 2023.

[3] Mathieu Blondel, Olivier Teboul, Quentin Berthet, and Josip Djolonga. Fast differentiable sorting and ranking. In International Conference on Machine Learning, pp. 950–959. PMLR, 2020.

[4] Quentin Berthet, Mathieu Blondel, Olivier Teboul, Marco Cuturi, Jean-Philippe Vert, and Francis Bach. Learning with differentiable pertubed optimizers. Advances in neural information processing systems, 33:9508–9519, 2020.

[5] Mathieu Blondel and Vincent Roulet. The elements of differentiable programming. arXiv preprint arXiv:2403.14606, 2024.

We thank the reviewer for their feedback and for taking the time to engage with our work. Below, we address the key points raised:

• On the scope of experiments: While we acknowledge the comment regarding limiting time to provide new experiments, our submission already includes a comprehensive set: gradient flows, point cloud registration, and a new experiment on detecting data from different domains. The added experiment specifically aligns with the setup in [1] that has been used to validate the advantage of computing the entire set of solutions (for all values of the hyper-parameter) in a similar context. This experiment uses data of dimension 784, demonstrating the applicability of our approach in high-dimensional settings.

• Regarding timings and $L_1$ norm: We respectfully disagree with the reviewer’s concerns about the $L_1$ norm. Its relevance in machine learning is well-established (e.g., WGAN) and it is often stated to be more robust (e.g. see [7] in the context of sliced Wasserstein). PAWL is particularly compelling because:

  • It provides all solutions across parameter values, which eliminates the need for repeated runs that may be required by competitors (e.g., during cross-validation).
  • Experiments demonstrate that PAWL (using the $L_1$ norm) is more than one order of magnitude faster than competitors. This efficiency remains significant even when extended to the $L_2$ norm.
  • Additionally, we emphasize that timing results are inherently implementation-dependent. As a consequence, the relative slope of the runtime curves across methods is more indicative of the scalability and theoretical efficiency of PAWL than the absolute values, which may vary due to implementation details. Looking at the slopes, PAWL $W_1$ is competitive with Fast-UOT and PAWL $W_2^2$ outperforms OPT.

Finally, if the reviewer thinks that it is useful, we can add a discussion about the choice of $L_1$ as a ground cost in the core of the paper, as well as a further analysis of the runtime.

[1] Laetitia Chapel, Rémi Flamary, Haoran Wu, Cédric Févotte, and Gilles Gasso. Unbalanced optimal transport through non-negative penalized linear regression. Advances in Neural Information Processing Systems, 34:23270–23282, 2021.

[7] Nietert, S., Goldfeld, Z., Sadhu, R., & Kato, K. (2022). Statistical, robustness, and computational guarantees for sliced wasserstein distances. Advances in Neural Information Processing Systems, 35, 28179-28193.