Reinforced dynamics for enhanced sampling in large atomic and molecular systems

I Introduction

Exploring the configuration space of large atomic and molecular systems is a problem of fundamental importance for many applications, including protein folding, materials design, and understanding chemical reactions, etc. There are several difficulties associated with these applications. The first is that the dimensionality of the configuration space is typically very high. The second is that there are often high energy barriers associated with the exploration. Both difficulties can be reduced by the introduction of collective variables (CVs) and the mapping of the problem to the CV space. The problem then becomes finding the free energy surface (FES) associated with the set of CVs, a problem that has attracted a great deal of interest in the last few decades kumar1992weighted ; kumar1995multidimensional ; voter1997hyperdynamics ; sugita1999replica ; vandevondele2000efficient ; earl2005parallel ; christ2008multiple ; gao2008self ; laio2002escaping ; barducci2008well ; maragliano2006temperature ; maragliano2008single ; abrams2008efficient ; yu2011temperature . One of the most effective techniques is metadynamics laio2002escaping , which computes a biasing potential by depositing Gaussian bases along the trajectory in the CV space. It is shown that the biasing potential converges to the inverted free energy at the end of the calculation barducci2008well . Also closely related to our work are the recent papers that propose to use machine learning methods to help parameterizing FES stecher2014free ; mones2016exploration ; galvelis2017neural ; schneider2017stochastic . In particular, the deep neural network (DNN) model has shown promise in effectively representing the FES defined on high dimensional CV space galvelis2017neural ; schneider2017stochastic .

In this work, we take metadynamics and machine learning methods one step further by making an analogy between reinforcement learning sutton1998reinforcement and the task of configuration space exploration and FES calculation. Classical reinforcement learning scheme involves a state space, an action space, and a reward function. The objective is to find the best policy function, which is a mapping from the state space to the action space, that optimizes the cumulative reward function. Our problem can be thought of as being a multi-scale reinforcement learning problem. We have a micro-state space, the configuration space of the detailed atomic system, and a macro-state space, the space of the CVs. The action space will be represented by the biasing potential in the biased molecular dynamics on the micro-state space. The optimal policy function is the inverted FES, defined on the macro-state space. The FES is parameterized by a carefully designed DNN model. Among other things, this allows us to handle cases with a large set of CVs. In the absence of an explicit reward function, we introduce an uncertainty indicator that can be used to quantify the accuracy of the FES representation. It is defined as the standard deviation of the predictions from an ensemble of DNN models, which are trained using the same dataset but different initialization of the model parameters. The bias is only adopted in regions where the uncertainty indicator is low, i.e. regions that are sufficiently explored, and thus the exploration in the insufficiently explored region is encouraged. We call this scheme the “reinforced dynamics”, to signal its analogy with reinforcement learning.

Roughly speaking, reinforced dynamics works as follows: The biasing potential, or the action, is initialized at 0 and is expected to converge to the inverted FES as the dynamics proceeds. Each step of the macro-iteration involves the following components. First, a biased MD is performed, in which the system is biased only in the regions where the uncertainty indicator is low. The biased simulation is likely to visit the CV regions never visited before or where the FES representation quality is poor. Next, a certain number of the newly visited CV values in regions where the uncertainty indicator is high are added to the training dataset. A restrained MD is performed to obtain the mean force, or the negative gradient of the FES, at each of the newly added CV values. Finally, the accumulated CV values and the mean forces are used as labels to train several network models, which give the current estimate of the biasing potential and the uncertainty indicator. This process is repeated iteratively until convergence is achieved, when the newly visited CV values all fall in the regions where the uncertainty indicator is low.

The quality of the free energy surface is determined by the quality of the CVs. Ideally we would like the FES to capture the structural and dynamic information of the system, such as the important metastable states and transitions between the metastable states. For many years, since our ability to accurately approximate the FES has been limited to systems with a small number of CVs, we have always faced the dilemma that choosing the right CVs is both critical and practically impossible. We believe that the ability of the reinforced dynamics to handle a large set of CVs will make the issue of choosing the right CVs much less critical.

