Progressive Entropic Optimal Transport Solvers

Summary

In this paper, the authors propose a new entropic optimal transport solver as an alternative to the commonly used Sinkhorn algorithm named ProgOT. This solver has three main properties: (i) It is less sensitive to the choice of the entropy-regularized parameter than the Sinkhorn algorithm; (ii) When computing couplings between point-clouds, the runtime of ProgOT is no longer than the Sinkhorn algorithm; (iii) The resulting optimal transport map estimator is consistent and stable. The authors provide both theoretical and empirical results to demonstrate the above claims.

Strengths

1. The proposed ProgOT is a new entropic optimal transport solver built based on the McCann-type interpolation. 2. The theoretical results in the paper are sound, and provided with rigorous proofs. 3. The paper is well written and organized. The background section is very useful, particularly for readers who are not familiar with optimal transport.

Weaknesses

1. Since the optimal transport map $T_0$ is unknown in practice, the assumption (A.2), which says that the inverse map $T^{-1}_0$ has at least three continuous derivatives, is quite strong. 2. In lines 33-36, when comparing the efficiency of the Sinkhorn algorithm to the linear programs for solving the EOT problems, the authors should make it more explicit by stating the computational complexity of both algorithms. 3. In lines 238-239, when initializing the entropy-regularized parameter $\varepsilon_0$, the authors should at least briefly present the intuition for setting it to be the average of the values in the cost matrix between sources and targets. 4. Minor issues: There are some undefined notations and grammatical mistakes: - In line 141, I think the term $S^{(1)}$ should be $S^{(0)}$. - In the inequality between lines 196-197, the notation $\lesssim_{\log(n),k}$ has not been defined yet. - In line 167, the notation $\alpha(k)$ has not been defined. Is it a constant depending on $k$? - In the assumption (A.3), what does the notation $D$ stand for? - In line 204, $\mu^{(k)}$ is corresponds a location --> grammatical mistake. 5. The authors should add a discussion about the limitations of the proposed method as well as future directions. For instance, can we generalize the ProgOT so that it applies for entropic unbalanced optimal transport. I believe that such discussion would make the paper more complete.

Questions

1. Are there any chances that the assumption (A.2) can be reduced? 2. What is the computational complexity of the ProgOT algorithm (Algorithm 2)? 3. Could the authors please explain more clearly why line 3 in Algorithm 2 helps improve the runtime? 4. Are there any instructions on how to choose the number of steps $K$?

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The limitations have not been discussed in the paper.

Summary

This paper introduces a new class of entropic optimal transport (EOT) solvers called PROGOT. This work aims to address the challenges of selecting entropic regularization strength $\epsilon$ for original EOT. As we know $\epsilon$ is significant to the performance of EOT like computation time and convergence rate. PROGOT utilizes a progressive framework that interpolates the whole entropic transportation process into multiple steps by using dynamic OT formulations. This work proposed algorithms to set the parameters throughout the interpolation process. As it claimed, the new framework enhances the robustness and performance of EOT, and avoid the headache of tuning the value of $\epsilon$. Experimental results show that PROGOT surpasses classic EOT in both synthetic and real dataset.

Strengths

- The major contribution of this work is the proposal of a novel framework to estimate entropic transport plans. This framework addressed the challenge of selecting an appropriate entropic regularization strength $\varepsilon$ in traditional EOT algorithms. - The mathematical notation and theoretical analysis are clear and easy to follow. And the theoretical analysis looks good to me. - The convergence of PROGOT map is supported both in theoretical proof and experimental results. - The methodology for selecting hyper-parameters (regularization and threshold schedule) and the corresponding justifications are well provided.

Weaknesses

