Progressive Entropic Optimal Transport Solvers

Thank you for taking the time to review our work. We have fixed the notation and typos issues, thank you for pointing them out.

Since the optimal transport map is unknown in practice, the assumption (A.2), which says that the inverse map has at least three continuous derivatives, is quite strong. Are there any chances that the assumption (A.2) can be reduced?

We highlight that requiring up to second-order derivatives on the inverse map is a standard requirement, see [1, 2, 3, 4]. The additional derivative comes from the assumptions in [4], which is, to date, the only work that demonstrates that the entropic OT map can statistically estimate the OT map. It is largely believed that the assumption is not superfluous, and that the extra derivative is required to make this statistically viable [5]. Though there are restricted scenarios under which the 3rd derivative is unnecessary, e.g., if $\nu$ is a discrete measure [5]. One can hope that ultimately these results can be bridged, but for now this is outside the scope of this work.

What is the computational complexity of the ProgOT algorithm (Algorithm 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.

The exact computational complexity depends on the problem geometry, since $\varepsilon_i$ is data-driven.

We can specify the worst-case computational complexity of ProgOT compared to Sinkhorn. The algorithm runs the sinkhorn subroutine $K$ times, with a decreasing sequence of regularizations $\varepsilon_i$. ProgOT also performs linear interpolations, yielding a total of $\mathcal{O}(nd + \sum_{i=1}^K C_{\mathrm{Sink}}(\varepsilon_i))$ operations, where $n$ is the size of the input domain, $d$ the dimension and $C_{\mathrm{Sink}}(\varepsilon_i)$ denotes the worst-case complexity of Sinkhorn with regularization $\varepsilon_i$. We added a paragraph to Appendix A discussing this matter.

In lines 238-239, when initializing the entropy-regularized parameter $\varepsilon$, 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.

We have updated that paragraph and included a brief explanation. This choice is not ours, but largely followed by the community. To avoid simple scaling effects (such as multiplying all features by a constant), cost matrices in entropic OT are rescaled, typically by their maximum value (as commonly done in POT [7]) or by mean-value (as done in OTT [6]). This rescaling is often absorbed into $\varepsilon$, and helps balance out cost/entropy in EOT optimization.

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.

This is indeed an interesting research direction. We have updated the conclusions section with a discussion on directions such as unbalanced OT. We note that such extension is far from being straightforward. Crucially, there is no equivalent of an “entropic potential map”, or McCann interpolation, that can be extended out of sample.

Could the authors please explain more clearly why line 3 in Algorithm 2 helps improve the runtime?

Line 3 implements a warm-start for the Sinkhorn algorithm. The learned potentials at iterate $k$ might still be relevant for the next Sinkhorn subroutine. Our update is essentially a back-of-the-hand envelope: since the distances between point clouds decrease at each iteration by roughly a factor of $(1-\alpha_k)$, we decrease the dual potentials by a similar amount. We observe that this simple update works well across various costs. Warm-starting Sinkhorn is an active area of work [6, 7, 8, 9].

Are there any instructions on how to choose the number of steps?

Intuitively, the number of steps should correlate positively with the distance between source and target, but for now, we do not have explicit instructions for this.

A simple solution can be choosing $K$ depending on the available compute. Choosing large $K$ typically does not affect performance negatively, but smaller $K$ may create a less stable algorithm (as we get back to Sinkhorn). In our experiments we have used K = 4, 8, 16, and noticed diminishing returns in increasing $K$.

References

[1] Hutter and Rigollet. "Minimax estimation of optimal transport maps", (2021)

[2] Manole, et al. "Plugin estimation of smooth optimal transport maps", (2021)

[3] Muzellec, et al. "Near-optimal estimation of smooth transport maps with kernel sums-of-squares", (2021)

[4] Pooladian, Aram-Alexandre, and Jonathan Niles-Weed. "Entropic estimation of optimal transport maps." arXiv preprint (2021)

[5] Pooladian et al. “Minimax estimation of discontinuous optimal transport maps: The semi-discrete case” ICML (2023)

[6] Cuturi, Marco, et al. "Optimal transport tools (ott): A jax toolbox for all things wasserstein." arXiv preprint (2022)

[7] Flamary, Rémi, et al. "Pot: Python optimal transport." Journal of Machine Learning Research (2021)

[8] Amos, B., Cohen, S., Luise, G., & Redko, I. Meta optimal transport. arXiv preprint (2022)

[9] Thornton, James, and Marco Cuturi. "Rethinking initialization of the sinkhorn algorithm." AISTATS (2023)

We will integrate the discussion in our final version and we thank you for the many comments you have made.

Between our code implementation and significantly larger experiment, we believe the paper has indeed improved further, and this is one of the merits of the rebuttal process.

We humbly ask if, as you mention that the discussions has consolidated your positive opinion of the paper, and your soundness / presentation / contribution scores al stand at "3:good", if you would consider increasing your score? The acceptance threshold at neurips this year is likely going to be around 5.5. Hence a score of 5 is, relatively to all other papers, negative in this context.

the authors

Many thanks for your careful review, we thank you for your positive score, encouraging comments and questions. We did our best to answer them.

… One may find this less favorable, as it replaces one hyper-parameter with several others.

While we certainly agree with this assessment, one message we want to convey is that having to choose a single variable $\varepsilon$ for everything (Sinkhorn, map estimation) holds too much sway over EOT. A "good" $\varepsilon$ can be difficult to nail right, notably for map estimation, and most practitioners want to bypass this. As it stands, this is likely one of the factors that limits usage of EOT.

Our goal is to "divide and conquer" $\varepsilon$ too, translating it into simpler quantities. In practice, a user would only need to only choose the number of steps $K$, according to the means of compute available to them, since we provide automatic routines for choosing the rest of the hyper-parameters: the step schedule is linear (a.k.a constant-speed), the threshold scaling log-linear, and the epsilon schedule is tuned as in Algo. 4.

The gradient of PROGOT is missed, which is important for machine learning application.

Differentiability is an important point. We would like to provide a satisfactory answer by splitting the various meanings this could take in the context of ProgOT.

ProgOT can return three objects that a user may want to differentiate:

• “Wasserstein-like” OT transport cost between two point-clouds, equal to $t(\mathbf{X},\mathbf{Y}):=\langle \mathbf{P}, [h(x_i-y_j)]_ij\rangle$ where $\mathbf{P}$ is the solution outputted by Algo.2. This is the quantity displayed in the x-axis of Fig. 5 (as well as in our new plots). We believe this is the quantity you have in mind when you mention a gradient. Differentiating this quantity w.r.t. $\mathbf{x}, \mathbf{y}$ is, in fact, very easy, since it amounts to using a “Danskin-like” linearization, i.e. use the chain rule to differentiate $t$ w.r.t. $\mathbf{X}$ or $\mathbf{Y}$ (or any other parameter), while keeping $\mathbf{P}$ fixed. We illustrate this in Figure H of the uploaded pdf, and also provided a colab link to the AC, who can share it with you.
• The Jacobian of the coupling matrix $\mathbf{P}$ returned by Algo.2. This can be achieved by differentiating chained Sinkhorn iterations. This is the strategy used, in a different context, when differentiating the optimal coupling obtained in Gromov-Wasserstein for instance. It is costly, but doable.
• The Jacobian of the ProgOT map (that returned by Algo. 3) at a given point $x$, w.r.t. Parameters. This is the same as above, and will require a chained differentiation of successive OT solutions.

Hence, it is very easy to use ProgOT to obtain descent directions to minimize OT-costs. Differentiating higher-order objects (couplings / map estimator) is more involved.

In Figure 5, could you provide the actual number of iterations instead of just indicating different marker sizes?

We have updated the figures in the paper to show the number of iterations, see Figure F for the updated version of Figure 5 in the paper. We have also added 4 new figures to the appendix in new data configurations. Fig G in the upload PDF shows an example. To answer your question, We have prepared Table E with the exact number of iterations corresponding to Figure F, (the updated variant of Figure 5).

**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.

Thank you for pointing this typo out, we have updated the text.

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 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.

Even if $K$ is small, the cost values might be different if the regularization parameter is different. We visualize Sinkhorn using different values $\beta$, and in many cases (c.f. Fig F and Fig G) ProgOT with $K=2$ lies close to a Sinkhorn instance. Perhaps your intuition is better reflected in the updated plots for this experiment.

We also highlight that the two axes of cost and entropy should be viewed together. One possible scenario is that by choosing a larger $K$, we target a smaller $\varepsilon$ in the last iteration, and gaining sharper entropy, at the price of having a larger cost.

Some of the results in Figure 5 in our initial draft were slightly off because we were using a maximum number of iterations when implementing Sinkhorn and ProgOT. While the overall message of these plots have not changed, this error sometimes contradicted with our claim that all methods were converging to target accuracy (L.322). We now set that maximum iterations to infinity, and therefore all solutions converge correctly within the desired threshold, irrespective of the number of iterations. See Fig F and Fig G in the uploaded PDF for an example.

We would like to thank you for your review, and for the many thought provoking questions you have asked. We did our best to answer all of them.

a crucial shortcoming of this method is scalability. […] severely limits its practical usability.

There might be a confusion: Theorem 3 proves a theoretical recovery guarantee for ProgOT as a map estimator. This is a worst-case statistical analysis. This worst-case study has little impact on prediction performance on real tasks (e.g. Table 1), and no impact on computational aspects (Figure 5) or coupling estimation, which are all practically governed by other factors ($K, \alpha_k, \varepsilon_k, \tau_k$ for instance).

Theorem 3 guarantees the soundness of the ProgOT map. In contrast, none of the NN-based methods (e.g. those compared in the [Korotin et al. 21] benchmark) have a theoretical statistical guarantee. The fact that ProgOT has such a guarantee (even if pessimistic) cannot therefore be seen as a weakness. Their estimation is hard, because of non-convexity. In contrast, ProgOT is recovered as a sequence of convex problems.

the authors do not test their approach using the provided benchmark pairs for $d$ larger than 64 […] it should be thoroughly tested on synthetic tasks.

Thanks for this suggestion. To demonstrate the scalability of ProgOT to larger $n$ and larger $d$, we have added two new experiments: MSE error on the Korotin benchmark for $d=128, 256$ (as you requested) and a new large scale experiment on all $n=60.000$ grayscale images ($d=1024$) of the CIFAR10 dataset.

GMM Experiment: We use the code from the [Korotin et al., 2021] benchmark to create high-$d$ synthetic problems, reaching the maximum dimension available $d=256$. Tab. D in the uploaded PDF presents the results. ProgOT (with default settings) dominates other baselines. This is remarkable, because: (1) this benchmark favors Neural OT solvers by design: the ground truth transport is itself generated as the gradient of an ICNN. (2) running ProgOT takes about 3 minutes, while Neural solvers take 10-20X longer (about 30’ for ICNN and 60’ for the Monge Gap), all on a A100 GPU.

CIFAR10 Experiment: We designed a novel benchmark task, to match images and their blurred versions. As described in our general response, the ground truth OT coupling should be the diagonal. We only compare ProgOT to Sinkhorn (ICNN and Monge Gaps do not return couplings). See Fig. A& Tab. B in the uploaded PDF for the results and refer to the general response for more details.

Note that a sample size $n=60,000$ is, to our knowledge, unheard of for entropic OT for such high-$d$ (we overwhelmingly see $n\leq 10,000$ in papers). We achieve this through a tight implementation of ProgOT within OTT-JAX, and JAX, getting the immediate benefits of sharding across GPUs.

You use entropic OT formulation with Euclidean quadratic cost due to […] Is there generalization of the method for arbitrary transport cost function?

This generalization is already in the paper: we defined ProgOT using arbitrary translation invariant (TI) costs, $c(x,y) = h(x-y)$ (see L.81). Alg. 2 and 3 both use $h$ explicitly. We use a a $p$-norm (with $p=1.5$) in the bottom of Fig. 5. We only need $h$ to be $\tfrac12\|\cdot\|^2$ when proving map guarantees in Section 3.2, see L.184.

To complete the empirical overview on general costs, we have added a new table to the appendix A, which benchmarks map estimation with general costs (similar to Table 1).

it is understandable that a sub-problem is easier and well-conditioned, than the initial problem […] could you provide a theoretical understanding why it is so?

Because the intermediate problems lie along the geodesic path, the distance/transport costs for them is smaller, L. 154~163. Calling Sinkhorn for problems with smaller transport costs is easier/more stable w.r.t. $\varepsilon$. This also allows warm-starting to recycle solutions from the previous step (L.3 in Algo. 2.)

Does the method need a pre-trained OT map?

No, ProgOT runs off-the-shelf using source/target data (see L.229, inputs to Algo.2) and is deterministic, as it relies on sub-Sinkhorn calls. Could you please clarify your question?

Is there intuition or methodology of picking a more appropriate step-size for McCann interpolator at step $k$?

We investigated this thoroughly. We found out that scheduling the step-size does not seem to play a key role, neither theoretically nor empirically (see Fig 7 in the Appendix). We stick to the simplest choice, the constant-speed, a.k.a. linear, schedule.

Why does the speed of convergence of the proposed algorithm compare to the speed of the Sinkhorn algorithm of EOT[…]

The convergence rates are comparable because ProgOT uses calls to local Sinkhorn problems as a subroutine. Theorem 3 guarantees theoretically that by using ProgOT we remain statistically sample efficient, while gaining a computational edge. Theorem 3 is illustrated empirically in Figure 4 (A), where we plot the MSE for various $n$, using the ground-truth OT map from [Korotin et al., 2021].

Could you provide practical evidence that obtained OT curves are really straight, at least synthetic Gaussian experiments?

Using Algo.2, curves can be either obtained by sampling straight lines from the final coupling matrix $\mathbf{P}$ returned by ProgOT, or using the map (Alg. 3) out of sample. Sampling from coupling yields straight lines, by construction. This might be used in, e.g., flow-matching estimation. The "straightness" of lines sampled with Alg. 3 will hinge on the various $\varepsilon_i$, as represented schematically in Fig. 3.

I suggest the authors to test their algorithm on benchmark with available ground-truth OT maps for bigger values of $d$

Many thanks for this suggestion. We hope our new results [Tab. D] assuage your concerns.

References

[1] Vacher, Adrien, and François-Xavier Vialard. "Parameter tuning and model selection in optimal transport with semi-dual Brenier formulation." Advances in Neural Information Processing Systems 35 (2022): 23098-23108.

[2] Van Assel, Hugues, et al. "Optimal Transport with Adaptive Regularisation." arXiv preprint arXiv:2310.02925 (2023).

[3] Scetbon, Meyer, Marco Cuturi, and Gabriel Peyré. "Low-rank sinkhorn factorization." International Conference on Machine Learning. PMLR, 2021.

[4] Pooladian, Aram-Alexandre, and Jonathan Niles-Weed. "Entropic estimation of optimal transport maps." arXiv preprint arXiv:2109.12004 (2021).

[5] Hutter and Rigollet. "Minimax estimation of optimal transport maps", 2021

[6] Manole, et al. "Plugin estimation of smooth optimal transport maps", 2021

[7] Muzellec, et al. "Near-optimal estimation of smooth transport maps with kernel sums-of-squares", 2021

[8] Bunne, Charlotte, et al. "Learning single-cell perturbation responses using neural optimal transport." Nature methods 20.11 (2023): 1759-1768.]