In this paper, we give a systematic presentation of the theoretical and practical aspects of reinforced dynamics. We first focus on methodology and introduce the theory and flowchart of the reinforced dynamics scheme. Then we use the classical example of alanine dipeptide and tripeptide with two and four CVs, respectively, as illustrations due to their intuitive appeal. The solvent effect is explicitly considered in both examples. The FESs constructed by the reinforced dynamics are compared with those constructed by long brute-force simulations (5.1 $\mu$s for alanine dipeptide and 47.7 $\mu$s for tripeptide) to demonstrate the accuracy and efficiency of the method. Finally, an application to the structural optimization of the polyalanine-10 system with 20 CVs is presented to demonstrate the practical promise of reinforced dynamics.

II Theory

We assume that the system we are studying has $N$ atoms, with their positions denoted by $\bm{r}=(\bm{r}_{1},\dots,\bm{r}_{N})$. The potential energy of the system is denoted by $U(\bm{r})$. Without loss of generality, we consider the system in a canonical ensemble. Given $M$ predefined CVs, denoted by $\bm{s}(\bm{r})=(s_{1}(\bm{r}),\dots,s_{M}(\bm{r}))$, the free energy defined on the CV space is

with $Z=\int e^{-\beta U(\bm{r})}\,d\bm{r}$ being the normalization factor. The brute-force way of computing the free energy (1) is to sample the CV space exhaustively and to approximate the probability distribution $p(\bm{s})$ by making a histogram of the CVs. This approach may easily become prohibitively expensive. In such a case, an alternative way of constructing the FES is to fit the mean forces acting on the CVs, i.e.,

Several ways of computing $\bm{F}(\bm{s})$ have been proposed ciccotti2005blue ; maragliano2006temperature ; abrams2008efficient . We will adopt the approach of restrained dynamics proposed in maragliano2006temperature . In this formulation, a new term is added to the potential of the system to represent the effect of the spring forces between the configuration variables and the CVs. It can be shown that the mean force is given by $F_{\alpha}(\bm{s})=\lim_{k_{\alpha}\rightarrow\infty}F_{\alpha}^{k}(\bm{s})$ for $\alpha=1,2,...,M$, where the $\alpha$-th component of $F^{k}$ is defined to be

Here $Z_{k}(\bm{s})=\int e^{-\beta U_{k}(\bm{r},\bm{s})}\,d\bm{r}$ is the normalization factor, $\{k_{\alpha}\,|\,\alpha=1,\dots,M\}$ are the spring constants for the harmonic restraining potentials, and $U_{k}(\bm{r},\bm{s})$ is defined by

In practice, the spring constants are chosen to be large enough to guarantee the convergence to the mean forces. The time duration for the restrained dynamics should be longer than the largest relaxation timescale of the fast modes of the system, in order for the ensemble average in Eq. (3) to be approximated adequately by the time average. In the rest of the paper, we do not explicitly distinguish $\bm{F}$ and $\bm{F}^{k}$.

II.2 Free energy representation

The free energy $A(\bm{s})$ will be represented by a deep neural network (DNN) model, in which the input CVs are first preprocessed, then passed through multiple fully connected hidden layers, and, in the end, mapped to the free energy. The structure of the DNN model is schematically illustrated in Fig. 1. Mathematically, a DNN representation with $N_{h}$ hidden layers is given by

is well defined since each layer of the construction (5) is differentiable, and hence the DNN representation of the free energy $\mathcal{A}(\bm{s})$ is also differentiable.

It should be noted that the design of the DNN model can be adapted to different kinds of problems. We use the fully-connected DNN model here for simplicity of discussion. For example, for some condensed systems, an alternative network model resembling the one used in the Deep Potential method should be preferred han2017deep ; zhang2017deep . We leave this to future work.

II.3 Training and uncertainty indicator

The DNN representation of the free energy is obtained by solving the following minimization problem

where $D$ denotes the set of training data and $|D|$ denotes the size of the dataset $D$. Here $\bm{\mathcal{F}}(\bm{s})$ comes from the DNN model, and $F(\bm{s})$ is the collected mean force for the data $s$. Precise ways of collecting the data will be discussed later. It should be noted that at the beginning of the training process, we have no data. Data is collected as the training process proceeds.

