Paper: Light Schrödinger Bridge (ICLR-2024)

Reviewer comment (Experiments):
The proposed method on image dataset requires pretrained latent space (512 dimension) that is already structurally informative.

Author response:
**(1) The proposed method on image dataset requires pretrained latent space (512 dimension) that is already structurally informative.**

We agree with the reviewer.

Reviewer comment (Novelty):
On the discussion of tractable / real-world SB given pairing, a few important references such as "Aligned SB" and "image-to-image SB" are missing.

Author response:
**(5) SB given pairing (...) references such as "Aligned SB" and "image-to-image SB" are missing**

Following your suggestion, we added the references to **Section 4** (related works) of the revised paper.

Reviewer comment (Writing):
All "Schrodinger" should be changed to Schrödinger.

Author response:
**(4) All "Schrodinger" should be changed to Schrödinger.**

Thanks for noting. Done.

Reviewer comment (Experiments):
In the single-cell experiments, the competitor results are taken from another paper [Tong et al., 2023]. As the evaluation is performed on a leave-one-domain out and training on the others, it is questionable that the training procedure was the same. I believe that some of these competitors should be reproduced by the authors to ensure that the training setting is similar, especially as the unpaired image-to-image experiment is a quantitative experiment without competitors.

Author response:
**(1) the single-cell experiments (...) relatively small dimensional setting (...) I am ready to reconsider my score if the authors can reproduce some of the competitor results and consider the highest dimensional setting of single-cell data experiments**

Following your question, we considered several new high-dimensional setups, provided the qualitative comparisons with several EOT/SB solvers and compared their training time. Please see **Q1 in the general answer** to all the reviewers.

Reviewer comment (Experiments):
The experiments were performed in a relatively small dimensional setting: the unpaired image-to-image experiments used a pre-trained feature extractor of dimension 512, the single-cell data experiment used the representation of the 5 whitened principal components (ie dimension of 5), and the EOT benchmark data were (at most) of dimension 128. Therefore, the question of the performance of the proposed method in high-dimensional data is legitimate, especially as the authors use a mixture of Gaussian representation.

Author response:
**(1) the single-cell experiments (...) relatively small dimensional setting (...) I am ready to reconsider my score if the authors can reproduce some of the competitor results and consider the highest dimensional setting of single-cell data experiments**

Following your question, we considered several new high-dimensional setups, provided the qualitative comparisons with several EOT/SB solvers and compared their training time. Please see **Q1 in the general answer** to all the reviewers.

Reviewer comment (Theory):
Little is said about the statistical estimation (with respect to the data dimension and sample size) of the proposed parametrization to the true entropic OT plan $\pi^\star$. It is known that it also suffers from the curse of dimension [2,3].

Author response:
**(2) Little is said about the statistical estimation (with respect to the data dimension and sample size) of the proposed parametrization to the true entropic OT plan (...)**

Following your question, we added a new **Appendix H** where we demonstrated that our solver experiences a usual **parametric** convergence rate in the statistical error  ($\sim \frac{1}{\sqrt{N}}+\frac{1}{\sqrt{M}}$, where $N,M$ are sample sizes), please see **Q2 in the general response** to all the reviewers. In the future, it would be interesting to further study the trade-off between the statistical and approximation errors, their dependence on $K$ (number of components) and $D$ (dimension).

Reviewer comment (Experiments):
How does your method depend on the dimensionality of the data? I think that this is a legitimate question, especially with the normalization constant being approximated by a mixture of Gaussians. I recommend adding more dimensionality in some experiments (like single-cell by considering a larger number of principal components. Maybe 5, 50, 200, 500, 1000?) to study how the proposed method performs as the dimension grows. I also recommend adding this discussion to the limitation paragraph in Appendix E.

Author response:
**(1) the single-cell experiments (...) relatively small dimensional setting (...) I am ready to reconsider my score if the authors can reproduce some of the competitor results and consider the highest dimensional setting of single-cell data experiments**

Following your question, we considered several new high-dimensional setups, provided the qualitative comparisons with several EOT/SB solvers and compared their training time. Please see **Q1 in the general answer** to all the reviewers.

Reviewer comment (Experiments):
On the single-cell experiment of the unpaired image-to-image experiment, could you show the training speed to reach convergence and the number of iterations it took? I checked the appendix and I did not find such plots. Maybe it would be interesting to compare to other OT solvers (sinkhorn or stochastic variants) on a simple problem to see the different behaviour. (I acknowledge that the Sinkhorn solvers are only usable in a discrete setting.)

Author response:
**(4) Comparison with Sinkhorn**

We have done comparison with the (discrete) Sinkhorn algorithm as a part of our additional comparison, please see again **Q2 in the general answer** to all the reviewers and Appendix C.4. We see that for small $\epsilon$ our solver yield even faster convergence than the Sinkhorn algorithm.