[9] Dessein, Arnaud, Nicolas Papadakis, and Jean-Luc Rouas. "Regularized optimal transport and the rot mover’s distance. arXiv e-prints." arXiv preprint arXiv:1610.06447 (2016).

Dear Reviewer Z1nb,

The discussion period is closing, and we will soon not be able to interact with you.

Still, we hope that you can consider our answers / additional material in coming days, during the reviewers-AC discussion.

Let us emphasize again that, while we have added results on the benchmark you requested, we have also designed and ran a large scale CIFAR-10 experiment ($n=60k, d=1024$) motivated in part by your comments on scalability and high dimensions. Such numbers are unheard of in the EOT literature. Hence, we are very grateful for your detailed feedback on the paper, which has triggered this exploration.

We have also provided a colab that illustrates that ProgOT is not only differentiable, but also ready to be used "off the shelf", as we claim in the paper.

You mentioned that you were open to adjusting your score based on the authors' answers. If these experiments seem convincing to you, we would of course be very grateful if you could so.

If you have any remaining questions, we might be able to answer them through the AC (in principle we should be able to communicate with the AC following the closing of the discussion period).

Respectfully, The Authors

Many thanks for your review and for your comments. Our response follows.

I think the contribution of this work is rather limited since the only benefit is to avoid the tuning of the parameter $\varepsilon$.

