Neural networks-based variationally enhanced sampling

TL;DR

This paper introduces a neural network-based approach to variationally enhanced sampling, improving the efficiency of sampling complex free energy surfaces in atomistic simulations by leveraging machine learning and reinforcement learning techniques.

Contribution

It develops a neural network implementation of the bias potential in variationally enhanced sampling, enabling better handling of complex free energy landscapes and multiple collective variables.

Findings

Accelerates sampling of complex free energy surfaces.

Reduces boundary effects artifacts.

Handles multiple collective variables more effectively.

Abstract

Sampling complex free energy surfaces is one of the main challenges of modern atomistic simulation methods. The presence of kinetic bottlenecks in such surfaces often renders a direct approach useless. A popular strategy is to identify a small number of key collective variables and to introduce a bias potential that is able to favor their fluctuations in order to accelerate sampling. Here we propose to use machine learning techniques in conjunction with the recent variationally enhanced sampling method [Valsson and Parrinello, Physical Review Letters 2014] to determine such potential. This is achieved by expressing the bias as a neural network. The parameters are determined in a reinforcement learning scheme aimed at minimizing the variationally enhanced sampling functional. This required the development of a new and more efficient minimization technique. The expressivity of neural networks allows accelerating sampling in systems with rapidly varying free energy surfaces, removing boundary effects artifacts, and making one more step towards being able to handle several collective variables.

Figures

Equations

$$P(\mathbf{s})=\int d\mathbf{R}\frac{e^{-\beta U(\mathbf{R})}}{Z}\delta \left(\mathbf{s}-\mathbf{s}(\mathbf{R})\right)$$

$$F(\mathbf{s})=-\frac{1}{\beta }\log P(\mathbf{s}).$$

$$\Omega [V]=\frac{1}{\beta }\log \frac{\int d\mathbf{s}{e}^{-\beta \left(F(\mathbf{s})+V(\mathbf{s})\right)}}{\int d\mathbf{s}{e}^{-\beta F(\mathbf{s})}}+\int d\mathbf{s}p(\mathbf{s})V(\mathbf{s})$$

$$F(\mathbf{s})=-V(\mathbf{s})-\frac{1}{\beta }\log p(\mathbf{s})$$

$$p_V(\mathbf{s})=p(\mathbf{s})$$

$$p_V(\mathbf{s})=\frac{e^{-\beta \left(F(\mathbf{s})+V(\mathbf{s})\right)}}{\int d\mathbf{s}{e}^{-\beta (F\left(\mathbf{s})+V(\mathbf{s})\right)}}$$

$$p(\mathbf{s})=\frac{e^{-\beta F(\mathbf{s})/\gamma }}{\int d\mathbf{s}{e}^{-\beta F(\mathbf{s})/\gamma }}\propto {\left[P(\mathbf{s})\right]}^{1/\gamma }$$

$${\mathbf{x}}^{l+1}=a({\mathbf{w}}^{l+1}{\mathbf{x}}^l+{\mathbf{b}}^{l+1})$$

$${\mathbf{F}}_i=-{\nabla }_{{\mathbf{R}}_i}V=-\sum _{j=1}^n\frac{\partial V}{\partial s_j}{\nabla }_{{\mathbf{R}}_i}{s}_j$$

$$\frac{\partial \Omega }{\partial \mathbf{w}}=-{⟨\frac{\partial V}{\partial \mathbf{w}}⟩}_V+{⟨\frac{\partial V}{\partial \mathbf{w}}⟩}_p$$

$$⟨O(\mathbf{R})⟩=\frac{{⟨O(\mathbf{R})e^{\beta V_s(s(\mathbf{R}))}⟩}_{V_s}}{{⟨e^{\beta V_s(s(\mathbf{R}))}⟩}_{V_s}}.$$

$$D_{KL}^{(n)}(p_V||p)=\sum _{\mathbf{s}}{p}_V^{(n)}(\mathbf{s})\log \frac{p_V^{(n)}(\mathbf{s})}{p^{(t)}(\mathbf{s})}$$

Full text

K****eywords Molecular dynamics $|$ Enhanced sampling $|$ Machine learning