Reviewer comment (Evaluation):
[Single-cell experiments] I understand the motivation for using a different metric than W_1 to compare the different approaches. However as it is not the standard metric, it is hard for me to compare your method with other standard single-cell trajectory methods. As the authors-reviewer discussion has ended (due to a late rebuttal submission), I would have appreciated seeing the Wasserstein 1 metric for the different approaches. Indeed, the novel dataset was first considered by [Tong et al., 2023] where they used the W_1 distance. As the authors did not reproduce their method with the considered metric, it is hard to compare with more standard approaches on single-cell datasets like ODE/SDE-based approaches [Tong et al., 2023]. Unfortunately, two of the considered competitive methods are not standard in the single-cell trajectory literature. Therefore, I still think that the experimental section lacks reproduced competitors especially as the metrics are different.

Author response:
Dear reviewer KWE5, we would like to add some clarifications regarding new experiments that we discussed in our previous response to you. 

For the new experiments with single cell data (Section 5.3) with dimensions: 50, 100, 1000, we consider the energy distance metric instead of Wasserstein-1 used in the previous single cell setup (Appendix B). Our motivation lies in the statistical properties of estimating these metrics by samples. The Wasserstein-1 metric has a bad sample complexity $O(\frac{1}{n^{\frac{1}{D}}})$, while the energy distance (a variant of the maximum mean discrepancy metric) has a sample complexity $O(\frac{1}{\sqrt{n}})$ and allows unbiased estimates from samples [1].

[1] Genevay A. et al. Sample complexity of sinkhorn divergences //The 22nd international conference on artificial intelligence and statistics. - PMLR, 2019. - С. 1574-1583.

Reviewer comment (Novelty):
**Novelty of the paper:** it seems like the paper is just a combination of the two previous works [1, 2]. Admittedly, the authors have clearly elaborated in their paper of such cases, but I do not see the clear novelty in the methodology part. Even on the proof of the universal theorem of the Gaussian mixture model, one can see that half of it is straightforward calculation from the two aforementioned paper.

Author response:
**(1) Novelty [...]: [...] the paper is just a combination of the two previous works [1, 2]. Admittedly, the authors have clearly elaborated in their paper of such cases, but I do not see the clear novelty in the  methodology part.**

In the mentioned paper [2], the authors propose a scheme that allows to construct, for a given distribution $p_0$, a new distribution $p_1$ with a known conditional entropy regularized optimal transport (entropic OT) plan $\pi^*(y|x)$. However, the authors **did not propose any algorithm** to solve entropic OT between two given distributions $p_0$ and $p_1$.

In the other paper [1], the authors propose a new **algorithm** based on the connection between entropic OT in dual form and energy based models. Due to the use of energy-based like objective, they need to use Langevin sampling at both training and inference steps, **which is very time consuming.**

In our paper, we propose **an lightweighted and fast algorithm that allows to directly minimize the KL-divergence KL** $(\pi^*||\pi\_{\theta})$ between the parametric approximation $\pi\_{\theta}$ and the solution to the entropic OT $\pi^*$ without knowing $\pi^*$ and even without needing to have samples of it (Proposition 3.1). Furthermore, we show that by using Gaussian mixture parameterization, this KL minimization can be performed by a standard gradient descent algorithm without the need to use time-consuming Langevin dynamics. Moreover, we proved that such parameterization has a universal approximation property (Theorem 3.4) similar to the Gaussian mixture model.

We believe this is a sufficient amount of novelty to demonstrate that our work is not just a combination of prior works.

**(2) Even on the proof of the universal theorem of the Gaussian mixture model, one can see that half of it is straightforward calculation from the two aforementioned paper.**

We respectfully disagree. The only thing where the proof may intersect with the above-mentioned papers [1,2] is the usage of the duality form from weak OT. This generic known result just serves an intermediate step to guarantee the existence of the nearly-optimal continuous dual variable (stage 1 of the proof). The main (*the trickiest and the longest*) part of the proof (stage 2) is to demonstrate that the dual form can be approximated arbitrarily well with the Gaussian mixtures. This is highly **non trivial** due to the necessity to simultaneously control the approximation error for the potential and its $C$-transform. To our knowledge, this principal analysis is completely novel.

*We emphasize that, to our knowledge, our Theorem 3.4 is the first ever universal approximation result for the SB.*

Reviewer comment (Experiments):
**Unclear benefit of the LightSB solver in realistic setting:** for the single cell dataset using W1 distance as a metric, LightSB's performance leaves some gaps with the two best methods, however I do not see author's comment about this. The other benchmark on unpaired image-to-image translation is very hard to judge, as for this task there is no quantitative metric to compare with other solvers. I do not know why the authors omit unconditional image generation tasks, as this is an important and popular benchmark that has FID as a standard metric. Moreover, there have already existed results on some of the neural EOT solvers or the diffusion SB solver (using iterative proportional fitting)/diffusion SB matching (using iterative Markovian fitting), or flow matching/rectified flow in this task, so it is easy to compare with the baseline.

Author response:
**(3) Unclear benefit the of LightSB solver in realistic setting: for the single cell dataset using W1 distance as a metric, LightSB's performance leaves some gaps with the two best methods, however I do not see author's comment about this.**

The main benefit of our solver is the simplicty and its speed. First, it is as simple as the well-celebrated Gaussian mxiture for the density estimation. Second, our solver converges on CPU in 1 minute while the rest methods require longer training times. As one may see, e.g., from Tables 3 and 4, despite its simplicity, *our method already beats many existing complex approaches.* Please also see **Q1 in the general response** to all the reviewers.