In light of your comment, we will delete the sentence “Setting $\varepsilon$ can be difficult, … and bias” in our abstract, so that readers do not get confused on what constitutes our contribution.

Our paper is not “only about avoiding tuning $\varepsilon$” in Sinkhorn. While this was the main focus of other recent works, such as [1, 2], L. 59, this is not our goal.

Our solver blends McCann type interpolation, local Sinkhorn computations and entropic map estimation to yield both a coupling and map estimator (L.12~18). While easing the burden to pick $\varepsilon$, our solver does a lot more than that, as listed in Contributions, L.65~76, and demonstrated in Fig. 5 (computations, highlighting very different outputs compared to Sinkhorn), or Table 1 (out-of-sample prediction).

Please also take a look at our new experimental results [Fig. A, Tab. B] on the entire CIFAR-10 dataset.

In my experience, tuning is usually not a big issue since Sinkhorn is rather robust.

Use cases for Sinkhorn vary a lot, and we do trust that you may not have run into such problems in the past. In practice, however, EOT solvers are widely used, for many tasks, impacting many user profiles.

In the simplest cases, e.g. when computing couplings on normalized data (e.g. embeddings on the sphere), we also observe that Sinkhorn's outputs can be fairly stable across various $\varepsilon$ values. On real data sources (such as the genomics point clouds), this is not necessarily the case. $\varepsilon$ has a considerable impact on stability and compute time.