To guarantee accuracy for this model, we require that the CV values in $D$ is an adequate sample of the CV space. This is made difficult due to the barriers on the energy landscape. The MD will tend to be stuck at metastable states without being able to escape. To help overcome this problem, we introduce a biased dynamics. Details of that will be discussed in the next subsection.

A key notion for reinforced dynamics is the uncertainty indicator. This quantity is important in the data collection step as well as in the biased dynamics step. Our intuition is that the DNN model should produce a reasonably accurate prediction of the free energy in regions that are adequately covered by $D$, but is much less so in regions that are covered poorly by $D$ (or have not been visited by the MD). To quantify this, we introduce a small ensemble of DNN models, where the only difference between these models is the random weights used to initialize them. We can then define the uncertainty indicator as $\mathcal{E}(\bm{s})$, the standard deviation of the force predictions, viz.

where the ensemble average $\langle\cdots\rangle$ is taken over this ensemble of models. One expects that this ensemble of models give rise to predictions of the mean forces $\bm{F}$ that are close to each other in regions well covered by $D$. In the regions that are covered poorly by $D$, the predictions will scatter much more. This is confirmed by our numerical results.

Finally, it is worth noting that the minimization problem (9) is solved by the stochastic gradient descent (SGD) method combined with the back-propagation algorithm lecun2012efficient . This has become the de facto standard algorithm for training DNN models. In all the test examples, we first adopt a random initialization procedure for the weights, where each component in $\bm{W}_{p}$ in Eq. (6) is initialized from a normal distribution with mean 0 and standard deviation ${1}/{\sqrt{|\bm{d}_{p-1}|+|\bm{d}_{p}|}}$, and each component in $\bm{b}_{p}$ is initialized from a normal distribution with mean 0 and standard deviation 1. Then at each training step, the weights are updated based on the evaluation of the loss function on a small batch, or subset $B$ of the training data $D$, i.e.,

II.4 Adaptive biasing

A way of encouraging the MD to overcome the barriers in the energy landscape and escape metastable regions is to add a bias to the potential. The force on the $i$-th atom then becomes:

Since the FES is the best approximation of the potential energy in the space of CVs, it is natural to use the current approximation of the FES, with a negative sign added, as the biasing potential, as is done in metadynamics laio2002escaping ; barducci2008well . We will adopt the same strategy but we propose to switch on the biasing potential only in regions where we have low uncertainty on the DNN representation of the FES:

where the biasing potential $\mathcal{A}(\bm{s}(\bm{r}))$ is the mean of the predefined ensemble of DNN models, and $\sigma(\cdot)$ is a smooth switching function defined by

Here $e_{0}$ and $e_{1}$ are two uncertainty levels for the accuracy of the DNN model. In regions where the uncertainty indicator $\mathcal{E}(\bm{s})$ is smaller than the level $e_{0}$, the accuracy of the DNN representation of $\mathcal{A}(\bm{s})$ is adequate, and hence the system will be biased by $\mathcal{A}(\bm{s})$. In the regions where $\mathcal{E}(\bm{s})$ is larger than level $e_{1}$, the accuracy of the DNN representation is inadequate, and the system will follow the original dynamics governed by the potential energy $U(\bm{r})$. In between $e_{0}$ and $e_{1}$, the DNN model is partially used to bias the system via a rescaled force term $-\sigma(\mathcal{E}(\bm{s}(\bm{r})))\,\nabla_{\bm{r}_{i}}\mathcal{A}(\bm{s}(\bm{r}))$.

II.5 Data collection

II.6 The reinforced dynamics scheme

Fig. 2 is a flowchart of the reinforced dynamics scheme. Given an initial guess of the FES represented by the DNN, a biased MD, i.e. Eq. (14), is performed to sample the CV space from an arbitrarily chosen starting point. If no a priori information on the FES is available, then a standard MD is carried out. The visited CV values are recorded at a certain time interval and tested by the uncertainty indicator to see whether they belongs to a region with high uncertainty in the CV space. If all the newly sampled CV values from the biased MD trajectory belong to the region with low uncertainty, it can be (1) the biased MD is not long enough, so parts of the CV space are not explored, (2) the interval for recording CV values along the biased MD is not small enough, so some visited CV values belonging to the region with high uncertainty are missed, or (3) the DNN representation for FES is fully converged, then the iteration should be stopped and one can output the DNN representation for the FES, namely the mean of the predefined ensemble of models. Case (1) can be excluded by systematically increasing the length of the biased simulation. Case (2) can be excluded by decreasing the recording interval.

