Stochastic Optimal Control for Collective Variable Free Sampling of Molecular Transition Paths

Summary

The paper proposes an interesting method called PIPS for sampling transition paths between metastable states in molecular systems. These transitions are difficult to sample using standard Molecular Dynamics (MD) simulations due to high energy barriers. Traditional methods augment MD with a bias potential based on Collective Variables (CVs), but selecting appropriate CVs can be challenging and limiting for larger systems. PIPS instead does not rely on CVs and utilizes stochastic optimal control with neural network policies to efficiently learn the optimal bias potential sample molecular transitions. The method demonstrates success in generating low-energy transitions for several molecular systems.

Strengths

1. Presentation: The technical part of the paper is very well-organized and easy to follow. Every component, like MD, optimal control, and their connections, are well explained to me. 2. Originality: The paper introduces a new ML method, PIPS, for sampling molecular transition paths without the need for Collective Variables. This offers a fresh perspective on solving the problem of high energy barriers in molecular dynamics simulations via stochastic optimal control. 3. Technical soundness: The paper establishes a formal relationship between the problem of sampling molecular transition paths and the stochastic optimal control. This theoretical foundation well motivates the credibility of the proposed method. 4. Generalizability: PIPS does not depend on CVs, making it applicable to many molecular systems where traditional methods may not be suitable,w without intense requirements for domain knowledge. This generalizability is particularly valuable as molecular systems are under exploration.

Weaknesses

1. Transferability across all molecular systems. I list this critical limitation in the limitation section below. 2. ML contribution: One minor concern for accepting the paper to NeurIPS is that the paper is its contribution as an ML paper. ML technical parts in the paper, such as optimal control and practical neural network implementations, are all existing things. Therefore the ML contribution is not significant for the ML community. 3. Domain-specific impact: Continue from point 1, apart from ML contribution, I fully understand the value of the application paper for addressing a critical real-world problem. Then my second minor concern is that even the application scenario is a little niche. Topics such as MD acceleration and protein folding are definitely highly impactful, but the specific problem "Free Sampling of Molecular Transition Paths" sounds niche to me. Maybe provide a little more real-world impact and successful application examples in the paper? 4. I have several questions listed below.

Questions

1. I may miss some details: is Schrodinger Bridge Problem really related to the method? I feel like the main method is just to use SOC to solve BPTP, and Sec 3.1 is kind of related but unnecessary. If so, I think it's unnecessary to incorporate these unnecessary contents that make the paper partially not coherent. 2. For the experiments, you train each model for a single molecular structure. Is it possible to train a model across several structures and make it generalizable to even other structures? This is the true pursuit of ML4Molecule studies. Just designing a method to solve a single structure's problem will be expensive and not widely useful in practice. 3. Why OpenMM is necessary for implementation? As the probabilistic formulation shown in equation 7, everything MD step is described as precise distributions. Then you can code the MD sampling and overcome the downstream mentioned in lines 237-242.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

The main limitation in my mind is transferability to all systems: While PIPS shows promise in generating transition paths when trained on a single molecular system, is it possible to transfer the effectiveness to other or even all molecular systems? Or, for a weaker setting, transferable to arbitrary predefined initial and terminal states? If not, for even solving one structure, we need to train PIPS with extensive MD samplings, which might be too expensive.

Summary

This paper addresses the task of sampling transition paths between metastable states in molecular systems. This problem is particularly pronounced when dealing with complex systems, where traditional methods often fall short due to the need for chemical intuition and the limitations in scalability. To overcome these obstacles, the authors propose a novel approach called path Integral stochastic optimal control (PIPS). Unlike previous methods, PIPS does not rely on the explicit selection of collective variables (CVs), making it applicable to larger systems without sacrificing accuracy.

Strengths

The paper provides a novel perspective and approach to sampling transition paths by leveraging stochastic optimal control theory and has already influenced a series of recent works. In particular, phrasing the problem as Schrödinger bridge (SB) problem is interesting and creative, as it presents a new application for optimal control as well as a clever parameterization of molecular dynamics frameworks.

Weaknesses