$\varepsilon$ also impacts disproportionately the predictive performance for map estimation. Not tuning it often results in poor outcomes, which is why all Sinkhorn baselines in our draft use cross-validation. To convince you, we have added a line with "untuned Sinkhorn" (i.e. using the default given by OTT-JAX) to our benchmarks [see Tab. D in the PDF], showing that it performs significantly worse than cross-validated Sinkhorn and ProgOT. Thank you for this suggestion.

In addition, the work does not discuss any computational tradeoff in doing so.

We did target this tradeoff explicitly in Figure 5, see L.314. We have updated appendix A with more experimental results to highlight that tradeoff.

Tuning the parameter only affects the accuracy of EOT,...

Can you clarify what you mean by accuracy of EOT? Do you mean proximity to the unregularized solution?

...which is rather minor since the main bottleneck of large-scale Optimal Transport is the number of support $n$

Each Sinkhorn iteration has indeed a cost in $n^2$. However, the number of iterations varies tremendously with $\varepsilon$, e.g. from 10 iterations for large $\varepsilon$ to 10k for small ones, as shown e.g. in our paper, and [Fig. F & Fig. G].

To our knowledge, the only linear time RegOT solver is the low-rank (LRSink) approximation from [3]. That solver does not yield a map estimator, as it does not solve the dual OT problem. Exploring hybridizing ProgOT with LRSink is an interesting direction.