If CV values belonging to the region with high uncertainty are discovered, they will be added to the training dataset $D$. The CV values that are already in the training dataset should be retained and serve as training data for later iterations. The mean forces at the added CV values are computed by the restrained dynamics Eq. (3). A new ensemble of DNN models for the FES are then trained, using different random initial guesses for $\{\bm{W},\bm{b}\}$. The standard deviation of the predictions from these models is again used to estimate the uncertainty indicator $\mathcal{E}(\bm{s})$. The iteration starts again using the biased MD simulation with the new DNN models.

Finally, it is worth noting that the restrained MD simulations for mean forces, which take over most of the computation time in the reinforced dynamics scheme, are embarrassingly parallelizable. The training of the ensemble of DNN models is also easily parallelizable. Several independent walkers can be set up simultaneously for a parallelized biased simulation, and this provides a more efficient exploration of the FES. These techniques can help accelerating the data collection process and benefit large-scale simulations for complex systems.

III Numerical examples: alanine dipeptide and tripeptide

We investigate the FES of the alanine dipeptide (ACE-ALA-NME) and alanine tripeptide (ACE-ALA-ALA-NME) modeled by the Amber99SB force field hornak2006comparison . The molecules are dissolved in 342 and 341 TIP3P jorgensen1983comparison water molecules, respectively, in a periodic simulation cell. All the MD simulations are performed using the package GROMACS 5.1.4 abraham2015gromacs . The cut-off radius of the van der Waals interaction is 0.9 nm. The dispersion correction due to the finite cut-off radius is applied to both energy and pressure calculations. The Coulomb interaction is treated with smooth particle mesh Ewald method essmann1995spm with a real space cut-off 0.9 nm and reciprocal space grid spacing 0.12 nm. The system is integrated with the leap-frog scheme at timestep 2 fs. The temperature of the system is set to 300 K by velocity-rescale thermostat bussi2007canonical with a relaxation time 0.2 ps. The solute and solvent are coupled to two independent thermostats to avoid the hot-solvent/cold-solute problem lingenheil2008hot . Parrinello-Rahman barostat parrinello1981polymorphic (GROMACS implementation) with a relaxation timescale 1.5 ps and compressibility $4.5\times 10^{-5}\,\textrm{Bar}^{-1}$ is coupled to the system to control the pressure to 1 Bar. For both the alanine dipeptide and tripeptide, any covalent bond that connects a hydrogen atom is constrained by the LINCS algorithm hess1997lincs . The H-O bond and H-O-H angle of water molecules are constrained by the SETTLE algorithm miyamoto2004settle .

III.2 Free energy surface construction

The information of the four-dimensional FES of the alanine tripeptide constructed by brute-force MD sampling and the reinforced dynamics is presented in Fig. 4, by projecting on the $(\varphi_{0},\psi_{0})$, $(\varphi_{1},\psi_{1})$ and $(\varphi_{1},\psi_{0})$ planes. For example, the projection onto the $(\varphi_{0},\varphi_{0})$ variables is defined by

where $C$ is a constant that is chosen to normalize the minimum value of $A(\varphi_{0},\psi_{0})$ to zero. Projected free energies $A(\varphi_{1},\psi_{1})$ and $A(\varphi_{1},\psi_{0})$ are defined analogously. The uncertainty levels of the reinforced dynamics are set to $e_{0}=3.0$ kJ/mol/rad and $e_{1}=3.5$ kJ/mol/rad. The biased MD simulation of the 72nd iteration does not find any CV value belonging to the region with high uncertainty, so the process stops. From the 0th to the 71st iteration, 1363 CV values are added to the training dataset $D$. The total biased MD simulation time is $73\times 0.14=10.22$ ns, while the total restrained MD simulation time is $1363\times 0.14=190.82$ ns. The total wall time of the restrained MD simulations is $6.2\times 10^{4}$s, while the total wall time for training the networks is $1.1\times 10^{5}$s. For comparison, we carried out 18 independent brute-force MD simulation, each of which is 2.65 $\mu$s long, so the total length of brute-force MD trajectories is 47.7 $\mu$s. Fig. 4 shows that the reinforced dynamics is able to reproduce the FES with satisfactory accuracy on all the projected planes. It is noted that the projected FESs on both the $(\varphi_{0},\psi_{0})$ and $(\varphi_{1},\psi_{1})$ variables are different from the FES of alanine dipeptide, which indicates the correlation of backbone atoms.

