Paper: Generative Entropic Neural Optimal Transport To Map Within and Across Space (ICLR-2024)

Reviewer comment (Experiments):
While in Proposition 3.1, the authors proved minimizing (7) results in the correct optimal plan, in practice in Algorithm 1,2, a **biased** loss is being optimized due to finite batch sizes. By biased, what I mean is that the expectation of the discrete estimator in Algorithm 1,2 is not necessarily the same as (7). One could expect the biasedness will decrease as the batch size increases. However, the effect of batch sizes is not discussed at all; it is taken to be 1024 throughout. I view this point as a central limitation of the current work. If theoretical analysis is difficult, empirical evidence to show how the results degrade as a function of batch size could be important. For instance, using a batch size of 1 would likely break the algorithm.

Author response:
➤ We have clarified in the updated pdf that Proposition 3.1 is given as a way to theoretically justify our loss, not as a proper justification of our algorithm. We have substantially changed the language in that part, and introduced a discussion on why, while not being a panacea, we believe that our mini-batch approach is particularly adapted to our goal of estimating entropic couplings (and *not* the elusive “cursed-by-dimension-true-OT-map”). 

➤  We added an analysis of the influence of the batch size in appendix E2. Conceptually it seems that a batch size of 1 would, in the balanced case, fall back on the independent sampling of flow matching.

Reviewer comment (Evaluation):
This paper lacks a quantitative evaluation regarding whether GENOT is a reasonable approach for estimating OT couplings (Only qualitative samples in Fig 1 and 4).

Author response:
➤ We believe the new experiments that we have added in Appendix E.2, paragraph 2, testing GENOT's ability to learn the known EOT coupling between Gaussians, for several dimensions $d$ and $\varepsilon$ values can provide such an answer.

Reviewer comment (Experiments):
This paper lacks a comparative evaluation in Sec 5 with previous approaches for various single-cell biology problems, such as [1, 2] (presented in Introduction Section).

Author response:
Regarding comparison with **[Schiebinger+2019]** (WOT, i.e. discrete unbalanced OT): because we perform a train-test split for the experiments, and hence it is not possible to directly compare with WOT, as WOT does not allow for out-of-sample estimation (WOT only outputs an in-sample coupling matrix, not a "Monge" map). Another contribution of WOT is the estimation of growth rates of cells, to adjust the marginal distribution. It would be straightforward to include this in the GENOT framework, but not really needed in this paper.

Regarding the comparison with **[Demetci+2022]** (SCOT, i.e. discrete Gromov-Wasserstein OT): **this is indeed the competing method we compare with**, and is denoted by GW-LR, which stands for “Gromov Wasserstein with Linear Regression”. Since SCOT is discrete, as WOT, and does not use parameterized maps, it is exclusively “in-sample”, and cannot be used “out-of-sample”. Therefore, we extend SCOT out-of-sample using simply a vector-valued linear regression. For the fused setting, which was not treated in SCOT, we do the same and denote this baseline by FGW-LR (“Fused Gromov Wasserstein with Linear Regression”) accordingly.

Reviewer comment (Theory):
Proposition 3.1 assumes $Y \sim \pi_{\epsilon}^{\star} (\cdot | x)$ in Eq 7. In practice, Algorithm 1 adopts the mini-batch estimate for $\hat{\pi}_{\epsilon}$. Can GENOT still recovers Optimal Conditional Generators by taking expectations over mini-batch estimate, e.g. [4]? In this respect, I am curious about the quantitative results for the OT coupling estimates from GENOT for Fig 1 and Fig 4.

Author response:
➤ We have changed the statement of our Proposition 3.1 in the paper, to make it clear that Proposition 31. is a theoretical justification for the loss and not supporting using mini-batch training.  We have also added a discussion of the use of mini-batches, and mentioned to other reviewers the fact that our goal to learn _entropic_ OT couplings (as opposed to unregularized) makes us fall in the better statistical regimes of entropic OT.

Reviewer comment (Reproducibility):
Which algorithms are used for this mini-batch estimate?