Of course, the theoretical foundation utilized in this work is well-established in other fields. While this can be considered a weakness, I feel this work still provides a valuable contribution to the machine learning community. The experiments are conducted on three molecular systems, i.e., alanine dipeptide, polyproline, and chignolin. While rather simple systems, I believe, these are sufficient illustrative examples to serve as an evaluation of this work. Could the authors still elaborate on how and why no larger molecules are considered? For me, a core limitation is the lack of baselines, i.e., how do classical SBs or diffusion SBs such as [1] and [2] in this framework? Is there any way to generalize to unseen molecules? Are there appropriate datasets available to tackle such tasks? What would be the necessary steps to achieve something like this? [1] Tianrong Chen, Guan-Horng Liu, and Evangelos A. Theodorou. "Likelihood training of Schrödinger bridge using forward-backward SDEs theory." International Conference on Learning Representations (ICLR), 2022. [2] Valentin De Bortoli, et al. "Diffusion Schrödinger bridge with applications to score-based generative modeling." Advances in Neural Information Processing Systems 34, 2021.

Questions

The paper is very well structured and presented. No further questions on my end. I would classify myself as an informed reader and thus cannot access properly how intuitive the introduction is for readers with little background in control theory.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

4 excellent

Limitations

See weaknesses.

Summary

This manuscript addresses the molecular transition pathway sampling problem, which is a fundamental and significant topic in computational chemistry/biology research. Specifically, the authors first made connections between transition path (TP) sampling, the Schrödinger Bridge Problem (SBP), and Stochastic Optimal Control (SOC), where the theoretical formulations are clearly described to facilitate understanding. Then, they proposed PIPS, a modified version of the PICE algorithm tailored to sample molecular transition paths without explicitly specifying the Collective Variables (CVs). It has been shown that one could learn a neural network bias potential (or force) policy to find the optimal transition paths in a few molecular systems, bypassing the need to use CVs or run long-time MD simulations. I find this work interesting and enjoyed reading it. However, please see below for some of my concerns and questions.

Strengths

- The research topic (i.e., molecular transition path sampling, rare-event sampling) is quite significant and has garnered much attention over the years. This work provides a novel perspective towards more effectively and efficiently identifying molecular conformation changes between metastable states. - The paper is clearly written and the main ideas are easy to follow. - I would say this work is inspirational and has the potential to lead to more follow-up work in this direction.

Weaknesses

- The theoretical part of this work has already been developed for a while (i.e., Kappen and Ruiz, 2016). This work is an application of existing theoretical tools (with some modifications) to solve a particularly important biological problem. It is not necessarily a *weakness* yet I just put this comment here. - While making connections between TP, BSP, and SOC, the transition into adopting the PISOC (line 177) is a little bit abrupt. - The policy network u has not been discussed much, i.e., is MLP+act_fn a good universal solution to any molecular system under the PIPS framework?

Questions

- For each molecular system, a new policy $\boldsymbol{u}$ needs to be trained. Therefore, is it correct that one needs to perform some hyperparameter tuning (i.e., Algorithm 1 in Appendix B) for each new system? - Follow-up, is it possible to develop a generalized policy network across different molecular systems? - While PIPS closely follow the PICE framework, I was wondering if there are other established methods also suitable for solving the TP sampling problem? It would be great if the authors could add some discussions about it. - What is the impact of temperature on PIPS performance? In Table 1, PIPS is only run at 300K. - The reported performance used MLP + ReLU activation, what is the impact of different network architectures? Some ablation studies would be appreciated. - Line 326, while 5000 fs is much shorter compared with 0.6 us, why is the final configuration still close to the ground truth? Under the current statement, what happens between 5000 fs and 0.6 us? - Minor, line 86, Canonical *or* NVT? - Minor, line 322, "from its low energy unfolded state", do you mean the unfolded state is metastable, yet it has higher energy than the folded conformation?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors discussed PIPS limitations in the final Discussion section.

Summary

Molecular dynamics simulations yield trajectories which show how a system develops from a starting configuration to a final configuration over time. Given a system with more than one (meta-)stable state, sampling trajectories which include transitions from one (meta-)stable state to another one might occur sparsely due to high energy barriers. A solution is biasing the trajectory sampling procedure towards the ones which include transitions. While classical approaches rely on Collective Variables (CV) to bias the sampling procedure, the authors present a CV-free method by learning an MLP, which takes the system configuration as inputs and biases the trajectory sampling procedure towards those who include transitions. The authors recap that the problem of sampling these paths can be solved by stochastic optimal control theory and that related work suggested the PICE algorithm to find a way to control the path sampling. They adapt the PICE algorithm and use the adaption - called PIPS - in their work. PIPS allows to control the path sampling by introducing a learnable bias module, which is the MLP mentioned above. The authors consider two versions of PIPS by modelling a) the bias potential and b) the bias force. The authors apply their method to three different molecular systems.

Strengths

