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

## Abstract

Abstract Understanding transition pathways between meta-stable states in molecular systems is crucial to advance material design and drug discovery. However, unbiased molecular dynamics simulations are computationally infeasible due to the high energy barriers separating these states. Although recent machine learning techniques offer potential solutions, they are often limited to simple systems or rely on collective variables (CVs) derived from costly domain expertise. In this paper, we introduce a novel approach that trains diffusion path samplers (DPS) for transition path sampling (TPS) without the need for CVs. We recast the problem as an amortized sampling of the target path measure, minimizing the log-variance divergence between the path measure induced by our DPS and the target path measure. To ensure scalability for high-dimensional tasks, we introduce (1) a new off-policy training objective based on learning control variates with replay buffers and (2) a scale-based equivariant parameterization of the bias forces. We evaluate our approach, coined TPS-DPS, on a synthetic double-well potential and three peptides: Alanine Dipeptide, Polyproline Helix, and Chignolin. Results show that our approach produces more realistic and diverse transition pathways compared to existing baselines. We also provide links to our project page and code .

## 1 Introduction

In material design and drug discovery, it is crucial to understand the mechanisms and kinetics of transitions between meta-stable states of molecular systems, such as protein folding and chemical reactions (Mulholland, ; Piana et al., ; Ahn et al., ; Spotte-Smith et al., ). Their comprehensive study requires sampling transition paths (Elber, ; Lee et al., ), which provides insight into mechanisms and energy landscapes. However, naïvely sampling transition paths by molecular dynamics (MD) simulations is often computationally intractable due to high energy barriers, which cause an exponential decay in probability to make a transition (Pechukas, ).

To address this problem, researchers have developed enhanced sampling approaches such as steered MD (SMD; Schlitter et al., ; Izrailev et al., ), umbrella sampling (Torrie & Valleau, ; Kästner, ), meta-dynamics (Ensing et al., ; Branduardi et al., ; Bussi & Branduardi, ), on-the-fly probability-enhanced sampling (OPES; Invernizzi & Parrinello, ), and adaptive biasing force (ABF; Comer et al., ) methods. These methods rely on bias forces to facilitate transitions across high energy barriers. They are mainly designed based on collective variables (CVs), which are functions of atomic coordinates that capture the slow modes of the transition. Although effective for some systems, the reliance on expensive domain knowledge limits the applicability of the methods to systems where CVs are less understood.

Recently, machine learning has emerged as a promising paradigm for CV-free transition path sampling (Das et al., ; Lelièvre et al., ; Holdijk et al., ). The key idea is to parameterize the bias force using a neural network and train the corresponding biased MD simulation to sample realistic transition paths. In particular, Lelièvre et al. () considered reinforcement learning to sample paths escaping meta-stable states. Das et al. () and Holdijk et al. () proposed to minimize the Kullback-Leibler (KL) divergence between the path measures induced by the neural network and the target path measure. However, minimizing the KL divergence leads to mode collapse, which is hazardous for sampling transition paths due to a lack of exploration. Furthermore, Das et al. () and Lelièvre et al. () limited their evaluation to low-dimensional synthetic systems. Designing machine learning algorithms for CV-free transition path sampling of large molecules remains an open problem.

Contribution. In this work, we propose to solve the transition path sampling problem with the diffusion path sampler (TPS-DPS). TPS-DPS trains the bias force to minimize a recently proposed log-variance divergence (Nüsken & Richter, ) between the path measure induced by the bias MD and the target path measure. To ensure scalability for high-dimensional tasks, e.g., Chignolin folding, we propose a new off-policy training scheme and a scale-based parameterization of the bias forces.

To be specific, we propose a new training objective to minimize the log-variance divergence, defined between a pair of path measures induced by the biased MD and the unbiased MD conditioned on the target meta-stable state. To minimize the variance of gradient estimators and improve sample efficiency, we propose to use (1) learnable control variates and (2) replay buffers with simulated annealing. Our choice of minimizing the log-variance is motivated by recent works (Richter et al., ; Nüsken & Richter, ; Richter & Berner, ) that demonstrate many desirable properties of the log-variance divergence such as a robust gradient estimator for training stability and avoidance of differentiation through the SDE solver. This allows the use of previously collected samples in a replay buffer, promoting sample efficiency in training and improving the diversity of the generated transition paths by preventing the mode collapse observed in KL divergence.

We also introduce a new $SE(3)$ equivariant scale-based parameterization for the bias force that guides the biased MD to sample meaningful paths more frequently, based on inductive bias. Our key idea is to parameterize the bias force to yield a positive inner product with a direction to decrease the distance between the molecular state and the target meta-stable state, which guarantees the bias force to (locally) decrease the distance. We also use the Kabsch algorithm (Kabsch, ) to align the molecule and the bias force to the target meta-stable state, guaranteeing $SE(3)$ equivariance for better generalization across molecular states that are identical up to roto-translations.

We extensively evaluate our method on the synthetic double-well potential and three peptides: Alanine Dipeptide, Polyproline Helix, and Chignolin. We compare TPS-DPS with prior ML approach (PIPS; Holdijk et al., ), as well as classical non-ML methods, e.g., steered MD (SMD; Izrailev et al., ). Our experiments demonstrate that TPS-DPS consistently generates realistic and diverse transition paths, similar to the ground truth ensemble. In addition, ablation studies of algorithmic components verify the effectiveness of our method.

## 2 Related work

Transition path sampling without ML. Metadynamics (Branduardi et al., ), on-the-fly probability-enhanced sampling (OPES; Invernizzi & Parrinello, ), adaptive biasing force (ABF; Comer et al., ), and steered molecular dynamics (SMD; Schlitter et al., ; Izrailev et al., )
were introduced to explore molecular conformations that are difficult to access by unbiased molecular dynamics (MD) within limited simulation times (Hénin et al., ). However, they mostly rely on collective variables (CVs) for high-dimensional problems and are inapplicable to systems with unknown CVs.
To sample transition paths without CVs, Dellago et al. () proposed shooting methods that use the Markov chain Monte Carlo (MCMC) procedure on path space.

Data driven ML approaches.
Recently, generative models have been trained to sample a new transition path given a dataset of transition paths. Petersen et al. (); Triplett & Lu () and Lelièvre et al. () applied diffusion probabilistic models (Ho et al., ) and variational auto-encoders (Kingma & Welling, ) for transition path sampling, respectively. However, these methods are limited to small systems. Concurrent to our work, Jing et al. () generated molecular trajectories by task-wise conditions given an initial state and viewed transition path sampling as an interpolation task between two molecular states.

We also note some related works that are not directly comparable to our method. Klein et al. () proposed to accelerate MD simulation through coarse-grained path generation that may include multiple meta-stable states. However, this work is incapable of targeting a specific meta-stable state and generating fine-grained transition paths. Duan et al. (); Kim et al. () use neural networks to generate transition states of a given chemical reaction, but cannot generate transition paths.

Data free ML approaches. Without a previously collected dataset, Das et al. (); Lelièvre et al. (); Hua et al. (); Holdijk et al. () trained generators to sample transition paths using biased MD. Lelièvre et al. () used reinforcement learning to train the generator but focused on escaping an initial meta-stable state rather than reaching the target meta-stable state. Next, Holdijk et al. () considered transition path sampling as minimizing the reverse KL divergence between path distribution from biased MD and transition path distribution. In this work, we mainly compare our method with Holdijk et al. () and provide quantitative analysis for real-world molecule systems. We also note that, concurrent to our work, Du et al. () also considered variational formulation of the TPS problem using Doob’s $h$-transform.