Author response:
➤ We briefly introduced all these solvers in the Background Sec. 2.  All our solvers are implemented in OTT-JAX, (https://ott-jax.readthedocs.io/en/latest/index.html). 

**GENOT-K** and **U-GENOT-K** use the entropic Kantorovich problem and its unbalanced counterpart, using the Sinkhorn algorithm [Cuturi+’2013] and its unbalanced generalized [Séjourné+’2023].

 **GENOT-(F)GW** and **U-GENOT-(F)GW** use entropic (Fused) Gromov-Wasserstein problem and its unbalanced counterpart, using [Peyré+’2016], and its unbalanced generalized [Séjourné+2022], with the trick outlined by [Scetbon+22] to reduce it to quadratic complexity. We use Anderson acceleration to speed up evaluations in the inner loops of Sinkhorn.

Reviewer comment (Theory):
The algorithm solves continuous entropy-regularized OT (GW) problems by distilling discrete solutions, but empirical sample-based plans are known to poorly approximate continuous solutions. The theoretical results do not clarify this issue.

Author response:
> **_“The algorithm provided by the authors solves the continuous entropy-regularized OT (GW) problem by distilling the solutions of the discrete one. This raises my major concern about the paper because plans estimated on empirical samples are known to be bad approximations of the continuous ones which was also mentioned by the authors (in Section 5.2). The provided theoretical results do not clarify this.”_**

➤ We agree. We believe, however, that two things should be considered:

First, the alternative is to use independent sampling of the marginals (corresponding to $\varepsilon = \infty$). We believe that even in high-dimensions and limited $n$, it is worth extracting more meaningful couplings from discrete solvers.

Second, the efficiency of discrete solvers is improving, notably for (F)GW solvers. Up to a few years ago, they were widely deemed cubic. There are now reliable off-the-shelf implementations, with suitable costs, that are quadratic. From a purely practical perspective, we don’t even think they should be run to convergence, since we mostly aim at recovering some stochasticity from them, in our sampling step.

We continue this discussion in the updated pdf (below Prop. 3.1)

Reviewer comment (Theory):
Proposition 1 may be misleading because achieving zero loss in equation (7) would recover the desired conditional plan only when the loss is calculated using samples from the ground truth conditional plan, not the discrete one as in the proposed algorithm. Therefore, even with a zero loss, it is unclear what is being learned.

Author response:
> **_“Furthermore, their Proposition 1 is a little bit misleading. It states that achieving zero value in the loss (7) one recovers the desired conditional plan. However, it is true when the loss is calculated using the samples from the ground truth conditional plan and not the discrete one as in the proposed algorithm.”_**

➤ We agree that Proposition 3.1 was not clear, and have rephrased its wording and its context, please see the updated pdf.

Reviewer comment (Theory):
This concern also applies to the unbalanced variant of the algorithm, with additional questions about the discrete approximation of scaling factors.

Author response:
> **_“Thus, even when the loss in this algorithm decreases to zero, it is not clear what is actually being learned. This concern extends to the unbalanced variant of the algorithm since it is based on the balanced one. Here additional questions arise about the discrete approximation of scaling factors.”_**

➤ We agree that recovering the unbalanced scalings is also yet another hurdle. Our intuition, from the Proof of Prop. 3.3, however, is that recovering a “good” scaling factor is not harder than recovering the “good” coupling.

Reviewer comment (Theory):
In the GW case, the problem is non-convex and may have multiple local minima, causing potential instability in optimization as the objective may change notably at each step, possibly making training results unreliable.

Author response:
> **“The situation is getting even more complex in the case of a GW problem, since it is non-convex and, as such, may have multiple local minimums. Thus, the optimization objective may notably change at every step depending on the calculated discrete OT plans, which may make the result of the training even more unreliable.”**

➤ We agree. We would like to add an important comment: recent theoretical works (e.g. [Rioux+23], added in the updated pdf) point out that the _larger_ entropic regularization, the “more convex” GW problems can become (to the extent that it does become convex above a certain $\varepsilon$).

We therefore think that your intuition that GW couplings cannot be used to guide FM should be nuanced by the fact that we target the _entropic_ GW coupling. In short, recovering, in the FGW case, a coupling $\pi_\varepsilon^\star$ is likely much easier and stable throughout iterations than $\pi_0^\star$, for $\varepsilon$ large enough.

Reviewer comment (Evaluation):
The evaluation is performed on low-dimensional 2D/3D experiments and several single-cell biological tasks, but does not demonstrate the nature of the learned solutions. Ground truth comparisons are done only in low dimensions, and in biological tasks, comparison is limited to marginal distributions, leaving it unclear what is learned in high dimensions.

Author response:
> **_“The provided algorithm is assessed in simulated low-dimensional 2D (3D) experiments and several real-world single-cell tasks. These experiments do not demonstrate the nature of the learned solutions. Indeed, the comparison with the ground truth plans is performed only in low dimensions, in biological one the comparison is reduced to marginal distributions.”_**

➤ We understand the concerns about the metrics in the single-cell data benchmarks. For **GENOT-K**, we follow the setup of previous neural OT papers which benchmark on single-cell data [Bunne+2022], [Uscidda+2023] which assess the learned Monge map only with respect to the fitting term. 

Yet, we would like to highlight that for **GENOT-GW** and **GENOT-FGW**, we _**do assess**_ the optimality of the learned transport plan using the FOSCTTM score (see Appendix C.2 for the definition of that metric), which involves the joint law of the coupling and not just its marginals.

Reviewer comment (Experiments):
The potential unreliability in high dimensions may impose severe risks when applying the algorithm to biological tasks. The reviewer suggests evaluating the method on a recent benchmark (Gushchin et al., 2023) for the non-GW case to better demonstrate its validity.

Author response:
> **_“Could you please evaluate your algorithm and compare with other entropy-regularized approaches in high dimensions using a recent benchmark (Gushchin et al., 2023)?”_**

➤ We agree with the reviewer that it is relevant to measure GENOT's ability to recover EOT couplings that are known analytically. Due to time constraints, we are not able to evaluate GENOT on the very recent [Guschin+’2023;a]’s benchmark. Instead, we evaluate GENOT's ability to learn the known EOT coupling between Gaussians [Janati+’2020], whose original code we use. We considered several dimensions $d$ and $\varepsilon$ values, as done in [Mokrov+2023], [Guschin+’2023;b], or [Shi’+2023]. These results paint GENOT under a favorable light.

We have **added these results to Appendix E2, paragraph**.

Reviewer comment (Reproducibility):
No code is provided for the method, which is problematic given that the biological setup seems not easy to reproduce.

Author response:
> **_“The authors do not provide code for their method which is important since the provided biological setup seems to be not easy to reproduce.”_**

➤ We have added code in the appendix. We are aware this is not ideal. We commit to implementing a readily usable GENOT implementation in a leading OT toolbox shortly, and will release code to reproduce experiments in a different, independent, repo.

Reviewer comment (Evaluation):
No clear assessment of the errors caused by discrete approximations of OT in high dimensions is provided (the reviewer's main concern).

Author response:
> **_“No clear assessment of the errors caused by discrete approximations of OT in high dimensions is provided (my main concern);”_**

➤ We feel this is a problem that all papers that use guidance from discrete solvers to explore a continuous formulation suffer from, including e.g. Genevay / Salimans early work, as well as, more recently, [Uscidda+23,Pooladian+23,Tong+23 a;b]’s approaches.

In fact, we argue in the updated pdf, below Prop. 3.1, that we suffer _less_ from that problem, because our “ground truth” is _precisely_ the _entropic_ coupling. All works above typically start their theoretical justification from $\pi_0^\star$, and use $\pi_\varepsilon^\star$ in practice. Our starting point is $\pi_\varepsilon^\star$.

Reviewer comment (Reproducibility):
No code is included in the supplementary material to reproduce the experiments.

Author response:
> **_“The authors do not provide code for their method which is important since the provided biological setup seems to be not easy to reproduce.”_**

➤ We have added code in the appendix. We are aware this is not ideal. We commit to implementing a readily usable GENOT implementation in a leading OT toolbox shortly, and will release code to reproduce experiments in a different, independent, repo.

Reviewer comment (Writing):
The significance from the biological point of view is unclear, at least to the reviewer, due to their lack of biological background.

Author response:
> ***“Not clear significance from the biological point of view (at least for me because I do not have a background in biology).”***

➤ We understand that the biological significance is hard to assess without a background in biology. Yet, an entire community in bioinformatics has been pushing for generative and foundation models in general, and neural OT, in particular. These tools are having a very large impact on very important tasks in single cell genomics, as documented by the many references we have included.

To help readers parse this, we have added additional references in the updated pdf, towards reviews or roadmap papers, in Section 5.1 and 5.2, to help the reader parse each of our experiments. 

We are open to adding more content if you feel this would help readability.

Reviewer comment (Experiments):
the authors did not investigate computational limits of the proposed approach: how do comp. efficiency and accuracy scale w.r.t. dimensionality? sample size?

Author response:
➤ We parse your question in two different blocks: compute complexity and statistical efficiency. We discuss statistical / minibatch aspects later in the response.

Just like other coupling-guided methods, estimating GENOT requires at each iteration a call to a discrete solver, to sample coupled points from two marginals. We have explored in this work (un)balanced {Kantorovich / linear, GW / quadratic, FusedGW / linear+quadratic} settings. These solvers’ complexities are all well understood, and not the crux of our contribution:

The Kantorovich problem is solved using the unbalanced [Frogner+15] formulation of the Sinkhorn algorithm *[Cuturi13]*, while the GW/fused uses *[Séjourné+22]*, itself an extension of a scheme proposed in *[Peyre+16]*

- For $n$ samples from two measures, and a cost function $c : \mathbb{R}^d \times  \mathbb{R}^d  \to \mathbb{R}$, Sinkhorn has quadratic $\mathcal{O}(n^2)$ time complexity. This does prevent using *very* large batches (above e.g. >20k). We just focused our efforts on running experiments with $n = 1024$. We could imagine exploring other low-rank alternatives, which are known to scale better, but this would change the overall messaging of our paper (since we emphasize entropy).

- For $n$ samples, the complexity of the implementations we use (e.g. in OTT-JAX) of (F)GW will depend on the intra-domain cost functions $c_\mathcal{X}$ and $c_\mathcal{Y}$, as well as extra-domain cost for fused. As shown in the first section of [Scetbon+22], this complexity is also quadratic if all of these cost functions have a low rank decomposition (which is typically the case in large data regime, where $n>d$, and $d$ is the dimensionality of points in $\mathcal{X}$ or $\mathcal{Y}$). This notably includes the inner product and the squared Euclidean costs. 

We added a very short comment in line 3 of the Algorithm on p.16. We can add a more detailed discussion in the appendix, along the lines of the paragraph above, if the reviewer thinks this would be helpful.

Reviewer comment (Reproducibility):
the code is not provided. However, the experimental protocol for working with biological data looks very complex. So it is not easy to replicate the results of the research

Author response:
➤ This is a good point. We have shared code in the appendix of the updated pdf. We are aware this is not ideal. We commit to implementing a readily usable GENOT implementation in a leading OT toolbox shortly, and will release code to reproduce experiments in a different, independent, repo.

Reviewer comment (Evaluation):
from a practical viewpoint it is difficult for a non-specialist to verify results of computational experiments in biological domain, whether they are significant or not. Moreover, in biological experiments the authors use only marginal distributions to estimate accuracy of their results. It is not clear to which OT plan the solution, delivered by the proposed algorithm, converges. Thus, it is important to estimate accuracy of the solution based on high-dimensional statistics (not only marginal distributions), although in biological use cases we may not have such possibility.

Author response:
➤ We understand that the biological significance is hard to assess without more background knowledge on these problems. We have added additional references in the updated pdf, towards reviews or roadmap papers, in Section 5.1 and 5.2, to help the reader parse each of our experiments. We are open to adding more content if you feel this would help readability.

➤ We understand the concerns about the metrics in the single-cell data benchmarks. For **GENOT-K**, we follow the setup of previous neural OT papers which benchmark on single-cell data  [Bunne+2022, Uscidda+2023] which assess the learned Monge map only with respect to the fitting term. Yet, we would like to highlight that for GENOT-GW and GENOT-FGW we do assess the optimality of the learnt transport plan: **The FOSCTTM score assesses how close a cell is to its true match** and was shown in **Figure 5** in our original submission (now Figure 6).

Reviewer comment (Experiments):
There exist several benchmarks for OT, e.g. https://arxiv.org/abs/2306.10161 (Building the Bridge of Schrödinger: A Continuous Entropic Optimal Transport Benchmark) and https://arxiv.org/abs/2106.01954 (Do Neural Optimal Transport Solvers Work? A Continuous Wasserstein-2 Benchmark). It is worth considering whether the efficiency of the proposed algorithm can be evaluated using high-dimensional test use cases from one of these benchmarks.

Author response:
➤ We agree with the reviewer that it is relevant to measure GENOT's ability to recover EOT couplings that are known analytically. Due to time constraints, we are not able to evaluate GENOT on the very recent [Guschin+’2023;a]’s benchmark. Instead, we evaluate GENOT's ability to learn the known EOT coupling between Gaussians [Janati+’2020]. We use [Janati+’2020]'s original code to create this ground-truth. We considered several dimensions $d$ and $\varepsilon$ values, as done in [Mokrov+2023], [Guschin+’2023;b], or [Shi’+2023]. Please look at these results in **Appendix E2, paragraph 2** of the updated pdf.

Note that the second benchmark that you suggested [Korotin'+2021] can only be used to evaluate models estimating deterministic OT maps (a.k.a. Monge maps) in the squared-Euclidean cost case. This is not so relevant for our probabilistic maps, in the entropy-regularized setup.

Reviewer comment (Writing):
page 30. What is "... choose alpha = 0.7m=, but ..."?

Author response:
➤Thanks for pointing out this typo.