W2-3. I recommend the authors more carefully break down and explain each of the components in eq (6) more clearly.

To resolve your concern, we provide two different interpretations.

Eq (6) in our submitted manuscript is as follows:

$

F _{v _{\theta}, \tilde{v}}(X)&= \frac{1}{2}\int ^{T} _{0}\lVert v _{\theta}(X _{t})\rVert ^{2}\mathrm{d}t-\int ^{T} _{0}(v _{\theta} \cdot \tilde{v})(X _{t})\mathrm{d}t\\ &\quad -\int ^{T} _{0} v _{\theta}(X _{t}) \cdot \mathrm{d}W _{t}+ \log 1 _{\mathcal{B}}(X)

$

The first interpretation is based on Radon-Nikodym derivatives. By applying the Girsanov theorem, eq (6) can be rewritten as

$

F _{v _{\theta}, \tilde{v}}(X)&= \frac{1}{2}\int^{T} _{0}\lVert v _{\theta}(X _{t})\rVert^{2}\mathrm{d}t-\int^{T} _{0}(v _{\theta}\cdot \tilde{v})(X _{t})\mathrm{d}t \\ &\quad -\int^{T} _{0}v _{\theta}(X _{t}) \cdot \mathrm{d}W _{t}+\log 1 _{\mathcal{B}}(X)\\ &= {\log\frac{\mathrm{d}\mathbb{P} _{0}}{\mathrm{d}\mathbb{P} _{v _{\theta}}}(X)}+{\log\frac{\mathrm{d}\tilde{\mathbb{Q}}}{\mathrm{d}\mathbb{P} _{0}}(X)},

$

here, $\tilde{\mathbb{Q}}$ is the (unnormalized) target path measure such that $\frac{\mathrm{d}\tilde{\mathbb{Q}}}{\mathrm{d}\mathbb{P} _{0}}(X)=1 _{\mathcal{B}}(X)$. The first three terms of RHS in eq (6) correspond to ${\log\frac{\mathrm{d}\mathbb{P} _{0}}{\mathrm{d}\mathbb{P} _{v _{\theta}}}(X)}$, the (log-likelihood) deviation of the path measure induced by biased MD from that by unbiased MD. The last term $1 _{\mathcal{B}}(X)$ corresponds to $\frac{\mathrm{d}\tilde{\mathbb{Q}}}{\mathrm{d}\mathbb{P} _{0}}(X)$, connecting path measure induced by unbiased MD and the (unnormalized) target path measure by reweighting the path measure depending on the success in reaching target.

The second interpretation is based on policies. The first term $\frac{1}{2}\int^{T} _{0}\lVert v _{\theta}(X _{t})\rVert^{2}\mathrm{d}t$ can be interpreted as regularization term for bias force. The second and third terms, rewritten as follows

$$\int^{T} _{0}v _{\theta}(X _{t})\cdot\left(\tilde{v}(X _{t}) + \frac{\mathrm{d}W _{t}}{\mathrm{d}t}\right)\mathrm{d}t=\left<v _\theta,\tilde{v}+\frac{\mathrm{d}W _{t}}{\mathrm{d}t}\right>,$$

can be interpreted as the similarity between current policy $v _\theta$ and noised policy of $\tilde{v}$ for the reference path measure.

---

W2-4. Personally, I think the logical flow of the paper would be enhanced if the training algorithm was placed in the main text, as it concisely encapsulates the training procedure.

We agree with this suggestion and moved the training algorithm to the main text in our updated manuscript.

---

W3. When evaluating the quality of transition paths, it would be useful to compare to the ground truth transition paths obtained from e.g. long, unbiased MD simulations. For example, for chignolin, these are obtainable from reference simulations performed by D.E. Shaw.

We tried this, but concluded that the trajectories are incomparable due to different simulation settings. D.E. Shaw [1] deployed MD trajectories with intervals of 200ps, which is too coarse compared to our step size of 1fs. Moreover, we experiment on implicit solvent settings, unlike D.E.Shaw.

---

W4. A citation of Sipka, et. al. would be well-placed as they also tackle TPS with ML.

To resolve your concern, in our updated manuscript, we place Sipka et al., [3] in related works. They used differentiable biased MD simulation to train bias potential and introduced partial back-propagation and graph mini-batching techniques to resolve computational issues in differentiable simulation.

Q1. Why are MLPs chosen to parameterize the bias forces for the molecular systems? Have you tried using simple GNNs?

We used MLP for the fair comparison to PIPS [2] which also used MLP. We tried to use EGNN [4] before and found that MLP seems to outperform EGNN in small systems. For larger systems, applying GNN seems to be more promising.

