Paper: Reflected Schr\"odinger Bridge for Constrained Generative Modeling (ICLR-2024)

Reviewer comment (Novelty):
This is a straight forward connection between reflected SDEs [1] and schrodinger bridge IPF [2,3], without a clear reason or motivation.

Author response:
**Straightforward connections between reflected SDEs and Schrodinger bridge (SB)**

While the SB algorithms (Valentin'21, Chen'22) have made significant strides in the study of dynamic optimal transport and generative modeling, our understanding of constrained generative modeling with non-linear transport and theoretical optimal transport guarantees remains incomplete, especially in light of real-world data often exhibiting bounded support. To address this challenge, we derive reflected forward-backward stochastic differential equations (reflected FB-SDEs) with Neumann and Robin boundary conditions, yielding a specific loss function for constrained generative modeling on arbitrary smooth domains. Furthermore, we establish connections between reflected FB-SDEs and Entropic Optimal Transport (EOT) on bounded domains. Notably, **our work opens avenues for analyzing the linear convergence of dynamic couplings on bounded domains, offering valuable insights into high-dimensional generative models.**

Reviewer comment (Evaluation):
Benefits of the method, in particular reduced number of diffusion steps or the benefit of entropic OT on constrained spaces are not really explored, it has not been made clear how this method is beneficial.

Author response:
**Empirical justification of Fewer NFEs**


We have included section D.5 in the appendix to study how optimal transport helps reduce NFEs in a simulation example. We note that optimizing the forward network plays a crucial role in training the backward network, especially when contrasted with the baseline where the forward network weights are set to 0. The well-optimized forward network leads to an improved sample quality even in cases where NFE is set to 10 and 12. We also wish to produce more examples in the final version.

Reviewer comment (Theory):
Likelihood based training of IPF appears incorrect. The first terms in algorithm 1 involve expected $\log y_T$ and $\log y_0$ starting from points from the alternate marginal distribution $x_0$ and $x_T$ respectively. My understanding from [3] is that these are the result of simulating the diffusion backwards and hence depend on the parameters of the networks. These log terms do not coincide with the log densities of the target measures unless the IPF procedure has converged, hence one cannot simply ignore them and use this as a loss without taking gradients through the diffusion simulation. I would argue this is not likelihood training and [3]'s argument is incorrect unless there's the assumption that the method has already converged, in which case there is no need for further iterations.

Author response:
**Log terms do not coincide with log densities unless the IPF procedure has converged and Why Chen'22's alternative training has no $E[\log \overleftarrow y_T]$ and $E[\log \overrightarrow y_0]$, while ours does.**

Although the likelihood in Proposition 1 is theoretically correct, we recognize the challenges in estimating $E[\log \overleftarrow y_T]$ and $E[\log \overrightarrow y_0]$ in Alg.1, particularly at non-differentiable terminal time points $t=0$ and $T$. A more pragmatic alternative involves omitting these terms following Chen'22, albeit at the cost of a slight increase in the variational gap. We have made the corresponding revision in Alg.1. We express our gratitude to reviewer vVua again for capturing the important details and improving the algorithmic demonstration.

Reviewer comment (Writing):
Section 1: Shi et al 2023 does not use IPF but IMF.

Author response:
While we maintain our belief that the GP-based SB exhibits scalability to high dimensions, we acknowledge that GP faces challenges in ultra-high-dimensional problems when compared to DNNs. To clarify, we wish to revise our statement: "Recent advances (De Bortoli et al., 2021; Chen et al., 2022b) have pushed the frontier of IPFs to ultra-high-dimensional generative models using DNN frameworks. Additionally, the GP-based SB (Vargas et al., 2021) demonstrates great promise and robustness in high-dimensional generations. Bridge matching methods (Shi et al.,
2023; Peluchetti, 2023) also offers promising alternatives for solving complex dynamic SB problems.

Reviewer comment (Experiments):
Vargas et al 2021 does not perform SB on high dimensional examples as claimed, maximum dimension is 4.

Author response:
While we maintain our belief that the GP-based SB exhibits scalability to high dimensions, we acknowledge that GP faces challenges in ultra-high-dimensional problems when compared to DNNs. To clarify, we wish to revise our statement: "Recent advances (De Bortoli et al., 2021; Chen et al., 2022b) have pushed the frontier of IPFs to ultra-high-dimensional generative models using DNN frameworks. Additionally, the GP-based SB (Vargas et al., 2021) demonstrates great promise and robustness in high-dimensional generations. Bridge matching methods (Shi et al.,
2023; Peluchetti, 2023) also offers promising alternatives for solving complex dynamic SB problems.

Reviewer comment (Novelty):
The authors mention other related work such as Riemannian SGM Bortoli et al 2022, it should be noted that even more relevant to this work is SB and hence OT has been performed on manifolds in Thornton et al Riemannian Diffusion Schrodinger Bridge, 2022.

Author response:
**Comments on the References**

We value your insightful suggestions regarding the references. In response, we have reintroduced the bridge matching method in the introduction section and incorporated the more pertinent reference [1] into our related works.

[1] Riemannian Diffusion Schrödinger Bridge.

Reviewer comment (Novelty):
My main concern is with the motivation of the paper. While it is certainly possible (as shown by the authors) to extend the Schrodinger Bridge framework to the reflected setting it is not clear why one would like to do this. I think that a contribution to a top tier venue such as ICLR requires more motivation (either theoretical, methodological or experimental), note that I am not talking about the novelty here (there is no doubt that the presented work is novel). It seems that the authors want to put forward that Schrodinger Bridge has the ability to reduce the length of the integration time in generative modeling (a reason first put forward in the work of [1]) "Notably, the forward process (1a) requires a long time T to approach the prior distribution, which inevitably leads to a slow inference". However there already exist many works focused on improving the speed of diffusion models (distillation, better samplers, etc.) which have proven to be more efficient than Schrodinger bridges.