III.3 Illustration of the adaptive feature

From the 5th to the 8th iteration, the DNN representation of the FES is of relatively good quality. The CV values added to the training dataset are those that sample the border of high energy peaks at $\varphi\approx 2$ rad and $\varphi\approx-0.5$ rad. At the 9th iteration, no CV value belonging to regions with high uncertainty is found because the pushing-back events happen so quickly that the CV values are not recorded by the biased MD trajectory with the 0.2 ps recording interval. However, if we reduce the CV recording interval from 0.2 ps to 0.04 ps, 19 CV values can still be identified to be in the regions with high uncertainty and used to start the next biasing-and-training iteration. Since the construction of high energy FES peaks is of less interest, for the sake of computational cost, we do not use the smaller recording interval in our result. This means that the we ignore the FES regions with sharp gradient so that the biased system can only stay for a time scale that is much shorter than the recording interval. Better stopping criteria that guarantee the representation quality of the important structures of FES and excludes the irrelevant energy peaks are left for future studies.

III.4 Remark on the choice of CVs

One important issue is to find the right set of CVs in order to capture the structural and dynamics information that we are interested in. However, this is a difficult problem and is not the topic of this work. Here, we will study how the enhanced sampling and free-energy estimation are affected when (unnecessary) additional CVs are included. We will see that the estimated free energy for the larger set of CVs is consistent with the one for the smaller set of CVs in the sense that after projecting the former onto the smaller set of CVs, one recovers the latter.

IV Application to polyalanine-10

In reinforced dynamics, both the neural network representation of the FES and the restrained simulation for mean forces are relatively insensitive to the dimensionality of the CV space. Thus it has the potential to be able to handle systems with a large set of CVs. As an illustrative example, we investigate the metastable conformations of a polyalanine-10 (ACE-(ALA)10-NME) molecule. In this example, rather than constructing an accurate free energy in the whole space of CVs, our goal is to efficiently search for the most stable structures in the conformational space of the system. We will demonstrate that reinforced dynamics allows us to explore very efficiently the most relevant metastable conformations of this molecule, including the $\alpha$-helix and $\beta$-strand conformations, and to provide estimates for the relative stability between different metastable states.

One technical remark is that for computational efficiency, we adopt a multi-walker scheme of reinforced dynamics for this relatively high-dimensional case. In each iteration of this scheme, different walkers undergo biased dynamics independently under the same biased potential. Next, a set of CV values with high uncertainty are selected and restrained simulations are performed to calculate the associated mean forces. Finally, the selected CV values and associated mean forces provided by all the walkers are merged and added to the dataset. An ensemble of new neural network models are then trained with this larger dataset. The multi-walker scheme improves the efficiency of the data collection step and it helps to accelerate the exploration procedure.

The system of polyalanine-10 (ACE-(ALA)10-NME) is modeled by the Amber96 forcefield kollman1996advances . The molecule is in the gas phase and is set in a $3.5~{}\textrm{nm}\times 3.5~{}\textrm{nm}\times 3.5~{}\textrm{nm}$ simulation region. To start with, we prepare misfolded initial configurations of the molecule in three stages. In the first stage, starting from an alpha-helix configuration, two ends of the molecule is pulled along the $z$ direction in an extended simulation region ($3.5~{}\textrm{nm}\times 3.5~{}\textrm{nm}\times 30~{}\textrm{nm}$) at rate 0.1 nm/ps for 100 ps. During this process, no thermostat is used for the system. At the end of this stage, the backbone of the molecule is fully extended, and the temperature of the system increases to 1155 K. In the second stage, the pulling force is removed and the molecule is equilibrated at 300 K for 200 ps, using the velocity-rescaling thermostat bussi2007canonical with 0.2 ps of relaxation time and an integration time step of 1 fs. In the third stage, an unbiased productive simulation of 200 ps is carried out at 300 K with a time step of 2 fs. 100 candidate configurations along the trajectory of this simulation are saved in every other 2 ps. Finally, 14 independent walkers are initialized with randomly chosen configurations from these candidates.