Q2 Is it possible to use the replay buffer to serve as positive examples for “virtual” goal states (e.g. Hindsight Experience Replay)? I.e, even if a path does not reach the target state, it does reach some other state, and it can be treated as a positive/expert demonstration for reaching that state.

Yes, Hindsight experience replay (HER) [5] can be used to augment target states. We tried it before and found that it doesn't seem to work well given only two meta-stable states of a molecule. If you are interested in training a single model for various meta-stable states and molecular systems, HER will be an effective design choice for the replay buffer.

Q3. As far as I could tell, ground truth MD data is not used anywhere in the method. I understand that unbiased MD rarely achieves transitions, but in the situations that it does, could you leverage these “expert demonstrations” to improve your method?

Yes, we tried to pre-train the model with small amounts of expert trajectories and then apply our approaches. While it seems to accelerate convergence, it does not improve final performances in the synthetic system.

References

[1] Lindorff-Larsen, Kresten, et al. "How fast-folding proteins fold." Science 334.6055 (2011): 517-520.

[2] Holdijk, Lars, et al. "Stochastic optimal control for collective variable free sampling of molecular transition paths." Advances in Neural Information Processing Systems 36 (2024).

[3] Sipka, Martin, et al. "Differentiable simulations for enhanced sampling of rare events." International Conference on Machine Learning. PMLR, 2023.

[4] Satorras, Vıctor Garcia, Emiel Hoogeboom, and Max Welling. "E (n) equivariant graph neural networks." International conference on machine learning. PMLR, 2021.

[5] Andrychowicz, Marcin, et al. "Hindsight experience replay." Advances in neural information processing systems 30 (2017).

---

W1-1. This could also be explored theoretically: How does the variance of the reweighting procedure grow as system dimensionality grows?

We would like to point out that the log-variance divergence doesn't require importance weight due to the freedom in the reference path measure. Thus, the variance of the reweighting procedure of our method is 0 regardless of system dimensionality.

---

W2-1. Further clarification on the motivation and benefits of the decomposition (scale-based parameterization) would be helpful. A figure illustrating the effect of this decomposition (scale-based parameterization) on sampling informative trajectories would be helpful.

First, we clarify that the bias force decomposition illustrates the effect of the positive correlation of bias force from scale-based parameterization with the direction to the target position. There is no explicit decomposition of bias force in our algorithm. To avoid confusion, we removed the explanation about force decomposition in our updated manuscript.

The motivation for scale-based parameterization is the failure of bias force and bias potential parameterizations in sampling informative paths during training. The key benefit of this parameterization is the guarantee of reducing distance to the target position for every MD step as proved in Appendix A.3 of our updated manuscript. This allows bias force to sample transition paths during training in larger systems.

To resolve your concern, we illustrate the effect of scale-based parameterization in Figure 2 of our updated manuscript. The figure describes bias forces positively correlated with the target directions in positive scaling parameterization.

---

W2-2. The KL divergence loss is mentioned several times as an alternative to the proposed log variance divergence objective, but is never explicitly written. As I understand it, the only difference between the loss function (eq 10) and a KL objective is that the expectation can be evaluated w.r.t an old policy (i.e off-policy learning). I think it would help to directly contrast the two objectives, perhaps with a side-by-side comparison.

To resolve your concern, we directly compare two objectives side-by-side in the following Table C and Appendix E. The discretized form of KL divergence (or relative entropy) and proposed log-variance divergence (Eqn 10) are explicitly written as follows:

$$\mathcal{L}_\text{KL}(\theta)=\mathbb{E} _{x _{0:L} \sim p _{v _{\theta}}(x _{0:L})}\left[\log\frac{p _{0}(x _{0:L})1 _{\mathcal{B}}(x _{0:L})}{p _{v _{\theta}}(x _{0:L})}\right],$$ $$$$

As shown in the following Table C, log-variance divergence (1) reuses samples from old policies with replay buffer, (2) uses an annealing technique to explore and reach the target well, (3) detaches gradients from the SDE solver, and (4) reduces the variance of gradient estimator with learnable control variate $w$. These benefits are impossible for on-policy loss, e.g. reverse KL divergence.

Table C. Side-by-side comparison of reverse KL and log-variance divergence (eq 10).

In Appendix E of our updated manuscript, we also provide the quantitative and qualitative results to demonstrate that the proposed loss leads to more accurate and diverse transition paths than reverse KL divergence.

