Accelerated Sampling of Rare Events using a Neural Network Bias Potential

Contemporary machine learning models suffer a substantial degradation in performance when confronted with long-tail events that are not represented well in collected data [1, 2, 3]. Several approaches, including the detection of out-of-distribution samples [4] and data resampling techniques [5, 6], can address this issue. In this paper, we propose a method that is specifically designed to efficiently sample these long-tail events, also referred to as rare events. Our method aims to efficiently collect rare events in simulation and increase their representation in datasets.

Specifically, in the domains of materials science and bio-chemistry, there is a pressing need to sample rare events that are associated with specific physical phenomena and estimate their probabilities. This is critical in advancing our understanding of various materials properties ranging from mechanical composites [7, 8, 9] to transport proteins [10], which are essential to aerospace and pharmaceutical industries. The characterization of rare event transitions between two stationary metastable states has been a focus of a big body of research in the past [11, 12, 13], and continues to stimulate the present research undertakings. In the same vein, we seek to study rare event transitions occurring in atomic systems with probabilities as low as $10^{-6}$. The molecules follow the Langevin dynamics governed by a potential energy function [14] which contains multiple minima corresponding to metastable states. During the dynamics, the molecular system spends most of the time performing random thermal motion near a metastable state before rarely escaping to the neighboring metastable state. Our goal is to sample such rare event transitions of the molecular system between two metastable states of the potential energy function. The traditional Monte Carlo sampling method is prohibitively expensive to capture such rare events and suffers from the curse of dimensionality in higher dimensions. Another line of research focuses on determining the committor function $-$ the probability of making the rare event transition from a given molecular configuration $-$ by solving a partial differential equation numerically with finite element method [15, 16]. However, such numerical determination of the committor function becomes exponentially difficult with the dimension of the problem. Thus, it is important to develop novel approaches to efficiently sample successful transitions and estimate the probability of rare events.

In this work, we present an approach that leverages deep neural networks (DNNs) to learn from and sample distributions of rare events in molecular dynamics. A holistic view of our method is presented in Fig. 1. The method of applying bias potential to enhance sampling rate is first introduced in [17, 18] and both works present promising results on 2D. The limitations of both methods lie in the construction of a trial importance function and the challenges with multiple degrees of freedom in higher dimensions. This work [19] first formulate the estimation of rare event probabilities as an optimization problem and employed importance sampling techniques to recover the unbiased probabilities. Our method adopts a similar optimization framework and provides a practical framework to train, test, and evaluate the sample quality. One novelty of our method is its ability to learn from past successful transitions, so that we can add humans in the loop to sketch possible transition paths, use collected transition paths in real-world experiments [20, 21] or we can remove partially the energy barrier first and then obtain some successful paths. The optimal DNN-based bias potential need to both minimize the KL divergence between the distribution of transition paths under the biased dynamics and unbiased dynamics and maximize the probability of rare events.

Our learning algorithm comprises three key steps: 1) The bias potential function is approximated by a DNN and the samples are generated by running the biased Langevin dynamics; 2) The DNN then learns to refine this sampled distribution and aligns it more closely to the unbiased distribution of rare event transitions; 3) After generating feasible rare events, we employ the importance sampling method to compute the actual unbiased probabilities of these rare events. Our paper makes several significant contributions:

1. C1

   We introduce an efficient algorithm for training a DNN to approximate the variance-free importance/bias potential function. This trained DNN significantly improves the efficiency of our rare event sampling in molecular dynamics. Our method can effectively scale to higher-dimensional energy functions since it does not require data pre-collection or solving any equations.

2. C2

   Our algorithm effectively learn from previously successfully sampled rare events which gives our algorithm the freedom and power to reuse and learn from any distribution of trajectories.

3. C3

   Since we are sampling the molecule’s trajectories, which exist in a space of more than 1000 dimensions, ensuring the statistical robustness of our estimator presents a significant challenge. To address this, we statistically assess the reliability of our estimator. This statistical measure enables us to compare the efficiency and accuracy of our estimator with other competing estimators in a fair and rigorous manner.

Optimal Control and Importance Sampling: Several papers have delved into computing an optimal bias potential and utilizing importance sampling to estimate probability of rare events. Among them, a number of works focus on theoretical results on the optimal control and on transition paths [22, 19, 23], and the others provide numerical results [24, 25, 26]. The main advantage of our work is the statistical guarantees of our estimator and the ability to learn from successful transitions. Additionally, our algorithm do not require data collection like in [25] and is scalable to higher dimension. In addition, [27] introduces the method to compute the bias potential to greatly enhance the rate of transitions, so that they can obtain the probability with much higher efficiency. However, their approach of computing the ground truth bias potential is restricted to dynamics on a grid. We have expanded upon this paper, adapting it to free dynamics in both 2D and higher dimensions.