Machine learning (ML) is changing the way in which modern science is conducted. Atomistic based computer simulations are no exceptions. Since the work of Behler and Parrinello [1] neural networks [2, 3] or Gaussian processes [4] are now almost routinely used to generate accurate potentials. More recently ML methods have been used to accelerate sampling, a crucial issue in molecular dynamics (MD) simulations, where standard methods allow only a very restricted range of time scales to be explored. An important family of enhanced sampling methods is based on the identifications of suitable collective variables (CVs) that are connected to the slowest relaxation modes of the system [5]. Sampling is then enhanced by constructing an external bias potential $V(\mathbf{s})$ which depends on the chosen collective variables $\mathbf{s}$. In this context machine learning has been applied in order to identify appropriate CVs [6, 7, 8, 9, 10] and to construct new methodologies [11, 12, 13, 14, 15]. From these early experiences, it is also clear that ML applications can in turn profit from enhanced sampling [16].

Here we shall focus on a relatively new method, called Variationally Enhanced Sampling (VES) [17]. In VES the bias is determined by minimizing a functional $\Omega=\Omega[V(\mathbf{s})]$. This functional is closely related to a Kullback Leibler divergence [18]. The bias that minimizes $\Omega$ is such that the probability distribution of $\mathbf{s}$ in the biased ensemble $P_{V}(\mathbf{s})$ is equal to a pre-assigned target distribution $p(\mathbf{s})$ . The method has been shown to be flexible [19, 20] and has great potential also for applications different from enhanced sampling. Examples of these heterodox applications are the estimation of the parameters of Ginzburg-Landau free energy models [18], the calculation of critical indexes in second-order phase transitions [21], and the sampling of multithermal-multibaric ensembles [22].

Although different approaches have been suggested [23, 24], the way in which VES is normally used is to expand $V(s)$ in a linear combination of orthonormal polynomials and use the expansion coefficients as variational parameters. In spite of its many successes, VES is not without problems. The choice of the basis set is often a matter of computational expediency and not grounded on physical motivations. Representing sharp features in the VES may require many terms in the basis set expansion. The number of variational parameters scales exponentially with the number of CVs and can become unmanageably large. Finally, non-optimal CVs may lead to very slow convergence.

In this paper, we use the expressivity [25] of neural networks (NN) to represent the bias potential and a stochastic steepest descent framework for the determination of the NN parameters. In so doing we have developed a more efficient stochastic optimization scheme that can also be profitably applied to more conventional VES applications.

Neural network based VES

Before illustrating our method, we recall some ideas of CV-based enhanced sampling methods and particularly VES.

Collective variables

It is often possible to reduce the description of the system to a restricted number of collective variables $\mathbf{s}=\mathbf{s}(\mathbf{R})$, functions of the atomic coordinates $\mathbf{R}$, whose fluctuations are critical for the process of interest to occur. We consider the equilibrium probability distribution of these CVs as:

[TABLE]

where Z is the partition function of the system, $U(\mathrm{R})$ its potential energy and $\beta=(k_{B}T)^{-1}$ the inverse temperature. We can define the associated free energy surface (FES) as the logarithm of this distribution:

[TABLE]

Then an external bias is built as a function of the chosen CVs in order to enhance sampling. In umbrella sampling [26] the bias is static, while in Metadynamics [27] is iteratively built as a sum of repulsive Gaussians centered on the points already sampled.

The variational principle

In VES a functional of the bias potential is introduced:

[TABLE]

where $p(\mathbf{s})$ is a chosen target probability distribution. The functional $\Omega$ is convex [17] and the bias that minimizes it is related to the free energy by the simple relation:

[TABLE]

At the minimum the distribution of the CVs in the biased ensemble is equal to the target distribution:

[TABLE]

where $p_{V}(\mathbf{s})$ is defined as:

[TABLE]

In other words, $p(\mathbf{s})$ is the distribution the CVs will follow when the $V(\mathbf{s})$ that minimizes $\Omega$ is taken as bias. This can be seen also from the perspective of the distance between the distribution in the biased ensemble and the target one. The functional can be indeed written as $\beta\Omega[V]=D_{KL}(p\ ||\ p_{V})-D_{KL}(p\ ||\ P)$ [18], where $D_{KL}$ denotes the Kullback-Leibler (KL) divergence.