Time complexity of training. The training and inference time complexity of our approach is $O(NMLJ)$ and $O(NML)$, respectively. ($N$ is the number of atoms, $M$ is the number of samples, $L$ is the number of MD steps, and $J$ is the number of rollouts.) To be specific, training consists of biased MD simulations with $O(NML)$ time complexity. Given the number of samples $M$, the time complexity of one biased MD step of TPS-DPS is $O(N)$.

To justify it, we note that the biased MD step consists of three stages: (1) calculating bias force, (2) calculating OpenMM force field, and (3) integrating the biased MD. Given the number of layers, and hidden units, MLP for bias force requires $O(N)$ and the Kabsch algorithm for equivariance requires $O(N)$. Calculating the force field with cut-off and integrating MD with VVVR integrator also requires $O(N)$.

Computational cost of experiments in the paper. To measure computational cost, we consider the number of energy evaluations which is proportional to the number of rollouts and runtime per rollout in training and inference time. As shown in the following Table A, TPS-DPS requires less energy evaluations than PIPS [2] in training time, producing informative paths within several rollouts. The inference cost of our approach is proportional to Steered MD (SMD) while ML methods require training for higher performance than SMD.

Table A. Comparision of computational costs with PIPS and SMD in training and inference time. EET and EEI denote the number of energy evaluations in training and inference time, respectively. RT and RI denote runtime per rollout in training and inference time, respectively.

TPS-DPS is a computationally efficient alternative to PIPS, with significantly reduced training costs while maintaining inference costs comparable to SMD. This efficiency makes it suitable for generating paths across both small and large molecular systems.

Limitation of our method. While we have provided some promising results, we emphasize that we are not declaring any victory since the experimented proteins (<50 amino acids) are still small. Our main bottleneck is the cost of running multiple MD simulations. On a single A5000 GPU, training our method for BBL with 47 amino acids requires around a week since it requires 560.95 (sec) per rollout as shown in Table A. Training our models for much larger proteins (> 500 amino acids) would require at least 10 weeks. Moreover, we do not consider explicit solvent since this would require much more computations for calculating the force fields than implicit solvent.

Dear reviewer aeFr,

We express our deep appreciation for your time and insightful comments. In our updated manuscript, we highlight the changes in Blue.

In what follows, we address your comments one by one.

---

W1: The authors mention lack of computational budget, but the paper lacks analysis of training cost or scaling behaviour. Investigating scalability of the method through theoretical or empirical investigation on higher-dimensional systems would be valuable.

We sincerely appreciate this feedback. We mentioned the lack of a computational budget due to our personal constraint at the time of submission, however, we fully agree with your comment that the thorough investigation would be valuable.

To this end, we investigate the scaling behavior of our method to larger proteins, analyze the computational complexity of our method, and report the training and inference cost. We conclude by pointing out the limitations of our method for much larger real-world problems, which is added in the limitation section of our updated manuscript.

Scaling behavior of our method to larger proteins. We investigate scaling behavior and observed that our method succeeds even in proteins such as Trpcage, BBA, and BBL from [1] that are larger than the previously considered Chignolin by two to five times. However, even with crude approximation, our method is not applicable to medium proteins (approximately at least 1700 GPU hours for medium proteins) or explicit solvent settings without investing a significant amount of computation.

We provide the quantitative results of scaling behavior of our method to Trpcage, BBA, and BBL in Table A and Appendix D of our updated manuscript. We also provide a video of the generated trajectory on the project page.

Table A. Benchmark scores on Trpcage, BBA and BBL at 400K. All metrics are averaged over 64 paths.

Surprisingly, one can observe that our method with scale-based parameterization still outperforms the considered baselines even for larger proteins. Even when inspecting the trajectories provided in our project page, we find the motions to be quite realistic. However, we are not declaring victory, as we point out later.

Dear reviewer ALy8,

We express our deep appreciation for your time and insightful comments. In our updated manuscript, we highlight the changes in Blue.

In what follows, we address your comments one by one.

---

W1. My main concern is with the baselines included in the experiments. Presently, only one ML baseline, PIPS, is included. The authors have mentioned several other baselines in the related work section, such as Jing et al. (2024), Das et al. (2021), Hua et al. (2024), and Du et al. (2024). Including these in the experiments would be beneficial.

To resolve your main concern, in our updated manuscript, we compare with two-way shooting (non-ML baseline) and Doob's Lagrangian [1] (ML baseline) in Alanine Dipeptide at 300K. As shown in the following Table A, our approach produces more plausible transition states than baselines and finds the transition states of various potential energies, unlike Doob's Lagrangian.

