Correlational Lagrangian Schrodinger Bridge: Learning Dynamics with Population-Level Regularization

Q. For L_ind the authors mention (l.192-194) that it is not suitable when time series data is not available and justify next introduced L as solutions that do not require individual trajectories as they get averaged. However, in section 3.4 L_ind is used in the final formulation (13). Does this imply that individual trajectory data is actually available? Are the introduced L effective by themselves when trajectory data is not available? This point seems unclear.

A. The reviewer has asked an important clarification question that is very relevant to our problem setup. We see that a confusion might be related to our presentation and thus clarify as follows. We also noticed many other questions are clarification type questions. We thank your questions that encourage us to further strengthen our presentation’s clarity and effectiveness.

• Individual trajectory data are not available but regularization is available, for instance, restricting motions of individual particles.
• $L_{ind}$ denotes individual regularization rather than trajectories. Trajectories are denoted as samples from distribution $p_{t_i, t_{i+1}}$ in formulation (3.1).
• The final formulation (13) unifies individual $L_{ind}$ and population regularization $L_{pop}$ into a general framework, weighted by hyperparameters. Details on hyperparameter tuning are provided in Appendix I.

---

Q. I also fail to understand how L_pop is an heterogeneous type regularizer. I get that it deals with a population state and not an individual action on single trajectory. However as E_\pi_t is involved in the definition isn't it still applied uniformly across the population?

A. Again, thank you for asking an important question that needs clarification as follows:

• Heterogeneity in systems suggests that states of individual particles can be very different from each other.
• Individual regularization is against heterogeneity – The same prior is uniformly applied to all individual particles
  • The prior is applied to individual trajectories first $prior(x) = h(x)$ and then rolled to all particles $E_{\pi_t} ( |\frac{d}{dt}prior(x)|^2)$.

For the points above we used gene expression as an example in Lines 79-80: “Gene expression covariance is only among a population of cells, while each individual cell has different/heterogeneous gene expressions. Uniformly restraining the states of individuals is thus oblivious to such knowledge”

• Population regularization better accommodates heterogeneity – The prior is applied on the population state:
  • The regularization of population states does NOT stress that individual states must behave similarly, even though the population state represents a summary of individuals via expectation (Analogy: a change in average velocity does not necessarily require the same change in individual velocities).
  • Specifically, the prior is first summarized on the population level $prior(x) = E_{\pi_t} (h(x))$ and then applied on the population-state trajectory $| \frac{d}{dt} prior(x) |^2$.
• For reasons above we wrote in Lines 84-86, “As the ensemble statistics of individual states, population states (i) respect the diversity (heterogeneity) of individual states, and importantly, (ii) can accommodate domain priors previously not utilized at the population level”. We also provided an illustration in Figure 1 left.

---

Q. The paper could benefit from a detailed explanation of how the final proposed formulation (13) helps with the "non convex and constrains wrt network parameters challenge" mentioned.

A. That is a good suggestion. We will strengthen the presentation with the following explanations:

• The original non-convex constrained optimization problem (6) is highly challenging to solve.
• We approximate the problem by reformulating it as an unconstrained optimization problem (12), which can be effectively handled using gradient-based approaches.
• This reformulation involves replacing the hard constraint with a penalty term in the objective function, effectively introducing it as a soft constraint.
• Such an approximation is also commonly adopted in other works, e.g. in [1] Section 3.2.

---

Q. The paper introduces multiple correlation-based Lagrangians, though only a few are used in the experiments.

A. The question might arise from some misunderstanding. We would like to clarify as follows:

• We have used all the three regularizers of covariance velocity (8), acceleration (9), and potential (12).
• As stated in the previous answer, The final formulation (13) unifies all regularizations into a general framework, weighted by hyperparameters.
• Details on hyperparameter tuning are provided in Appendix I.

---

[1] Neural Lagrangian Schrödinger Bridge: Diffusion Modeling for Population Dynamics

Q. There are some grammatical errors, such as "overwrote" in line 175, "we have" in line 319. "We achieves..." in line 299.

A. We thank the reviewer for pointing them out and we have corrected them in the updated pdf.

---

Q. A lot of variables are introduced in passing. I like the table that reminds what the difference between sub/superscript is and think this could be extended to definitions of variables xxx.

A. We sincerely thank the reviewer for the comments. We have provided an extended version of the notation Table 1 in the updated pdf.

---

Q. No code to validate.

A. We have uploaded our anonymous code in the supplementary materials.

---

Q. The presentation of the results is a little confusing. Figure 2 uses PCs for the gene expression trajectories but it would make more sense to present the way Tong et al. 2023 did in this paper in Figure 3. This would make the value of the method more immediate to computational biologists. The current method showing PCs of individual genes is unintelligible. This change would greatly help clarity for a more diverse audience. Additionally, I'd recommend using this compressed representation in the appendix. Just showing pairwise plots of seemingly random genes and expressed vs. predicted is not very helpful, unless I'm missing something related to the covariance that this method proposes. If there's something there that makes those plots more interpretable, you should highlight that in the main paper.

A. We appreciate your great comment on improving the visualization.