Learning Committor Function: Another popular line of research involves using neural networks to approximate the committor function, which requires solving a partial differential equation with neural networks [28, 29, 30, 31]. One notable limitation of this method is the significant growth in the size of the datasets and computational costs as the dimension grows. Specifically, it mandates the preparation of a dataset prior to the training. In contrast, our method can actively generate samples and learn from the samples simultaneously.

Traditional Sampling Method: There have been various sampling method to increase the efficiency of sampling low-probability events in molecular dynamics [32], for example, the army ants algorithm [33] and adaptive importance sampling method [17]. One shortcoming mentioned in [17] is difficulty of finding a suitable trial function for importance sampling that is successful in higher dimension and involving multiple degrees of freedom. Many books [34] and review papers give a comprehensive summary to all the methods.

We study a class of rare events associated with transition paths in molecular dynamics of chemical reaction networks. The dynamics follow overdamped Langevin dynamics [35, 14]:

Here $\mathbf{x}\in\mathbb{R}^{d}$, $k_{B}=8.617\times 10^{-5}$ is Boltzmann’s constant, $T$ is the absolute temperature, $m$ is the mass of the particle, $\gamma$ is the damping ratio, and $B(t)$ is a Brownian motion. For further simplification, we define ${\epsilon}=2k_{B}T/(m\gamma)$. To introduce the discretized dynamics with timestep of $\Delta t$, we first let $\omega=(\delta_{0},\delta_{1},...,\delta_{N})$ to be a sequence of i.i.d noises, where $\delta_{i}\sim\mathcal{N}(0,\Delta t),\Delta t>0$. As a result, the discretized dynamics follows the equation:

In this work, we focus on the 2D domain, but our method can be extended to general higher dimensions. We let $m=1,\gamma=1$, and an energy function $U$ on position $(x,y)$, defined as

Figure 1 illustrates the energy function that has two potential wells separated by a potential barrier. In this figure, A and B represent the two minima of the energy function ¹¹1In the scenario where energy minima are unknown, we can use linear search and gradient descent method to determine energy minima.. Our objective is to calculate the probability of a particle starting near A and reaching a region close to B before a set deadline. When the temperature is low, escaping the energy well around A becomes highly challenging, resulting in a very low success rate.

We aim to calculate the probability that a particle starts near A, escapes the energy basin, and moves to point B within a predetermined number of steps, denoted as $N\in\mathbb{N^{+}}$. We consider the particle to be close to B if its distance to B is within $\delta$. The particle forms an $N$-step trajectory $\nu=(\mathbf{x}_{0},\mathbf{x}_{1},...,\mathbf{x}_{N})$ as it follows the Langevin dynamics. In addition, we denote the step that the particle is closest to $B$ as $\tau_{B}$, that is

The simulation stops when the particle reaches B or the deadline arrives. In this way, the rare event is defined as:

From the original Langevin dynamics in eq. 1, we aim to learn a neural network with parameters $\theta$ to parameterize the bias potential function so the modified dynamics equation becomes

The optimal modified dynamics significantly increases the probability of rare event $S$ and the successful trajectories are similar to the trajectories sampled with the original potential. Lastly, we employ importance sampling to recover the probability of the rare event with the original energy function $U$.

Dimension of the problem: Although the molecule’s movement is restricted to a 2D plane, we aim to sample its trajectories. For each step in these trajectories, its position ($x$ and $y$) is recorded, resulting in a distribution with a dimensionality of $2N$, where $N$ is the number of steps in a trajectory. In our experiments, $N=500$, so the distribution is 1000-dimensional.

Choice on $\tau$: Generally, the time taken for a transition to occur is of the order $\mathcal{O}\left(\exp\left(\frac{1}{{\epsilon}}\right)\right)$, but the particle spends a substantial proportion of this time around A without crossing the energy barrier. Thus, if we want the particle to start from A and move to B without going back to A, we can select a deadline $\tau=\mathcal{O}\left(\frac{1}{{\epsilon}^{r}}\right)$ for some suitably chosen $r$, so that the particle has sufficient time to travel to B. More precisely, we choose $r$ such that $\tau>\mathbb{E}[\tau_{B}\;|\;\|\mathbf{x}_{0}-A\|<\delta,\tau_{B}<\tau_{A}]$.

