Transition Path Sampling with Improved Off-Policy Training of Diffusion Path Samplers

Authors: Kiyoung Seong, Seonghyun Park, Seonghwan Kim, Woo Youn Kim, Sungsoo Ahn
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: molecular dynamics, transition path sampling

Abstract:
Understanding transition pathways between two meta-stable states of a molecular system is crucial to advance drug discovery and material design. However, unbiased molecular dynamics (MD) simulations are computationally infeasible because of the high energy barriers that separate these states. Although recent machine learning techniques are proposed to sample rare events, they are often limited to simple systems and rely on collective variables (CVs) derived from costly domain expertise. In this paper, we introduce a novel approach that trains diffusion path samplers (DPS) to address the transition path sampling (TPS) problem without requiring CVs. We reformulate the problem as an amortized sampling from the transition path distribution by minimizing the log-variance divergence between the path distribution induced by DPS and the transition path distribution. Based on the log-variance divergence, we propose learnable control variates to reduce the variance of gradient estimators and the off-policy training objective with replay buffers and simulated annealing techniques to improve sample efficiency and diversity. We also propose a scale-based equivariant parameterization of the bias forces to ensure scalability for large systems. We extensively evaluate our approach, termed TPS-DPS, on a synthetic system, small peptide, and challenging fast-folding proteins, demonstrating that it produces more realistic and diverse transition pathways than existing baselines. We also provide links to [project page](https://kiyoung98.github.io/tps-dps/) and [code](https://github.com/kiyoung98/tps-dps).

Decision: accept (poster)
Reviewer scores: 6, 8, 6, 6

Reviews and discussion:
<thread>
creation date:
1729383357227

thread_type:
official_review

<thread_object>
creator: reviewer_aefr

<content>
summary:
The authors present Transition Path Sampling with Diffusion Path Sampling (TPS-DPS) an approach to accelerate the sampling of transition pathways in molecular systems. As is somewhat standard, this is formulated as learning a potential which biases the original (Langevin) dynamics in order to facilitate the transition. The authors formulate the problem as minimizing a log-variance divergence between the target path measure (i.e transition path ensemble of unbiased MD simulations) and the path measure induced by the biased dynamics. This is claimed to be better than a KL divergence objective in terms of variance reduction. The core contribution of the paper is to formulate an off-policy, RL-style training workflow in which trajectories are sampled from the biased dynamics and added to a growing replay buffer, and  the log variance divergence is continuously minimized along with a learnable control variate. A special bias parameterization is chosen in order to promote more frequent transitions. Results are demonstrated on several model systems, including a double well potential, alanine dipeptide, the polyproline helix, and chignolin folding. The authors demonstrate that they outperform non-ML and ML baselines in terms of the diversity and quality of the sampled transition paths.

strengths:
1. The authors are tackling an important and challenging problem that has received less attention from the ML community compared to topics like machine learning force fields, etc, although that is now changing.
2. The core method is fairly simple and straightforward to implement, and the results on the model systems are strong relative to the baselines.

weaknesses:
My primary concerns center around the lack of clarity in presentation, and not considering higher-dimensional systems (other than a brief mention at the end of the paper). I will elaborate further.

Higher-Dimensional Systems:

I will preface this critique with the following disclaimer: I am aware that the systems considered in this work (especially alanine dipeptide and chignolin) are widely used as models to study transition states (including in the PIPS paper which the authors directly compare to), and I understand why the authors chose to demonstrate their methods on these systems. However, given the growing set of ML methods in this area which all benchmark on these same tasks, standout methods are those which demonstrate potential to generalize to more realistic, higher-dimensional systems. At the end of this paper, the authors acknowledge that “we did not evaluate our algorithm for even larger proteins due to the lack of computational budget. In addition, our work does not generalize across unseen pairs of meta-stable states or different molecular systems”. To me, this is not sufficient.

The authors claim that they don’t scale up due to a lack of computational budget. However, it’s hard to place this statement in context with the rest of the paper, since as far as I could tell, no training cost/runtime/FLOP analysis was provided for the existing systems. How does the training cost scale with the dimensionality of the system and/or the height of energy barriers between metastable states? Is it truly just a scale problem, or will the method fundamentally break down for larger systems due to the difficulty in sampling informative trajectories for off-policy learning? This could also be explored theoretically: how does the variance of the reweighting procedure grow as system dimensionality grows and transition states become harder to find? I would like to see an investigation of these questions on a higher-dimensional system (perhaps one of the larger fast folding proteins from Lindorf-Larssen et. al. [1]). Even if the method breaks down, a failure mode analysis would be a valuable addition to the literature on ML for transition path sampling, which has so far been characterized by many papers “declaring victory” on simple systems. This would help the community more realistically assess the promise of ML-based TPS. If this is done convincingly and well, I would be willing to reconsider my evaluation.

Lack of Clarity of Presentation

Several aspects of the paper were confusing to follow.
1. I did not understand why the decomposition of the bias force into a component aligned with the direction of the target state and a component orthogonal to it is useful (Section 3.3). Do you only retain the component which is in the direction of the target state? Further clarification on the motivation and benefits of this decomposition would be helpful. I believe removing this decomposition is tested in the Ablations, but only in terms of the final loss, not in terms of the effect it has on sampling informative trajectories during training. A figure illustrating the effect of this decomposition could make this a lot more clear.

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 (Eqn 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.

3. Quoting from the paper, "The first three terms in Equation (6) correspond to the deviation of the biased MD from the unbiased
MD integrated over the path sampled from Pv ". I interpret the first term as a regularization of the bias force to be small, which makes sense. However, the second and third terms are not clear to me. In the second term, should this be interpreted as a dot product between the forces from the old and new policy? And the interpretation for integrating the third term against the Brownian motion is not clear to me. As the quantity in Equation 6 appears many times throughout the paper, I recommend the authors more carefully break down and explain each of the components more clearly.

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 is concisely encapsulates the training procedure.

A Few Other Miscellaneous Points:

1. 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 [1]. 

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

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

2. Is it possible to use the replay buffer to serve as positive examples for “virtual” goal states (e.g. Hindsight Experience Replay [3])? 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.

3. 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?

[1] Lindorff-Larsen, K., Piana, S., Dror, R. O., & Shaw, D. E. (2011). How fast-folding proteins fold. Science, 334(6055), 517-520.

[2] Sipka, M., Dietschreit, J. C., Grajciar, L., & Gómez-Bombarelli, R. (2023, July). Differentiable simulations for enhanced sampling of rare events. In International Conference on Machine Learning (pp. 31990-32007). PMLR.

[3] Andrychowicz, M., Wolski, F., Ray, A., Schneider, J., Fong, R., Welinder, P., ... & Zaremba, W. (2017). Hindsight experience replay. Advances in neural information processing systems, 30.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730667687700

thread_type:
official_review

<thread_object>
creator: reviewer_aly8

<content>
summary:
This paper introduces an approach that trains diffusion path samplers for transition path sampling without requiring collective variables. To ensure scalability in high-dimensional tasks, the authors propose an off-policy training objective that incorporates learning control variates with replay buffers and a scale-based equivariant parameterization of bias forces. The proposed TPS-DPS method has been evaluated on a synthetic double-well potential and three peptides.

strengths:
The paper is clearly written and contributes to the important drug discovery task of understanding transition pathways between meta-stable states in molecular systems. The design of CV-free diffusion path sampler using log-variance divergence with the learnable control variate is novel to me. For the experimental part, the visualizations of the sampling paths effectively demonstrate the accuracy and diversity of the sampled transition paths. The ablation study thoroughly illustrates the effectiveness of each design component.

weaknesses:
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.

questions:
In the introduction section, a reference is needed for the claim that "minimizing the KL divergence leads to mode collapse, ...".
</content>
</thread_object>

</thread>
<thread>
creation date:
1730753605318

thread_type:
official_review

<thread_object>
creator: reviewer_r5xj

<content>
summary:
The paper proposes a set of improvements to existing data-free machine learning methods for sampling transition paths. They show that avoiding KL divergence loss, using a replay buffer, and other changes all produce positive results over previous methods on small molecular test systems.

strengths:
The lop-variance minimization is a clear improvement over KL-based minimization. Similarly, the off-policy training scheme with a replay buffer also improves things. For evaluating these contributions, figure 9 is clear and strong.

The results on all of the toy systems evaluated are strong.

weaknesses:
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. Beyond this, there is not much negative to say about the method. The results for chignolin in figure 6 are a bit lacking. It is difficult to draw conclusions from only 16 paths, but I certainly understand the computational limitations here.

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)

questions:
What is the computational cost compared to SMD?
What limitations prevent you from applying this method to larger systems?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730760741526

thread_type:
official_review

<thread_object>
creator: reviewer_cja8

<content>
summary:
This paper addresses the transition path sampling problem in molecular systems, a critical challenge in materials design and drug discovery. Traditional unbiased molecular dynamics (MD) simulations are computationally infeasible, while existing machine learning approaches often rely heavily on domain-specific knowledge. In response, the authors propose a diffusion path sampler that trains a neural network by minimizing the log-variance divergence between the path measure induced by biased MD and the target path measure.

strengths:
1. The paper is well-written and easy to follow.

2. The use of log-variance as the training objective provides a more robust gradient estimator, eliminating the need to differentiate through the SDE solver. This approach also mitigates issues with the traditional KL-divergence objective, such as mode collapse.

3. A thorough empirical study demonstrates the effectiveness of the proposed method, and the comprehensive ablation study justifies each design choice.

4. Code is provided, enhancing the reproducibility of the results

weaknesses:
1. 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.

2. 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.

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

questions:
Please see the weakness section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732337610342

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
# General response

We sincerely thank all reviewers for their valuable and insightful comments. We highlight the updates of manuscript in blue. The key updates are as follows:

### Summary of key updates
- In Appendix C and our updated [project page](https://anonymous.4open.science/w/tps-dps-0941/), we provide quantitative and qualitative experiments on three fast folding proteins larger than Chingolin: Trpcage (Tryptophan-cage), BBA (Beta-beta-alpha) and BBL (a protein domain of biotin ligase enzyme). We also provide the model parameters in the updated [code](https://anonymous.4open.science/r/tps-dps-0941/).
- In Appendix D, we analyze the time complexity of our approach. We provide the number of energy evaluations and runtime spent in training and inference.
- In Appendix A.3, we prove that the bias force of positive scaling parameterization guarantees to decrease in the distance to the target state for a small MD step size.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes a novel method for Transition Path Sampling, i.e. given the stochastic dynamic that defines the evolution of the system; the task is to sample the trajectories between endpoints (the initial and the final states of the system) that are distributed according to the corresponding measure. The main contribution of the paper is the optimization of the path measure (via the parameterized drift function) via the log-variance divergence, which (unlike the KL-divergence) allows for various Monte Carlo estimates. The proposed algorithm is supported by an empirical study.