IV.2 Structure optimization

To find different metastable states and their relative stability, we combine the exploration stage, provided by the adaptively biasing procedure in reinforced dynamics, with an optimization stage, which can be viewed as a postprocessing of the explored configurations. In the exploration stage, due to the complexity of the 20-dimensional FES, we do not wait for the reinforced dynamics to stop by itself. Instead, we stop the process at the 210th iteration. The outputs of the biased MD simulations in each iterations, in total $14\times 211=2954$ configurations, are thus selected for the next stage. We remark that basins associated to important metastable conformations may not be visited during the 210 iterations. This seems to be a common issue of algorithms for conformation space exploration, no matter by enhanced sampling or by brute-force simulation. However, reinforced dynamics drastically accelerates the efficiency of exploration and, due to the biasing procedure, new low-energy states are more likely to be explored in earlier iterations. Although we stop the process at a certain number of iteration, further tests based on the accumulated dataset and restarted from the simulation can always be performed to check the results. In the optimization stage, the 2954 configurations are first relaxed by brute-force MD for 200 ps. Then the CV values corresponding to the relaxed configurations are taken as initial guesses for the unconstrained minimization on the DNN represented FES, which is solved by the Broyden-Fletcher-Goldfarb-Shanno (known as BFGS) method fletcher1987practical , and the solutions are local minima of the FES. The configurations are further relaxed with a restrained MD simulation centered at the corresponding local minima for 100 ps at a time step of 1 fs.

The native conformation (C004) and five metastable conformations with the lowest free energies are presented in Fig. 7. Their relative stability with respect to the native state and the standard deviations of the free energy predictions are also presented in the figure. The metastable conformation C000 corresponds to the $\beta$-strand conformation, while the metastable conformations C008, C009, C012 and C027 are misfolded conformations. The predicted free energies of the metastable conformations are very close, thus considering the uncertainties in these free energies, we can not tell whether one metastable state is more stable than another from the current reinforced dynamics simulation.

We also computed the transition paths from the native state to the five metastable state using the string method e2002string ; e2007simplified . The strings are discretized by 224 nodes. At each node a restrained MD of length 1600 ps is performed, and the CV values are recorded every 0.01 ps to compute the mean force by Eq. (3). The free energies are then computed by using thermodynamic integration along the string (see the green lines in Fig. 8). As a comparison, the free energies predicted by the reinforced dynamics along the same paths are plotted as the red lines, the standard deviations in the free energy model predictions are presented by the red shadows. The free energy predicted by the reinforced dynamics is in satisfactory agreement with the thermodynamic integration for the transitions C004$\rightarrow$C000, C004$\rightarrow$C009 and C004$\rightarrow$C012. The computation of the transition paths C004$\rightarrow$C009 and C004$\rightarrow$C012 are easier, because the $\alpha$-helical segments in the conformations C009 and C012 make them closer to the native state. It is also observed that the free energy barriers in transitions C004$\rightarrow$C009 and C004$\rightarrow$C012 are lower than others. Along the paths C004$\rightarrow$C008 and C004$\rightarrow$C027, the reinforced dynamics is quite accurate near the native and the metastable states. However, in the middle section of the paths, there are clear differences from the result of the thermodynamic integration. Many factors may contribute to this: Between the native and a metastable state, there may exist multiple transition paths; The path computed by the string method may not be the most probable path; Some conformations along the path may not be well sampled by the reinforced dynamics.

V Conclusion and perspective

In summary, reinforced dynamics is a promising tool for exploring the configuration space and calculating the free energy of atomistic systems. Even though we only presented examples of bio-molecules, it should be clear that the same strategy should also be applicable to many different tasks like studying the phase diagrams of condensed systems. In particular, due to the ability of the deep neural networks in representing high dimensional functions han2017deep ; zhang2017deep ; schneider2017stochastic ; lecun2015deep , we expect the reinforced dynamics to be particularly powerful when the dimensionality of the CV space is high. In addition, one should be able to couple it with optimization algorithms in order to perform structural optimization.

References