Neural networks-based variationally enhanced sampling

Significance

Atomistic-based simulations are one of the most widely used tools in contemporary science. However, in the presence of kinetic bottlenecks, their power is severely curtailed. In order to mitigate this problem, many enhanced sampling techniques have been proposed. Here, we show that by combining a variational approach with deep learning, much progress can be made in extending the scope of such simulations. Our development bridges the fields of enhanced sampling and machine learning and allows us to benefit from the rapidly growing advances in this area.

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 [O. Valsson, M. Parrinello, Phys. Rev. Lett. 113, 090601 (2014)] in order to determine such potential. This is achieved by expressing the bias as a neural network. The parameters are determined in a variational learning scheme aimed at minimizing an appropriate functional. This required the development of a more efficient minimization technique. The expressivity of neural networks allows representing rapidly varying free-energy surfaces, removes boundary effects artifacts, and allows several collective variables to be handled.

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 (NNs) (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\left(\mathbf{s}\right)$, which depends on the chosen CVs $\mathbf{s}$. In this context, ML has been applied in order to identify appropriate CVs (6–10) and to construct new methodologies (11–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 $\mathrm{\Omega }=\mathrm{\Omega }\left[V\left(\mathbf{s}\right)\right]$. This functional is closely related to a Kullback–Leibler (KL) divergence (18). The bias that minimizes $\mathrm{\Omega }$ is such that the probability distribution of $\mathbf{s}$ in the biased ensemble $P_V\left(\mathbf{s}\right)$ is equal to a preassigned target distribution $p\left(\mathbf{s}\right)$. 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\left(s\right)$ in a linear combination of orthonormal polynomials and use the expansion coefficients as variational parameters. Despite 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, nonoptimal CVs may lead to very slow convergence.

In this paper, we use the expressivity (25) of NNs 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 Networks-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 CVs $\mathbf{s}=\mathbf{s}\left(\mathbf{R}\right)$, 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:

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

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

$$F\left(\mathbf{s}\right)=-\frac{1}{\beta }\log P\left(\mathbf{s}\right).$$ [2]

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), it 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:

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

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

$$F\left(\mathbf{s}\right)=-V\left(\mathbf{s}\right)-\frac{1}{\beta }\log p\left(\mathbf{s}\right).$$ [4]

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

$$p_V\left(\mathbf{s}\right)=p\left(\mathbf{s}\right),$$ [5]

where $p_V\left(\mathbf{s}\right)$ is defined as:

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

In other words, $p\left(\mathbf{s}\right)$ is the distribution the CVs will follow when the $V\left(\mathbf{s}\right)$ that minimizes $\mathrm{\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 \mathrm{\Omega }\left[V\right]=D_{KL}\left(p\Vert p_V\right)-D_{KL}\left(p\Vert P\right)$ (18), where $D_{KL}$ denotes the KL divergence.

The Target Distribution.

In VES, an important role is played by the target distribution $p\left(\mathbf{s}\right)$. A careful choice of $p\left(\mathbf{s}\right)$ 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:

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

where $P\left(\mathbf{s}\right)$ 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\left(\mathbf{s}\right)$ has proven to be highly efficient (19). Since at the beginning of the simulation, $F\left(\mathbf{s}\right)$ is not known, $p\left(\mathbf{s}\right)$ 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.

NN Representation of the Bias.

The standard practice of VES has been so far of expanding linearly $V\left(\mathbf{s}\right)$ 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 NN, as shown in Fig. 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 nonlinear function $a$ (25):

$${\mathbf{x}}^{l+1}=a\left({\mathbf{w}}^{l+1}{\mathbf{x}}^l+{\mathbf{b}}^{l+1}\right).$$ [8]

Here, the nonlinear 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.

Fig. 1.

NN representation of the bias. The inputs are the chosen CVs, whose values are propagated across the network in order to get the bias. The parameters are optimized according to the variational principle of Eq. 3.

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

$${\mathbf{F}}_i=-{\nabla }_{{\mathbf{R}}_i}\hspace{0.17em}V=-\sum _{j=1}^n\frac{\partial V}{\partial s_j}{\nabla }_{{\mathbf{R}}_i}{s}_j,$$ [9]

where the first term is efficiently computed via back-propagation. The coefficients $\left\{w^i\right\}$ and $\left\{b^i\right\}$ that we lump in a single vector $\mathbf{w}$ will be our variational coefficients. With this bias representation, the functional $\mathrm{\Omega }\left[V\right]$ 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 NN applications, training plays a crucial role. The functional $\mathrm{\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 $\mathrm{\Omega }$ derivatives:

$$\frac{\partial \mathrm{\Omega }}{\partial \mathbf{w}}=-{⟨\frac{\partial V}{\partial \mathbf{w}}⟩}_V+{⟨\frac{\partial V}{\partial \mathbf{w}}⟩}_p,$$ [10]

where the first average is performed over the system biased by $V\left(\mathbf{s}\right)$ and the second over the target distribution $p\left(\mathbf{s}\right)$. 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 NN 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\left(\mathbf{s}\right)$ 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 of 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), which 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 NN community. Namely, we do not aim at reaching the minimum, rather at determining a value of the bias $V_s\left(s\right)$ that is close enough to the optimum value to be useful. In the present context, this means that a run biased by $V_s\left(s\right)$ can be used to compute the Boltzmann equilibrium averages from the well-known umbrella sampling-like formula:

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

In order to measure the progression of the minimization toward Eq. 5, we monitor at iteration step $n$ an approximate KL divergence between the running estimates of the biased probability $p_V^{\left(n\right)}\left(s\right)$ and the target $p^{\left(n\right)}\left(s\right)$:

$$D_{KL}^{\left(n\right)}\left(p_V\Vert p\right)=\sum _{\mathbf{s}}{p}_V^{\left(n\right)}\left(\mathbf{s}\right)\log \frac{p_V^{\left(n\right)}\left(\mathbf{s}\right)}{p^{\left(t\right)}\left(\mathbf{s}\right)}.$$ [12]

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 3 parts. In the first one, the ADAM optimizer (34) is used, until a preassigned threshold value $ϵ$ for the distance $D_{KL}\left(p_V\Vert p\right)$ is reached. From then on, the learning rate is exponentially brought to zero provided that $D_{KL}$ remains below $ϵ$. 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 3 needs 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 2D Wolfe–Quapp potential, rotated as in ref. 24. This is shown in Fig. 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 suboptimal 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).

Fig. 2.

Results for the 2D model. (A) Potential energy surface of the model. (B, Upper) CV evolution as a function of the number of iterations. The points are colored according to their energy value. (B, Lower) Evolution of the KL divergence between the bias and the target distribution (green) and learning rate scaling factor (gray). When the KL divergence is lowered below the threshold value, the learning rate is decreased exponentially until it is practically zero. (C) Free-energy profiles obtained from the NN and with the reweighting procedure, compared with the reference obtained by integrating the model. For the DEEP-VES curve, also an estimate of the error is given by averaging the results of 8 different simulations. The lack of accuracy on the left shoulder is a consequence of the combination of the suboptimal character of the chosen variable and the strength of the bias, which creates new pathways other than the one of minimum energy. These artifacts are removed by performing the reweight with the static bias potential $V_s$. (D) RMSE on the FES computed in the regions of 10 $k_{B}T$ from the reference minimum.

We use a 3-layer NN with [48,24,12] nodes, resulting in 1,585 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 $ϵ$ 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 (SI Appendix).

In the upper right panel (Fig. 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 suboptimal 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 suboptimal 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\left(\mathbf{s}\right)$ (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 $ϵ$; 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 2 small peptides, alanine dipeptide (Fig. 4) and alanine tetrapeptide (Fig. 5) 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 ${\varphi }_i$ and ${\psi }_i$, where the first ones are connected to the slowest kinetic processes. The smaller Ala2 has only 1 pair of such dihedral angles, which we will denote as $\left\{\varphi ,\psi \right\}$, while Ala4 has 3 pairs of backbone angles denoted by $\left\{{\varphi }_i,{\psi }_i\right\}$ with $i=\mathrm{1,2,3}$.

Fig. 4.

Alanine dipeptide.

Fig. 5.

Alanine tetrapeptide.

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 dihedral angles $\left\{\varphi ,\psi \right\}$ 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: $\left\{\varphi ,\psi \right\}\to \left\{\cos \left(\varphi \right),\sin \left(\varphi \right),\cos \left(\psi \right),\sin \left(\psi \right)\right\}$.

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 (SI Appendix). In Fig. 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 20 kJ/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.

Fig. 3.

Alanine dipeptide free-energy results. (A) DEEP-VES representation of the bias. (B) FES profile obtained with reweighting. (C) Reference from a 100-ns metadynamics simulation. The main features of the FES are captured by the NN, which allows for an efficient enhanced sampling of the conformations space. Finer details can be easily recovered with the reweighting procedure.

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 6 dihedral angles $\left\{{\varphi }_1,{\psi }_1,{\varphi }_2,{\psi }_2,{\varphi }_3,{\psi }_3\right\}$ (Fig. 6). 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 6D 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\left(\mathbf{s}\right)$ 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 of 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 Appendix).

Fig. 6.

Two-dimensional FESs of Ala4 obtained with the reweighting, as a function of the 3 dihedral angles $\left\{{\varphi }_1,{\varphi }_2,{\varphi }_3\right\}$. The FES in terms of other pairs, as well the comparison with the references, is reported in SI Appendix.

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 CV. 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 NNs, which we use here to represent rapidly varying features of the FES and to avoid boundaries artifacts. In Fig. 7, we show the free-energy profile learned 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 2 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 30 ns.

Fig. 7.

FES of silicon crystallization, in terms of the number of cubic diamonds atoms in the system. Snapshots of the 2 minima, as well as the transition state, are also shown. The reweight is obtained using $V_s$.

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 3 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 nonoptimal 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 learned and the time needed to converge for the Wolfe–Quapp potential.

Conclusions

In this work, we have shown how the flexibility of NNs 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 6D 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 NN 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 KL divergence between the current and the target probability distribution. We think that conventional VES can also benefit from our 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 CVs on-the-fly. This work also allows tapping into the immense literature on ML and NNs 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). All data and input files required to reproduce the results reported in this paper are available on PLUMED-NEST (www.plumed-nest.org), the public repository of the PLUMED consortium (39), as plumID:19.060 (40). We plan to release the code also in the open-source PLUMED package. We take care of the construction and the optimization of the NN 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 of the examples reported in the paper.

Wolfe–Quapp Potential.

The Wolfe–Quapp potential is a fourth-order polynomial: $U\left(x,y\right)=x^4+y^4-2\hspace{0.17em}{x}^2-4\hspace{0.17em}{y}^2+xy+0.3\hspace{0.17em}x+0.1y$. We rotated it by an angle $\theta =-3/20\hspace{0.17em}\pi$ in order to change the direction of the path connecting the 2 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. An NN 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 $ϵ=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 (41) patched with PLUMED. The peptides are simulated in the canonical ensemble using the Amber99-SB force field (42) with a time step of 2 fs. The target temperature of 300 K is controlled with the velocity rescaling thermostat (43). For Ala2, we use the following parameters: decay constant equal to $10^3$ and threshold equal to $ϵ=0.5$. A grid of 50 × 50 bins is used. For Ala4, the Metropolis algorithm is used to generate a set of points according to the target distribution. At every iteration, 25,000 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 bias factor is equal to 10.

Silicon.

For the silicon simulations, we use Large-Scale Atomic/Molecular Massively Parallel Simulator (44) patched with PLUMED, employing the Stillinger and Weber potential (45). A 3 × 3 × 3 supercell (216 atoms) is simulated in the isothermal-isobaric ensemble with a timestep of 2 fs. The temperature of the thermostat (43) is set to 1,700 K with a relaxation time of 100 fs, while the values for the barostat (46) are 1 atm and 1 ps. The CV used is the number of cubic diamond atoms in the system, defined according to ref. 47. The decay time for the learning rate is $2\cdot 10^3$, and a grid of 216 bins is used. The bias factor used is equal to 100.