**(4) The other benchmark on unpaired image-to-image translation is very hard to judge (...). I do not know why the authors omit unconditional image generation tasks (...).**

In our paper, we emphasize that we develop a solver for SB applications in moderate-dimensional spaces, not for large-scale generative modeling. Just like the basic Gaussian mixture model is **not used** for generative modeling, our solver is **not aimed** for learning large-scale generative models. Still the basic Gaussian mixture is a very easy-to-use baseline for many data analysis tasks, so we do believe that our solver will play an analogous role in the field of SB. Specifically, we already see that despite its simplicity it outperforms many methods.

Reviewer comment (Theory):
**Questionable theoretical result:** in the proof of the universality of the Gaussian mixture parameterization for SB (theorem 3.4), in the paragraph below equation (29), I fail to understand why the authors wrote  
> "Besides, it also has scalar covariances of its components because multiplier $exp(− |x_1|^2 /2\epsilon )$’s covariance is scalar itself"  
but what I understood in the paper's settings is that $x_1 \sim \pi_1$ is an unknown distribution, without assumptions on its parametric form. This is a key argument, as without it the factorization failed to be what the authors claimed in the statement of the theorem, I hope the authors could clarify this to me, otherwise it would be a hole in the proof of an important analysis of the paper.

Author response:
**(5) Questionable (...) universality of the Gaussian mixture parameterization (...) I fail to understand why the authors wrote "Besides, it also has scalar covariances of its components because multiplier $exp(-|x_1|^2/2\epsilon)$'s’s covariance's scalar itself", but what I understood in the paper's settings is that  is an unknown distribution, without assumptions on its parametric form. This is a key argument, as without it the factorization failed to be what the authors claimed in the statement of the theorem, (...).**

We think there is a misunderstanding of our Theorem 3.4 from your side and are happy to give clarifications below.

**We do not make any assumptions on the form of distributions $p_0,p_1$ and the OT plan $\pi^{*}$**. We only require $p_0,p_1$ to be supported on the compact sets. What we do is we prove the existence of a parametric approximator $v_{\theta}$ for the adjusted Schrodinger potential which yields the approximation $\pi_{\theta}\approx \pi^{*}$ of the OT plan. This **approximator has a specific parametric form**, namely, $v_{\theta}$ is an unnormalized Gaussian mixture (GM) with scalar covariances of components.

In the proof of our universal approximation theorem, we first pick a GM approximator $v_{\widetilde{\theta}}$ (eq. 28) with scalar covariance for the function $\exp(\frac{\widetilde{f}(x_1)}{\epsilon})$ which is introduced during the proof. Then we define a new function by $v_{\theta}(x_1)=v_{\widetilde{\theta}}(x_1)\cdot \exp(-\frac{\|x_1\|^{2}}{2\epsilon})$. Hence, the statement *"it ($v_{\theta}$) also has scalar covariances of its components because multiplier $exp(-|x_1|^2/2\epsilon)$'s covariance's scalar itself*" is related exclusively to the parametric approximator $v_{\theta}$, not to $p_0,p_1$ or $\pi^{*}$.

*As per your request, in the revised paper, we added the detailed derivation of this statement to the proof of this our Theorem 3.4 (Appendix A).*

**Concluding remarks**. 
We would be grateful if you could let us know if the explanations we gave have been satisfactory in addressing your concerns about our work. If so, we kindly ask that you consider increasing your rating. We are also open to discussing any other questions you may have.

**Additional references.**

[1] Petr Mokrov, Alexander Korotin, and Evgeny Burnaev. Energy-guided entropic neural optimal transport. arXiv preprint arXiv:2304.06094, 2023.

[2] Nikita Gushchin, Alexander Kolesov, Alexander Korotin, Dmitry Vetrov, and Evgeny Burnaev. Entropic neural optimal transport via diffusion processes. In Advances in Neural Information Processing Systems, 2023a.

Reviewer comment (Theory):
Equation (5) seems incomplete or not well-defined. The equation is a crucial part of the methodology, and its clarity is essential for both understanding the method and replicating the results.

Author response:
It seems like we do not completely understand your comment. Equation (5) just explains that we would like to find the best approximation (in terms of KL divergence) of optimal plan $\pi^{*}$ in a given set of distributions $\pi_{\theta}$ indexed by a parameter $\theta$. The particular parameterization is introduced in the following narrative (equations 7 and 9).

Reviewer comment (Experiments):
It is pointed regarding other SB solvers that `they expectedly require time-consuming training/inference procedures.` The authors state that their method avoids `time-consuming max-min optimization, simulation of the full process trajectories, iterative learning, and MCMC techniques,` which are commonly employed in existing solvers. While this claim holds conceptual interest, empirical evidence to substantiate this would significantly strengthen the paper.

Author response:
Following your question, we considered several new high-dimensional setups, provided the qualitative comparisons with several EOT/SB solvers and compared their training time. Please see **Q1 in the general answer** to all the reviewers.