...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.

To prove guarantees for estimators, the strength of assumptions should naturally be commensurate with the difficulty of the task.

Proving theoretical recovery of OT maps is hard, and to our knowledge, the entropic map estimator [4] that we build upon is the only tractable estimator that has such guarantees (L. 126~129). We follow their assumptions, and we do not see any alternatives now. The fact that the inverse mapping must have three continuous derivatives is a strong assumption, but we believe this is a fundamental issue, and not an artefact or our analysis. Other approaches (e.g., [5, 6, 7]) share many of the same regularity assumptions as we do.

Note that many OT map estimators are routinely used without any theoretical guarantees (e.g. all neural-OT approaches, flow-matching like formulations), even in biological settings [8]. Aside from the entropic map, ProgOT is the only other tractable map estimator that comes with convergence guarantees.

To summarize, apart from the entropic map estimator, ProgOT is the only currently known tractable OT map estimator with guarantees, and we show that it performs better in practice.

How would $\nu_{\min}, \nu_{\max}$ affect the theoretical analysis?

$\nu_{\min}$ and $\nu_{\max}$ appear as constants in the bounds, as established by [4], which are just meant to be multiplicative constants in the bounds. The role is analogous in the other papers mentioned above

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.

This is an interesting research direction. However, extending our work to the settings you mention (quadratic regularization/ unbalanced settings) is far from being straightforward. Crucially, there is no equivalent of an “entropic potential map”, that can be extended out of sample, for the dual solutions outputted from quadratic solvers (first proposed in [DPR16]), nor for the unbalanced case.

We have added references that you mentioned to the conclusion section, to highlight possible future work around ProgOT.

We are very thankful for your various comments. With these clarifications and the first (to our knowledge) demonstration that ProgOT can run an EOT based solver at large scales (60,000 points, d=1024), and still return a meaningful coupling, we believe we have a significant and convincing contribution.

We are also grateful for your score increase.

While there is little time left in the dicussion, we remain available to answer any other questions you may have.

the authors