Using the bias potential, we want to see the likelihood of rare events significantly increases, and the distribution of successful trajectories remains unchanged. In our experiments, we use the energy function 1 as the loss function and batch stochastic gradient descent method [36] to train our DNN. After many simulation runs, the bias potential $U_{B}$ allow us to sample from a distribution of trajectories $Q$ that minimizes the objective function (7). Algorithm 1 describes the training process. In line 12, we minimize two components in the loss function: the first part is the KL divergence between $Q$ and $P$, which tries to match the distribution of successful trajectories under the modified and the original energy, and the second part penalizes unsuccessful trajectories.

After training the neural network, we employ importance sampling to compute the rare event’s probability under the original distribution $P$. We run simulations using the biased dynamics as described in eq. 4. As in eq. 7, we define $Q$ to only model the successful trajectories. After computing $\omega_{p}$, the probability of $\nu\in S$ under $P$ is given by:

Here $P$ represents the distribution of $\nu$ under the original potential and $Q$ is the distribution of $\nu$ under the modified potential. The density ratio $\frac{p(\omega_{p})}{p(\omega)}$ can be computed analogously to eq. 6.

To achieve an efficient estimation, we want the density ratio of successful trajectories $W_{i}=\frac{dP(\nu)}{dQ(\nu)}$ to be at the same scale with every sample, so that the estimation is not described by a limited subset of samples. Therefore, we use a metric called Effective Sample Size (ESS) [37] and coefficient of variance (CV) to measure the efficiency of importance sampling:

Additionally, we utilize the confidence interval and standard deviation of the estimation to measure the reliability of our estimators.

The neural network architecture consists of 4 hidden layers and $\tanh$ as the activation function (except the output layer). It maps a molecular position to the bias energy: $\mathbb{R}^{2}\rightarrow\mathbb{R}.$ We find that adding two control variables $\exp(\|x-A\|_{2})$ and $\exp(\|x-B\|_{2})$ allows the DNN to better utilize the information of the distance between $x$ and A or B. The project is implemented using the PyTorch library. The gradient of bias energy $\nabla U_{B}(\mathbf{x})$ is computed by the PyTorch autograd module. All the experiments are performed on one AWS c5.4xlarge instance.

Our method is a generative approach, so it does not require a standard train/val/test split. Every trajectory is sampled with random Gaussian noises. During the test, we freeze the DNN-based bias potential and use it in eq. (4) to sample trajectories. We train our method with a batch size of 512 and train 300 steps. The comparisons of the confidence interval, success rate, ESS, and time consumption under two temperature settings are demonstrated in Table 1 and Table 2. We use a statistical significance at a P-value of 0.05 throughout the tests. We showcase two functionalities of our method:

In this work, we focus on challenges of efficiently sampling rare events in molecular dynamics, which underpins the technologically important properties of materials and molecules. We propose to utilize deep neural networks as a bias/importnace potential function. We successfully formulate the probability as an optimization problem which is used to train the neural network. Our proposed approach can be scaled to high-dimensional cases and also provides robust statistical guarantees on the accuracy of the estimated probabilities. The ability of our algorithm to learn from successful samples makes our method versatile and marks an improvement over existing methodologies. We compare our method with traditional Monte Carlo sampling and numerical FEM method under different temperatures to measure the efficacy of our approach. Our estimator exhibits a smaller variance than the traditional Monte Carlo and achieves about 5x speedup comparing the training and test time combined, and more than 44x speedup if only comparing the test time. We also test the robustness and scalability of our bias potential. With more test samples increasing, the variance decreases in the order of $\mathcal{O}(1/\sqrt{n})$, and effective sample size and time increase in the order of $\mathcal{O}(n)$.

The immediate future direction is applying our method to higher dimensional models, similar to [39]. In higher dimensions, we anticipate our method to offer even greater speed advantages over the naive Monte Carlo approach, which suffers from the curse of dimensionality. Another intriguing avenue to explore is the potential for training neural networks using existing molecular dynamics datasets, as described in [40, 21]. One limitation of our method lies in the variance of likelihood ratio of importance sampling. We plan to test the stability and sensitivity of the importance sampling method by adding artificial noises to the bias potential and measuring how much the results change. We will also compare against another line of research [28, 31], which solves the committor function and tries to use it to enhance the sampling rate.

We would like to express our gratitude to the anonymous reviewers for their insightful comments. Our appreciation extends to Professor Masha (Maria) K. Cameron for sharing her FEM code with us, and to Jiaxin (Margot) Yuan for valuable discussions. We are also thankful to Dr. Tianyi Shi for meticulously proofreading this manuscript and for meaningful discussions. The research is supported by the Air Force Office of Scientific Research under award number FA9550-20-1-0397. Additional support is gratefully acknowledged from NSF 1915967, 2118199, 2229012, 2312204.