## 3 Transition path sampling with diffusion path samplers

Figure: Figure 1: Problem setup. Potential energy landscape of Alanine Dipeptide, where the sampled path from state $\mathcal{A}$ to state $\mathcal{B}$ is highlighted in yellow dotted lines. We visualize snapshots in the transition path highlighted in white circles and the transition state highlighted in the white star.
Refer to caption: x1.png

In this section, we introduce our method, coined transition path sampling with diffusion path sampling (TPS-DPS). Our main idea is to formulate the transition path sampling (TPS) problem as a minimization of log-variance divergence (Nüsken & Richter, ) between two path measures: the path measure of transition paths and that of biased MD simulations parameterized by a neural network. Our main methodological contribution is twofold: (1) a new training algorithm that minimizes the log-variance divergence with a learned control variate and replay buffer, and (2) a $SE(3)$ equivariant scale-based parameterization of the biased MD that provides inductive bias for dense training signals in high-dimensional problems.

### 3.1 Problem setup

Our goal is to sample transition paths from one meta-stable state to another meta-stable state given a molecule system. We provide an example of the problem for Alanine Dipeptide in Figure 1. We view this as a task to sample paths from an unbiased molecular dynamics (MD), described by a stochastic differential equation (SDE), conditioned on its starting and ending points. To solve this task, we train the bias force parameterized by a neural network to amortize the sampling procedure.

Molecular dynamics. We consider a MD simulation on time interval $[0,T]$ that describes motion of a molecular state $\bm{X}_{t}=(\bm{R}_{t},\bm{V}_{t})\in\mathbb{R}^{6N}$ at time $t$ where $N$ is the number of atoms, $\bm{R}_{t}\in\mathbb{R}^{3N}$ is the atom-wise positions and $\bm{V}_{t}\in\mathbb{R}^{3N}$ is the atom-wise velocities. In particular, we adopt Langevin dynamics (Bussi & Parrinello, ) defined as the following SDE:

$$ $\mathrm{d}\bm{X}_{t}=\bm{u}(\bm{X}_{t})\mathrm{d}t+\Sigma\mathrm{d}\bm{W}_{t}, ~{}~{}\bm{u}(\bm{X}_{t})=\left(\bm{V}_{t},-\frac{\nabla U(\bm{R}_{t})}{\bm{m}} -\gamma\bm{V}_{t}\right),~{}~{}\Sigma=\text{diag}\left(\bm{\zeta},\sqrt{\frac{ 2\gamma k_{B}\lambda}{\bm{m}}}\right)$ (1) $$

where $U$, $\bm{m}$, $\gamma$, $k_{B}$, $\lambda$, and $\bm{W}_{t}$ denote the potential energy function, the atom-wise masses, the friction term, the Boltzmann constant, the absolute temperature, and the Brownian motion, respectively, and $\bm{\zeta}\in\mathbb{R}^{3N}$ is a vector of positive infinitesimal values. We let $\mathbb{P}_{0}$ denote the path measure induced by Equation 1 in the space $\mathcal{C}([0,T];\mathbb{R}^{6N})$ of continuous functions $\bm{X}:[0,T]\rightarrow\mathbb{R}^{6N}$.

Transition path sampling. One of the challenges in sampling transition paths through unbiased MD simulations is the meta-stability: a state remains trapped for a long time in the initial meta-stable state $\mathcal{A}\subseteq\mathbb{R}^{3N}$ before transitioning into a distinct meta-stable state $\mathcal{B}\subseteq\mathbb{R}^{3N}$. To capture the rare event where transition from $\mathcal{A}$ to $\mathcal{B}$ occurs, we constrain paths $\bm{X}=(\bm{X}_{t})_{0\leq t\leq T}$ sampled from unbiased MD to satisfy $\bm{R}_{0}\in\mathcal{A}$, $\bm{R}_{T}\in\mathcal{B}$ for a fixed time $T$. Since the meta-stable state $\mathcal{A}$ and $\mathcal{B}$ are not well-specified for many molecular systems, we simplify this task by (1) fixing a local minima $\bm{R}_{\mathcal{A}},\bm{R}_{\mathcal{B}}$ of the potential energy function in the meta-stable states $\mathcal{A},\mathcal{B}$ and (2) sampling a transition path $\bm{X}$ that starts from the state $\bm{R}_{0}=\bm{R}_{\mathcal{A}}$ and ends at the vicinity of $\bm{R}_{\mathcal{B}}$.

To be specific, we aim to sample from the target path measure $\mathbb{Q}$ which reweights the path measure $\mathbb{P}_{0}$ via the Radon–Nikodym derivative, defined as follows:

$$ $\frac{\mathrm{d}\mathbb{Q}}{\mathrm{d}\mathbb{P}_{0}}(\bm{X})=\frac{1_{ \mathcal{B}}(\bm{X})}{\mathbb{E}_{\mathbb{P}_{0}}\left[1_{\mathcal{B}}(\bm{X}) \right]},\quad 1_{\mathcal{B}}(\bm{X})=\begin{cases}1&\text{if}~{}\lVert\bm{R} _{T}-\rho_{T}\cdot\bm{R}_{\mathcal{B}}\rVert\leq\delta,\\ 0&\text{otherwise},\end{cases}$ (2) $$

where $\cdot$ denotes group action associated with the $SE(3)$ space and $\rho_{T}\cdot\bm{R}_{\mathcal{B}}$ is the aligned position by the optimal roto-translation $\rho_{T}\in SE(3)$ to minimize its Euclidean distance to $\bm{R}_{T}$, i.e., $\rho_{T}=\operatorname*{arg\,min}_{\rho\in SE(3)}\lVert\bm{R}_{T}-\rho\cdot\bm
{R}_{\mathcal{B}}\rVert$. Such a transformation can be obtained from the Kabsch algorithm in $O(N)$ complexity (Kabsch, ).

Note that one may consider naïve rejection sampling to sample transition paths, based on running unbiased MD to sample a path $\bm{X}$ from the path measure $\mathbb{P}_{0}$ and accepting if the path $\bm{X}$ arrives at the neighborhood of the position $\bm{R}_{\mathcal{B}}$ with the radius $\delta$. However, this method does not scale to high-dimensional or low-temperature problems, since the sampled path by unbiased MD rarely reaches the target states due to the high energy barriers, i.e., the rejection ratio is too high.

### 3.2 Log-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 $\mathbb{Q}$, using the log-variance divergence (Nüsken & Richter, ) between the path measures. We propose a new training scheme to minimize the log-variance divergence based on learning the control variate of its gradient and a replay buffer to improve sample efficiency and diversity.

Amortizing transition path sampling with log-variance divergence. To match the target path measure $\mathbb{Q}$, we consider a biased MD defined by a policy $\bm{v}$ (or bias force $\bm{b}$) as the following SDE:

$$ $\mathrm{d}\bm{X}_{t}=(\bm{u}(\bm{X}_{t})+\Sigma\bm{v}(\bm{X}_{t}))\mathrm{d}t+ \Sigma\mathrm{d}\bm{W}_{t},\quad\bm{v}(\bm{X}_{t})=\Sigma^{-1}\left(\bm{0}, \frac{\bm{b}(\bm{X}_{t})}{\bm{m}}\right).$ (3) $$