Table A. Comparision with variable length two-way shooting and Doob's Lagrangian with internal coordinate and 2 mixtures in Alanine Dipeptide at 300K.

Other baselines are not comparable to our method. Jing et al [2]., generated time-coarsened transition paths with step size 10ps which is too coarsened to capture transition states while we use small steps size 1fs. They also required datasets generated by long unbiased MD simulation to train the model while we can train the model without datasets collected in advance. Das et al [3]., and Hua et al [4]., limited their evaluation to low-dimensional synthetic systems and did not consider real molecular systems.

---

Q1. In the introduction section, a reference is needed for the claim that "minimizing the KL divergence leads to mode collapse, ...".

To resolve your concern, in the introduction section of our updated manuscript, we provide the reference [5, 6] to claim that reverse KL divergence suffers from mode collapse.

---

References

[1] Du, Yuanqi, et al. "Doob's Lagrangian: A Sample-Efficient Variational Approach to Transition Path Sampling." The Thirty-eighth Annual Conference on Neural Information Processing Systems.

[2] Jing, Bowen, et al. "Generative Modeling of Molecular Dynamics Trajectories." The Thirty-eighth Annual Conference on Neural Information Processing Systems.

[3] Das, Avishek, et al. "Reinforcement learning of rare diffusive dynamics." The Journal of Chemical Physics 155.13 (2021).

[4] Hua, Xinru, et al. "Accelerated Sampling of Rare Events using a Neural Network Bias Potential." arXiv preprint arXiv:2401.06936 (2024).

[5] Vargas, Francisco, Will Sussman Grathwohl, and Arnaud Doucet. "Denoising Diffusion Samplers." The Eleventh International Conference on Learning Representations.

[6] Richter, Lorenz, and Julius Berner. "Improved sampling via learned diffusions." arXiv preprint arXiv:2307.01198 (2023).

Dear Reviewer CJA8,

We express our deep appreciation for your time and insightful comments. In our updated manuscript, we highlight the changes in Blue.

In what follows, we address your comments one by one.

---

W1. The paper's novelty is somewhat limited, primarily focusing on the use of the log-variance objective. A deeper theoretical analysis would strengthen the contribution.

We would like to gently point out that our novelty is not limited to simply using the existing log-variance divergence. We modify the original log-variance divergence [1] with learnable control variate and off-policy training to improve sample-efficiency and diversity. We also propose a novel scale-based parameterization of bias force which guarantees to decrease the distance to the target state given a small MD step size, sampling informative trajectories in large molecules.

To further alleviate your concern about the lack of theoretical analysis, in Appendix A.3, we provide a theoretical analysis of the proposed scale-based parameterization. We show that (1) bias force with scale-based parameterization is positively correlated with the direction to the target position and (2) there exists a small MD step size such that the bias force reduces the distance to the target position.

---

W2. The term "diffusion path sampler" may imply a link to diffusion models, which could mislead readers. It would be helpful to clarify if there is a relationship between the proposed method and diffusion models, as the term "diffusion" only appears twice in the manuscript.

Thank you for pointing out the possible confusion. We clarify the term diffusion path sampler in the footnote of our updated introduction as follows:

We coin our method diffusion path sampler, since it samples paths using diffusion SDE. Our notation closely aligns with diffusion samplers [2, 3] that samples from the final state (instead of the entire path) using diffusion SDE.

---

W3. More discussion of concurrent work, particularly Doob’s Lagrangian: A Sample-Efficient Variational Approach to Transition Path Sampling, would provide valuable context.

To incorporate your comment, in the related work section of our updated manuscript, we discuss more about the Doob’s Lagrangian [4] as follows:

Doob's Lagrangian considered the TPS problem as finding Doob's $h$ transformation via the Lagrangian method with boundary constraints. They parameterized marginal distribution as (mixture) Gaussian path distribution to satisfy the boundary constraints without relying on simulation in training time and sampled transition paths with the bias force derived from the Fokker-Planck equation in inference time.

In the updated experiments, we also compare with Doob's Lagrangian in Alanine Dipeptide at 300K. As shown in the following Table A, our approach produces more plausible transition states than baseline and finds the transition states of various potential energy.

Table A. Comparision with Doob's Lagrangian with internal coordinate and 2 mixtures of Gaussian in Alanine Dipeptide at 300K.

---

References

[1] Nüsken, Nikolas, and Lorenz Richter. "Solving high-dimensional Hamilton–Jacobi–Bellman PDEs using neural networks: perspectives from the theory of controlled diffusions and measures on path space." Partial differential equations and applications 2.4 (2021): 48.

