Transition Path Sampling with Improved Off-Policy Training of Diffusion Path Samplers
Authors: Kiyoung Seong
Seonghyun Park
Seonghwan Kim
Woo Youn Kim
Sungsoo Ahn
Venue: ICLR 2025 Poster (2025)
Keywords: molecular dynamics
transition path sampling
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We introduce TPS-DPS, a novel diffusion path sampler for collective variable free transition path sampling, improved with off-policy training and scale-based equivariant parameterization for high-dimensional systems.

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

Published as a conference paper at ICLR 2025
TRANSITION PATH SAMPLING WITH IMPROVED OFF- POLICY TRAINING OF DIFFUSION PATH SAMPLERS
Kiyoung Seongi, Seonghyun Park Seonghwan Kim, Woo Youn Kim, Sungsoo Ahn
KAIST (kyseong98, hyun26, dmdt.ka0,, wocyoun, sungsoo. ahn) @kaist.ea ac.kr
ABSTRACT
Understanding transition pathways between meta-stable states of 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. 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 reduce the variance of gradient estimators and the off-policg training objective with replay buffers and simulated annealing techniques improve sample efficiency and diversity. We also propose a scale-based equivariant paramcterization of the bias forces to ensure scalability for large systems. We extensively evaluate our approach, termed TPS- DPS, on synthetic system, small peptide, and challenging fast-folding proteins, demonstrating that produces more realistic and diverse transition pathways than existing baselines. We provide links to our project page and code.
INTRODUCTION
In drug discovery and material design, it is crucial to understand the mechanisms of transitions between meta-stable states of molecular systems, such as protein folding (Salsbury Jr, 2010; Piana et al., 2012), chemical reaction (Mulholland, 2005; Ahn et al., 2019), and nucleation (Thanh et al., 2014; Beaupere et al., 2018). Molecular dynamics (MD) simulations have become one of the most widely used tools for sampling these transitions. However, sampling transition paths with unbiased MD simulations computationally expensive due to high energy barriers, which cause an exponential decay in the probability of making transition (Pechukas, 1981).
To address this problem, researchers have developed enhanced sampling approaches such as steered MD (SMD; Schlitter al., 1994; Izrailev et al., 1999), umbrella sampling (Torrie & Valleau, 1977; Kästner, 2011), meta-dynamics (Ensing et al., 2006; Branduardi et al., 2012; Bussi & Branduardi, 2015). adaptive biasing force (ABF; Comer et al., 2015), on-the-fly probability-enhancee sampling (OPES; Invernizzi & Parrinello, 2020) methods. These methods utilize bias forces to facilitate transitions across high energy barriers. They require collective variables (CVs), functions of atomic coordinates designed to capture the slow degree of freedom. Although effective for simple systems, the reliance on expensive domain knowledge limits the applicability of the methods to complex systems where CVs are less understood.
Recently, machine learning has emerged as a promising paradigm for CV-free transition path sampling (TPS) (Das et al., 2021; Lelièvre et al., 2023; Holdijk et al., 2024). The key idea is to parameterize the bias force using a neural network and train it to sample transition paths directly via the biased MD simulation. In particular, Lelièvre et al. (2023) considered reinforcement learning to sample paths
Work done in POSTECH

---

Published as a conference paper at ICLR 2025
escaping meta-stable states. Das et al. (2021); Hua et al. (2024); Holdijk al. (2024) considered TPS problem as minimizing the reverse Kullback- induced by the biased MD and the target path measure. However, minimizing the reverse KL divergence suffers from mode collapse, capturing only a subset of modes of the target distribution (Vargas et al., 2023; Richter & Berner, 2024). Furthermore, Das al. (2021); Lelièvre al. (2023); Hua et al. (2024); Holdijk et al. (2024) limited their evaluation to low-dimensional synthetic systems or small peptides. Developing machine learning algorithms that generate accurate and diverse transition pathways for complex molecular systems remains an open challenge.
(KL) divergence between the path measures
Contribution. In this work, we propose the diffussion path sampler (DPS) to solve the transition path sampling problem.' Our approach, coined TPS-DPS, (1) trains the bias force by minimizing a recently proposed log-variance divergence (Nüsken & Richter, 2021) between the path measure induced by the biased MID and the target path measure, and (2) uses scale-based parameterization of the bias force to handle large systems including fast-folding proteins. Specifically, to leverage desirable properties of the log-variance divergence, such as robustness of gradient estimator and degree of freedom in reference path measure, we propose to learn a control variate for reducing the variance of gradient estimators and employ off-policy training scheme with replay buffer and simulated annealing to improve sample efficiency and prevent the mode collapse.
We also introduce new scale-based equivariant parameterization for the bias force to frequently sample meaningful paths in training. Our key idea is to predict the atom-wise positive scaling factor of displacement from current molecular states the target meta-stable state. This guarantees the bias force to decrease the distance between them for every MD step. We also use the Kabsch algorithm (Kabsch, 1976) align the current molecular states with the target meta-stable state, guaranteeing S (3) equivariance of bias force for better generalization across the states.
We extensively evaluate our method on the synthetic double-well potential with dual channels, Alanine Dipeptide, and four fast-folding proteins: Chigmolin, Trp-cage, BBA, and BBL Lindorff-Larsee al., 2011). We compare TPS-DPS with the ML approach (PIPS; Holdijk et al., 2024), as well as classical non-ML. methods, e.g., steered MD (SMD; Schlitter et al., 1994; Izrailev et al., 1999). Our experiments demonstrate that TPS-DPS consistently generates realistic and diverse transition paths, similar to the ground truth ensemble. In addition, we do ablation studies of the proposed components.
2 RELATED WORK
Transition path sampling without ML. Metadynamics (Branduardi al., 2012), on- on-the-fly probability-enhance sampling (OPES; Invernizzi & Parrinello, 2020), adaptive biasing force (ABF; Comer et al., 2015), and steered molecular dynamics (SMD; Schlitter et al., 1994; Izrailev et al., 1999) were introduced to explore molecular conformations that are difficult to access by unbiased molecular dynamics (MD) within limited simulation times (Hénin al., 2022). However, they mostly rely on collective variables (CVs) for nigh-ddmensiioal problems and are inapplicable to systems with unknown CVs. To sample transition paths without CVs, Dellago et al. (1998) proposed shooting methods that use the Markov chain Monte Carlo (MCMC) procedure on path space. However, these methods often suffer from long mixing times in path space, which hinders the exploration of diverse transition paths. Additionally, high rejection rates can further reduce their efficiency.
Data driven MI. approaches. Recently, generative models have been used to sample new transition paths given a dataset of transition paths. Petersen et al. (2023); Triplett & Lu (2023) and Lelièvre et al. (2023) applied diffusion probabilistic models (Ho et al., 2020) and variational auto-encoders (Kingma& & Welling, 2013) for transition path sampling, respectively. However, these methods are limited to small systems and struggle to collect data using unbiased MD simulations due to the high energy barriers. Klein et al. (2024); Schreiner et al. (2024); Jing al. (2024) proposed accelerate MD using time-coarsened dynamics, but the time-coarsened dynamics cannot capture the details of the transition, e.g., the transition states. Duan et al. (2023); Kim et al. (2024) used neural networks to generate transition states of a given chemical reaction, but cannot generate transition paths.
Data free MIL approaches. Without previously collected dataset, Das et al. (2021); Lelièvre et al. (2023); Sipka et al. (2023); Hua et al. (2024); Holdijk et al. (2024) trained the bias forces to
We coin our method diffusion path sampler since it samples paths using diffusion SDE, similar to diffusion samplers (Zhang& & Chen, 2022; Vargas et al., 2023) that use diffusion SDEs for sampling the final state.
2

---