We also let $\mathbb{P}_{\bm{v}}$ denote the path measure induced by the SDE. To amortize transition path sampling, we match the path measure $\mathbb{P}_{\bm{v}_{\theta}}$ of a parameterized policy $\bm{v}_{\theta}$ with the target path measure $\mathbb{Q}$ by minimizing the log-variance divergence:

$$ $D_{\text{LV}}^{\mathbb{P}}(\mathbb{P}_{\bm{v}_{\theta}}\|\mathbb{Q})=\mathbb{V }_{\mathbb{P}}\left[\log\frac{\mathrm{d}\mathbb{Q}}{\mathrm{d}\mathbb{P}_{\bm{ v}_{\theta}}}\right]=\mathbb{E}_{\mathbb{P}}\left[\left(\log\frac{\mathrm{d} \mathbb{Q}}{\mathrm{d}\mathbb{P}_{\bm{v}_{\theta}}}-\mathbb{E}_{\mathbb{P}} \left[\log\frac{\mathrm{d}\mathbb{Q}}{\mathrm{d}\mathbb{P}_{\bm{v}_{\theta}}} \right]\right)^{2}\right],$ (4) $$

where $\mathbb{P}$ is an arbitrary reference path measure with $\mathbb{E}_{\mathbb{P}}[\log({\mathrm{d}\mathbb{Q}}/{\mathrm{d}\mathbb{P}_{\bm
{v}_{\theta}}})]<\infty$. To express the log-variance divergence in detail, we let $\mathbb{P}=\mathbb{P}_{\tilde{\bm{v}}}$ for some policy $\tilde{\bm{v}}$ and apply the Girsanov’s theorem to Equation 4, deriving the following formulation:

$$ $D_{\text{LV}}^{\mathbb{P}_{\tilde{\bm{v}}}}(\mathbb{P}_{\bm{v}_{\theta}}\| \mathbb{Q})=\mathbb{E}_{\mathbb{P}_{\tilde{\bm{v}}}}\left[(F_{\bm{v}_{\theta}, \tilde{\bm{v}}}-\mathbb{E}_{\mathbb{P}_{\tilde{\bm{v}}}}[F_{\bm{v}_{\theta}, \tilde{\bm{v}}}])^{2}\right]\,,$ (5) $$