The target distribution

In VES an important role is played by the target distribution $p(\mathbf{s})$. A careful choice of $p(\mathbf{s})$ may focus sampling in relevant regions of the CVs space and in general accelerate convergence [28]. This freedom has been taken advantage of in the so called well-tempered VES [19]. In this variant, one takes inspiration from well-tempered metadynamics [29] and targets the distribution:

[TABLE]

where $P(\mathbf{s})$ is the distribution in the unbiased system and $\gamma>1$ is a parameter that regulates the amplitude of the $\mathbf{s}$ fluctuations. This choice of $p(\mathbf{s})$ has proven to be highly efficient [19]. Since at the beginning of the simulation $F(\mathbf{s})$ is not known, $p(\mathbf{s})$ is determined in a self-consistent way. Thus also the target evolves during the optimization procedure. In this work we update the target distribution every iteration, although less frequent updates are also possible

Neural network representation of the bias

The standard practice of VES has been so far of expanding linearly $V(\mathbf{s})$ on a set of basis functions and use the expansion coefficients as variational parameters. Here this procedure is circumvented as the bias is expressed as a deep neural network as shown in figure 1. We call this variant DEEP-VES. The inputs of the network are the chosen CVs and this information is propagated to the next layers through a linear combination followed by the application of a non-linear function $a$ [25]:

[TABLE]

Here the non linear activation function is taken to be a rectified linear unit. In the last layer only a linear combination is done, and the output of the network is the bias potential.

We are employing NN since they are smooth interpolators. Indeed the NN representation ensures that the bias is continuous and differentiable. The external force acting on the i-th atom can be then recovered as:

[TABLE]

where the first term is efficiently computed via back-propagation. The coefficients $\{w^{i}\}$ and $\{b^{i}\}$ that we lump in a single vector $\mathbf{w}$ will be our variational coefficients. With this bias representation, the functional $\Omega[V]$ becomes a function of the parameters $\mathbf{w}$. Care must be taken to preserve the symmetry of the CVs, such as the periodicity. In order to accelerate convergence, we also standardize the input to have mean zero and variance one [30].

The optimization scheme

As in all neural network applications, training plays a crucial role. The functional $\Omega$ can be considered a scalar loss function with respect to the set of parameters $\mathbf{w}$. We shall evolve the parameters following the direction of the $\Omega$ derivatives:

[TABLE]

where the first average is performed over the system biased by $V(\mathbf{s})$ and the second over the target distribution $p(\mathbf{s})$. It is worth noting that the form of the gradients is analogous to unsupervised schemes in energy-based models, where the network is optimized to sample from a desired distribution [31].

At every iteration, the sampled configurations are used to compute the first term of the gradient. The second term, which involves an average over the target distribution, can be computed numerically when the number of CVs is small or using a Monte Carlo scheme when the number of CV is large. In doing so, the exponential growth of the computational cost with respect to the number of CV can be circumvented. This scheme allows the neural network to be optimized on-the-fly, without the need to stop the simulation and fit the network to a given dataset. In this respect, the approach used here resembles a reinforcement learning algorithm. The potential $V(\mathbf{s})$ can be regarded indeed as a stochastic policy that drives the system to sample the target distribution in the CVs space. This is a type of variational learning akin to what is done in applications of neural networks to quantum physics [32].

Since the calculation the gradients requires performing statistical averages, a stochastic optimization method is necessary. In standard VES applications the method of choice has so far been the one of Bach and Moulines [33], that allows reaching with high accuracy the minimum, provided that the CVs are of good quality. Applying the same algorithm to NNs is rather complex. Therefore we espouse the prevailing attitude in the neural networks community. Namely, we do not aim at reaching the minimum rather at determining a value of the bias $V_{s}(s)$ that is close enough to the optimum value to be useful. In the present context this means that a run biased by $V_{s}(s)$ can be used to compute the Boltzmann equilibrium averages from the well-known umbrella sampling-like formula:

[TABLE]