- While the effort in setting the hyper-parameters and the justifications is acknowledged, from a broader perspective, to replace the selection of $\varepsilon$, PROGOT introduces a series of new hyper-parameters (step/regularization schedules and the length of the schedules, $K$) and theoretic assumptions on the inputs. One may find this less favorable, as it replaces one hyper-parameter with several others. - Some of the experiments would be benefit from including of the real OT cost (or map) as a reference. For example, adding the real OT map in Figure 1 and adding the real OT cost data point in Figure 5 will be helpful. - The gradient of PROGOT is missed, which is important for machine learning application. - In Figure 5, could you provide the actual number of iterations instead of just indicating different marker sizes? For example, comparing the size of markers for $\beta=0.08$ vs. $K=8$ is difficult. Additionally, since all the sub-figures use the same legend, it should be clarified whether they run the same number of iterations for the same configuration.

Questions

- Line 141: do you mean $S^{(0)}$ instead of $S^{(1)}$ in the equation $\mu^{(1)}=S^{(1)}\mu$? The same issue in line 143. - Figure 5, why does PROGOT with larger $K$ value has larger cost and appear closer to the original Sinkhorn point? Please correct my intuition if wrong: PROGOT with smaller $K$ runs fewer interpolation steps, leading to fewer number of iterations and closer to the original Sinkhorn cost. When $K=0$, PROGOT should be the same as EOT. In figure 5, the number of iteration results align with my intuition, but the "distances" to Sinkhorn results looks inconsistent.

Rating

7

Confidence

3

Soundness

4

Presentation

3

Contribution

3

Limitations

None

Summary

The choice of appropriate entropic trade-off term is one of the main headaches for finding maps between data distributions with sample access when considering Optimal Transport (OT) with entropic regularization. While the selection of sufficiently small regularization terms leads to unstable learning, the picking of large ones causes biased solutions. Considering the EOT problem with Euclidean quadratic cost, the authors propose to divide the entropic problem to small sub-problems between intermediate distributions with their own regularization terms. The final entropic map between initial data distributions is a composition of maps from sub-problems. The authors provide the theoretical guarantees that the constructed map does not differ much from the true one. Moreover, the authors offer the algorithm for selecting appropriate regularization terms for each step as well as demonstrate the method’s performance on synthetic data and single-cell experiments.

Strengths

- The paper provides a new methodology for solving OT problems with entropy regularization which does not combine the ideas of prior works. The paper is well-written, well-organized and clear. The storyline is perfect and understandable through the paper. - Since the method alleviates the tuning process for entropic regularization terms and allows to avoid unstable learning as well as finding primitive OT maps, this work might be interesting to the ML community. I am sure that other researchers might apply this methodology for Flow matching or Schroedinger bridge methods that aim to build interpolation from one dataset to another one. - The ProgOT alleviates the tuning process for entropic regularization terms, thereby it does not suffer from unstable learning as well as primitive solutions.

Weaknesses

Although the paper offers a unique theoretical approach, a crucial shortcoming of this method is **scalability**. Indeed, the estimation, which is provided by Theorem 3, demonstrates poor scalability of the method since the rate of convergence to the true OT maps depends on the dimension $d$. From my point of view, this issue of the approach could not be ignored since it severely limits its practical usability. This my concern is supported by the set of experiments considered in the paper. It seems that the authors avoid empirical evaluation of the method in high-dimensional problems (which might be useful for mitigating the concern). The performance of the method is tested only in low-dimensional ($d<=64$) synthetic data experiments using the benchmark with available ground-truth OT maps (Korotin et al, 2021), as well as biological experiments with no sufficiently large dimensionality ($d<=256$) of data. I am especially confused by the fact that the authors do not test their approach using the provided benchmark pairs for $d$ larger than 64 and at the same time apply the method to an important biological task in case of $d=256$. However, I believe that prior to adapting the algorithm to the tasks from single-cell biology (especially as important as prediction of cancer cells responses to drug perturbation), it should be thoroughly tested on synthetic tasks. **Overall**, although the proposed approach is as fast as the Sinkhorn algorithm and demonstrates better performance in low dimensional problems, there is no guarantees and understanding of the method’s behavior in high dimensional problems as well as there are no comparisons with other EOT solvers in this case. This is the main reason of my current score. However, I am open to adjust the score based on the authors' answers.