$$ $F_{\bm{v}_{\theta},\tilde{\bm{v}}}(\bm{X})=\frac{1}{2}\int^{T}_{0}\lVert\bm{v} _{\theta}(\bm{X}_{t})\rVert^{2}\mathrm{d}t-\int^{T}_{0}(\bm{v}_{\theta}\cdot \tilde{\bm{v}})(\bm{X}_{t})\mathrm{d}t-\int^{T}_{0}\bm{v}_{\theta}(\bm{X}_{t}) \cdot\mathrm{d}\bm{W}_{t}+\log 1_{\mathcal{B}}(\bm{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 $\mathbb{P}_{\tilde{\bm{v}}}$. The last term reweights the unbiased MD to the target path measure $\mathbb{Q}$. As a result, minimizing Equation 5 could be thought as minimizing the variation between $\mathbb{P}_{\bm{v}_{\theta}}$ and $\mathbb{Q}$. We provide the full derivation in Section A.1. Compared to KL divergence, the log-variance divergence provides a robust gradient estimator and avoids differentiating through the SDE solver. (Richter et al., ; Nüsken & Richter, ).

Minimizing with learnable control variate.
To minimize the log-variance divergence, we consider the following loss function that replaces the estimation of $\mathbb{E}_{\mathbb{P}_{\bm{v}_{\theta}}}[F_{\bm{v}_{\theta},\bm{v}_{\theta}}]$ by learning a parameter $w$:

$$ $\mathcal{L}(\theta,w)=\mathbb{E}_{\mathbb{P}_{\bm{v}_{\theta}}}\left[(F_{\bm{v }_{\theta},\bm{v}_{\theta}}-w)^{2}\right],$ (7) $$

where $w$ is a control variate that controls the variance of the gradient estimator of $\nabla_{\theta}\mathcal{L}(\theta,w)$ without changing the gradient. Note that we set $\tilde{\bm{v}}=\bm{v}_{\theta}$ in Equation 6, which implies that the gradient of Equation 7 coincides with the KL divergence (Richter et al., ; Nüsken & Richter, ). When optimized, the control variate $w$ estimates the expectation $\mathbb{E}_{\mathbb{P}_{\bm{v}_{\theta}}}[F_{\bm{v}_{\theta},\bm{v}_{\theta}}]$ since $\operatornamewithlimits{argmin}_{w}\mathcal{L}(\theta,w)=\mathbb{E}_{\mathbb{P
}_{\bm{v}_{\theta}}}[F_{\bm{v}_{\theta},\bm{v}_{\theta}}]$. Thus, jointly optimizing $(\theta,w)$ with the gradient step can be interpreted as jointly minimizing log-variance divergence and estimating $\mathbb{E}_{\mathbb{P}_{\bm{v}_{\theta}}}[F_{\bm{v}_{\theta},\bm{v}_{\theta}}]$ using $w$.

Off-policy training with replay buffer.
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-policy training, which is widely used in discrete-time reinforcement learning (Mnih et al., ; Bengio et al., ). For the sample efficiency, we reuse the samples with a replay buffer $\mathcal{D}$ which stores path samples from the path measure $\mathbb{P}_{\bm{v}_{\bar{\theta}}}$ associated with previous policies $\bm{v}_{\bar{\theta}}$. Our modified loss function $\mathcal{L}^{\mathcal{D}}$ with replay buffer $\mathcal{D}$ is defined as follows:

$$ $\mathcal{L}^{\mathcal{D}}(\theta,w)=\mathbb{E}_{(\bm{v}_{\bar{\theta}},\bm{X}) \sim\mathcal{D}}[(F_{\bm{v}_{\theta},\bm{v}_{\bar{\theta}}}(\bm{X})-w)^{2}].$ (8) $$

Using the replay buffer also prevents mode collapse, using diverse paths from different path measures. Similar to other off-policy training algorithms (Malkin et al., ; Kim et al., ), we use simulated annealing to collect diverse paths that cross high energy barriers.

Discretization.
To implement the algorithm, we discretize Equation 8. Given a discretization step size $\Delta t$, we consider the discretized paths $\bm{x}_{0:L}=(\bm{x}_{0},\bm{x}_{1},\ldots,\bm{x}_{L})$ of $\bm{X}$ from MD simulations where $L=T/\Delta t$ and $\bm{x}_{\ell}=\bm{X}(\ell\Delta t)$. In discrete cases, the discretized paths $\bm{x}_{0:L}$ from previous policies $\bm{v}_{\bar{\theta}}$ and their (gradient-detached) policy values $(\bm{v}_{\bar{\theta}}(\bm{x}_{0}),...,\bm{v}_{\bar{\theta}}(\bm{x}_{L}))$ are used to approximate the value $F_{\bm{v}_{\theta},\bm{v}_{\bar{\theta}}}(\bm{X})$ in Equation 6 as follows:

$$ $\hat{F}_{\bm{v}_{\theta},\bm{v}_{\bar{\theta}}}(\bm{x}_{0:L})=\frac{1}{2}\sum_ {\ell=0}^{L-1}\lVert\bm{v}_{\theta}(\bm{x}_{\ell})\rVert^{2}\Delta t-\sum_{ \ell=0}^{L-1}(\bm{v}_{\theta}\cdot\bm{v}_{\bar{\theta}})(\bm{x}_{\ell})\Delta t -\sum_{\ell=0}^{L-1}\bm{v}_{\theta}(\bm{x}_{\ell})\cdot\bm{{\epsilon}}_{\ell}+ \log 1_{\mathcal{B}}(\bm{x}_{0:L})\,,$ (9) $$

where the noise $\bm{{\epsilon}}_{\ell}=\Sigma^{-1}(\bm{x}_{\ell+1}-\bm{x}_{\ell}-(\bm{u}(\bm{x
}_{\ell})+\Sigma\bm{v}_{\bar{\theta}}(\bm{x}_{\ell}))\Delta t)$ is the discretized Brownian motions of the Langevin dynamics with policy $\bm{v}_{\bar{\theta}}$. For implementation, we further derive a simple discretized loss of Equation 8 from Equation 9 as follows:

$$ $\mathbb{E}_{\bm{x}_{0:L}\sim\hat{\mathcal{D}}}\left[\left(\log\frac{p_{0}(\bm{ x}_{0:L})1_{\mathcal{B}}(\bm{x}_{0:L})}{p_{\bm{v}_{\theta}}(\bm{x}_{0:L})}-w \right)^{2}\right],$ (10) $$

where the buffer $\hat{\mathcal{D}}$ stores paths $\bm{x}_{0:L}$ sampled from the previous policies, and $p_{0}$ and $p_{\bm{v}_{\theta}}$ denote discrete time transition probability induced by Equations 1 and 3, respectively. We provide a formal derivation of the discretized loss in Section A.2. Note that the same objective was derived in the name of relative trajectory balance by Venkatraman et al. ().

We describe our training algorithm in Section A.3. Overall, our off-policy training algorithm iterates through four steps: (1) sampling paths from the biased MD simulation with current policy $\bm{v}_{\theta}$ at high temperature, (2) storing sampled paths in the replay buffer $\hat{\mathcal{D}}$, (3) sampling a batch of the paths from the replay buffer, and (4) training current policy $\bm{v}_{\theta}$ by minimizing the loss in Equation 10. After minimization, biased MD simulation can directly sample transition paths by amortizing inference in the target path measure.

### 3.3 Parameterization for high dimensional tasks

In this section, we introduce new parameterizations of the bias force and the indicator function for high-dimensional tasks. Our parameterization is designed around alleviating the problem of sparse training signal, where the model is unable to learn anything due to the failure of collecting meaningful paths that end near the meta-stable state. This problem is especially severe in high-dimensional problems such as large molecules.

Bias force parameterization. To frequently sample the meaningful paths, we parameterize the bias force to guarantee (locally) decreasing the distance between the current molecular state and the target meta-stable state. This is done by parameterizing the bias force as positive atom-wise scaling of the force minimizing the distance, which implies that the bias force can be decomposed into two terms: (1) the direction to decrease the Euclidean distance and (2) the direction that leaves the Euclidean distance invariant. The bias force also satisfies roto-translational equivariance to the current state $\bm{X}_{t}$, aligning with the symmetry of the transition path sampling problem for better generalization.

To be specific, we use a roto-translation $\rho_{t}\in SE(3)$ to align $\bm{R}_{\mathcal{B}}$ with $\bm{R}_{t}$, as we do for indicator function in Equation 2. To achieve $SE(3)$ equivariant, we parameterize the bias force as follows:

$$ $\bm{b}(\bm{X}_{t})=\text{diag}(\bm{s}_{\theta}(\rho_{t}^{-1}\cdot\bm{X}_{t}))( \rho_{t}\cdot{\bm{R}}_{\mathcal{B}}-\bm{R}_{t}),$ (11) $$

where $\bm{s}_{\theta}(\cdot)\in\mathbb{R}_{+}^{3N}$ is a neural network constrained to have positive output elements. One can observe the positive inner product between the bias force and the direction to minimize the Euclidean distance, i.e., $(\rho_{t}\cdot\bm{R}_{\mathcal{B}}-\bm{R}_{t})$, as follows:

$$ $\bm{b}(\bm{X}_{t})^{\top}(\rho_{t}\cdot\bm{R}_{\mathcal{B}}-\bm{R}_{t})=(\rho_ {t}\cdot\bm{R}_{\mathcal{B}}-\bm{R}_{t})^{T}\text{diag}(\bm{s}_{\theta}(\rho_{ t}^{-1}\cdot\bm{X}_{t}))(\rho_{t}\cdot\bm{R}_{\mathcal{B}}-\bm{R}_{t})\geq 0,$ (12) $$

which implies that the bias force can be decomposed into a vector positively aligning with $(\rho_{t}\cdot\bm{R}_{\mathcal{B}}-\bm{R}_{t})$ and a vector orthogonal to $(\rho_{t}\cdot\bm{R}_{\mathcal{B}}-\bm{R}_{t})$.

We also note that the bias force is $SE(3)$ equivariant to the current state $\bm{X}_{t}$ since it is the product of an invariant term $\text{diag}(\bm{s}_{\theta}(\rho_{t}^{-1}\cdot\bm{X}_{t}))$ and an equivariant term $(\rho_{t}\cdot\bm{R}_{\mathcal{B}}-\bm{R}_{t})$. This stems from the fact that the Kabsch algorithm is equivariant to the current state $\bm{X}_{t}$, i.e., roto-translating position $\bm{R}_{t}\mapsto\bar{\rho}\cdot\bm{R}_{t}$ leads to roto-translating output of the Kabsch algorithm $\rho_{t}\mapsto\bar{\rho}\cdot\rho_{t}$.

In the experiments, we also consider other equivariant parameterizations that are less constrained: (1) directly predicting the bias force by $\rho_{t}\cdot\bm{b}_{\theta}(\rho_{t}^{-1}\cdot\bm{X}_{t})\in\mathbb{R}^{3N}$ and (2) predicting the force as gradient of the invariant bias potential $\nabla b_{\theta}(\rho_{t}^{-1}\cdot\bm{X}_{t})\in\mathbb{R}$. We observe the parameterizations to be useful for low-dimensional tasks, but suffer from the sparse training signal on high-dimensional tasks.

Indicator function parameterization. We also propose to relax the indicator function $1_{\mathcal{B}}$ as a radial basis function (RBF) kernel $\tilde{1}_{\mathcal{B}}(\bm{X})=k(\bm{R}_{T},\rho^{-1}\cdot{\bm{R}}_{\mathcal{
B}};\sigma^{2})$ which measures the similarity between two positions and $\sigma>0$ controls the degree of relaxation. The range of RBF kernel $k$ is bounded by the interval $(0,1]$ so that $\log\tilde{1}_{\mathcal{B}}(\bm{X})$ is well-defined and $\tilde{1}_{\mathcal{B}}(\bm{X})$ represents the binary indicator function smoothly. To capture high training signal, we focus on the signal from subtrajectories of sampled paths, taking maximum over RBF kernel values of all intermediate states by $\tilde{1}_{\mathcal{B}}^{\text{max}}(\bm{X})=\max_{t\in[0,T]}k(\bm{R}_{t},
\tilde{\bm{R}}_{\mathcal{B}};\sigma^{2})$. 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. Also, the relaxed indicator function is $SE(3)$ invariant to $\bm{R}_{t}$ due to the Kabsch algorithm.

## 4 Experiment

**Table 1: Benchmark scores on the double-well system and three real-world molecular systems: Alanine Dipeptide, Polyproline Helix, and Chignolin. All metrics are averaged over 1024 paths for the double-well system, and 64 paths for real-world molecular systems. ETS is computed for paths that hit the target meta-stable state, and the best results are highlighted in bold. TPS-DPS predicting the bias force, potential, and atom-wise scaling are denoted by (F), (P), and (S), respectively. UMD ($\lambda$) denotes unbiased MD with temperature $\lambda$ 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 real-world molecules.**
| Method | RMSD ($\downarrow$) | THP ($\uparrow$) | ETS ($\downarrow$) |
| --- | --- | --- | --- |
| Å | % | $\mathrm{kJmol}^{-1}$ |  |
| Double-well |  |  |  |
| UMD (1200K) | 2.21 $\pm$ 0.10 | 0.00 | - |
| UMD (2400K) | 2.11 $\pm$ 0.38 | 3.03 | 1.69 $\pm$ 0.31 |
| UMD (3600K) | 1.85 $\pm$ 0.68 | 12.60 | 2.12 $\pm$ 0.41 |
| UMD (4800K) | 1.54 $\pm$ 0.81 | 21.58 | 2.77 $\pm$ 0.69 |
| SMD (0.5) | 0.98 $\pm$ 0.90 | 52.15 | 1.54 $\pm$ 0.21 |
| SMD (1) | 0.14 $\pm$ 0.08 | 99.80 | 1.85 $\pm$ 0.16 |
| TPS-DPS (F, Ours) | 0.01 $\pm$ 0.02 | 99.90 | 1.38 $\pm$ 0.16 |
| TPS-DPS (P, Ours) | 0.01 $\pm$ 0.03 | 99.71 | 1.36 $\pm$ 0.15 |
| TPS-DPS (S, Ours) | 0.01 $\pm$ 0.03 | 99.80 | 1.73 $\pm$ 0.20 |
| Polyproline Helix |  |  |  |
| UMD (300K) | 2.22 $\pm$ 0.11 | 0.00 | - |
| UMD (1200K) | 1.38 $\pm$ 0.45 | 10.94 | 1010.28 $\pm$ 38.44 |
| SMD (5k) | 1.68 $\pm$ 0.17 | 54.69 | 350.58 $\pm$ 14.36 |
| SMD (10k) | 1.26 $\pm$ 0.06 | 100.00 | 355.62 $\pm$ 14.83 |
| PIPS (F) | 2.64 $\pm$ 0.15 | 0.00 | - |
| PIPS (P) | 1.85 $\pm$ 0.13 | 93.75 | 574.66 $\pm$ 20.49 |
| TPS-DPS (F, Ours) | 1.53 $\pm$ 0.12 | 98.44 | 418.17 $\pm$ 45.54 |
| TPS-DPS (P, Ours) | 1.35 $\pm$ 0.12 | 100.00 | 345.00 $\pm$ 32.58 |
| TPS-DPS (S, Ours) | 1.17 $\pm$ 0.02 | 100.00 | 342.00 $\pm$ 20.28 |

In this section, we compare our method, called TPS-DPS, with both classical and ML approaches, assessing the accuracy and diversity of sampled transition paths. We begin with a synthetic double-well system at 1200K, followed by three real-world molecular systems with various numbers of amino acids: Alanine Dipeptide, Polyproline Helix, and Chignolin. Additionally, we conduct ablation studies to validate the effectiveness of each component in our method. All real-world molecular systems are simulated using the OpenMM library (Eastman et al., ). Details on OpenMM simulation and model configurations are provided in Sections B.1 and B.2, respectively.

Evaluation Metrics.
We consider three metrics for evaluation: RMSD, THP, and ETS. The root mean square distance (RMSD) calculates the Euclidean distance between the atoms in the final position of the sampled path and the target position $\bm{R}_{\mathcal{B}}$, with the final position aligned to the target. The target hit percentage (THP) measures the number of paths that successfully arrive at the target meta-stable state $\mathcal{B}$. Finally, the energy of the transition state (ETS) measures the ability to identify probable transition states. For further details, refer to Section B.3.

Baselines.
We compare TPS-DPS with both non-ML and ML baselines. For non-ML baselines, we consider unbiased MD (UMD) in various temperatures and steered MD (SMD; Izrailev et al., ) with various force constants $k$ and collective variables. For ML baselines, we consider a CV-free transition path sampling method, path integral path sampling (PIPS; Holdijk et al., ) which also trains a bias force by minimizing the KL divergence between path measures induced by the biased MD and the target path measure. For simplicity, we denote parameterizations for predicting force, potential, and atom-wise scaling as (F), (P) and (S), respectively.

Figure: (a) Double-well energy landscapes
Refer to caption: x2.png

### 4.1 Double-well system

Figure: (a) UMD
Refer to caption: x4.png

We begin by evaluating our method on a two-dimensional synthetic system, i.e., a double-well potential at 1200K. This system has two global minima (white circles) representing the meta-stable states, and two reaction pathways via saddle points (white stars).
We sample transition paths from the left meta-stable state $\bm{R}_{\mathcal{A}}$ to the right meta-stable states $\mathcal{B}=\{\bm{R}~{}|~{}\lVert\bm{R}-\bm{R}_{\mathcal{B}}\rVert<0.5\}$, using over-damped Langevin dynamics. We collect ground truth path ensembles by rejection sampling which proposes paths sampled from the unbiased MD simulations and accepts if the final states are in the target meta-stable states $\mathcal{B}$. We provide more details on the system in Section B.4.

In Table 1, TPS-DPS shows superior performance for the double-well potential regardless of the bias force design. In Figure 2, the neural bias potential accelerates the transition by increasing the potential energy near the initial meta-stable state while decreasing the potential energy near the two energy barriers. Moreover, the distribution of energy and $y$ coordinates of the transition states from TPS-DPS is closest to the ground truth compared with other baselines, successfully capturing two reaction channels. In Figure 3, unbiased MD at 1200K fails to escape the initial state while SMD struggles to pass the saddle points. Our approach generates transition paths more similar to the ground truth than other methods.

Figure: (a) UMD
Refer to caption: x10.png

Figure: Figure 5: Snapshots of transition path for Alanine Dipeptide (left) and Polyproline Helix (right).
Refer to caption: x16.png

Figure: (a) UMD
Refer to caption: extracted/5906922/resources/images/chig_unbiased.png

### 4.2 Alanine Dipeptide conformation change

We first consider Alanine Dipeptide for a real-world molecule consisting of two alanine residues, sampling transition paths from the $C5$ (upper left) to the $C7ax$ (lower right) as seen in Figure 4. The target meta-stable states are defined as $\mathcal{B}=\{\bm{R}~{}|~{}\lVert\xi(\bm{R})-\xi(\bm{R}_{\mathcal{B}})\rVert<0
.75\}$, where $\xi(\bm{R})=(\phi,\psi)$ is a well-known collective variable which consists of two backbone dihedral angles. Also, Alanine Dipeptide has two reaction channels between the $C5$ and $C7ax$ state passing through the saddle points marked as white stars.

In Table 1 and Figure 4, unbiased MD at 300K fails to escape the initial state, while SMD with the collective variable $\xi$ generates transition paths with less probable transition states. PIPS generates transition paths of only one reaction channel, suffering from mode collapse. Compared to the baselines, our method successfully generates diverse transition paths that pass two reaction channels and show superior performance regardless of the bias force design. For qualitative analysis, we plot the snapshots of the transition path sample from TPS-DPS in Figure 5.

### 4.3 Polyproline Helix isomerization

Next, we consider a more complex molecule, Polyproline Helix consisting of three proline residues. We sample transition paths from the left-handed state (PPII) in the cis-configuration to the right-handed state (PPI) in the trans-configuration. For simplicity, we define the target meta-stable state based on handedness $H$ as $\mathcal{B}=\{\bm{R}~{}|~{}H(\bm{R})>0\}$. For the formal definition of handedness, refer to Section B.4.
As seen in Table 1, our models consistently outperform all baselines: unbiased MD, SMD, and PIPS. In Figure 7(a), a transition path sampled from TPS-DPS crosses the energy barrier around 0.4ps and changes handedness for 1.5ps. Additionally in Figure 5, we plot the snapshots of the sampled transition path.

Figure: (a) Potential energy and CVs over time.
Refer to caption: x17.png

### 4.4 Chignolin folding

Finally, we consider a challenging molecule, Chignolin, an artificial protein consisting of 10 amino acids (Honda et al., ), which folds into a $\beta$-hairpin structure by hydrogen bonds. We sample transition paths from the unfolded state (right white circle) to the folded state (left white circle) in Figure 6. We define the target meta-stable state $\mathcal{B}=\{\bm{R}~{}|~{}\lVert\xi(\bm{R})-\xi(\bm{R}_{\mathcal{B}})\rVert<0
.75\}$ where $\xi$ consists of the top two components of time-lagged independent component analysis (TICA; Pérez-Hernández et al., ). We further describe TICA in Section B.4.

In Table 1 and Figure 6, unbiased MD, PIPS, TPS-DPS (F), and TPS-DPS (P) fail to hit the target meta-stable state. While SMD hits the target meta-stable, its transition paths do not pass probable transition states. Only TPS-DPS (S) successfully samples transition paths that pass probable transition states.
In Figure 7, we further validate the sampled paths using the potential energy and donor-acceptance distance of the two key hydrogen bonds. The sampled path forms two key hydrogen bonds, ASP3OD-THR6OG and ASP3N-THR8O, reducing the donor-acceptance distance below the threshold $3.5$Å. Additionally, the transition time of the sampled path by TPS-DPS is 5ps, where the average folding time of Chignolin is known to be 0.6$\mu$s (Lindorff-Larsen et al., ).

### 4.5 Ablation study

Effectiveness of algorithmic components. First, we conduct ablation studies to verify the effectiveness of the five proposed components: log-variance loss, learnable control variate, replay buffer, simulated annealing, and maximum over RBF values for various path lengths. To be specific, we (1) replace our loss with the KL divergence, (2) replace the learnable control variate with the local control variate used in Nüsken & Richter (), (3) remove the replay buffer and use data only from the current policy, (4) use only one temperature $\lambda$, and (5) remove maximum operation over RBF values using only the final state.

Figure: Figure 8: Trade-off between THP and ETS for various $k$ in SMD. The further left and higher up ($\nwarrow$), the better.
Refer to caption: x19.png

As seen in Figure 9, 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 KL 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.

Figure: (a) Component-wise performance on Alanine Dipeptide.
Refer to caption: x20.png

Comparison to steered MD with varying force constant $k$.
The performance of SMD differs by the force constant $k$. As seen in Figure 8, one can see the trade-off between THP and ETS for various $k$ values. TPS-DPS outperforms SMD in THP and ETS without showing a trade-off between them.

## 5 Conclusion

In this work, we introduced a novel CV-free diffusion path sampler, called TPS-DPS, to amortize transition path sampling, using log-variance divergence with the learnable control variate, and the replay buffer with simulated annealing for off-policy training. We also propose new $SE(3)$ equivariant scale-based parameterization for bias force and relaxed indicator function for frequent training signals. Evaluations on synthetic double-well and real-molecule systems such as Alanine Dipeptide, Polyproline Helix, and Chignolin demonstrate superior accuracy and diversity of our model compared to both classical and ML approaches.

#### Limitation.

While our algorithm successfully samples transition paths for small proteins, 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. These points to an interesting venue for future research, which would be more appealing for practical applications in drug discovery or material design.

## Appendix 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 (, Appendix A.1) for the derivation in more general settings.

Our goal is to derive that

$$ $\mathbb{E}_{\mathbb{P}_{\tilde{\bm{v}}}}\left[\left(\log\frac{\mathrm{d} \mathbb{Q}}{\mathrm{d}\mathbb{P}_{\bm{v}_{\theta}}}-\mathbb{E}_{\mathbb{P}_{ \tilde{\bm{v}}}}\left[\log\frac{\mathrm{d}\mathbb{Q}}{\mathrm{d}\mathbb{P}_{ \bm{v}_{\theta}}}\right]\right)^{2}\right]=\mathbb{E}_{\mathbb{P}_{\tilde{\bm{ v}}}}\left[(F_{\bm{v}_{\theta},\tilde{\bm{v}}}-\mathbb{E}_{\mathbb{P}_{\tilde{ \bm{v}}}}[F_{\bm{v}_{\theta},\tilde{\bm{v}}}])^{2}\right],$ (13) $$

To this end, we focus on calculating $\log\frac{\mathrm{d}\mathbb{Q}}{\mathrm{d}\mathbb{P}_{\bm{v}_{\theta}}}(\bm{X})$ when $\bm{X}\sim\mathbb{P}_{\tilde{\bm{v}}}$.
Following (Nüsken & Richter, , Lemma A.1), we apply Girsanov’s Theorem to calculate the Radon-Nikodym derivative $\frac{\mathrm{d}\mathbb{P}_{\bm{v}_{\theta}}}{\mathrm{d}\mathbb{P}_{0}}$ as follows:

$$ $\frac{\mathrm{d}\mathbb{P}_{\bm{v}_{\theta}}}{\mathrm{d}\mathbb{P}_{0}}(\bm{X} )=\exp\left(\int^{T}_{0}(\bm{v}_{\theta}^{T}\Sigma^{-1})(\bm{X}_{t})\cdot \mathrm{d}\bm{X}_{t}-\int^{T}_{0}(\Sigma^{-1}\bm{u}\cdot\bm{v}_{\theta})(\bm{X }_{t})\mathrm{d}t-\frac{1}{2}\int^{T}_{0}\lVert\bm{v}_{\theta}(\bm{X}_{t}) \rVert^{2}\mathrm{d}t\right)$ (14) $$

since the state $\bm{X}_{t}$ follows the SDE $\mathrm{d}\bm{X}_{t}=(\bm{u}(\bm{X}_{t})+\Sigma\tilde{\bm{v}}(\bm{X}_{t}))
\mathrm{d}t+\Sigma\mathrm{d}\bm{W}_{t}$. We plug it into Equation 14 and utilize the definition of the target path measure $\mathbb{Q}$ in Equation 2 to compute $\log\frac{\mathrm{d}\mathbb{Q}}{\mathrm{d}\mathbb{P}_{0}}$ as follows:

$$ $\displaystyle\log\frac{\mathrm{d}\mathbb{Q}}{\mathrm{d}\mathbb{P}_{\bm{v}_{ \theta}}}(\bm{X})$ $\displaystyle=\log\frac{\mathrm{d}\mathbb{Q}}{\mathrm{d}\mathbb{P}_{0}}\frac{ \mathrm{d}\mathbb{P}_{0}}{\mathrm{d}\mathbb{P}_{\bm{v}_{\theta}}}(\bm{X})$ (15) $\displaystyle=\log 1_{\mathcal{B}}(\bm{X})-\log Z-\int^{T}_{0}(\bm{v}_{\theta} ^{T}\Sigma^{-1})(\bm{X}_{t})\cdot\mathrm{d}\bm{X}_{t}$ (16) $\displaystyle\quad+\int^{T}_{0}(\Sigma^{-1}\bm{u}\cdot\bm{v}_{\theta})(\bm{X}_ {t})\mathrm{d}t+\frac{1}{2}\int^{T}_{0}\lVert\bm{v}_{\theta}(\bm{X}_{t})\rVert ^{2}\mathrm{d}t$ (17) $\displaystyle=\log 1_{\mathcal{B}}(\bm{X})-\log Z-\int^{T}_{0}(\bm{v}_{\theta} \cdot\tilde{\bm{v}})(\bm{X}_{t})\mathrm{d}t$ (18) $\displaystyle\quad-\int^{T}_{0}\bm{v}_{\theta}(\bm{X}_{t})\cdot\mathrm{d}\bm{W }_{t}+\frac{1}{2}\int^{T}_{0}\lVert\bm{v}_{\theta}(\bm{X}_{t})\rVert^{2} \mathrm{d}t$ (19) $\displaystyle=F_{\bm{v}_{\theta},\tilde{\bm{v}}}(\bm{X})-\log Z$ (20) $$

Since $\log Z$ is the constant, it is canceled out in the log-variance divergence as follows:

$$ $\mathbb{E}_{\mathbb{P}_{\tilde{\bm{v}}}}\left[\left(\log\frac{\mathrm{d} \mathbb{Q}}{\mathrm{d}\mathbb{P}_{\bm{v}_{\theta}}}-\mathbb{E}_{\mathbb{P}_{ \tilde{\bm{v}}}}\left[\log\frac{\mathrm{d}\mathbb{Q}}{\mathrm{d}\mathbb{P}_{ \bm{v}_{\theta}}}\right]\right)^{2}\right]=\mathbb{E}_{\mathbb{P}_{\tilde{\bm{ v}}}}\left[(F_{\bm{v}_{\theta},\tilde{\bm{v}}}-\mathbb{E}_{\mathbb{P}_{\tilde{ \bm{v}}}}[F_{\bm{v}_{\theta},\tilde{\bm{v}}}])^{2}\right],$ (21) $$

### A.2 Connection to existing loss functions on discrete-time domain

In this section, we connect our discretized loss of Equation 8 to the loss function, called relative trajectory balance (Venkatraman et al., , RTB). Like our methods, RTB also amortized inference in target path distribution by training forward distribution on 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 $\bm{x}_{0:L}$ sampled from the path measure $\mathbb{P}_{\bm{v}_{\bar{\theta}}}$,

$$ $(\hat{F}_{\bm{v}_{\theta},\bm{v}_{\bar{\theta}}}(\bm{x}_{0:L})-w)^{2}=\left( \log\frac{p_{0}(\bm{x}_{0:L})1_{\mathcal{B}}(\bm{x}_{0:L})}{Z_{\theta}p_{\bm{v }_{\theta}}(\bm{x}_{0:L})}\right)^{2},$ (22) $$

where $w=\log Z_{\theta}$ is a learnable scalar parameter, and path distribution $p_{\bm{v}}(\bm{x}_{0:L})=\prod_{\ell=0}^{L-1}p_{\bm{v}}(\bm{x}_{l+1}|\bm{x}_{l})$ is Markovian, and its transition kernel $p_{\bm{v}}(\bm{x}_{l+1}|\bm{x}_{l})$ are derived from Euler-Maruyama discretization of the SDE in Equation 3 as follows:

$$ $\bm{x}_{l+1}=\bm{x}_{l}+\bm{u}(\bm{x}_{l})\Delta t+\Sigma\bm{v}(\bm{x}_{l}) \Delta t+\Sigma\bm{{\epsilon}}_{l},$ (23) $$

where $\bm{{\epsilon}}_{l}\sim\mathcal{N}(\bm{0},\Delta t)$. To this end, we can calculate as follows:

$$ $\displaystyle\log p_{0}(\bm{x}_{0:L})-\log p_{\bm{v}_{\theta}}(\bm{x}_{0:L})$ (24) $\displaystyle=\sum_{\ell=0}^{L-1}\log p_{\bm{v}_{\theta}}(\bm{x}_{l+1}|\bm{x}_ {l})-\sum_{\ell=0}^{L-1}\log p_{0}(\bm{x}_{l+1}|\bm{x}_{l})$ (25) $\displaystyle=\frac{1}{2}\sum_{\ell=0}^{L-1}(\Sigma\bm{v}_{\bar{\theta}}\Delta t +\Sigma\bm{{\epsilon}}_{l}-\Sigma\bm{v}_{\theta}\Delta t)^{T}(\Sigma^{T}\Sigma \Delta t)^{-1}(\Sigma\bm{v}_{\bar{\theta}}\Delta t+\Sigma\bm{{\epsilon}}_{l}- \Sigma\bm{v}_{\theta}\Delta t)$ (26) $\displaystyle-\frac{1}{2}\sum_{\ell=0}^{L-1}(\Sigma\bm{v}_{\bar{\theta}}\Delta t +\Sigma\bm{{\epsilon}}_{l})^{T}(\Sigma^{T}\Sigma\Delta t)^{-1}(\Sigma\bm{v}_{ \bar{\theta}}\Delta t+\Sigma\bm{{\epsilon}}_{l})$ (27) $\displaystyle=\frac{1}{2\Delta t}\sum_{\ell=0}^{L-1}(\lVert\bm{v}_{\bar{\theta }}\Delta t+\bm{{\epsilon}}_{l}-\bm{v}_{\theta}\Delta t\rVert^{2}-\lVert\bm{v}_ {\bar{\theta}}\Delta t+\bm{{\epsilon}}_{l}\rVert^{2})$ (28) $\displaystyle=\frac{1}{2}\sum_{\ell=0}^{L-1}\lVert\bm{v}_{\theta}(\bm{x}_{\ell })\rVert^{2}\Delta t-\sum_{\ell=0}^{L-1}(\bm{v}_{\theta}\cdot\bm{v}_{\bar{ \theta}})(\bm{x}_{\ell})\Delta t-\sum_{\ell=0}^{L-1}\bm{v}_{\theta}(\bm{x}_{ \ell})\cdot\bm{{\epsilon}}_{\ell}$ (29) $\displaystyle=\hat{F}_{\bm{v}_{\theta},\bm{v}_{\bar{\theta}}}(\bm{x}_{0:L})- \log 1_{\mathcal{B}}(\bm{x}_{0:L})$ (30) $$

which implies

$$ $\hat{F}_{\bm{v}_{\theta},\bm{v}_{\bar{\theta}}}(\bm{x}_{0:L})=\log\frac{p_{0}( \bm{x}_{0:L})1_{\mathcal{B}}(\bm{x}_{0:L})}{p_{\bm{v}_{\theta}}(\bm{x}_{0:L})},$ (31) $$

by subtracting $w$ and squaring both sides, we have

$$ $(\hat{F}_{\bm{v}_{\theta},\bm{v}_{\bar{\theta}}}(\bm{x}_{0:L})-w)^{2}=\left( \log\frac{p_{0}(\bm{x}_{0:L})1_{\mathcal{B}}(\bm{x}_{0:L})}{Z_{\theta}p_{\bm{v }_{\theta}}(\bm{x}_{0:L})}\right)^{2}$ (32) $$

We can view $p_{0}(\bm{x}_{0:L})1_{\mathcal{B}}(\bm{x}_{0:L})$ as the unnormalized target distribution discretized from the target path measure $\mathbb{Q}$, $Z_{\theta}$ as the estimator for normalizing constant $Z=\int p_{0}(\bm{x}_{0:L})1_{\mathcal{B}}(\bm{x}_{0:L})d\bm{x}_{0:L}$, and
$p_{\bm{v}_{\theta}}(\bm{x}_{0:L})$ as forward probability distribution to amortize inference in the target distribution. Based on these results, we provide our training algorithm in Section A.3.

### A.3 Training algorithm

Figure: Algorithm 1 Training

## Appendix B Experiment details

### B.1 OpenMM configurations

For real-world molecules, Alanine Dipeptide, Polyproline Helix, and Chingolin, we use the VVVR integrator (Sivak et al., ) with the step size $\Delta t=1~{}\mathrm{fs}$ and the friction term $\gamma=1~{}\mathrm{ps}^{-1}$. In the TPS-DPS training algorithm, we start simulations at a temperature $\lambda_{\text{start}}=600$K, and end at a temperature $\lambda_{\text{end}}=300$K. Other OpenMM configurations are shown in Table 2.

**Table 2: OpenMM configurations for three real-world molecular systems.**
| Molecule | Simulation time ($T$) | Force field | Solvent |
| --- | --- | --- | --- |
| Alanine Dipeptide | $1000$ | amber99sbildn (Lindorff-Larsen et al., ) | vaccum |
| Polyproline Helix | $10000$ | amber/protein.ff14SBonlysc (Maier et al., ) | implicit/gbn2 |
| Chignolin | $5000$ | amber/protein.ff14SBonlysc | implicit/gbn2 |

### B.2 Model configurations

We use a 3-layer MLP for the 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 (Zheng et al., ) to the MLP output. As an input to the neural network, we concatenate the current position $(\bm{R}_{t})_{i}$ of the $i$-th atom with its distance to the target position $d_{i}=\lVert(\tilde{\bm{R}}_{\mathcal{B}})_{i}-(\bm{R}_{t})_{i}\rVert_{2}$. For real-world molecules, we apply the Kabsch algorithm (Kabsch, ) for heavy atoms to align $\bm{R}_{\mathcal{B}}$ with $\bm{R}_{t}$. We update the parameters of the neural network with a learning rate of $0.0001$, while the scalar parameter $w$ is updated with a learning rate of $0.001$. We clip the gradient norm with 1 to prevent loss from exploding. we train $J=1000$ times per rollout. We report other model configurations in Table 3. For PIPS, we use the model configurations reported by Holdijk et al. (). For CVs of SMD, we use backbone dihedral angles $(\phi,\psi)$ for Alanine Dipeptide and RMSD for Polyproline Helix and Chignolin.

**Table 3: Model configurations of TPS-DPS.**
| System | # of rollouts ($I$) | # of samples ($M$) | Batch size ($K$) | Buffer size | Relaxation ($\sigma$) |
| --- | --- | --- | --- | --- | --- |
| Double-well | 20 | 512 | 512 | 10000 | 3 |
| Alanine Dipeptide | 1000 | 16 | 16 | 1000 | 0.1 |
| Polyproline Helix | 1000 | 16 | 4 | 200 | 0.2 |
| Chignolin | 1000 | 16 | 4 | 200 | 0.5 |

### B.3 Evaluation metrics

#### Root mean square distance (RMSD).

We use the Kabsch algorithm (Kabsch, ) for heavy atoms to align the final position with the target position $\bm{R}_{\mathcal{B}}$, 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 $\bm{R}_{\mathcal{B}}$.

#### Target hit percentage (THP).

THP measures the success rate of paths arriving at the target meta-stable state $\mathcal{B}$ in a binary manner. Formally, given the final positions $\{R^{(i)}\}^{M}_{i=1}$ of $M$ paths, THP is defined as follows:

$$ $\operatorname{THP}=\frac{|\{i:\bm{R}^{(i)}\in\mathcal{B}\}|}{M}$ (33) $$

#### 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 a transition path. Formally, given a transition path $\bm{x}_{0:L}$ of length
$L$ that reaches the target meta-stable state i.e., $\bm{R}_{L}\in\mathcal{B}$, ETS is defined as follows:

$$ $\operatorname{ETS}(\bm{x}_{0:L})=\max_{\ell\in[0,L]}U(\bm{R}_{\ell})$ (34) $$

### B.4 System details

#### Double-well potential and dynamics.

Double-well system follows the overdamped Langevin dynamics defined as follows:

$$ $\mathrm{d}\bm{R}_{t}=\frac{-\nabla U(\bm{R}_{t})}{\bm{m}}\mathrm{d}t+\sqrt{ \frac{2\gamma k_{B}\lambda}{\bm{m}}}\mathrm{d}\bm{W}_{t}\,.$ (35) $$

For simplicity, we let $\bm{R}=(x,y)\in\mathbb{R}^{2}$, $\bm{m}=I,\gamma=1,\Delta=0.01,T=10$, and $\lambda=1200$K. To evaluate the ability to find diverse transition paths, we consider the following double-well potential (Hua et al., ):

$$ $U(x,y)=\frac{1}{6}\left(4(1-x^{2}-y^{2})^{2}+2(x^{2}-2)^{2}+[(x+y)^{2}-1]^{2}+ [(x-y)^{2}-1]^{2}-2\right)\,.$ (36) $$

This potential has global minima and two saddle points, having two meta-stable states and two reaction channels.

#### Handedness.

For four points $A,B,C,D$, we define a handedness (Moradi et al., ) as follows:

$$ $H_{ABCD}=\frac{\overrightarrow{EF}\cdot(\overrightarrow{CD}\times \overrightarrow{AB})}{|\overrightarrow{EF}|\cdot|\overrightarrow{CD}|\cdot| \overrightarrow{AB}|}.$ (37) $$

Here, $E$ and $F$ are the midpoints of the vectors $\overrightarrow{AB}$ and $\overrightarrow{CD}$, respectively. With the backbone atoms $X_{1},X_{2},...,X_{N}$ of a Polyproline Helix, we can define the handedness for $N$ atoms as $H=\sum_{i=1}^{N-3}H_{X_{i}X_{i+1}X_{i+2}X_{i+3}}$. In our experiments, we take the alpha carbon $X_{1}$, the carbonyl carbon $X_{2}$ of the first proline residue, the alpha carbon $X_{3}$ of the second residue and alpha carbon $X_{4}$ of third residues to compute the handedness.

#### Time-lagged independent components (TICA).

To extract the collective variable (CV) for Chignolin, we consider components of time-lagged independent component analysis (TICA; Pérez-Hernández et al., ). We run 1$\mu$s unbiased MD simulations and get states per 1ps to collect MD trajectories, using the OpenMM library with the same Chignolin system configuration as in Section B.1. For the top two TICA components, we use PyEMMA library (Scherer et al., ) with a time lag $\tau=500$ps.

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

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: Please see the weakness section.
soundness: 3
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
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.
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.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: What is the computational cost compared to SMD?
What limitations prevent you from applying this method to larger systems?
soundness: 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.

The results on all of the toy systems evaluated are strong.
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.
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)

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: In the introduction section, a reference is needed for the claim that "minimizing the KL divergence leads to mode collapse, ...".
soundness: 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.
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.
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.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
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.
soundness: 3
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.
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.
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.

Meta-review: 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.

Decision: Accept (Poster)