• We have added the new visualization as suggested in Figure 22.
• Some clarification on gene selection: They are not randomly selected. The selected visualized gene pairs are the most correlated ones in the dataset, providing a “local” view of the dynamics.
• For PC: Additionally, we offer a “global” view by visualizing the dynamics in the PC space, where the PC is computed based on all genes.
• We will follow your suggestion to add more visualizations in the revision.

---

Q. Do you use the product over the probability densities in equation 6 because of the assumption of independence? This wasn't explained well and this lack of definition makes the whole paper hard to follow since it is used extensively. it's hard for me to validate the proofs missing interpretation about what the product is that defines the correlation structure in the first place.

A. The question might arise from some misunderstanding.

• We didn’t make assumptions on independence and conduct factorization on joint distributions.
• Product is operated across variables within the same time stamp $t$, and the expectation is computed on the joint distribution $\pi_t$.
• Such a formulation is to describe the correlation across variables within the same time stamp.

Q. One thing that wasn't completely clear to me (this isn't an area I am particularly familiar with) was what the main reasons for the audience of this conference to be interested in this work are. It seems like a nice contribution to learning for particular settings, in which a population-level view is reasonable but beyond the rather specific numerical examples in a cellular context it wasn't clear to a non-expert what those contexts might be. Are the authors able to explain this a little more explicitly? Are there general settings in which the authors envisage the proposed approach being effective? Are there contexts of broad interest to the ICLR community which involve cross-sectional observations in which population-level regularization would be reasonable?

A. We are grateful to the reviewer for the useful comments. We would like to clarify for this question (and also other questions) that our focus on methodology and applications is general rather than restricted to a specific domain.

• We focus on the task of learning population dynamics from cross-sectional data, where individual trajectories cannot be tracked, and only samples of different individuals are available at various time points. In fact, the individuals across different time points are different (thus cannot be tracked) but drawn from the same time-dependent distribution.
• This study is motivated by a biological problem involving single-cell data, where measurements are destructive, making it impossible to observe the same cell at multiple time points.
• However, this setting is not limited to the biological domain; it extends to societal applications as well. For example, it can be used to model the evolution of population statistics [1] or even public opinion [2], where it is not guaranteed that data for the same individual will be available across different time periods.
• Our experiments are also not confined to biological datasets; for instance, we explore opinion depolarization in Appendix I.

---

Q. The underlying motivations wasn't articulated as clearly as it might have been in my view -- beyond leading to a more tractable optimisation problem, when/why would one choose to view the problem of inferring dynamics in this population-level way? (The fact the individuals may behave heterogeneously and one may not be able easily to relate observations to individuals does not in and of itself seem to mean that a population-level regularisation is the natural way to proceed.)

A.

• Clarification: Our end goal is still to predict the dynamics of individuals while we implement the regularization at the population level.
• Specifically, we optimize the generative model of the form of $v_t(x), \Sigma_t(x)$ which is responsible for individual evolution.
• The motivation to implement regularization at the population level is to respect the heterogeneity of individual behavior, where individual states could be very different in certain systems (e.g. cells differentiate into diverse cell types, or people’s opinions diverge into various focal points), the population state is a more robust metric to regularize.

---

Q. The numerical evaluation appears to be limited to a single setting which is quite difficult for a non-expert to interpret, although there are numerical quantifications provide which appear to demonstrate good performance in that setting. It might be more persuasive if there were numerical results illustrating the proposed approach in several settings illustrating good performance across a wider range of scenarios particularly if the paper's contribution is to be viewed as general methodology rather than an application to the setting of the numerical example. There are some additional experiments in Appendix I, but these are primarily very closely related to the original application with on exception on "opinion depolarisation" but the presentation here is so condensed that it is not completely obvious what is shown.

A.

• Our experiments encompass the biological field of single-cell analysis, the sociological study of opinion polarization, and the creation of synthetic datasets.
• We have added more background discussions for the opinion polarization experiment in Appendix I.

---

Q. There are places where more proofreading would help, e.g. in the references line 544 has `Schr"odinger'

A. We have updated the pdf to correct the typo.

Q. The authors should perhaps give some idea of the computational costs of their methods. I am also not sure why in Figure 2 the cyan dots (in the lower panel) have much smaller variance than the red ones.

A. Thank you very much for the great comments.

• The computational complexity depends on the specific form of the regularizer.
• For covariance velocity (8), it is $O(NTD^2)$, where $N$ is the number of sampled trajectories, $T$ is the number of sampled time point, and $D$ is the number of variables (dimensionality).
• The variance can be either larger or smaller depending on many reasons, as shown in more visualizations in Appendix I.

---

Q. What happens if the covariance structure undergoes a change over time? Can the method adapt to this scenario? Have you considered things like normalizing flows as a comparison?

A. This is a wonderful idea.

• The method can be incorporated into a dynamically evolving covariance structure.
• To incorporate it, we need only specify that $L_pop$​ in equation (13) is also dependent on the index $t$.
• In our paper, the normalizing flow is compared under a different name, for example, OT-Flow in Table 2.
• In the future, we can also combine normalizing flow and population regularization for a potentially more effective dynamics modeling.