Questions

- From intuition’s point of view, it is understandable that a sub-problem is easier and well-conditioned, than the initial problem. Anyway, could you provide a theoretical understanding why it is so? It seems that this fact depends on the step-size of the McCann interpolator. - You probably use entropic OT formulation with Euclidean quadratic cost due to the convenient use of linear McCann interpolation between distributions. Is there generalization of the method for arbitrary transport cost function? - Does the method’s algorithm need a pre-trained OT map between initial distributions that approximately matches one distribution to another? - Is there intuition or methodology of picking a more appropriate step-size for McCann interpolator at step k? - Why does the speed of convergence of the proposed algorithm compare to the speed of the Sinkhorn algorithm of EOT problem between initial distributions? Could you provide a theoretical explanation and plot of convergence, at least in a synthetic Gaussian example? - In accordance with the developed methodology, the geodesic curves between initial distributions should be straight. Could you provide practical evidence that obtained OT curves are really straight, at least synthetic Gaussian experiments? I suggest the authors to test their algorithm on benchmark with available ground-truth OT maps for bigger values of $d$. **References.** A. Korotin, L. Li, A. Genevay, J. M. Solomon, A. Filippov, and E. Burnaev. Do neural optimal transport solvers work? A continuous Wasserstein-2 benchmark. Advances in neural information processing systems, 2021.

Rating

5

Confidence

4

Soundness

2

Presentation

3

Contribution

3

Limitations

In accordance with the theorem 3 of the main text, the main limitation of the method is scalability. However, it seems that the authors did not mention this fact.

Summary

The authors proposed ProgOT, a method to solve a sequence of EOT problems so that practitioners do not have to tune the entropic regularizer parameter $\varepsilon$ and strike a good balance between computational and statistical complexity.

Strengths

The manuscript is well-written and it was easy to understand.

Weaknesses

I have several concerns about the contributions and the assumptions of this work: - First of all, I think the contribution of this work is rather limited since the only benefit is to avoid the tuning of the parameter $\varepsilon$. In my experience, tuning $\varepsilon$ is usually not a big issue since Sinkhorn is rather robust. In addition, the work does not discuss any computational tradeoff in doing so. Tuning the parameter $\varepsilon$ only affects the accuracy of EOT, which is rather minor since the main bottleneck of large-scale Optimal Transport is the number of support $n$. - Secondly, this work uses a lot of strong assumptions to establish theoretical guarantees. For example, it assumes Euclidean cost, convex, and compact condition and the inverse mapping has at least three continuous derivatives. These assumptions are very restrictive and would limit the applicability of the theoretical analysis.

Questions

- How would $\nu_{min}, \nu_{max}$ affect the theoretical analysis? - I would recommend the authors to consider quadratic regularized OT/UOT formulations [1], [2] and entropic-UOT [3]. - In addition, I would recommend the authors to cite and discuss the following works as well. References: [1] Smooth and Sparse Optimal Transport. Mathieu Blondel, Vivien Seguy, Antoine Rolet. In Proc. of AISTATS 2018. https://arxiv.org/abs/1710.06276 [2] "On Unbalanced Optimal Transport: Gradient Methods, Sparsity and Approximation Error". Quang Minh Nguyen, Hoang Huy Nguyen, Lam Minh Nguyen, Yi Zhou. Journal of Machine Learning Research, (JMLR), 2023 [3] Pham, K., Le, K., Ho, N., Pham, T., & Bui, H. (13--18 Jul 2020). On Unbalanced Optimal Transport: An Analysis of Sinkhorn Algorithm. In H. D. Iii & A. Singh (Eds.), Proceedings of the 37th International Conference on Machine Learning (pp. 7673–7682). Retrieved from https://proceedings.mlr.press/v119/pham20a.html

Rating

4

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

See above.

Response to the authors

Thank you very much for your detailed and insightful responses. Though I'm still not very convinced of the significance of the work, the responses have addressed many of my concerns. I will increase the score.