Published as a conference paper at ICLR 2025
sample transition paths using the biased MD. Lelièvre et al. (2023) used reinforcement learning to train the bias forces but focused on escaping an initial meta-stable state rather than targeting given meta-stable state. Sipka et al. (2023) used differentiable biased MD simulation to train bias potential and introduce partial back-propagation and graph mini-batching techniques to resolve computational issues in differentiable simulation. Das et al. (2021); Hua et al. (2024); Holdijk et al. (2024) considered the TPS problem as minimizing the reverse KL divergence between path distribution from biased MD and transition path distribution. Das al. (2021); Hua al. (2024) limited their evaluation to low- -ww-dimeniiona synthetic systems. In this work, we mainly compare our method with (PIPS; Holdijk al., 2024). Concurrent to our work, Du et (2024) considered the TPS problem as minimizing Doob's s Lagrangian objective with boundary constraints. They parameterized marginal distribution as (mixture) Gaussian path distribution to satisfy the boundary constraints without simulation in training time and sampled transition paths with the bias force derived from the Fokker- Planck equation in inference time.
3 TRANSITION PATH SAMPLING WITH DIFFUSION PATH SAMPLERS
In this section, we introduce our method, coined transition path sampling with diffusion path sampler (TPS-DPS). Our main idea is formulate the transition path sampling (TPS) problem as minimization of the log-variance divergence (Nüsken & Richter, 2021) between two path measures: the path measure induced by DPS and that of transition paths. Our main methodological contribution is twofold: (1) a new off- off-policy training algorithm that minimizes the Dogvvarincc divergence with the learnable control variate, replay buffer, and simulated annealing (2) a SF(33 equivariant scale-based parameterization of the bias force that has an inductive bias for dense training signals in large systems.
3.1 PROBLEM SETUP
Our goal is to sample transition paths from one meta- stable state to another meta-stable state given a molec- ular system. We provide an example of the problem for Alanine Dipeptide in Figure We view this as task to sample paths from unbiased molecular dynamics (MD) in Equation (1) conditioned on its starting and ending points of initial and target meta- stable states, respectively. To solve this task, we train the bias force parameterized by a neural network amortize the sampling procedure.
State
State
Molecular dynamics. We consider MD simulation on time interval [0, T" i.c., the motion of molecular state X, (R,, V,) R&N at time / where N is the number of atoms, R, E R3N is the dlom-wisee positions and V, E R is the atom-wise velocities. In particular, we adopt Langevin dynamics (Bussi & Parrinello, 2007) defined as the following SDE:
Figure Problem setup. The sampled tran- sition path (yellow dotted lines) from the state A the state B the free energy landscape of Alanine Dipeptide. We visualize the snap- shots (white circles) of the transition path and the transition state (white star).
dx;, -0:==(XX,,, u(X,)dt EdW,, u(X,) = (V, (U
vv.). Σ ddig (c,, 2
(5.V2rhhe))
(1)
m
m
where U, m. m, kB. X, and W, denote the potential energy function, the atom-wiss wise masses, the friction term, the Boltzmann constant, the absolute temperature, and the Brownian motion, respectively, and < IR3N is a vector of positive infinitesimal values. MD in Equation (1) induces the path measure, denoted by Po. which refers to the positive measure defined on measurable subsets of the path space C[0 T]; R6N) consisting of continuous functions
[0,T] + R The path (probability) measure Po induced by MD assigns high probability to set of the probable paths when solving MD.
Transition path sampling, One of the challenges in sampling transition paths through unbiased MD simulations is the meta-stability: state remains trapped for long time in the initial meta-stable state A C R before transitioning into a distinct meta-stable state C R3N, To capture the rare event where transition from A to B occurs, constrain paths X (Xt)osesT sampled from
3

---

Published as a conference paper at ICLR 2025
unbiased MD satisfy Ro E A, Rr B for fixed time T. Since the meta-stable state A and B are not well-specified for many molecular systems, we simplify this task by (1) fixing a local minima R. Rg of the potential energy function in the meta-stable states A, B and (2) sampling a transition path that starts from the state Ro = R and ends at the vicinity of Rg-
To be specifie, aim sample from the target path measure Q. which is obtained by reweighting the path measure Po with the (normalized) indicator function. The indicator function assigns zero weight paths that do not reach the vicinity of the target position Rg. Formally, the reweighting function is called the Radon-Nikodyn derivative defined as follows:
dQ
(
1s( (x( Z
1 1x(X)
{o
if ll Rs Rrll 6 7 = Br [1b(X))
dIF 0
otherwise,
where denotes group action associated with the SE(3) space and pr Rg is the aligned target position by the optimal roto-translation pr SE) minimize its Euclidean distance Rr, i.e., Pr = argmin argmin SE)(3) Rtr - P Rs|l- Such a transformation can be obtained from the Kabsch algorithm in ((N) complexity (Kabsch, 1976).
Note that one may consider naive rejection sampling to sample transition paths, based on running unbiased MD to sample a path X from the path measure Po and accepting if the path X arrives at the neighborhood of the position Rg with the radius 6. However, this method struggles to sample transition paths of systems with high energy barriers since the sampled path by unbiased MD rarely reaches the target states, i.e., the rejection ratio too high.
3.2 VAR-VARIANCE MINIMIZATION
In this section, we propose our algorithm to amortize transition path sampling. Our key idea is to train a neural network to induce a path measure that matches the target path measure Q, using the oog-variecce divergence between the path measures. We propose the log-variance divergence based on learning the control variate of its gradient and a replay buffer to improve sample efficiency and diversity.
training scheme to minimize
Amortizing transition path sampling with log-variance divergence. To match the target path measure Q we consider biased MD defined by policy v (or bias force b) the following SDE:
(0. b b(x))
dx, (u (X) (x(X.))))))) EdW,, v v(X.) =-1
m
We also let F. denote the path measure induced by the SDE. To amortize transition path sampling, we match the path measure Pe,, of a parameterized policy vo with the target path measure Q by minimizing the hor-variance divergence:
log dQ
(105 [ogg
dQ dIp Do
dQ dIf
1)
Div (P on [Q) Vp
Ep
ES
log
(4)
dPP
where P is arbitrary reference path measure with Ez [log(dQ/dIe m..)] variance divergence in detail, we let IP = I for some policy ö and apply the Girsanov's 's theorem to Equation (4), deriving the following formulation:
co. To express the log-
DES (Pe,,
[[() = Er, [(For. - Es, (Fo..)))))))
(5)
Fo.. (X) = _ |000 (x (x) f- (vx))) /' no (x)) dW, log (X).
(6)
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 IP.. The last term reweights the unbiased MD the target path measure Q. As a result, minimizing Equation (5) could be thought as minimizing the variation between P., and Q. We provide the full derivation in Appendix A. Compared KL divergence, the lor-variance divergence provides a robust gradient estimator and avoids differentiating through the SDE solver (Richter et al., 2020; Nüsken & Richter, 2021).
Minimizing with learnable control variate. To minimize the log-variance divergence, we consider the following loss that replaces the estimation of Er.,, [Fr.,.l by learning scalar parameter w::
L(0, w) Er,, [(Fro.us w) )
4

---

Published as a conference paper at ICLR 2025
Algorithm 1 Training
1: Initialize an empty replay buffer D, an policy ug. a scalar parameter w. the number of rollout and training per rollout J, and an annealing schedule Aatan
X Aend.
2: for i = I,. do
3: Generate M paths (=04)) Jm=1 from the biased MD simulations with by at temperature X
(m))
4: Update the replay buffer D + DU ( (m)
Jm=
5: for j = J.....JJJ do
6: Sample K data f0 R=1 from D.
7: Update 0 and with the gradient of _ Ek
(log
pa(mn 2. (1 (1 W
Pug (20
8: end for
9:
end for
where w is a control variate that controls the variance of the gradient estimator of VoC(o, w) without changing the gradient. Note that we set i uo in Equation (6), which implies that the gradient of Equation (7) coincides with the KI. divergence (Richter et al., 2020; Nüsken & Richter, 2021). When optimized, the control variate w estimates the expectation Ep., (L.g.to) since argmin,,. C( w) Er., vol. Thus, jointly optimizing (0, w) with the gradient step can be interpreted as jointly minimizing log- agrvaciaece divergence and estimating Ep.., (Fee.o)) utilizing w.
Off-policy training with replay buffer and simulated annealing, To leverage the degree of freedom in reference path measure for the log variance divergence, we allow discrepancy between reference path measure and current path measure, called off aff-policy training, which is widely used in discrete- time reinforcement learning (Mnih et al., 2013; Bengio et al., 2021). For the sample efficiency, we reuse the samples with a replay buffer DD which stores path samples from the path measure Pe, associated with previous policies ug. Our modified loss function _ with D is defined as follows:
[P w) = x)=p..) - w)
(8)
The replay buffer also prevents mode collapse, using diverse paths from different path measures. Additionally, in line with other off-policy training algorithms (Malkin al., 2022; Kim et al., 2023), we utilize the simulated annealing technique to sample diverse paths that cross high-energy barriers.
Discretization. To implement the algorithm, we discretize Equation (8). Given discretization step size At, we consider the discretized paths To.L = (æo, "i....,,.1 xt.) of X from MD simulations where L T/At and æ; = X ((At). In discrete cases, the discretized paths æo.L from previous policies uy and their (gradient-detanl detached) policy values (vg (Zo); ..., uy (x))) are used to approximate the value Fu..up (X) in Equation (6) follows:
L-1
L Σ 2=0
Fen.op (Zo:L)
Σ
1o (x))
2
(vo 2y) (Ix)
E vx(x)) et log 1 (x(oL)), (9)
= 2
8=
where the noise c2 = -- of the Langevin dynamics with policy vg. For implementation, we further derive a simple discretized loss of Equation (8) from Equation (9) as follows:
(u(xz) + Eng ( x At) is the discretized Brownian motions
(10g Po (
27
E LDD
w) )
(10)
Pep (wo.L)
where the buffer D stores paths To.L sampled from the previous policies, and Po and Pe, denote discrete time transition probability induced by Equations (1) and (3), respectively. We provide formal derivation of the discretized loss Appendix A.2. Note that the same objective was derived in the name of relative trajectory balance by Venkatraman et al. (2024).
We describe our training algorithm Algorithm 1. Overall, our off-policy training algorithm iterates through four steps: (1) sampling paths from the biased MD simulation with current policy de at high temperature, (2) storing sampled paths in the replay buffer D, (3) sampling batch of the paths from the replay buffer, and (4) training current policy Vn by minimizing the loss in Equation (10). After training, the biased MD simulation can directly sample transition paths from the target path measure.
5

---

Published as a conference paper at ICLR 2025
3.3 PARAMETERIZATION FOR LARGE SYSTEMS
In this section, we introduce new parameterizations of the bias force and the indicator function to sample transition paths of large systems. Our parameterization is designed to alleviate the problem of sparse training signal, where the model struggles to collect meaningful paths that end at the vicinity of the target meta-stable state in training. This problem is especially severe in large systems.
Bias force parameterization. To frequently sample the meaningful paths, aim parameterize the bias force which guarantees to reduce the distance between the current molecular state and the target meta-stable state for every MD step. This is achieved by predicting the atom-wise positive scaling factor of the direction to the aligned target meta-stable state rather than predicting force or potential directly. Moreover, current molecular state input X,, aligning with the symmetry of the transition path sampling problem for better generalization and parameter efficiency.
design the bias force satisfy roto-tran ito-trraslational equivariance the
To be specific, we achieve the inductive bias for dense training signals and the SF(3) equivariance to the current molecular state input X, by parameterizing the bias force follows:
b xx)) diag(so(p5* X.) Rg - R₂),
(11)
where (.)+) E &R is neural network constrained to have positive output elements and predicts atom-wise scaling factors and the optimal roto-translation Pr argmin
R2 - p R3 which aligns Rg with RG, as in Equation (2). We note that the bias force (divided by atom-wise masses) is positively correlated with the direction to target state, i.e., (b( (b(xx)/m /)mm (p: Rs -RR)) > 0.
SE))))) argmin
To formalize the benefit of the positive correlation between the bias force and the direction to the target state, one can prove that there always exists a small enough step size At that decreases the distance between the current state R, and the aligned target state Pi+Ae" RB, i.e.
ll''+a: Rg R'+sell < llpe Rs R.
(12)
where R = R. + b(X, At /// is the position updated by the bias force with step size At and Pl+ = aigminpeesEnn A - p- -.Rlll We provide the proof of Equation (12) in Appendix A.3.
In the experiments, we also consider other equivariant parameterizations that are less constrained: (1) directly predicting the equivariant bias force b(X,) = Pt bo(p,' .X,) E IRN and (2) predicting the invariant bias potential (p X,) E TR to obtain the bias force b (X,)) = (bg X,) E R3N We observe these two parameterizations be useful for low-dimensional tasks but struggle to produce meaningful paths in large systems during training. As shown in Figure 2, bias forces with the positive scaling parameterization are positively correlated with the direction to the target position regardless of network parameters, unlike direct force parameterizations.
Indicator function parameterization. We propose to relax the indicator function 1s into radial ba- sis function (RBF) kernel 1 T (X) A<RT,P"
Rs; which measures the similarity between two positions where a controls the degree of relax- ation. The range of RBF kernel A is bounded by the interval (0, 1] so that log (xx is well-defined and (xx represents the binary indicator function smoothly. To capture a high training signal from subtrajectories of sampled paths, we propose to take maximum over RBF kernel values of all intermediate states by mm (X) = max4e (0,17 (R,, Rg; o²). To extract the subtrajectory with a high training signal, we can truncate the paths at the time that maximizes RBF kernel values. allowing variable path lengths. Notably, the relaxed indicator function is SE(3) in- variant to R. because of the Kabsch algorithm.
(a) Direct prediction
(b) Positive scaling
Figure 2: Visualization of the bias force
fields of two different bias force parameter- izations with initialized neural networks. (a) directly predicting the bias force and (b) predicting the positive scaling factors of di- rection to the target position (white circle).

---

Published as a conference paper at ICLR 2025
Table I: Benchmark scores on the double-well system and Alanine Dipeptide. All metries are averaged over 1024 paths for the double-well system, and 64 paths for Alanine Dipeptide. ETS is computed for paths that hit the target meta-stable state, and the best results are highlighted in bold. Predicting the bias force, potential, and atom-wis wise positive scaling are denoted by (F), (P), and (S), respectively. UMD (A) denotes unbiased MD with temperature À and SMD (k) denotes steered MD with the force constant k: Unless otherwise specified, paths are generated by MD simulation at 1200K for double-well and 300K for Alanine Dipeptide.
RMSD (L)
THP (t) %
ETS (L) IJmmol
RMSD (L)
THP (f)
ETS (L)
Method
Method
%
kJmoo
Double-well
Alanine Dipeptide
(1200K)
2.21
0.00
(300K)
0.00
UMD (2400K)
2,11 + 0.38
3,03
1.69
0.31
UMD (3600K)
1.19 + 0.32
6.25
812.47 + 148.80
UMD (3600K)
1,85
12.60
2.12
(10)
0.21
7.81
33.15 + 78.40 + 12.76
UMD (4800K)
1.54 0.81
21.58
2.77
0.69
SMD (20)
0.56 +-0.27
54.69
SMD (0.5)
0.98 0.90
$2.15
1.54
0.21
PIPS (F)
0.66 : 0.15
43.75
28.17 + 10.86
SMD
0.14 0.08
99.80
1.85
0.16
PIPS
1,66 0.03
0.00
TPS-DPS (F, Ours)
0.01 0,02
99.90
1.38 + 0.16
TPS-DPS (F, Ours)
0.16 + 0.06
92.19
19.82 + 15.88
TPS-DPS (P. Ours)
0.01 0.03
99.71
1.36 = 0.15
TPS-DPS (P. Ours)
0.16 0.10
87.50
18.37 + 10.86
TPS-DPS (S, Ours)
0.01 0.03
99.80
1.73 0.20
TPS-DPS (S, Ours)
0,25 + 0.20
76.00
22.79 = 13.57
U
bp
U + by
5M0
Energy Ikymoo
(a) Potential energy landscapes
(b) Distributions of the energy and
Figure 3: Visualization of potential energy landscapes and distributions double-well potential. (a) Visualization of the learned bias potential by of TPS-DPS (P). (b) Distributions of the potential energy and y coordinate of transition states from 1024 transition paths sampled by each method.
EXPERIMENT
In this section, we compare method, termed TPS-DPS, with both classical non-ML and ML approaches, assessing the accuracy and diversity of sampled transition paths. We begin with the double-well system and Alanine Dipeptide followed by the four fast-folding proteins: Chignolin, Trp-cage, BBA, and BBL. Additionally, we conduct ablation studies to validate the effectiveness of each component in method. All real-world molecular systems
simulated using the OpenMM library (Eastman et al., 2023). Details on OpenMM simulation and model configurations are provided in Appendices B. and B.2, respectively. Appendix C, we analyze the time complexity of TPS-DPS and evaluate the number of energy evaluations and runtime in training and inference time.
Evaluation Metrics. We consider three metrics to evaluate models: (Kabsch) RMSD, THP, and ETS. The root mean square distance (RMSD) measures the ability of the model to produce final positions of paths close to the target position Rg. The target hit percentage (THP) measures the ability of the model to produce final positions of paths that successfully arrive at the target meta-stable state B. Finally, the energy of the transition state (ETS) measures the ability of the model to identify probable transition states. For further details, refer to Appendix B.3.
Baselines. We compare TPS-DPS with both non-ML and ML baselines. For non-ML baselines, we consider unbiased MD (UMD) with various temperatures and steered MD (SMD; Schlitter et al., 1994; Izrailev et al., 1999) with various force constants and collective variables (CVs). For ML baselines, consider CV-free transition path sampling method, path integral path sampling (PIPS; Holdijk et al., 2024) which also trains a bias force by minimizing the KL divergence between path measures induced by the biased MD and the target path measure.
7

---

Published as a conference paper at ICLR 2025
UMD
SMD (0.5)
SMD (1)
IPS-DPS (F)
TPS-DPS (P)
Ground Truth
UMD
SMD (10)
SMD (20)
PIPS (F)
TPS-DPP SPS-DPSIIN (F) TPS-DPS (P)
Figure 4: Visualization of sampled paths on energy landscapes. For the double-well system, we aim sample transition paths from the left meta-stable state the right on the potential energy landscape (top). For Alanine Dipeptide, we aim to sample conformational changes from the C5 (upper left) to the C7aa (lower right) on the Ramachandran plot (bottom). White circles and stars indicate meta-stale stable states and saddle points, respectively.
4.1 DOUBLE-WELD WELL SYSTEM
We begin by evaluating our method on a synthetic two-dimensional double-well system with dual channels. This system has two global minima representing the meta-stable states and two reaction pathways with the corresponding saddle points. We aim to sample transition paths from the left state R. to the right meta-stable states B = ( R Rg < 0.5). We collect ground truth path ensembles by rejection sampling which proposes paths sampled from the unbiased MD simulations and accepts the final states are in the target meta-stable states B. We provide more details on the double-well system in Appendix B.4.
As shown in Table 1, Figure 3, and Figure 4, TPS-DPS outperforms baselines regardless of bias force parameterizations and generates more similar transition paths to the ground truth than baselines. In Figure 3, the bias potential accelerates the transition by increasing the potential energy near the initial meta-stable state while decreasing the potential energy the two energy barriers. Moreover, the distribution of energy and y coordinate of the transition states from TPS-DPS is closest the ground truth compared with other baselines, successfully capturing two reaction channels. In Figure 4, TPS-DPS (F) and (P) generate similar transition paths to the ground truth while UMD at 1200K fails escape the initial state and SMD struggles pass the saddle points.
4.2 ALANINE DIPEPTIDE
For real-world molecules, we first consider Alanine Dipeptide consisting of two alanine residues and aim sample conformational changes from the meta-stable state C5 C'7ax as shown in Figure 4. The target meta-stable states are defined as B = (R (I)( E( (Rs) < 0.75]. where
{(R)= (d) v a well-known collective variable which consists of two backbone dihedral angles. Alanine Dipeptide has two reaction channels with the corresponding saddle points.
In Table and Figure 4, our method shows superior performance regardless of the bias force parameterizations and successfully generates diverse transition paths that capture two reaction channels. Compared to our method, UMD at 300K fails to escape the initial state, SMD with the two backbone torsion CV generates transition paths with less probable transition states, and two-way shooting struggles to find plausible transition states. PIPS generates transition paths with only one reaction channel, suffering from mode collapse.
8

---

Published as a conference paper at ICLR 2025
Table 2: Benchmark scores on fast-folding proteins. All metries are averaged over 64 paths. Unless otherwise specified, paths are generated at 300K for Chignolin and 400K for the others.
RMSD (L)
THP <t>
ETS (1)
RMSD (.)
THP (
ETS (
Method
Method
kJmol 1
Chignolin
Tip-cage
(300K)
7,98
(400K)
UMD (1200K)
7.23 + 0.93
1.56
388 17
UMD (1 (000))
8.27 2: 1.13
0.00
SMD (10k)
1,26 1: 0.31
6.25
-527.95 = 93.58
SMD (10K)
1.88 1: 0.23
3,12
-312.54 1: 20.67
SMD (206)
138.87
(20k)
PIPS (F)
4,66 0.17
0.00
PIPS (F)
7.47 0.11
0.00
PIPS
4.67 0.32
PIPS (P)
6.07 4-.26
0.00
TPS-DPS (F. Ours)
4,41
TPS-DPS (P. Ours)
6,35 0.31
0.00
IPS-DPS IP. Ours)
3,87 1: 0.42
0.00
TPS- DPS IP. Ours)
3.15 1: 0.52
12,50
-512.97 1: 56.89
TPS-DPS Outs)
1.17
-780.18
TPS-DPS (S, Ours)
-317.61 140.89
BBA
BBL
UMD (400K)
10.03 1: 0.39
UMD (400K)
18.48 0.63
0.00
(120K)
(1200K)
18,90 1.16
SMD (10k)
2,89 +: 0.32
0.00
SMD (10k)
3,67 +.2-2:0.22
0.00
SMD
1,66 0.30
-3104.95 97.57
2.33
-1738.57 386811
PIPS (F)
9.84 0.18
0.00
PIPS (F)
17,92
0.00
PIPS (P)
9.09 1: 0.36
0.00
PIPS (P)
12.67 1:07:::00 0.331
0.00
9.45
TPS-DPS (E. Ous))
10.15
TPS-DPS (P. Ours)
3,89 +: 0.35
0.00
TPS- DPS (P. Ours)
6.45 0.26
0.00
TPS-DPS (S. Ours)
1,21 + 0.09
84,38
-3801.68
139,38
TPS-DPS iS, Ours)
1.60 :: 0.19
43.75
-3616.32 213.66
66
0.0
0.9
3.0
4.0
4.5
Potential energy
Distance (Rad)
2.4-
28
00
03
03
5.0
(a) Alanine Dipeptide
Potential energy
31000
-50
Distance (A)
33.0
AGROUD TIREOO
Apran THABO
4.5
0.9
THESOC
RCEO
1250
2.5-
TC1
(b) Chignolin
(c) Visualization of Chignolin
Figure 5: Qualitative evaluation on transition path sampled from TPS-DPS. (a) Potential energy
and the two backbone dihedral angle distances between the current and target states. (b) Potential energy and the two hydrogen bond distances between the current and the target state. (c) Visualization of hydrogen bond formation in Chignolin. We highlight each hydrogen bond in green and yellow.
4.3 FAST-FOLDING PROTEINS
Finally, for challenging molecules, we consider four fast-folding proteins (Lindorff- et al., 2011): Chignolin, Trp-cage, BBA, and BBL. with 10, 20, 28, and 47 amino acids, respectively. Proteins fold into stable structures by forming networks of hydrogen bonds. We aim to sample protein folding processes as seen in Figure 7. We define the target meta-stable state as = {R [|É(R) - (Rt)) < 0.75J where & consists of the top two components of time-lagged independent component analysis (TICA; FérezzHerránne
et al., 2013). We further describe TICA in Appendix B.4.
As shown in Table 2 and Figure 7, only TPS-DPS (S) successfully samples transition paths that pass probable transition states while TPS-DPS (F) and TPS-DPS (P) fail to hit the target meta-stable state due to the lack of meaningful training signal. While SMD hits the target meta-stable, its transition paths pass less probable transition states. UMD and PIPS fail to the target meta-stable state. In Figure 5, we validate the sampled paths using the potential energy and accor-acceptor distance of the two key hydrogen bonds in Chignolin folding, ASP3OD-THR600 and ASP3N-THR used in (Yang al., 2024). The sampled path of reduces the donor-acceptor distance below the threshold 3.5A. For 3D videos of transition paths of fold-ng proteins, we refer project page.

---

Published as a conference paper at ICLR 2025
RMSD (L)
THP (t)
ETS
10a1
14
Method
Á
kJmol
Ours
KL
0.16 + 0.06
0.43 = 0.34
53.12
92.19
27.88 + 14.38
19.82 + 15.88
Loss
10-1
RMSD
0.8
1.0
local
0.24 0.15
73.44
22.53
10-1
0.6
w/o replay
0.33 0.27
64.06
24.38 1 12.31
wlo annealing
0.67 + 0.21
9.38
69.86 + 30.15
10-1
wo various len
0.23 = 0.11
75.00
29.49 + 14.13
1000
500
w/o equivariance
1 0.17
56.25
22.12 1 16.96
Number of Rollouts
Number of Rollouts
(a) Component-wise performance on Alanine Dipeptide.
(b) Loss and RMSD curves over rollouts.
Figure 6: Ablation studies on each proposed component of TPS-DPS (F) in Alanine Dipeptide.
(a) Benchmark scores on Alanine Dipeptide. (b) Loss and RMSD curves averaged over seeds.
4.4 ABLATION STUDY
We conduct ablation studies to verify the effectiveness the six proposed components: log-variance loss, learnable control variate, replay buffer, simulated annealing, maximum over RBF values for various path lengths, and SE(3) equivariance. To be specific, we (1) replace the log- log-variance divergence with the KL divergence, (2) replace the learnable control variate w IR with the local control variate which is Monte-Carlo estimator used Nüsken & Richter (2021), (3) remove the replay buffer and rely solely on data from the current policy, (4) use only one temperature À 300K in sampling, (5) remove maximum operation over RBF kernel values using only the final state, (6) skip the alignment of neural network input states with the target position by the Kabsch algorithm.
As seen in Figure 6, all the proposed components improve performance. Our loss is smaller than the KL divergence by more than two orders of magnitude and significantly improves performance. Learning the control variate slightly improves performance, showing that utilizing data from previous policies is effective. The replay buffer significantly improves training efficiency, and shows that the large performance gap between our loss and KI. divergence comes from the replay buffer. Simulated annealing for biased MD simulation is critical to finding transition paths. RMSD does not decrease without simulated annealing while loss decreases significantly. For the relaxed indicator function, maximum operation accelerates convergence and improves performance with frequent training signals from the subtrajectories. Leveraging the symmetry of the bias force with the Kabsch algorithm improves performance. We further compare with reverse KL divergence Appendix
CONCLUSION
In this work, we introduced a novel CV-free diffusion path sampler, called TPS-DPS, amortize the cost of sampling transition paths. We propose the log-variance divergence with the learnable control variate and off off-policy training with the replay buffer and simulated annealing. We propose new scale-based equivariant parameterization of bias force and relaxed indicator function for reaching target meta-stable states and frequent training signals, particularly in large molecules. Evaluations on double-well, Alanine Dipeptide, and four fast-folding proteins demonstrate that TPS-DPS is superior in accuracy and diversity compared to non-ML and ML approaches.
Limitation. While experiments show promise, they are limited to small fast-folding proteins (up to 50 amino acids). The applicability of our method to real-world proteins with more than 500 amino acids remains unexplored. Additionally, integrating various libraries, such as PLUMED (plu, 2019) and DMFF (Wang et al., 2023), for MD with neural network bias force, has not yet been investigated.
Our method does not generalize across unseen pairs of meta-stable states different molecular systems. These points to an interesting venue for future research, which would be more appealing for practical applications in drug discovery material design.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
This work partly was supported by Institute for Information & communications Technology Technol- ogy Planning & (Valuation(IITP) grant funded by the Korea (Faram((KAST)) and the National Research Foun-
Artificial Intelligence Graduate School Support dation of Korea(NRF) grant funded by the Ministry of Science and ICT(MSIT) (No. RS-2022- NR072184), and GRDC(Global Research Development Center) Cooperative Hub Program through the National Research Foundation of Korea(NRF) grant funded by the Ministry of Science and (CS(M (No. RS-2024-004361155)
We sincerely thank Soojung Yang for providing helpful feedback and suggestions in preparing the early version of the manuscript.

---

Published as a conference paper at ICLR 2025
REFERENCES
Promoting transparency and reproducibility in enhanced molecular simulations. Nature methods, 16 (8):670-673, 2019.
Seihwan Ahn, Mannkyu Hong, Mahesh Sundararajan, Daniel H Ess, and Mu-Hyun Baik. Design and optimization of catalysts based on mechanistic insights derived from quantum chemical reaction modeling. Chemical reviews, 196(1))559966560 2019.
Noé Beaupere, U Soupremanien, and Laurent Zalewski. Nucleation triggering methods in supercooled phase change materials (pem), a review. Therrmcchmica Acta, 670:184 177::84-210 184-201, 2018.
Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, and Yoshua Bengio. Flow network based generative models for non-iterative diverse candidate generation. Advances in Neural Information Processing Systems, 34:27381-27394 2021.
Davide Branduardi, Giovanni Bussi, and Michele Parrinello. Metadynamics with adaptive gaussians. Journal of chemical theory and computation, 8(7):2247-2244, 2012.
Giovanni Bussi and Davide Branduardi. Free-energy calculations with metadynamics: Theory and practice. Reviews in Computational Chemistry Volume 28, pp. 1-49, 2015.
Giovanni Bussi and Michele Parrinello. Accurate sampling using langevin dynamics. Physical Review E, 75(5):056707, 2007.
Jeffrey Comer, James C Gumbart, Jérôme Hénin, Tony Lelièvre, Andrew Pohorille, and Christophe Chipot. The adaptive biasing force method: Everything you always wanted know but were afraid to ask. The Journal of Physical Chemistry B, (193)))111-1111111
Avishek Das, Dominic C Rose, Juan Garrahan, and David T Limmer. Reinforcement learning of rare diffusive dynamics. The Journal of Chemicul Physics, 155(13), 2021.
Christoph Dellago, Peter G Bolhuis, and David Chandler. Efficient transition path sampling: Applica- tion to lennard-jones cluster rearrangements. The Journal of chemical physics, (18(2)):236-92455 1998.
Yuanqi Du, Michael Plainer, Rob Brekelmans, Chenru Duan, Frank Noe, Carla P Gomes, Alan Aspuru-Guzik, and Kirill Neklyudov. Doob's lagrangian: A sample-efficient variational approach to transition path sampling. In ICML 2024 Al for Science Workshop, 2024.
Chenru Duan, Yuanqi Du, Haojun Jia, and Heather Kulik. Accurate transition state generation with an object-aware equivariant elementary reaction diffusion model. Nature Computational Science, 2023.
Peter Eastman, Raimondas Galvelis, Raúl P Peláez, Charlles RA Abreu, Stephen Farr, Emilio Gallicchio, Anton Gorenko, Michael M Henry, Frank Hu, Jing Huang, et al. Openmm 8: Molecular dynamics simulation with machine learning potentials. The Journal of Physical Chemistry B, 128 (1):109-116, 2023.
Bernd Ensing, Marco De Vivo, Zhiwei Liu, Preston Moore, and Michael Klein. Metadynamics as a tool for exploring free energy landscapes of chemical reactions. Accounts of chemical research, 39 (2):73-81, 2006.
Jérôme Hénin, Tony Lelièvre, Michael R Shirts, Omar Valsson, and Lucie Delemotte. Enhanced sam- pling methods for molecular dynamics simulations [article v1. 0]. Living Journal of Computational Molecular Science, 4(1):1583-153)) 15833 2022.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 18:6840-6851 2020.
Lars Holdijk. Yuanqi Du, Ferry Hooft, Priyank Jaini, Berend Ensing, and Max Welling, Stochastic optimal control for collective variable free sampling of molecular transition paths. Advances in Neural Information Processing Systems, 36, 2024.
12

---

Published as a conference paper at ICLR 2025
Xinru Hua, Rasool Ahmad, Jose Blanchet, and Wei Cai. Accelerated sampling of rare events using neural network bias potential. arXiv preprint arXiv:2401.069933 2024.
Michele Invernizzi and Michele Parrinello. Rethinking metadynamics: from bias potentials to probability distributions. The journal of physical chemistry letters, 11(7):233-2236, 2020.
Sergei Izrailev, Sergey Stepaniants, Barry Isralewitz, Dorina Kosztin, Hui Lu, Ferenc Molnar, Willy Wriggers, and Klaus Schulten. Steered molecular dynamics. In Computational Molecular Dynamics: Challenges, Methods, Ideas: Proceedings of the 2nd International Symposium on Algorithms for Macromolecular Modelling, Berlin, May 21-24, 1997, pp. 39-65. Springer, 1999.
Bowen Jing, Hannes Stark, Tommi Jaakkola, and Bonnie Berger. Generative modeling of molecular dynamics trajectories. ICML'24 '2 Workshop for Life and Material Science: From Theory to Industry Applications, 2024.
Wolfgang Kabsch. A solution for the best rotation relate sets of vectors. Acta Crystallographica Section A: Crystal Physics, Diffraction, Theoretical and General Crystallography, 32(5):922-923, 1976.
Johannes Kästnnr. Umbrella sampling. Wiley Interdisciplinary Reviews: Computational Molecular Science, ((6):932-942) 2011.
Minsu Kim, Joohwan Ko, Dinghuai Zhang, Ling Pan, Taeyoung Yun, Woo Chang Kim, Jinkyoo Park, and Yoshua Bengio. Learning to scale logits for tenperature-condiiional gflownets. In NeurIPS 2023 Al for Science Workshop, 2023.
Seonghwan Kim, Jeheon Woo, and Woo Youn Kim. Diffusion- generative ai for exploring transition states from 2d molecular graphs. Nature Communications, 15(1):341, 2024.
Diederik P Kingma and Max Welling, Auto-encoding variational bayes. In The Second International Conference on Learning Representations, 2013.
Leon Klein, Andrew Foong, Tor Fjelde, Bruno Mlodozeniec, Mare Brockschmidt, Sehastian Nowozin, Frank Noé, and Ryota Tomioka. Timewarp: Transferable acceleration of molecular dynamics by learning time-coarsened dynamics. Advances in Neural Information Processing Systems, 36, 2024.
Tony Lelièvre, Geneviève Robin, Innas Sckkat, Gabriel Stoltz, and Gabriel Victorino Cardoso. Generative methods for sampling transition paths in molecular dynamics. ESAIM: Proceedings and Surveys, 73:238-255 2566, 2023.
Kresten Lindorff-Larsen Stefano Piana, Kim Palmo, Paul Maragakis, John L Klepeis, Ron O Dror, and David Shaw. Improved side-chain torsion potentials for the amber ff99sb protein force field. Proteins, 18:1950-1958, 2010.
Kresten Lindorff-Larseen Stefano Piana, Ron O Dror, and David E Shaw. How fast-folding proteins fold. Science, 33446605))517-5200 2011.
James A Maier, Carmenza Martinez, Koushik Kasavajhala, Lauren Wickstrom, Kevin Hauser, and Carlos Simmerling. ff14sb: improving the accuracy of protein side chain and backbone parameters from ff99sb. Journal of chemical theory and computation, 11(8):3696-33 077366666773,,,0 2015.
Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio. Trajectory balance: Improved credit assignment in gflownets. Advances in Neural Information Processing Systems, 35: 5955-5967, 2022.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning, 2013.
Adrian Mulholland. Modelling enzyme reaction mechanisms, specificity and catalysis. Drug discovery today,
2005.
Hai Nguyen, Daniel R Roe, and Carlos Simmerling. Improved generalized born solvent model parameters for protein simulations. Journal of chemical theory and computation, 94):2020-2034 2013.
13