In order to measure the progression of the minimization towards eq. 5 we monitor at iteration step $n$ an approximate KL divergence between the running estimates of the biased probability $p_{V}^{(n)}(s)$ and the target $p^{(n)}(s)$:

[TABLE]

These quantities are estimated as exponentially decaying averages. In such a way only a limited number of configurations contribute to the running KL divergence.

The simulation is thus divided into three parts. In the first one the ADAM optimizer [34] is used, until a pre-assigned threshold value $\epsilon$ for the distance $D_{KL}(\ p_{V}\ ||\ p\ )$ is reached. From then on the learning rate is exponentially brought to zero provided that $D_{KL}$ remains below $\epsilon$. From this point on, the network is no longer updated and statistics is accumulated using eq. 11. The longer the second part, the better the bias potential $V_{s}$, and the shorter phase three need to be. Faster decay times instead need to be followed by longer statistics accumulation runs. However, in judging the relative merits of these strategies it must be taken into account that the third phase in which the bias is kept constant involves a minor number of operations and it is therefore much faster.

Results

Wolfe-Quapp potential

We first focus on a toy-model, namely the two-dimensional Wolfe-Quapp potential, rotated as in ref. [24]. This is shown in figure 2a. The reason behind this choice is that the dominant fluctuations that lead from one state to the other are in an oblique direction with respect to $x$ and $y$. We choose on purpose to use only $x$ as a CV in order to exemplify the case of a sub-optimal CV. This is representative of what happens in the practice when, more often than not, some slow degrees of freedom are not fully accounted for [24].

We use a three-layer neural network with [48,24,12] nodes, resulting in 1585 variational parameters, which are updated every 500 steps. The KL divergence between the biased and the target distribution is computed with an exponentially decaying average with a time constant of $5\cdot 10^{4}$ iterations. The threshold $\epsilon$ is set to 0.5 and the learning rate decay time to $5\cdot 10^{3}$ iterations. The results are robust with respect to the choice of these parameters, see the Supporting Information (SI).

In the upper right panel (figure 2b) we show the evolution of the CV as a function of the number of iterations. At the beginning, the bias changes very rapidly with large fluctuations that help to explore the configuration space. It should be noted that the combination of a sub-optimal CV and a rapidly varying potential might lead the system to pathways different from the lowest energy one. For this reason, it is important to slow down the optimization in the second phase, where the learning rate is exponentially lowered until it is practically zero. When this limit is reached the frequency of transitions becomes lower, reflecting the sub-optimal character of the CV [24], as can be seen in fig. 2b.

Although the bias itself is not yet fully converged, still the main features of the FES are captured by $V_{s}(\mathbf{s})$ (fig. 2c). This is ensured by the fact that while decreasing the learning rate the running estimate of the KL divergence must stay below the threshold $\epsilon$, otherwise the optimization proceeds with a constant learning rate. This means that the final potential is able to promote efficiently transitions between the metastable states. The successive static reweighting refines the FES and leads to a very accurate result (fig. 2d), removing also the artifacts caused in the first part by the rapidly varying potential.

Alanine dipeptide and tetrapeptide

As a second example we consider the case of two small peptides, Alanine Dipeptide and Alanine Tetrapeptide in vacuum, which are often used as a benchmark for enhanced sampling methods. We will refer to them as Ala2 and Ala4 respectively. Their conformations can be described in terms of the Ramachandran angles $\phi_{i}$ and $\psi_{i}$, where the first ones are connected to the slowest kinetic processes. The smaller Ala2 has only one pair of such dihedral angles, which we will denote as $\{\phi,\psi\}$, while Ala4 has three pairs of backbone angles denoted by $\{\phi_{i},\psi_{i}\}$ with $i=1,2,3$.

We want to show here the usefulness of the flexibility provided by DEEP-VES in different systems. For this purpose, we use the same architecture and optimization scheme as in the previous example. We only decrease the decay time of the learning rate in the second phase since the diedral angles $\{\phi,\psi\}$ are known to be a good set of CVs. In order to enforce the periodicity of bias potential, the angles are first transformed in their sines and cosines: $\{\phi,\psi\}\rightarrow\{\cos(\phi),\sin(\phi),\cos(\psi),\sin(\psi)\}$.