[2] Zhang, Qinsheng, and Yongxin Chen. "Path Integral Sampler: A Stochastic Control Approach For Sampling." International Conference on Learning Representations.

[3] Vargas, Francisco, Will Sussman Grathwohl, and Arnaud Doucet. "Denoising Diffusion Samplers." The Eleventh International Conference on Learning Representations.

[4] Du, Yuanqi, et al. "Doob's Lagrangian: A Sample-Efficient Variational Approach to Transition Path Sampling." The Thirty-eighth Annual Conference on Neural Information Processing Systems.

Q1. What is the computational cost compared to SMD?

Time complexity. The training time complexity of our approach is $O(NMLJ)$ while the time complexity of Steered MD (SMD) is $O(NML)$. ($N$ is the number of atoms, $M$ is the number of samples, $L$ is the number of MD steps, and $J$ is the number of rollouts.) To be specific, training consists of biased MD simulations with $O(NML)$ time complexity. Given the number of samples $M$, the time complexity of one biased MD step of TPS-DPS is $O(N)$

To justify it, we note that the biased MD step consists of three stages: (1) calculating bias force, (2) calculating OpenMM force field, and (3) integrating the biased MD. Given the number of layers, and hidden units, MLP for bias force requires $O(N)$ and the Kabsch algorithm for equivariance requires $O(N)$. Calculating the force field with cut-off and integrating MD with VVVR integrator also requires $O(N)$.

Empirical cost. To measure computational cost, we consider the number of energy evaluations and runtime per rollout in training and inference time. As shown in the following Table B, the inference cost of our approach is proportional to SMD. The additional runtime in inference is from predicting bias forces and assigning them to external force objects of OpenMM Library. Note that our method requires training for higher performance than SMD.

Table B. Comparision of computational costs with SMD in training and inference time. EET and EEI denote the number of energy evaluation in training and inference time, respectively. RT and RI denote runtime per rollout in training and inference time, respectively.

In Appendix D of our updated manuscript, we further compare the number of energy evaluations and runtime with the other ML baselines.

---

Q2. What limitations prevent you from applying this method to larger systems?

Our computational limitation is MD simulation with $O(NML)$ time complexity. Given the number of samples $M=16$, $L=5000$ steps, $J=1000$ rollouts, and a single A5000 GPU, training our method for BBL protein with 47 amino acids requires around a week since it requires 560.95 (sec) per rollout. Training our models for much larger proteins (> 500 amino acids) requires at least 10 weeks.

Dear Reviewer r5xj,

We express our deep appreciation for your time and insightful comments. In our updated manuscript, we highlight the changes in Blue.

In what follows, we address your comments one by one.

---

W1. While the paper provides clear improvements over previous methods, it is not clear that the improvements enable the method to tackle the grand challenge of finding plausible transition paths for real biomolecular systems.

We agree that, at the moment, our work does not solve the grand challenge of finding plausible transition paths for real biomolecular systems, such as medium-sized proteins in explicit solvent settings with all-atom simulations. This would require significant algorithmic developments and inverestment of computational resources that are currently unavailable in most academic setting.

Nevertheless, to further emphasize the promise of our method on scaling to real-world biomolecular systems, we experiment on fast proteins of larger size in Appendix C of our updated manuscript. We experiment on Trpcage, BBA, and BBL with 20, 28, and 47 amino acids, respectively, which are significantly larger than the previously considered Chignolin with 10 amino acids. As shown in the following Table A, our approach with scale-based parameterization outperforms the considered baselines.

Table A. Benchmark scores on Trpcage, BBA and BBL at 400K. All metrics are averaged over 64 paths.

We also plot 64 sampled paths on the top two TICA components with free energy surface in Appendix C of our updated manuscript and provide 3D videos of the sampled transition paths in our updated project page. For experimental settings, refer to Appendix B.1 of our updated manuscript.

---

W2. It is difficult to draw conclusions from only 16 paths, but I certainly understand the computational limitations here.

This was not due to computational limitations but for better visualization. To resolve your concern, in Appendix C of our updated manuscript, we plot 64 sampled paths of fast-folding proteins on the top two TICA components with free energy surfaces.

---

W3. The paper could be clearer in laying out the overall problem. Without reading the previous papers in this area, it was difficult to follow. (eg. common terms like "path measure" could be more clearly defined)

To clarify the mathematical term, in the method section of our updated manuscript, we provide additional descriptions for path measure and Radon-Nikodym derivative.