Venue: ICLR 2023 Conference

Reviews:

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This work tackles the problem of learning dynamics from population observations over time. Existing works have used Continuous normalizing flows (CNFs) or schrodinger bridges (SBs) to tackle this problem. This work (NLSB) combines adhoc regularizations in Tong et al. 2020 and more recent SDE formulations to learn a principled and regularized SDE to model cell populations and incorporate additional regularizations. The method is tested on a synthetic OU SDE and a single-cell RNA-seq dataset measuring embryoid body development over time. 
strength_and_weaknesses: Strengths:

- Clearly explained motivation formalizing previously adhoc regularizations in Tong et al. 2020
- Adds ideas from cellular modeling to SDE and Schrodinger bridge models for improved modeling of a cellular system.
- Good tradeoff of time vs. accuracy in figure 5 across dimensionalities.
- Interesting architecture adaptation for fast gradient computation.

Weaknesses:

- Limited evaluation (2 settings) with relatively minor improvements in performance.
- The network is not as easy to invert as other models which use ODE dynamics (Tong et al. 2020) or train both a forward and backwards model. I see that a reverse time simulation was performed in E.4, but it is unclear to me how well this works in practice (quantitatively). I have had numerical difficulties training this backwards network in the past.

Comments / questions:

There is something strange in these comparisons… what are the differences with NeuralSDE?

Figure 10, are the captions correct? It looks to me that (e) and (f) are switched.

“The weight coefficients $\lambda_e$ and $\lambda_h$ are tuned for each interval $[t_{k-1}, t_k]$ respectively”. Can the authors elaborate here? This seems concerning as there may be potential for data leakage.

It might be interesting to visualize the potential landscape (as is done in Bunne et al. 2022). Biologists are quite interested in what is termed “Waddington’s landscape” for developing systems. 

Would it be possible to include an additional evaluation metric (besides the 2-Wasserstein distance)? The 2-Wasserstein distance mostly cares about outliers, that the distribution is “mostly right” which might miss some of the other differences between the predicted and true distributions. Common choices I believe are 1-Wasserstein distance (Tong et al. 2020, Vargas et al. 2021) and Maximum mean discrepancy (Huguet et al. 2022). Since you directly optimize for the 2-Wasserstein distance this may be an overoptimistic metric. If this is not possible it would at least be good to note the limitations of the current evaluation.
clarity,_quality,_novelty_and_reproducibility: Writing is clear. Formulation of the regularizations and architecture are novel to my knowledge. Architecture and setup is clear and reproducible. I have some questions as to the exact hyperparameter tuning setup, including “We searched all weight coefficients of regularization terms in [0.0, 0.5]”, this can’t possibly be true right? Also is EMD-L2 the same as Wasserstein-2? these are not the same to me as Wasserstein-2 generally implies squared L2 cost. 
summary_of_the_review: This paper introduces a new architecture and training procedure for neural Lagrangian Schrodinger bridges. This provides a more principled approach to cellular systems than Tong et al., and includes additional prior information not present in the general Schrodinger bridge formulations. This is an incremental but solid contribution to the literature on modelling population dynamics with additional constraints.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The paper aims to learn population dynamics using neural stochastic differential equations.  The paper formulates the Lagrangian Schrodinger bridge (LSB) of optimal transport. Experiments are conducted on synthetic data sets and single-cell embryoid body scRNA-seq data.

strength_and_weaknesses: Strength:
- To my knowledge, the approach using the LSB is interesting and novel.
 

Weaknesses:
- The paper claims that one of the main contributions is the model architecture of potential functions. It seems that the paper adopts such an architecture from OT-Flow.
- The paper does not highlight or motivate the reason for using Lagrangian Schrodinger bridge.
- The empirical results, unfortunately, only give marginal improvements compared to models (i.e. IPML).

clarity,_quality,_novelty_and_reproducibility: The Appendix A to explain Eq. (7,8,9) is provided but is not self-contained. Some more discussions on HJB solutions (based on OT-Fflow) are needed to arrive the final formulation in Eq. (7,8,9).

There are some minor gramatical errors which can be improved.
summary_of_the_review: The paper presents an interesting approach to learn neural SDEs for the problems of population dynamic modelling. However, the paper is not well-motivated and the technical contributions seem incremental.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The authors formulated the Lagrangian Schrodinger bridge problem and proposed to solve it approximately by the advection-diffusion process with regularized neural SDE. The expensive trace computation operation was also alleviated by adopting a model architecture motivated by OT-Flow. A few experiments were conducted on population dynamics and showed the efficiency in modeling stochastic behavior.
strength_and_weaknesses: **Pros:** 

1. An interesting extension of Schr\"{o}dinger bridge by including additional regularizers, which ends up with an interesting Lagrangian Schr\"{o}dinger bridge that generalized the quadratic transport cost function.