The simulation of Ala2 reached the KL divergence threshold after 3 ns, and from there on the learning rate was exponentially decreased. The NN bias was no longer updated after 12 ns. At variance with the first example, when the learning is slowed down and even when it is stopped the transition rate is not affected: this is a signal that our set of CVs is good (see SI). In figure 3 we show the free energy profiles obtained from the NN bias following eq. 4 and the one recovered with the reweighting procedure of eq. 11. We compute the root mean square error (RMSE) on the FES up to 20kJ/mol as in ref. [19]. The errors are 1.5 and 0.45 kJ/mol, corresponding to 0.6 and 0.2 $k_{B}T$, both well below the threshold of chemical accuracy.

In order to exemplify the ability of DEEP-VES to represent functions of several variables we study the FES of Ala4 in terms of its six dihedral angles $\{\phi_{1},\psi_{1},\phi_{2},\psi_{2},\phi_{3},\psi_{3}\}$ (fig. 5). In the case of a small number of variables the second integral on the right-hand side of eq. 10 is calculated numerically on a grid. This is clearly not possible in a six-dimensional space, thus we resort to a Monte Carlo estimation based on the Metropolis algorithm [35]. Convergence is helped by the fact that the target $p(\mathbf{s})$ is broadened by the well tempered factor $\gamma$ (eq. 7). In this case, we cannot use eq. 12 to monitor convergence. However, as the number of CV treated increases the probability of the space spanned by the CVs of covering all the relevant slow modes is very high. Thus less attention needs to be devoted to the optimization schedule. The transition from the constant learning rate phase to the static bias region can be initiated after multiple recrossings between metastable states are observed. In order to verify the accuracy of these results, we compared the free energy profiles obtained from the reweight with the results from metadynamics and its Parallel Bias version [36] (SI).

Silicon crystallization

In the last example, the phase transition of silicon from liquid to solid at ambient pressure is studied. This is a complex phenomenon, characterized by a high barrier around 80 $k_{B}T$, and several CVs have been proposed to enhance this process. The crystallization and the melting process are marked by a local character. As a consequence, many CVs are defined in terms of per-atom crystallinity measures. Then the number of solid-like atoms is used as collective variable. In terms of these variables, the free energy profile is characterized by very narrow minima, around 0 and $N_{atoms}$. In order to enhance in an efficient way the fluctuations in such space one often needs to employ variants of metadynamics [37] or a high order basis set expansion of the VES.

This can be easily addressed using the flexibility of neural networks, which we use here to represent rapidly varying features of the FES and to avoid boundaries artifacts. In figure 6 we show the free energy profile learnt by the NN as well the one recovered by reweighting. Even in the presence of very sharp boundaries, the NN bias is able to learn a good $V_{s}$, which allows for going back and forth efficiently between the two states and for recovering a good estimate of the FES in the subsequent reweight. The KL divergence threshold is reached in 10 ns, and the bias is no longer updated after 30ns.

Choice of the parameters

After presenting the results, we would like to spend a few words on the choice of parameters. This procedure requires setting three parameters, namely the time scale for calculating the KL divergence, the KL threshold and the decay time for the learning rate. Our experience is that if the chosen set of CVs is good this choice has a limited impact on the simulation result. However, in the case of non-optimal CVs, more attention is needed.

The KL divergence should be calculated on a time scale in which the system samples the target distribution, so the greater the relaxation time of the neglected variables, the greater this scale should be. However, this is only used to monitor convergence and therefore it is possible to choose a higher value without compromising speed. A similar argument applies to the decay constant of the learning rate, but in the spirit of reaching a constant potential quickly, one tries to keep it lower (in the order of thousands of iterations). Finally, we found that the protocol is robust with respect to the epsilon parameter of the KL divergence threshold, provided that it is chosen in a range of values around 0.5. In the case of FES with larger dimensionality, it may be appropriate to increase this parameter in order to reach rapidly a good estimate of $V_{s}$. In the supplementary information, we report a study of the influence of these parameters in the accuracy of the bias learnt and the time needed to converge for the Wolfe-Quapp potential.

Conclusions