Author response:
**Major Novelty and Motivation of the paper**

Our major novelty comes from the constrained generation on **arbitrary smooth** bounded domains with **theoretical support**. To achieve this goal, we developed the reflected Schrodinger bridge algorithm (rSB), which provides the first work of generative models that extends to arbitrary smooth geometries with optimal transport guarantees. Theoretically, we establish the elegant connections between rSB and the linear convergence of optimal transport on bounded domains. This not only enhances our understanding of the algorithm but also lays the foundation for exploring the theoretical properties of dynamic Schrödinger bridges—a domain that remains less explored. This novel linkage contributes significantly to the field, providing a deeper comprehension of the algorithm and paving the way for the study of dynamic Schrödinger bridges.

Reviewer comment (Theory):
The authors claim that Proposition 1 is a valid ELBO, however this result is only true at equilibrium. It is no longer true when the vector fields are parametric functions.

Author response:
**Proposition 1 is not a valid ELBO when vector fields are parametric functions**

We respectfully **disagree** with this argument. Firstly, we did not claim that Proposition 1 constitutes an Evidence Lower Bound (ELBO), as it is not parameterized by $\overrightarrow z_t^{\theta}$ and $\overleftarrow z_t^{\omega}$; rather, it serves as the precise likelihood /loss function. Secondly, the statement that the parametrized loss bears resemblance to the ELBO is initially posited in Theorems 1 and 2 in [1], and this statement holds true in our context as well.

[1] Song, Maximum Likelihood Training of Score-Based Diffusion Models. 2021

Reviewer comment (Evaluation):
I am not very convinced by the significance of the empirical results. As I emphasized earlier, I appreciate the fact that the authors provide an extensive image investigation, however the results are not very convincing. The quality seems to be worse than [2]. If the authors claim that the reflected SB method improves the speed of the generation then it would have been useful to compare NFE.

Author response:
**Empirical justification of Fewer NFEs**

Experimentally, we acknowledge that the class of Schrodinger bridge algorithms is expensive and may not provide compelling large-scale empirical results compared to techniques such as distillation. We have included section D.5 in the appendix to study the reduced NFEs in a simulated example, which shows that a well-optimized forward network significantly contributes to training the backward network compared to the baseline with the forward network weights being 0, which leads to an improved sample quality even in cases where NFE is set to 10 and 12. We wish to produce more examples in the final version.

Reviewer comment (Novelty):
It is not clear at all for me what are the differences between the theoretical section and the work of [3]. The work of [3] is cited but not comparison or discussion is provided.

Author response:
**Difference between the theoretical section and the work of [3]**

We want to emphasize that the convergence of optimal transport in **general domains and bounded domains are fundamentally different**. For instance, [1] delved into linear convergence on bounded domains, whereas [2] investigated sublinear convergence in general domains.

Concerning the convergence of optimal transport with practically approximated marginals, [3] explored the **primal** Iterative Proportional Fitting (IPF) in a **general** space with a **sublinear** rate. In contrast, our work examines the convergence of the **dual** IPF in a **bounded** space with a **linear** rate. It is noteworthy that our study on the reflected Schrödinger bridge also reveals the linear convergence of optimal transport on bounded domains, offering a more practical and **tangible impact and guiding the hyperparameter tuning of $\varepsilon$**.

[1] Guillaume Carlier (2022). On the Linear Convergence of the Multi-Marginal Sinkhorn Algorithm. 

[2] Ghosal and Nutz (2022) On the Convergence Rate of Sinkhorn’s Algorithm.

[3] Chen (2023) Provably Convergent Schrodinger Bridge with Applications to Probabilistic Time Series Imputation.

Reviewer comment (Novelty):
In the introduction [5] should be cited when citing [6] as the works are concurrent.

Author response:
**missing one literature on bridge matching**

Thank you for bringing attention to the noteworthy work on bridge matching studies. We have incorporated this relevant contribution into the revised version.

Reviewer comment (Experiments):
It would be very interesting to understand if the current setting can accommodate for an understanding of the dynamic thresholding procedure used in [8].

Author response:
**If the algorithm accommodates dynamic thresholding**

Yes. Similar to the study of Corollary A.11 in [1], dynamic thresholding approximates the projection operator on hypercubes given a fine enough discretization and inner pointing drifts.

[1] Lou (2023) Reflected Diffusion Models.

Reviewer comment (Novelty):
I wouldn't say that [1] doesn't necessarily propose approaches to deal with general geometries. In particular, they propose a diffeomorphic mapping scheme that should help generalization to other spaces although there are of course analytic blowups at edges/corners.

Author response:
We appreciate the reviewer for the insightful suggestions.

**Reflected diffusion model may also tackle general domains.**

We have refined our description of reflected diffusion models in the abstract and conclusion sections. For example, instead of saying "reflected diffusion models lack the flexibility to adapt to diverse domains", we wrote "reflected diffusion models may not easily adapt to diverse domains without the derivation of proper diffeomorphic mappings" to emphasize the flexible constrained generation ability when diffeomorphic mappings are available.

Reviewer comment (Experiments):
There could be more experimental results. In particular, the current results mostly show that the method works and can scale to high dimensions. Importantly, a more involved example (e.g. high dimensional simplex) would be needed to show that the method retains its niceties, as compared with [1].

Author response:
We appreciate the reviewer for the insightful suggestions.

**More experimental results, including a high dimensional simplex result.**

We have included section D.4 to study the generation in the simplex domain, which obtained a similar promising performance compared with Lou'23; we also observe that a well-optimized forward network helps in reducing the number of function evaluations in a simulation example and results in a better sample quality in section D.5