2. avoids the expensive computation of the divergence suffered by the Hutchinson estimator; the adoption of architecture in OT-Flow is interesting.
clarity,_quality,_novelty_and_reproducibility: 
1) In page 5, when $L=\frac{1}{2} ||u||^2$, $\boldsymbol f=[\nabla_u L]^{-1} (-\nabla_x \Phi)\neq -\nabla_x \Phi$, the transition from LSB to SB fails, any comments on that?

2) Any comments on m in section B of the appendix is greatly appreciated. What is the benefit to improve $m$.

3) Could the authors draw more connections between LSB and OT-Flow? If there is indeed a close connection, the more mathematical, the better.

4) Is this algorithm appliable to generative tasks such as CIFAR10 simulation? If not, is there a scalability issue?


Minor: 
2) page 14 $\mathcal{L}(p, m, \Phi)$ and $\mathcal{L}(\Phi, p, m)$ is not consistent.
summary_of_the_review: An interesting extension of SB with speed-ups on divergence estimation.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: This paper proposes the Lagrangian Schrodinger bridge (LSB) problem, a framework to recover the population dynamics from temporal data. LSB is a stochastic optimal control problem, which enables the authors to leverage the structure and propose the training objective (7). This method is evaluated on synthetic and single-cell population datasets.
strength_and_weaknesses: Strength:
The proposed framework is quite flexible. Moreover, it is simple, and it seems to work in experiments. The application of sing-cell population dynamics is interesting.

Weaknesses:
The major issue of the paper for me is that the contributions of the submission are not clear: 
- The objective (7) is not really new. The idea of relaxing the endpoint constraint by a softer one ($\mathcal{W}_2$ in (7)) was exploited by the prior work of Lavenant et al. 2021 in a similar context, although without the neural SDE; the authors seem have missed this reference. The idea of incorporating the Hamiltonian term (9) was also present in existing literature (Yang & Karniadakis 2020). As for the Lagrangian term $\mathcal{R}_e$, the authors simply reuse the regularization of Tong et al. 2020 without designing new ones.

- There is also little contribution from the algorithmic side, since the authors essentially combined various known methods to treat each term separately in (7). 

- I'm not sure why the authors mentioned various SB related training but did not compare to them in the experiments. For instance, what is IPML (NN) in Table 1? Is it from Vargas et al. 2021 or De Bortoli et al. 2021? Is there any reason of including these methods, instead of, say, Chen et al. 2021a?

- The $\mathcal{W}_2$ in the proposed objective (7) suddenly turns into the entropy-regularized $\mathcal{W}_\epsilon$ in (11). Why not propose $\mathcal{W}_\epsilon$ to begin with?

-----
Reference
Lavenant et al. 2021, Towards a mathematical theory of trajectory inference.
clarity,_quality,_novelty_and_reproducibility: Please see the Strength and Weaknesses above. I have concerns regarding the novelty of the submission. I also think the presentation can be substantially improved to highlight the contributions.
summary_of_the_review: This paper proposed an interesting framework for population dynamics. However, in its present state, it seems like gluing together several pieces (SB, stochastic optimal control, regularization from previous work, etc.) without any core contribution. I'm therefore inclined to rejection.

I suggest the following improvements:

1. Revise the presentation to highlight the contribution, focusing especially on the benefits of the neural SDE component in (6). 

2. Expand Section 3.3 to include more examples, even designing new ones, to reflect the generality of the proposed LSB. Otherwise, if the authors merely included one previous example (the cellular system), then the applicability of LSB is not really convincing to me.

3. Since the theoretical contribution is fairly limited, the authors should perform more extensive experiments. In particular, for the single-cell population dynamics, the authors might consider including the dataset of Schiebinger et al. 2019. In addition, following the point above, experiments on more applications, instead of just synthetic and single-cell population dynamics, are recommended.
review_title: null

Meta-review: metareview:_summary,_strengths_and_weaknesses: The authors propose to model a diffusion between two known marginal distributions. They formulate the Lagrangian Schrödinger bridge (LSB) problem, a generalization of the Schrödinger bridge problem using a Lagrangian to quantify displacement costs (see Sec.3.3). They solve this process using samples through regularized neural SDEs.

The paper was initially criticized for lacking on the experimental side, and for failing in some aspects w.r.t. the accuracy of its bibliographic references. Most of these concerns have been solved following a lengthy rebuttal process. I am therefore happy to suggest acceptance, with the important caveat that the authors will need to put all of the observations in their final version, including comparisons w.r.t. other baselines that have been requested.
justification_for_why_not_higher_score: I think the paper underwent several modifications, which do not make me feel confident to support an oral, but I think the content is very timely and will be of interest to the diffusion crowd.
justification_for_why_not_lower_score: Not reject because the rebuttal convinced a very knowledgeable reviewer to switch from a low to very high (8) score. I prefer having this as a spotlight because I think the paper is great and addresses a hot topic.
Decision: Accept: notable-top-25%