In this work, we have shown how the flexibility of neural networks can be used to represent the bias potential and the FES as well, in the context of the VES method. Using the same architecture and similar optimization parameters we were able to deal with different physical systems and FES dimensionalities. This includes also the case in which some important degree of freedom is left out from the set of enhanced CVs. The case of Alanine Tetrapeptide with a 6-dimensional bias already shows the capability of DEEP-VES of dealing with high-dimensional FES. We plan to extend this to even higher dimensional landscapes, where the power of NNs can be fully exploited.

Our work is an example of a variational learning scheme, where the neural network is optimized following a variational principle. Also, the target distribution allows an efficient sampling of the relevant configurational space, which is particularly important in the optics of sampling high dimensional FES.

In the process, we have developed a minimization procedure alternative to that of ref. [17], which globally tempers the bias potential based on a Kullback Leibler divergence between the current and the target probability distribution. We think that also conventional VES can benefit from our new approach and that we have made another step in the process of sampling complex free energy landscapes. Future goals might include developing a more efficient scheme to exploit the variational principle in the context of NNs, as well as learning not only the bias but also the collective variables on-the-fly. This work allows also tapping into the immense literature on machine learning and neural networks for the purpose of improving enhanced sampling.

Materials and methods

The VES-NN is implemented on a modified version of PLUMED2 [38], linked against LibTorch (PyTorch C++ library). We plan to release the code in the open-source PLUMED package in the future. We take care of the construction and the optimization of the neural network with LibTorch. The gradients of the functional with respect to the parameters are computed inside PLUMED according to eq. 10. The first expectation value is computed by sampling, while the second is obtained by numerical integration over a grid or with Monte Carlo techniques. In the following we report the setup for the VES-NN and the simulations for all the examples reported in the paper.

Wolfe-Quapp potential

The Wolfe-Quapp potential is a fourth-order polynomial: $U(x,y)=x^{4}+y^{4}-2\ x^{2}-4\ y^{2}+xy+0.3\ x+0.1y$. We rotated it by an angle $\theta=-3/20\ \pi$ in order to change the direction of the path connecting the two minima. A langevin dynamics is run with PLUMED using a timestep of 0.005, a target temperature of 1 and a friction parameter equal to 10 (in terms of natural units). The biasfactor for the well-tempered distribution is equal to 10. A neural network composed by [48,24,12] nodes is used to represent the bias. The learning rate is equal to 0.001. An optimization step of the NN is performed every 500 timesteps. The running KL divergence is computed on a timescale of $5\cdot 10^{4}$ iterations. The threshold is set to $\epsilon=0.5$ and the decay constant for the learning rate is set to $5\cdot 10^{3}$ iterations. The grid for computing the target distribution integrals has 100 bins. The parameters of the NN are the kept the same also for the following examples, unless otherwise stated.

Peptides

For the Alanine dipeptide (Ace-Ala-Nme) and tetrapeptide (Ace-Ala3-Nme) simulations we use GROMACS [39] patched with PLUMED. The peptides are simulated in the NVT ensemble using the Amber99-SB force field [40] with a time step of 2 fs. The target temperature of 300K is controlled with the velocity rescaling thermostat [41]. For Ala2 we use the following parameters: decay constant equal to $10^{3}$ and threshold equal to $\epsilon=0.5$. A grid of 50x50 bins is used. For Ala4 the Metropolis algorithm is used to generate a set of points according to the target distribution. At every iteration 25000 points are sampled. The learning rate is kept constant for the first 10 ns, and then decreased with a decay time of $2\cdot 10^{3}$ iterations. In both cases the biasfactor is equal to 10.

Silicon

For the silicon simulations we use LAMMPS [42] patched with PLUMED, employing the Stillinger and Weber potential [43]. A 3x3x3 supercell (216 atoms) is simulated in the NPT ensemble with a timestep of 2 fs. The temperature of the thermostat [41] is set to 1700K with a relaxation time of 100 fs, while the values for the barostat [44] are 1 atm and 1 ps. The CV used is the number of cubic diamond atoms in the system, defined according to ref. [45].The decay time for the learning rate is $2\cdot 10^{3}$, and a grid of 216 bins is used. The biasfactor used is equal to 100.