---

Published as a conference paper at ICLR 2025
Nikolas Nüsken and Lorenz Richter. Solving high-dimensional caniittonjjccbbl jaccobi-bellman bellman pdes using neural networks: perspectives from the theory of controlled diffusions and measures on path space. Partial differential equations and applications, 2(4):48, 2021.
Philip Pechukas. Transition state theory. Annual Review of Physical Chemistry, 32(1):159-177) 1981.
Guillermo Pérez-Hernándee Heránndzz Fabian Paul, Toni Giorgino, Gianni De Fabritiis, and Frank Noé. Identification of slow molecular order parameters for markov model construction. The Journal of chemical physics, 139(1); 2013.
Magnus Petersen, Gemma Roig, and Roberto Covino. Dynamicsdiffusion: Generating and rare event sampling of molecular dynamic trajectories using diffusion models. In NeurlPS 2023 Al for Sciienee Workshop, 2023.
Stefano Piana, Kresten Lindorff-L and David E Shaw. Protein folding kinetics and thermody- namics from atomistic simulation. Proceedings of the National Academy of Sciences, 109(44): 17845-17850, 2012.
Lorenz Richter and Julius Berner. Improved sampling via learned diffusions. In The Twelfth International Conference on Learning Representations, 2024. URL htttps://openzevien net/Eorum: iid h4Noh-0.0000000....
Lorenz Richter, Ayman Boustati, Nikolas Nüsken, Francisco Ruiz, and Omer Deniz Akyildiz. Vargrad: low-variance gradient estimator for variational inference. Advances in Neural Information Processing Systems, 33:13481- 13492, 2020.
Freddie R Salsbury Jr. Molecular dynamics simulations of protein dynamics and their relevance to drug discovery. Current opinion in pharmacology. 10(6):738-744, 2010.
Martin K Scherer, Benjamin
Saaeeelkam--Schhheee Fabian Paul, Guillermo Srrzz-Hernández, Moritz Hoffmann, Nuria Plattner, Christoph Wehmeyer, Jan- lan-Hendrik Prinz, and Frank Noé. Pyemma 2: software package for estimation, validation, and analysis of markov models. Journal of chemical theory and computation, 11(11):5525- 5542, 2015.
Jürgen Schlitter, Michael Engels, and Peter Krüger. Targeted molecular dynamics: a new approach for searching pathways of conformational transitions. Journal of molecular graphics, 12(2):84 89, 1994.
Mathias Schreiner, Ole Winther, and Simon Olsson. Implicit transfer operator learning: Multiple time-resolution models for molecular dynamics. Advances in Neural Information Processing Systems, 36, 2024.
Martin Sipka, Johannes CB Dietschreit, Lukáš Grajciar, and Rafael Gómez-Bombarelli Differentiable simulations for enhanced sampling of rare events. In International Conference on Machine Learning, pp. -20007 PMLR, 2023.
David A Sivak, John Chodera, and Gavin E Crooks. Time step rescaling recovers continuous-time dynamical properties for discrete-time time langevin integration of nonequilibrium systems. The Journal of Physical Chemistry B, (1824)):6666-477) 2014.
Nguyen TK Thanh, N Maclean, and S Mahiddine. Mechanisms of nucleation and growth of nanoparticles in solution. Chemical reviews, 114(15)7610-7630, 2014.
Glenn M Torric and John P Valleau. Nonphysical sampling distributions in monte carlo free-energy
estimation: Umbrella sampling. Journal of computational physics, 23(2):187-199, 1977.
Luke Triplett and Jianfeng Lu. Diffusion methods for generating transition paths. arXiv preprint arXiv:2309.10221 2023.
Francisco Vargas, Will Sussman Grathwohl, and Arnaud Doucet. Denoising diffusion samplers. In The Eleventh International Conference on Learning Representations, 2023. URL https:
//oppereviiccc ew net wwwpt/eetttteiiooeoo
forum? dtt.e/orrmm....aammfrmattl
14