(S1) **significance**: Being able to sample trajectories of molecular systems which include transitions is a significant problem in physics. Also, exploration of CV-free methods is a significant research direction for molecular dynamics related fields. (S2) **originality**: Applying the adaption of PICE to molecular systems is novel. Also the comparison of the two versions of PIPS (a) and b)) includes some novelty. The smoothing of the loss function is introduced newly and intuitively seems to be an important detail for succesfully ML model training. (S3) **quality**: - Code is given. The illustration notebooks - e.g. ExampleAlanine.ipynb - helps to understand what was done and give insights to model training and path sampling. - The section in which connections between transition path sampling, the Schrödinger bridge problem, and stochastic control theory are summarized and described is very informative and well understandable. (S4) **clarity**: The problem setting and the introduction to molecular dynamics and potentially occurring problems in sampling transition paths under the presence of high energy barriers are written clearly and very well. Section 3 is well structured and explains very well why SOC theory is useful for the given problem setting.

Weaknesses

(W1) **significance**: - Applying ML models to molecular systems to sample rare events is already done here [1]. Notably, the input to the model described in [1] is CV-free since Cartesian coordinates of the atoms are fed into the model. - While significance is given from a physics point of view, there is hardly any core ML-related novelty. The ML-related part of the manuscript is learning the bias function with a standard MLP. (W2) **originality**: While (S2) is given, the connection between biasing path sampling and SOC is already described by Kappen and Ruiz. The adaption of PICE to derive the new method PIPS just seems to include a smoothing of the loss function. So it seems like the major part of the method is re-using PICE. (W3) **quality**: The experimental section is flawed because of missing baselines and missing error bars and statistical tests: - All performance values in Table 1 are reported without error bars. Also statistical tests are missing. - Variation across training reruns is missing - [1] is neglected - CV-based baselines are not included but would be desired since it is unclear how PIPS performs compared to methods which use CVs. (W4) **clarity**: - Sparse method description (PIPS): Having described very well where PICE comes from and that PICE is able to provide a solution to the transition path sampling problem, the manuscript lacks information about the presented method (PIPS). E.g., pseudo code for training and sampling would have been helpful. - Additional confusion comes from the fact that the authors describe first that the PICE algorithm helps to find a suitable choice of the control $\boldsymbol{u}_\theta$ but later they describe that the learnable component of PIPS is $b$ (or $\boldsymbol b$). Even though the relation between the control $\boldsymbol u$ and the bias potential is mentioned it remains unclear what the MLP exaclty models and how the model outcome exactly influnences the path sampling. - Even though some of the information might already be spread in the manuscript, the authors should think about providing (again) relevant information for PIPS in section 4. This also includes information about the exact loss function, training setup, and hyperparameter selection. - Theorem 3.2 shows that SOC solves the BPTP problem with the "standard" choice of $\varphi$ but the authors choose a different one for PIPS. The authors missed to include mathematical implication of changing $\varphi$ in terms of whether it is still guaranteed find a BPTP problem solution with SOC. - Mathematical notation and formulas: The clarity is diminished due to unclear notation, e.g., $\pi _0$ (l. 171) vs $\pi^0$ (l. 191) and $\pi _{\boldsymbol u^*}$ (l. 201) vs $\pi ^{\boldsymbol u^*}$. - Section 3 - contribution vs related work: It remains unclear which parts are taken from related work and which parts are newly developed. E.g, [2] includes already a lot of the findings presented in equations (9) - (14). Minor points: - l. 174 - typo: $\boldsymbol u(\boldsymbol x_t)$ - l. 184 - typo: (eq. (5)) [1] Sun, Lixin, et al. "Multitask machine learning of collective variables for enhanced sampling of rare events." Journal of Chemical Theory and Computation 18.4 (2022): 2341-2353. [2] Kappen, Hilbert Johan, and Hans Christian Ruiz. "Adaptive importance sampling for control and inference." Journal of Statistical Physics 162 (2016): 1244-1266.

Questions

1. The authors should include training reruns, report error bars, and perform statistical test for quantitative experiments 2. The authors should include SOTA CV-based methods and [1] as baselines and compare PIPS to them. 3. The author should include additional information about PIPS, e.g. a) what does the MLP predict exactly, b) what is the loss function, c) how does the training and sampling pipeline look like. 4. The authors should clarify unclear notation. 5. The authors should clarify which parts of section 3 are taken from related work and which parts are newly developed.

Rating

3: Reject: For instance, a paper with technical flaws, weak evaluation, inadequate reproducibility and incompletely addressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

-

Thank You for the Rebuttal

Thank you for you rebuttal. I will keep my score.