---

Published as a conference paper at ICLR 2025
Siddarth Venkatraman, Moksh Jain, Luca Scimeca, Minsu Kim, Marcin Sendera, Mohsin Hasan, Luke Rowe, Sarthak Mittal, Pablo Lemos, Emmanuel Bengio, et al. Amortizing intractable inference in diffusion models for vision, language, and control. arXiv preprint arXiv:2405.20977, 2024.
Xinyan Wang, Jichen Li, Lan Yang, Feiyang Chen, Yingze Wang, Junhan Chang, Junmin Chen, Wei Feng, Linfeng Zhang, and Kuang Yu. Dmff: an open-source automatic differentiable platform for molecular force field development and molecular dynamics simulation. Journal of Chemical Theory and Computation, 19(17::89775500, 2023.
Soojung Yang. Juno Nam, Johannes CB Dietschreit, and Rafael Gómez-Bombarelli Learning collec- tive variables for protein folding with labeled data augmentation through geodesic interpolation. arXiv preprint arXiv:2402.01542, 2024.
Qinsheng Zhang and Yongxin Chen. Path integral sampler: A stochastic control approach for sampling. In International Conference on Learning Representations, 2022. URL https:/ openreview. net/ forum? iidddid=ddceee _uCb2ynRu7y.
15

---

Published as a conference paper at ICLR 2025
A METHOD DETAILS
A.1 LOG VARIANCE FORMULATION
In this section, we derive Equation (5) from Equation (4) to get the explicit expression for log-variance divergence in terms of SDE in Equation (1) and Equation (3). We refer to Nüsken & Richter (2021, Appendix A.I) for the derivation in more general settings.
Our goal is to derive that
EF. (l0g dQ dP.,,
dQ dP ms
1)
2 = Epp. [[Fvo, - Epp [Fe.. +))],
Bp log
(13)
To this end, we focus on calculating log aQ aFve
(X) when X Pa. Following (Nüsken & Richter,
2021, Lemma A.1), apply Girsanov's Theorem to calculate the Radon-Nikodym derivative decs dPo
as follows:
dlp, ve (X) exp (f dP
dx, - S/¹
vo) (X,) 1 2 / lox
(14)
Since the state X, follows the SDE dx, (u(X,) + E + EdW, We plug it into Equation (14) and utilize the definition of the target path measure Q in Equation (2) to compute dQ as follows:
log aP
dQ
dQ dFu
log dP
(X) = log
(X)
(15)
dFo dP...
log 1s(X) - log 2 - / (wb' dd.
(16)
1 ( (--1u u v0) (X,)dt 1 /-' |l0)
(17)
log 1x(X)= - -log 2 /'' (vo ()
(18)
Vol X,) dW, + 1 //' [| (I,
(19)
F (X) - log Z
(20)
Since log Z is the constant, it is canceled out in the log- divergence as follows:
log dQ
dQ dF
1)
27
Ex.
Ep.
log
Ep. [(Fv.. - Ep, (Fo...)))] =))?],
(21)
dP 06
A.2 CONNECTION TO EXISTING LOSS FUNCTIONS ON DISCRETE-TIME DOMAIN
In this section, connect discretized loss of Equation (8) the loss function, called relative trajectory balance (Venkatraman et al., 2024, RTB). Like our methods, RTB also amortized inference in target path distribution by training forward distribution discrete-time domains such as vision, language, and control tasks. When discretized, our loss function is equivalent to the RTB objective.
Our goal is to show that for every paths Lo.L sampled from the path measure IP P.y.
(f vo vg (ZO.L) w)² [og Po (1(
(22)
Zopes (XO.L)
where w = log Za is a learnable scalar parameter, and path distribution (= IHL- (x(++1/))) is Markovian, and its transition kernel pv pr (ax)) are derived from Euler- Maruyama discretization of the SDE in Equation (3) as follows:
1 क + (a At Ev(x:) + Eei,
At
(23)
16

---

Published as a conference paper at ICLR 2025
where E NV At). To this end, we can caleulate follows:
log po (2o.L) - log Poo (ZO:L)
(24)
log Pvo (0014 (211)
log po (2141 Σ
(25)
t=
Σ (EngAt+ + Ee
(ETEE (E"EEA)"EEOAAA Eg Eng
(26)
EvgAAt + Del [ "S (SegAt+
(2
2 2=
U u
10y + eil1²)
(28)
2
4-1
L-1
Uo (are) (1 112
Σ
(up (g)=v))(x;))4)
Σ
vo( EE.
(29)
6=
= Fon. (ZO.L) log (0(0)
(30)
which implies
F Fn (2 = log
(31)
Poo (10u1)
by subtracting w and squaring both sides, we have
(F (FO.L) w)2² (log Po (o((ooL)) 2 ZoP (Fo.L)
(32)
We can view Po ( g ( as the unnormalized target distribution discretized from the target path measure Q, and Zo as the estimator for normalizing constant Z = f po ( (.o.....d.... and Pes (L forward probability distribution to amortize inference in the target distribution. Based on these results, we provide our training algorithm in Algorithm 1.
A.3 PROOF OF SCALE-BASED PARAMETERIZATION
In this section, we prove that our scale-based parameterization of bias force strictly decreases the distance to the (aligned) target position for small step sizes, improving the ability to find informative paths in large molecules.
Proposition 1. Consider the molecular state R² at the t-th time step and the next state RY+ RR b (X) At/ m updated by step size At and the bias force 6(X.) diag( (tag(sa(pp" .x)))) Rg R.). Then there always exists a small enough At that strictly decreases the distance towards the target state Rg:
lIt+A RB R''sell < lipe Rs - RR
(33)
where P44 argminpeSE(()) [|p- R5 R2 and we assume that there does not exist a rotation that exactly aligns the current molecular state to the target state, i.e., ||p: Rs R, || > 0.
Proof. The proof consists of two steps. We first show the (strictly) positive correlation between the bias force and the direction from the 1-th state R2 the target state Rs. Next, we show that the positive correlation gaurantees strict decrease in distance between the states, i.c., lp: Rg Rell. given that the distance was not already zero.
Step 1:
First, we show that the bias force (divided by atom-wiss masses) is positively correlated with the direction to the target position, i.e., (b(X,)/ m) (pt Rs R2) > O. This follows from:
(6(x))mm) T(p. Rg Rτ) = (pt Rg - R₂) diag(8n 80 (pp" 1 X.)) (p: Rg R.)
(34)
m
3N
8;
(pr Rg R1) > 0,
(35)
mi
17

---

Published as a conference paper at ICLR 2025
where 8; 0 is the i-th element of so(pr Xi) and (pe Rs Ri); is the i-th element of the direction to the target position.
Step 2: Next, show that the positive correlation ensures distance reduction for small enough step size. Consider the squared distance between the target position Pt Rg and updated position by bias force
1p: Rs R
|lpe Rg - OR, + b( b(x,)4t/m))
(36)
1ptt Rs Rτ) b(X,)AI ml|
(37)
|l+. Rs R (2,()( (p. Rs R.) + (A/))(I(X.))mmll
(38)
Due to step 1, i.e. (b(X,)/m) (pr Rg R₂) > O, there exists step size A satisfying:
At ((b ()) (p: Rg Re)
(39)
66 [6(X,)/ml|
With this choice of AI, multiplying Ai||b(X ()) /
m|2 leads to the following inequaliity:
(4)))|6((X)//mll < 2 (pr Rs - R))-
(40)
By subtracting the right-hand side from both sides and adding llp: RB R2 112 both sides, have the following inequality:
pl RB Bi
< llp: Rg Rella
(41)
Taking the square root of both sides, we have the following inequality:
|r++ Rs R' < ||p: Rs RR'ssell ll Rs Rell
(42)
where the first inequality follows from the definition of PLAs argmin lp Rs R4 This completes the proof.
18

---

Published as a conference paper at ICLR 2025
B EXPERIMENT DETAILS
B.I OPENMM CONFIGURATIONS
For real- molecules, we use the VVVR integrator (Sivak et al., 2014) with the step size A/ = fs and the friction term 7 1 ps- In the TPS-DPS training algorithm, simulate MD with 7 1ps for Alanine Dipeptide and 7' 5ps for the fast-folding proteins. We start simulations at temperature Astart 600K, and end
temperature Xend 300K for Alanine Dipeptide and Chignolin and Aend = 400K for the others. We use the amber sbi ldn (Lindorff- Larsen et al., 2010) force field for Alanine Dipeptide in vacuum and the f 4SBon ysc (Maier et al., 2015) force field for the fast-folding proteins with the gbn2 implicit solvation model (Nguyen et al., 2013).
B.2 MODEL CONTIGURATIONS
We use a 3-layer MLP for double-well system, and a 6-layer MLP for real-world molecules with ReLU activation functions for neural bias force, potential, and scale. To constrain the output of the neural bias scale parameterization to a positive value, we apply Softplus to the MLP output. As input the neural network, concatenate the current position (R,); of the i-th atom with its distance the target position d, [| [RB); (Rt); /|2. For real-world molecules, we apply the Kabsch algorithm (Kabsch, 1976) for heavy atoms to align Rs with RG We update the parameters updated with
of the neural network with a learning rate of 0.0001, while the scalar parameter learning rate of 0.001. We clip the gradient norm with to prevent loss from exploding. we train 1000 times per rollout. We report other model configurations in Table 3. For PIPS, use the model configurations reported by Holdijk et al. (2024). For CVs of SMD, we use backhone dihedral angles (ø, 1)) for Alanine Dipeptide and RMSD for fast-folding proteins.
Table 3: Model configurations of TPS-DPS.
System
rollouts (I)
of samples (M)
Batch size (K)
Buffer size
Relaxation (o)
Double-well
20
512
$12
10000
3
Alanine Dipeptide
1000
16
1000
0.1
Chignolin
100
16
200
0.5
Trpcage
100
16
100
0.5
BBA
100
16
4
100
0.5
BBL
100
16
100
0.5
B.3
EVALUATION METRICS
Root mean square distance (RMSD). We use the Kabsch algorithm (Kabsch, 1976) for heavy atoms to align the final position with the target position Rg, using the optimal (proper) rotation and translation to superimpose two heavy atom positions. We calculate RMSD between heavy atoms of the final position and the target position Rg.
Target hit percentage (THP). THP measures the success rate of paths arriving at the target meta- stable state in a binary manner. Formally, given the final positions { ( J of M paths, THP defined as follows:
: R E BB
THP
(43)
M
Energy of transition state (ETS). ETS measures the ability of the method to find probable transition states when crossing the energy barrier. ETS refers to the maximum potential energy among states in transition path. Formally, given transition path æo. of length L that reaches the target meta-stable state i.e., Rt, E 13, ETS is defined as follows:
ETS( = max U(R.)
(44)
EE(O.L]
19

---

Published as a conference paper at ICLR 2025
B.4 SYSTEM DETAILS
Systhetic double-well system. Double-well system follows the overdamped Langevin dynamics defined as follows:
dR, (V,(R)) dt +
2nkBA dW.
(45)
For simplicity, we let R = (z, y) E TR? m ,, 1,A = 0.01, 7 10, and = 1200K. To evaluate the ability to find diverse transition paths, consider the following double-well potential (Hua et al., 2024):
Uzz (4(1 (4(1-22-y)) +2(2² - 2)² [ -) 1]2 [x -
(46)
This potential has global minima and two saddle points, having meta-stable states and two reaction channels.
Time-lagged independent components (TICA). To extract the collective variable (CV) for the four fast-folding proteins, we consider components of time-lageed independent component analysis (TICA; Péréz-Hernández et al., 2013). We J unbiased MD simulations with 2fs step size and record states per 2ps to collect MD trajectories, using the OpenMM library with the same configuration as in Appendix B.1. For the top two TICA components, we use PyEMMA library (Scherer et al., 2015) with time lag r = 500ps for Chignolin T = 200ps for the others.
Reproducibility. We describe experiment details in Appendix B, including detailed simulation configuration and hyper-parameters. In the anonymous link, we provide the code for TPS-DPS.
20

---

Published as a conference paper at ICLR 2025
C COMPUTATIONAL COST
In this section, we analyze the time complexity TPS-DPS and provide the number of energy evaluations and runtime in training and inference time for real molecules.
The training and inference time complexity of TPS-DPS is Q ML) and MLM,L), respectively, where 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 M. the total complexity one biased MD step is (
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 layers, and hidden units, MLP for bias force requires (NN) and the Kabsch algorithm for equivariance requires O(N). Calculating force with cut-off and integrating MD with VVVR integrator also requires (
To measure computational cost, we consider the number of energy evaluations and runtime per rollout training and inference time. shown in Table 4, the inference cost of TPS-DPS proportional to UMD and SMD which have time complexity O TPS-DPS requires less energy evaluations than PIPS training since TPS-DPS finds transition paths faster than PIPS by utilizing the replay buffer and simulated annealing.
Table 4: Cost comparison. EET and EEI refer to the number of energy evaluations by the OpenMM library in training and inference, respectively. RT and RI denote runtime (second) per rollout in training and inference for Alanine Dipeptide, 7 = 5ps for other systems, and A/ 1fs. The number of samples for each training rollout and inference is 16 and 64, respectively. Other configurations of TPS-DPS are provided in Appendix B.2 and those of PIPS can be found in (Holdijk et al., 2024)
single RTX A5000 GPU. MD simulations are conducted with T 1ps
Molecule
Method
(L)
(4)
UMD
64K
29.49
SMD
64K
47.45
PIPS (T)
240M
64K
44.22
71.05
Alanine Dipeptide
TPS DPS (F. Ours)
16M
64K
24.93
70.50
TPS-DPS Ours)
27.25
78.83
TPS-DPS (S. Ours)
16M
64K
25.11
73.04
SMD
320K
380.37
277.29
Chignolin
PIPS (P)
40M
320K
317.08
622.87
IPS-DPS (F, Ours)
8M
320K
209.29
562.90
TPS-DPS DPS (P. Ours)
RM
320K
224.36
623.63
TPS (5, Qurs)
8M
320K
215.18
581.26
UMD
320K
258.29
320K
323.52
PIPS
40M
320K
360.55
652.59
Trp-Cage
PIPS (P) Ours)
40M
320K
417.93
718.05
22
TPS DPS (P. Ours)
SM
320K
301.76
699, 44
(S. Ours)
293.51
UMD
320K
395.1 12
PIPS (I)
40M
320K
506.57
1080.62
(P) TPS-DPS (F. Ours)
8M
320K
422.23
1042.81
IPS-DPS (P. Ours)
8M
320K
430.24
1091.97
TPS DPS (S. Ours)
SM
320K
426.48
1068.68
SMD
320K
853.77
PIPS (F)
40M
320K
622.22
1558.55
BBL
PIPS
40M
320K
660.02
1612.78
TPS-DPS (F. Ours)
8M
320K
560.95
1520.05
IPS-DPS Ours)
572.77
1607.62
TPS DPS (S. Ours)
SM
320K
563.45
1553,89
21

---

Published as a conference paper at ICLR 2025
VISUALIZATION OF SAMPLED PATHS OF FAST FOLDING PROTEINS D
mea
ICC-2
TNC
TIC
11C
11C
1CC1
T.c.
TC1
TC1
TCC
TC1
TC
TC1
(a) UMD
(b) SMD (20k)
(c) PIPS (P)
(d) TPS-DPS (F) TPS-DPS (P) (i) TPS-DPS (S)
Figure 7: Visualization of sampled paths of four fast-folding proteins on free energy landscapes for top two TICA components. We aim sample folding processes for the four fast folding proteins: Chignolin, Trpcage, BBA, and BBL (from top to bottom rows).
22

---

Published as a conference paper at ICLR 2025
E COMPARISION WITH REVERSE KL DIVERGENCE
Table 5: Benchmark scores of reverse KL divergence and TPS-DPS on Alanine Dipeptide system.
Metrics averaged over 64 paths. TPS-DPS consistently outperforms reverse KI. divergence all metrics regardless of predicting bias force or potential.
RMSD (4)
THP (t)
ETS (4)
Method
Á
%
kJmol
Reverse KL (F)
0.43 + 0.34
53.12
27.88 1 14.38
Reverse KL (P)
0.58 0.34
48,43
21.61
+ 11.76
TPS-DPS (F, Ours)
0.16 0.06
92.19
19.82
15.88
TPS-DPS (P, Ours)
0.16 - 0.10
87.50
18.37
+ 10.86
(a) Reverse KL (F)
(b) Reverse KL (P)
TPS-DPS (F)
(d) TPS-DPS (P)
Figure 8: Visualization of sampled paths of Alanine Dipeptide on the Ramachandran plot. The reverse KI. divergence struggles to find diverse reaction channels suffering from mode collapse issues while the log-vau log-variance divergence of our method can capture two reaction channels and reach the target states better.
23


Reviews:
[{"review_id": "iAdInHC8k8", "reviewer": "Reviewer_CJA8", "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.", "questions": "Please see the weakness section.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper is well-written and easy to follow.\n\n2. 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.\n\n3. A thorough empirical study demonstrates the effectiveness of the proposed method, and the comprehensive ablation study justifies each design choice.\n\n4. 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.\n\n2. 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.\n\n3. More discussion of concurrent work, particularly Doob’s Lagrangian: A Sample-Efficient Variational Approach to Transition Path Sampling, would provide valuable context.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "j80cwJvirW", "reviewer": "Reviewer_r5xj", "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.", "questions": "What is the computational cost compared to SMD?\nWhat limitations prevent you from applying this method to larger systems?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n\nThe 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.\n\nThe 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)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9H7diXk3KQ", "reviewer": "Reviewer_ALy8", "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.", "questions": "In the introduction section, a reference is needed for the claim that \"minimizing the KL divergence leads to mode collapse, ...\".", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pQrID4wwBY", "reviewer": "Reviewer_aeFr", "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.", "questions": "1. Why are MLPs chosen to parameterize the bias forces for the molecular systems? Have you tried using simple GNNs?\n\n2. 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.\n\n3. 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?\n\n[1] Lindorff-Larsen, K., Piana, S., Dror, R. O., & Shaw, D. E. (2011). How fast-folding proteins fold. Science, 334(6055), 517-520.\n\n[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.\n\n[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.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "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.\n2. 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.\n\nHigher-Dimensional Systems:\n\nI 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.\n\nThe 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.\n\nLack of Clarity of Presentation\n\nSeveral aspects of the paper were confusing to follow.\n1. 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.\n\n2. 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.\n\n3. Quoting from the paper, \"The first three terms in Equation (6) correspond to the deviation of the biased MD from the unbiased\nMD 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.\n\n4. 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.\n\nA Few Other Miscellaneous Points:\n\n1. 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]. \n\n2. A citation of Sipka, et. al. [2] would be well-placed as they also tackle TPS with ML.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)