Accurate and Efficient Multi-Level Free Energy Simulations with Neural Network-Assisted Enhanced Sampling

Abstract

Free energy differences $\left(\Delta F\right)$ are essential to a quantitative characterization and understanding of chemical and biological processes. Their direct estimation with an accurate quantum mechanical potential is of great interest and yet impractical due to high computational cost and incompatibility with typical alchemical free energy protocols. One promising solution is the multi-level free energy simulation in which the estimate of $\Delta F$ at an inexpensive low level of theory is combined with the correction toward a higher level of theory. The poor configurational overlap generally expected between the two levels of theory, however, presents a major challenge. We overcome this challenge using a deep neural network model and enhanced sampling simulations. An adversarial autoencoder is used to identify a low-dimensional (latent) space that compactly represents the degrees of freedom that encode the distinct distributions at the two levels of theory. Enhanced sampling in this latent space is then used to drive the sampling of configurations that predominantly contribute to the free energy correction. Results for both gas phase and condensed phase systems demonstrate that this data-driven approach offers high accuracy and efficiency with great potential of scalability to complex systems.

Graphical Abstract

1. Introduction

The reliable estimation of free energy differences $\left(\Delta F\right)$ between two thermodynamic states of a (bio)molecular system has been an active research area for decades.^1–4 $\Delta F$ is not only closely tied to various important quantities such as binding constant, ${\text{pK}}_a$, and solubility, but also able to provide compelling insights into processes for which direct experimental information is difficult or even impossible to obtain.^5,6 As a result, free energy simulations (FES) have been applied to diverse areas including drug discovery⁷ and protein design/engineering.⁸ Although conventional methods based on molecular mechanical (MM) force fields (FF) and molecular dynamics (MD) simulations have been widely employed,^7,9,10 FES with quantum mechanical (QM) methods have increasingly gained interest as they are more accurate than MM models for the description of intramolecular conformations and intermolecular interactions,^11–15 and they generally require little to none empirical parameterizations. Moreover, by treating reactive groups with a QM method and the surrounding environment with MM FF, the combined QM/MM method has become a routine tool in the study of covalent inhibition and enzyme catalysis.^16,17 However, the high computational cost of QM methods significantly limits the practical application of QM-based FES; another technical challenge is that QM methods are not readily compatible with the typical alchemical free energy protocols¹⁸. Consequently, it is imperative to further develop free energy methods that can benefit from both the accuracy of QM methods and the efficiency of MM models.

One promising framework along this line is the multi-level free energy simulation (Figure 1a),^19–25 which is also referred to as “indirect cycle free energy simulation”,²⁴ “bookending method”,²⁶ or “end state correction scheme”²⁷ in the literature. Using a thermodynamic cycle, to obtain the free energy difference between two states A and B with an accurate but expensive method (the high level of theory, H), one can perform extensive configurational sampling with an efficient but less accurate method (the low level of theory, L), and then apply the low-to-high free energy correction, $\Delta F_{A/B}^{L\to H}$, to the end states (eq. 1):

$$\Delta F_{A\to B}^H=-\Delta F_A^{L\to H}+\Delta F_B^{L\to H}+\Delta F_{A\to B}^L.$$ (1)

Figure 1: The intermediate model potentially improves the convergence of the correction term in multi-level free energy simulation.

a A thermodynamic cycle is employed to compute the free energy difference between states A and B at the high level (H) by extensively sampling the low level (L) and improving the accuracy to the high level through correction terms (vertical legs). An intermediate model (M) is introduced to improve the convergence of the correction terms. Note that only the intermediate model for state B is shown explicitly for simplification (denoted as $B^M$ in the yellow rectangle). b The potential energy difference $\left(\Delta U\right)$ distributions of the low and high levels $\left(P_L\left(\Delta U\right)\right.$ and $P_H\left(\Delta U\right)$, respectively) often barely overlap for realistic molecular systems, which limits the reliability of methods that rely only on samples from the low level. The distribution of the intermediate model $\left(P_M\left(\Delta U\right)\right)$ bridges the distributions of the two levels of theory and therefore can significantly increase the computational efficiency and accuracy.

To completely avoid sampling at the high level, the single-sided free energy estimator, free energy perturbation (FEP),²⁸ is the most rudimental method for computing $\Delta F_{A/B}^{L\to H}$. A well-known limitation of FEP is that it is reliable only when a sufficient amount of configurations sampled at the low level are representative of those highly populated at the high level^29,30. A particularly relevant and quantitative measure of the similarity between the two configurational ensembles concerns the distribution of potential energy difference $\left(\Delta U\right)$ between the two levels, the central quantity that enters FEP.^31,32 Reliable estimates can be obtained when the distribution of $\Delta U$ in the forward perturbation (low-to-high) overlaps significantly with that in the backward perturbation (high-to-low). In practice, the two levels of theory often feature substantial differences in certain degrees of freedom (DOFs), which lead to a poor $\Delta U$ distribution overlap³³ and introduce large uncertainties into FEP results. Accordingly, numerous methodologies have been proposed to improve the evaluation of the correction term, $\Delta F_{A/B}^{L\to H}$.

In general, two levels of theory may differ in both stiff (e.g., bonds and angles) and soft (e.g., dihedral angles and collective motions) degrees of freedoms. For small differences in the stiff DOFs, corrections to the free energy can be estimated semi-analytically with either harmonic or anharmonic models.³⁴ For larger differences, the nonequilibrium-work (NEW) method in which simulations are conducted by rapidly switching between low and high level potentials can be effective.²⁴ Since the switching simulations require the evaluation of the high level potential at each MD step, the NEW trajectories need to be short and have the lengths of a few picoseconds. Although such a switching timescale is sufficient to relax the stiff DOFs, differences in the soft DOFs are difficult to capture.³⁵ To consider more complex discrepancies between different levels of theory, the low level potential can be modified to better mimic the high level potential,³⁶ using, for example, force matching.^21,26 With sufficient improvement of the low level potential function, FEP can then be used to better estimate the free energy correction toward the high level; the improved low-level method can also be used to compute forces at inner time steps in a multiple time step framework,³⁷ while high level forces are computed only at outer time steps. Alternatively, one may improve the low level of theory using machine learning based techniques; i.e., one parameterizes the difference $\left(\Delta \right)$ between the two levels of theory with neural networks based on samples from low-level MD simulations and uses such $\Delta$-learning to iteratively improve the low level method towards the high level.^38–42 In general, a robust improvement of the low-level model requires a substantial number of snapshots, which limits the size of systems amenable to the $\Delta$-learning approach. Finally, approaches based on the normalizing flow model⁴³ have been developed to improve the convergence of multi-level FES^44,45. Normalizing flow uses a chain of probability transformations to map a low-level distribution to a high-level distribution. The transformations are exact so in the ideal scenario the mapped distribution can fully overlap with the high level distribution, and hence FEP calculations can be efficient and accurate. Unfortunately, the normalizing flow method has to model the entire system, which means the modeled distribution can be extremely high-dimensional for condensed phase systems, making the model training almost impossible. As a result, the normalizing flow-based approaches have only been demonstrated to be effective in the gas phase or with an implicit solvent model.

In a different framework, without explicitly modifying the low-level potential, one may introduce an intermediate model (denoted as M in Figure 1a) that features good configurational overlaps with the high level,²⁵ and the total free energy correction can be evaluated in a staged fashion, i.e.

$$\Delta F_{A/B}^{L\to H}=\Delta F_{A/B}^{L\to M}+\Delta F_{A/B}^{M\to H}.$$ (2)

In some cases, the intermediate model differs only from the low level in terms of charge distributions.^25,46 Alternatively, when ensembles at the two levels of theory show substantial structural differences, the intermediate model can be derived by performing biased simulations at the low level toward configurations favored by the high level of theory. Such a staged approach is particularly attractive because the biased simulations required to establish the intermediate model can follow collective coordinates that compactly describe the “problematic” DOFs, rather than specifying the latter individually as required in methods that explicitly modify the low-level potential function. On the other hand, the major challenge is then the automated identification of such collective coordinates that dictate the configurational overlap between different levels of theory.

In recent years, the incorporation of the autoencoder⁴⁷ (AE), a neural network framework, and its variants into physics-based methods has resulted in great success in the collective variable (CV) or reaction coordinate (RC) discovery,^48–52 where CVs/RCs are low-dimensional representations that capture the most essential information regarding processes occurring in a high dimensional space. AEs map the input data into a low-dimensional latent space, which is the neural network-predicted RC, with little-to-none prior knowledge about the underlying process. Depending on the choice of input features and model architectures, the latent space usually captures the slow^52,53 or large-magnitude^48,49 motion of the system. The success of these AE-based methods suggests that a similar approach can be valuable to the identification of DOFs essential to the bridging of configurational distributions at different levels of theory in multi-level FES. We surmise that, even with limited sampling at the high level, if the neural network is trained on data from both levels, the optimal RCs are able to capture the nature of the problematic DOFs. Furthermore, the RCs are learned via a chain of continuous and differentiable functions of the input features and therefore can be conveniently integrated into enhanced sampling simulations^52–55 to generate the intermediate model (Figure 1b) for computing the free energy correction term in the staged-transformation framework.²⁵

Specifically, we develop a neural network based method to learn the optimal RCs so that enhanced sampling simulations in this compact space at the low level can be performed to preferentially sample the configurations that contribute significantly to the free energy correction toward the high level. The neural network model is based on the adversarial autoencoder (AAE),⁵⁶ a modern variant of AE, with key modifications that maximize the likelihood of success for our approach. While MD simulations at both levels are conducted to generate the training data, enhanced sampling simulations are conducted at the low level and the high level simulations are short and only required at the end states. The significant increase in the accuracy and efficiency compared to direct FEP or straightforward Bennett acceptance ratio (BAR)⁵⁷ calculations is demonstrated with both gas phase and solvated systems, showing great promise of the method for more complex condensed phase applications. Although the use of neural network in this study is directly inspired by the previous machine learning CV/RC discovery and enhanced sampling methods, our approach differs from those in several ways. First, rather than capturing the slow or large-magnitude motion(s), the latent space in our approach represents the most important structural differences that dictate the configurational overlap between two levels of theory, making the approach particularly suited for the application of multi-level FES. Second, as further elaborated in the following sections, the use of AAE allows for an explicit control of the latent space distribution, which makes the subsequent enhanced sampling simulations more straightforward and interpretable compared with previous multi-level methods.

2. Methods and Computational Setup

2.1. Typical Architecture of Adversarial Autoencoder

To provide the general background of our development, the typical architecture of AAE is shown in Figure S1a. The upper part is the same as AE (Figure S1b), in which an encoder maps a high-dimensional dataset, $X$, into a low-dimensional representation (the latent space), $z$. Then a decoder reconstructs the input data from $z$. The training process minimizes the reconstruction error, so after training, the learned latent space, $z$, captures the most essential information about $X$ while filtering out most of the noise. It’s therefore efficient to sample the latent space and generate new data that resemble the original data. The key improvement of AAE over AE is that the distribution of the latent space is matched to a prior distribution, whose functional form and parameters are determined before training. The match of the latent space to the prior distribution is realized by the adversarial training, an idea borrowed from generative adversarial networks (GAN).⁵⁸ In GAN, two models, a generator and a discriminator, are trained adversarially. The generator produces fake data that mimics real data while the discriminator attempts to distinguish fake and real data (Figure S1c). In AAE, the encoder part of the autoencoder acts as the generator. Fake data are the latent space variables generated by the encoder, while real data are sampled from the prior distribution. The outcome of the adversarial training in AAE is that the distribution of the latent space variables, $z$, matches the prior distribution (Figure S1a). Compared with AE, a better latent space can be learned in AAE.⁵⁶

2.2. An AAE based Workflow for Multi-level Free Energy Simulation

Now we present the proposed workflow for an efficient evaluation of the free energy correction $\left(\Delta F^{L\to H}\right)$ between different levels of theory in the multi-level FES framework. Short MD simulations at low and high levels are conducted and pairwise distances between the heavy atoms are computed from the MD trajectories. These features are commonly used in ML-assisted free energy calculations as they are translationally and rotationally invariant. The pairwise distances of the two levels are combined together as the training dataset, to ensure that the model can learn the difference between the two distributions. We take advantage of the flexibility of AAE to maximize the likelihood of success of our workflow by adopting two strategies. First, during the adversarial training stage of the discriminator, each data entry is assigned a label, making the model a supervised one. The label informs the discriminator whether the data belong to the low level or the high level. As suggested in the original AAE paper,⁵⁶ the labeled data from different classes can be mapped onto different regions of the prior distribution. Second, considering that the potential energy difference between the low and high levels of theory $\left(\Delta U\right)$ is critical to the evaluation of $\Delta F^{L\to H}$, the output of the decoder is not the reconstructed input, but scaled $\Delta U$ (see below) associated with the corresponding input data; this helps ensure that the latent space is optimized for capturing the information most critical to the evaluation of $\Delta F^{L\to H}$.

The above effort on maximizing the model performance results in an optimal latent space that can be used as the RCs in the subsequent umbrella sampling (US) simulations.⁵⁹ Thanks to the learned relationship between the configurational space and the latent space, we can bias the simulations at the low level so that the sampled ensembles are increasingly representative of the high level. Moreover, due to the use of labels during training, the latent space is aware of the information about the $\Delta U$ distributions at different levels of theory; thus we expect that the $\Delta U$ distributions during US sampling will also progressively shift toward the $\Delta U$ distribution of the high level. As a result, an intermediate model (M) can be readily established whose $\Delta U$ distribution overlaps with that of the high level. As a result, $\Delta F^{M\to H}$ can be robustly estimated with rapid convergence, leading to an efficient computation of $\Delta F^{L\to H}$ following eq. 2.

2.3. Computational Setup of the Test Systems

Two benchmark systems are tested in this study: alanine dipeptide (ALA) in the gas phase and bis-2-chloro diethyl ether (2CLE) in a water box. MD simulations are conducted at the low and high levels, respectively, in the CHARMM package.⁶⁰ For ALA in the gas phase, the high level is the Density Functional Tight Binding⁶¹ up to third Order (DFTB3)⁶² with the 3OB parameter set^63,64 (DFTB3/3OB). The low level is the CHARMM36m⁶⁵ force field. Non-bonded interactions are computed using a standard group based cutoff scheme⁶⁶. The Langevin integrator with a 2-fs timestep is used and a friction coefficient of 2 ps⁻¹ is applied to all atoms. The temperature is kept at 303.15 K.

In the case of 2CLE, the solute is solvated in a cubic water box with a 10.0 Å edge distance under periodic boundary condition (PBC). The low level is a MM model, in which the parameters for the solute were generated in the previous study²² from the CHARMM general force field (CGenFF)⁶⁷ and refined with the general automated atomic model parameterization (GAAMP) method.⁶⁸ The high level is the QM/MM potential where the solute is treated by DFTB3/3OB, while the water molecules are treated by the CHARMM-modified TIP3P water model.⁶⁹ The electrostatic interactions are calculated by the Particle-mesh Ewald (PME) summation.⁷⁰ Van der Waals (vdW) interactions are treated by a cutoff distance of 12 Å and a switching distance of 10 Å. NPT (constant particle number, pressure, and temperature) simulations at 303.15 K and 1 bar are carried out with an integration time step of 2 fs. Snapshots are saved every 2 ps.

For all simulations of both systems, a weak harmonic restraint (force constant: 10 kcal mol⁻¹ Å⁻²) is applied to the center-of-mass of ALA or 2CLE to keep them around the origin. The SHAKE algorithm is used to constrain bonds involving hydrogen atoms to the respective optimal bond distances at the low and high levels of theory. By constraining these bonds, we do not explicitly consider their vibrational contributions to the free energy correction, which can be estimated based on harmonic approximations.³⁴ This choice is made based on the consideration that for many processes of interest, vibrations of these stiff degrees of freedom do not make the dominant contribution to the free energy differences.^71,72

2.4. Dataset Preparation

Each snapshot from MD simulations is represented by a 1D list of pairwise distances between the non-hydrogen (heavy) atoms of ALA or 2CLE. Distances from both levels are combined as the input data for the neural network model training. Two training datasets containing 250 or 1250 conformations from each level of theory are created, corresponding to simulation lengths of 0.5 and 2.5 ns, respectively. Test datasets are generated from separate MD simulations. To mimic the realistic scenario where the high level sampling is limited, only 0.5 ns simulations are performed at both levels to generate a small dataset for the evaluation of model performance, which corresponds to 250 conformations at each level. A larger test dataset, which contains 40,000 snapshots collected from 80 ns of trajectory at each level of theory, is also generated for better examining the latent space mapping but not used in any model training/evaluation steps. Regardless of the level of the theory used for the configurational sampling, the potential energy difference is defined as $\Delta U=U_H-U_L$ for each snapshot, where $U_H$ and $U_L$ are the potential energies evaluated at the high and low levels, respectively. To avoid large numerical values for $\Delta U$, unscaled $\Delta U$ values are offset by 16,400 (ALA) or 12,200 kcal/mol (2CLE) before any computation or visualization. MinMaxScaler in scikit-learn⁷³ is applied to scale the pairwise distances and $\Delta U$ values to the range of [−1, 1]. The prior distribution is taken to be a mixture of two two-dimensional Gaussians. By design, the component centered in the first quadrant represents the low-level distribution while the other component centered in the third quadrant represents the high-level distribution. A two-dimensional prior model is used here because it strikes the balance between providing sufficient flexibility for the latent space optimization and being amenable to efficient sampling. For example, sampling in a three-dimensional latent space is substantially more costly.

2.5. Neural Network Training

All neural networks are implemented with PyTorch.⁷⁴ The encoder, decoder, and discriminator each consists of two fully-connected layers. The first layer has 64 nodes and the second layer has 32 nodes. The Leaky Rectified Linear Unit activation function (LeakyReLU)⁷⁵ is applied in each hidden layer. The functional form of LeakyReLU is as follows:

$$y=\{\begin{array}{ll}x\hfill & \text{if }x\ge 0\hfill \\ ax\hfill & \text{if }x<0\hfill \end{array},$$ (3)

where $a=0.2$ in our neural network models. A hyperbolic tangent function (tanh) is applied to normalize the output of the decoder to the range of [−1, 1]. The Mean Square Error loss function (MSELoss), $MSE=\frac{1}{n}{\sum }_{i=1}^n {\left(y_{i,\mathit{\text{original}}}-y_{i,\mathit{\text{predicted}}}\right)}^2$, is used for both the decoder and the discriminator. The Adam optimizer⁷⁶ is used with default parameters. The learning rate is 0.00025 for the encoder and the decoder while it is set to 0.0001 for the discriminator. Initial weights are generated from a normal distribution, $𝒩\left(\mathrm{0,0.01}\right)$. Five-fold cross-validation⁷⁷ is used to evaluate the performance of models.

Kullback–Leibler (KL) divergence is a non-negative number that measures the overlap between two probability distributions. A smaller KL divergence value indicates a larger overlap. KL divergence between the distributions of the reconstructed scaled $\Delta \hat{U}$ and the original scaled $\Delta U,D_{KL}(\Delta U\parallel \Delta \hat{U})={\int }_{-\infty }^{\infty } p\left(\Delta U\right)\text{log}\left(\frac{p\left(\Delta U\right)}{q\left(\Delta \hat{U}\right)}\right)d\Delta U$, is computed. As the KL divergence calculation is a part of the model evaluation, the small test dataset described in the above subsection is used.

2.6. Umbrella Sampling Simulations

Umbrella sampling (US) simulations are conducted in OpenMM 7.6⁷⁸ interfaced with the open-source, community-developed PLUMED library,⁷⁹ version 2.7.⁸⁰ The reaction coordinates are the learned latent space from neural network training. While the work was in process, the PLUMED library did not directly recognize files in the PyTorch format. Therefore, all weights and biases in the neural network layers are extracted after training, and the transformations from the input pairwise distances to the latent space are reconstructed in a PLUMED input file by using linear and non-linear functionalities available in PLUMED. These operations are numerically stable and efficient as they are all arithmetic computations and do not involve any gradient calculations. By the time our work was finished, the PYTORCH module was added to a later version of the PLUMED library (version 2.9), which allows users to directly load models defined in PyTorch.

The sampling of the $i$-th umbrella window is restrained to a narrow region around its equilibrium position on the latent space, $\left(z_1^i,z_2^i\right)$, through a harmonic restraining potential, $U^i=\frac{1}{2}k{\left(z_1-z_1^i\right)}^2+k{\left(z_2-z_2^i\right)}^2$, where $k=1.5\text{\hspace{0.17em}kcal}/\text{mol}$ is the force constant. The centers of the US windows $\left(z_1^i,z_2^i\right)$ are taken to be along the straight line that connects the two centers of the prior distributions (indicated as black dots in Figure 5a and summarized numerically in Table S1). During US simulations, at each integration step, the latent space variables are computed on-the-fly by the PLUMED library using the method described above. All other set-up and simulation parameters are the same as the unbiased low level simulations. 10 ns trajectory (5000 snapshots) is collected for each window. With the sufficiently large dataset, we then compute the free energy difference as a function of the simulation length (see below) to evaluate the convergence of our AAE-based approach.

Figure 5: Free energy calculation results for ALA in the gas phase.

a The underneath gray contour is the prior distribution. The black dots represent the centers of each umbrella window. The rainbow contours represent the sampled regions on the latent space for each window. b The orange and blue dotted curves represent the $\Delta U$ distributions of the ensemble sampled at the high and low levels, respectively. The two distributions only have a negligible overlap while the $\Delta U$ distribution of the intermediate state has decent overlap with both levels, especially with the high level. MM = CHARMM36m, DFTB = DFTB3/3OB. c The free energy difference between the low and high levels, $\Delta F^{L\to H}$, as a function of the simulation length. The shaded areas represent the standard error of the mean (SEM) computed from five replicas. Note that the model is trained on the first 2.5 ns trajectories from both levels, as indicated by the vertical dotted line in c, but extended simulations are conducted to compute the reference FEP and BAR values.

2.7. Free Energy Calculations

Following eq. 2, the correction term is computed as $\Delta F_{A/B}^{L\to H}=\Delta F_{A/B}^{L\to M}+\Delta F_{A/B}^{M\to H}$, where M denotes the intermediate model. In the current work, we focus on a single end state and thus drop the subscript notation in the free energy differences hereafter.

The WHAM package of the Grossfield Lab⁸¹ is used to calculate the free energy difference between umbrella windows. The original low level simulation is treated as one umbrella window without the restraining potential (denoted as the 0-th window). 9 umbrella windows are simulated in both test cases (1-st to 9-th windows), so $\Delta F^{L\to M}=F^9-F^0$. The free energy difference between the intermediate model and the high level, $\Delta F^{M\to H}$, as well as the direct BAR results are calculated with the BAR estimators implemented in the pymbar package,⁸² while the direct FEP results are calculated with in-house codes. The reported free energy differences are the average values of 5 independent replicas and statistical errors are estimated as the standard error of the mean (SEM) across the 5 replicas. The SEM of $\Delta F^{L\to M},s_{L\to M}$, and the SEM of $\Delta F^{M\to H},s_{M\to H}$, are propagated to obtain the SEM of $\Delta F^{L\to H},s_{L\to H}$, using the formula $s_{L\to H}=\sqrt{s_{L\to M}^2+s_{M\to H}^2}$. The computed free energy differences (Figures 5c, 6b,d and 7d), $\Delta F^{L\to M}/\Delta F^{M\to H}$ components (Figures S6 and S11) and the statistical errors (Figures S7 and S12) are shown as functions of the simulation length (i.e., by including an increasing number of snapshots in the free energy calculations), which is taken to be the same for the low level, high level and umbrella sampling simulations.

Figure 6: The KL divergence between the reconstructed and the original $\Delta U$ distributions is a good indicator of whether the sampling is sufficient for training.

a When a much smaller training dataset (250 conformations from each level) is used, the model with the optimal architecture results in a poor overlap between the reconstructed $\Delta \hat{U}$ and the original $\Delta U$ distributions for the test dataset. The KL divergence between the two distributions is 4.20 ± 0.54. b The free energy differences also deviate significantly from the average BAR value, especially in the first a few nanoseconds. c A smaller KL divergence of 2.46 ± 0.32 and therefore a good overlap between the reconstructed $\Delta \hat{U}$ and original $\Delta U$ distributions are obtained in another model with slightly different architecture but trained on 1250 conformations from each level. d The free energy estimates obtained in this model show much quicker and better convergence than the model in b and is comparable with the optimal model discussed in the text. Note that the model is trained on the first 0.5 ns or 2.5 ns trajectories from both levels, as indicated by the vertical dotted lines in b and d, respectively, but extended simulations are conducted to compute the reference values. The shaded areas represent the SEM from 5 replicas.

Figure 7: The AAE-assisted free energy calculation results on the solvated bis-2-chloro diethyl ether (2CLE) system.

a The solute, 2CLE, is solvated in a water box. b The prior distribution and the centers of umbrella windows follow the similar definitions as the ALA example in Figure 5a. c The $\Delta U$ distribution of the intermediate model (the solid green curve), i.e. the last umbrella window, has significant overlap with the $\Delta U$ distributions of both the low (the dotted blue curve) and high (the dotted orange curve) levels. d The free energy difference as a function of the simulation length shows the similar pattern as in the ALA case. The vertical dotted line represents the simulation length used for model training. The shaded areas represent the SEM from 5 replicas.

3. Results and Discussion

3.1. Training of the Optimal Model

The performance of our approach is first evaluated with an alanine dipeptide (ALA) in the gas phase (Figure 3a), a well-studied benchmark system to test free energy calculation methods.⁸³ Although it is a small molecule, some DOFs such as backbone and sidechain torsional angles are likely to feature different distributions at different levels of theory. Here, the low level is the MM force field⁶⁵ (MM), while the high level is a semi-empirical QM method, the third-order density functional tight binding (DFTB3/3OB) method.^61–64 We chose this semi-empirical QM method as the high level instead of an ab initio QM approach because DFTB3/3OB is more efficient with a reasonable accuracy, making it possible to obtain ‘exact’ reference results by running brute-force sampling, which can be used to evaluate the reliability and convergence of our multi-level FES approach. As stated earlier, the pairwise distances between the heavy atoms from MD trajectories at the MM and DFTB3/3OB levels are combined as the training dataset. The $\Delta U$ distributions at the two levels have large widths of 3.46 and 4.28 kcal/mol, respectively, and barely overlap, which highlights the difficulty in obtaining accurate free energy results with conventional methods^29,30 even for a gas phase system (Figure S2).

Figure 3: Training of the optimal model using the gas-phase alanine dipeptide (ALA) system as an example.

a The three-dimensional structure of ALA. b The blue solid curve and the green dotted curve represent the generator losses for the training and validation datasets, respectively. The orange solid curve and the red dotted curve represent the discriminator losses for the training and validation datasets, respectively. The shaded areas represent the one standard deviation of losses computed from the five-fold cross-validation. c After training, the test dataset is used to evaluate the performance of the model. The reconstructed $\Delta \hat{U}$ distribution overlaps well with the original distribution. The decoder predicts scaled $\Delta U$ from both levels and therefore, to make a direct comparison, the distribution of the original scaled $\Delta U$ in c also includes data from both levels. The KL divergence between the two distributions is 2.79 ± 1.12. The corresponding comparison for the training set is shown in Figure S4, in which the distributions include data from both levels as well. d The mapping of another larger test dataset to the learned latent space. The colorbar represents the scaled $\Delta U$ value associated with each latent space variables. The arrow points from the regions representing the dominant populations of the low level to regions representing the dominant populations of the high level (i.e. from 1 to −1).

The training and validation loss curves demonstrate the advantage of our AAE architecture (Figure 3b). By implementing the two strategies mentioned before, the discriminator loss quickly converges to the equilibrium value with very small variation. The same trend is observed for the generator loss, which decays sharply to a small and stable value. Moreover, the validation loss agrees well with the training loss, indicating neither over-fitting nor under-fitting. For comparison, two additional architectures are also implemented. One is an unsupervised and symmetric AAE, i.e. the label information is not provided to the discriminator and the decoder reconstructs the pairwise distances. The other is a supervised and symmetric AAE, i.e. the label information is provided to the discriminator and the decoder still reconstructs the pairwise distances (i.e. not $\Delta U$). The discriminator loss generally has much larger variation and does not converge to the expected value; the generator loss is also significantly larger than the best model (Figures 4a and b). In the supervised and symmetric AAE, the validation loss of the discriminator also deviates markedly from the training loss (Figure 4b). These observations clearly show that, without the two strategies we adopted, the optimal model will not be learned.

Figure 4: Training outcomes of the ALA test system for neural network models trained with non-optimal architectures.

a Loss curves of the unsupervised and symmetric AAE model, i.e. without label information and the decoder reconstructs the input data. b Loss curves of the supervised and symmetric AAE. c Mapping of a separate test dataset onto the learned latent space of the unsupervised and symmetric AAE. d Mapping of a separate test dataset onto the learned latent space of the supervised and symmetric AAE. All models are trained with the dataset containing 1250 conformations from each level and other hyper-parameters are the same as the optimal architecture.

After training, the distribution of the reconstructed $\Delta \hat{U}$ for the training set overlaps well with the original $\Delta U$ distribution (Figure S4). To evaluate the model performance on data not used for training, a separate test dataset is used to obtain the distribution of reconstructed potential energy difference $\left(\Delta \hat{U}\right)$ after the model training (see Figure 3c). We deliberately used a small test set (250 conformations) to mimic the realistic scenario where one aims to minimize the computational cost by limiting high level samplings to the generation of the training dataset. Since the $\Delta \hat{U}$ values are generated by decoding latent space variables, the significant overlap between the distributions of the reconstructed $\Delta \hat{U}$ and the original $\Delta U$ indicates that the learned latent space is indeed able to capture the essential DOFs that encode the difference in the $\Delta U$ distributions and therefore potentially facilitate the free energy calculations.

Mapping of the test dataset onto the learned latent space can strongly suggest whether the encoder has learned the representation of the input data distribution. Here, to better visualize the effect of the mapping, we used a much larger test dataset. In models with non-optimal architectures (i.e. the unsupervised and symmetric AAE or the supervised and symmetric AAE), the mapping of the test dataset only partially overlaps with the prior distribution (Figures 4c and 4d), leaving a large portion of the prior distribution uncovered. The uncovered regions represent the lack of a well-defined relationship between the prior distribution and the learned latent space in models with non-optimal architectures, preventing a clear interpretation of the mapping of the test dataset. By contrast, with the optimal architecture, the mapping of the test dataset overlaps well with the prior distribution, indicating that the learned latent space indeed highly resembles the desired prior distribution (Figure 3d). The optimal mapping provides further evidence that our model has successfully learned the relationship between the configurational space and the latent space, including the information regarding the $\Delta U$ distributions. First, the input data from low and high levels are mapped onto different regions in the latent space with the supervised discriminator (Figure S5). Second, the fact that the learned latent space captures the essential information about $\Delta U$ distributions can be better demonstrated by coloring each latent space variable with its associated $\Delta U$ value. Data points with similar $\Delta U$ values are clustered to the nearby region in the latent space. This naturally forms a pathway from regions representing the dominant populations of the low level to regions representing the dominant populations of the high level. Thus if we sample along this pathway in the latent space, we are able to collect configurations that are increasingly representative of the dominant populations of the high level method.

3.2. Construction of the Intermediate State Through Umbrella Sampling

With the optimal latent space, umbrella windows are constructed to follow the pathway that represents the continuous transition of the latent space mapping. Each umbrella window overlaps with its neighbors to a notable extent, which ensures that their $\Delta U$ distributions also have decent overlap (Figure 5a). As expected, their $\Delta U$ distributions progressively shift toward that of the high level. The last umbrella window features the largest overlap with the high level $\Delta U$ distribution and is chosen as the intermediate model (Figure 5b), for which $\Delta F^{M\to H}$ can be efficiently and accurately calculated through the two-sided BAR estimator.⁵⁷ The free energy difference between the original (unbiased) low level simulation and the intermediate model, $\Delta F^{L\to M}$, is readily calculated through the weighted histogram analysis method (WHAM).⁸⁴ The sum of these two terms gives the total free energy difference between the low level and the high level, $\Delta F^{L\to H}=\Delta F^{L\to M}+\Delta F^{M\to H}$, which is the correction term of interest in multi-level FES.

To demonstrate the advantage of our approach, we compare the free energy results from our AAE-assisted staged transformation approach with the direct calculations of $\Delta F^{L\to H}$ using FEP²⁸ and BAR.⁵⁷ Although FEP is an efficient estimator since it requires sampling only at the low level, the reliability of FEP calculations depends critically on the $\Delta U$ distributions overlap between levels. BAR in principal has a looser dependence on the configurational overlap although it requires explicit sampling at the high level. With the 2.5-ns MD trajectories used to train the AAE model, the FEP results (the dotted blue curve in Figure 5c) deviate significantly from the reference value and this trend persists even when more data are included (data points after the vertical dotted line). Moreover, large fluctuations are observed for FEP results (the blue shaded area in Figure 5c). Using the same 2.5-ns MD trajectories, the average values of BAR are stable (the dashed black line in Figure 5c). Nevertheless, the large statistical uncertainty for the first few nanoseconds (the grey shaded area in Figure 5c) indicates that the BAR results are reliable only if the sampling is long enough to develop sufficient overlap in $\Delta U$ distributions (for example, longer than 10 ns; see Figure 5c). This behavior makes the straightforward BAR calculation less practical in multi-level FES due to the high computational cost of extensive sampling at the high level. Nevertheless, the better convergence of BAR than FEP highlights that some amount of sampling at the high level is necessary to ensure the reliability of the free energy estimate.

In sharp contrast to the free energy differences computed with FEP and BAR, our AAE-based approach captures the free energy difference with high accuracy and efficiency (the orange curve in Figure 5c; for components and statistical errors, see Figures S6 and S7). The first term, $\Delta F^{L\to M}$, is efficiently calculated using the original low level trajectory and the intermediate model trajectory through the US simulations and WHAM as discussed above. The significant configurational overlap between the intermediate model and the high level greatly reduces the uncertainty in the calculation of the second term, $\Delta F^{M\to H}$. Altogether, the $\Delta F^{L\to H}$ computed from our approach quickly converges to the reference value when only the 2.5-ns high level data are included in the free energy calculations. Within this timescale, the results from our approach are more reliable than other methods, as demonstrated by the much smaller statistical uncertainty. To further support the observation, we also calculated free energy differences using the same model trained on the 2.5-ns trajectories but with more data for calculating free energies (see the orange curve in Figure 5c after the vertical dotted line). These additional calculations show comparable average results and similar level of statistical uncertainty as results for the first 2.5 ns, again supporting that a modest amount of sampling at the high level is sufficient to obtain converged results with our approach.

3.3. Proper Size of Training Dataset

In our method, an issue of major importance is the determination of the length of training simulations, especially sampling at the high level of theory. The expectation is that although the size of the training dataset at the high level is not large enough for computing accurate free energy differences, the learned latent space is still able to guide the construction of the intermediate model so that a reliable estimate of the free energy difference can eventually be obtained. In this regard, we surmise that the KL divergence between the reconstructed scaled $\Delta \hat{U}$ and the original scaled $\Delta U$ distributions is a proper metric for the overlap between the two distributions and therefore an indicator of whether the amount of training data is sufficient.

The model trained on the 2.5-ns dataset (1250 conformations from each level) shows a small KL divergence (Figure 5c) while models trained on a 0.5-ns dataset (250 conformations from each level) show large KL divergence (Figure 6a). Correspondingly, free energy differences with models trained on the smaller dataset show a large deviation from the reference value, suggesting that the training dataset has not captured the major distinctions between the low and high levels, leading to a systematic error in the AAE-assisted latent space discovery and subsequent free energy estimates (Figure 6b); nevertheless, we note that even with the smaller dataset, the results are substantially better than FEP. Furthermore, the model with a different set of hyper-parameters but similar KL divergence (Figure 6c) as the “optimal” model discussed above is able to maintain a similar level of accuracy in free energy calculations (Figure 6d). These observations support that the KL divergence between the reconstructed and the original $\Delta U$ distributions can be a proper indicator for determining the required length of the training simulations.

3.4. Free Energy Difference in a Solvated System

Having demonstrated the success in a gas phase system, we now apply the computational framework to a condensed phase system, which includes a small solute, bis-2-chloro diethyl ether (2CLE), in a water box (Figure 7a). Here, the low level is still the MM FF, and the high level potential is a combined QM/MM method where the solute is treated by DFTB3/3OB, and the solvent is treated by the MM FF. In addition to intra-solute DOFs, solute-solvent interactions also vary with the level of theory. Thus the low and high level distributions are broad (with a width of 3.22 and 4.70 kcal/mol, respectively) and likely to have a poor overlap, making conventional methods even more difficult to converge.³⁵ For the optimization of the latent space, to avoid modeling the large number of water molecules, we aim to include minimal information about the solvent in the training dataset. In this particular example, only pairwise distances between the heavy atoms of the solute are used for training. In this way, we expect that the solvent effect is implicitly reflected by the distributions of the solute-only features.

For the AAE-based training, the two modifications of the architecture discussed earlier are also proven important in the solvated case. As shown in Figure S8, the loss curves in the solvated case capture the similar behavior as observed in the gas phase example. With the optimal architecture, the training loss and the validation loss agree well with each other (Figure S8a). The discriminator loss quickly converges to the equilibrated value with a small variation and the generator loss also decays rapidly to a small and stable value. As in the gas phase example, models with non-optimal architectures generally show large variations in the discriminator loss and/or poor overlap between training and validation losses (Figures S8b and c). KL divergence between the reconstructed $\Delta \hat{U}$ and the original $\Delta U$ distributions is also used to determine the length of the training dataset (Figure S9). We found that the models trained on the dataset containing 2.5 ns trajectories from each level show favorable overlap between two distributions and a small KL divergence (Figure S9a), which indicates that this dataset is sufficient for training a reliable model, as also demonstrated by the loss curves in Figure S8a.

With the optimal latent space, US simulations are conducted to establish the intermediate model (Figure 7b). As in the gas phase example, the $\Delta U$ distribution of the last umbrella window has a large overlap with the high level distribution and therefore chosen as the intermediate model (Figure 7c). Figure 7d shows the free energy difference as a function of the simulation length using FEP, BAR, and our approach. The FEP results never converge to the reference value and exhibit large statistical uncertainties. The BAR results slowly converge to the reference value but show large statistical uncertainties for all simulation lengths, suggesting that the results are less reliable even if we run much longer simulations at both levels. By contrast, the free energy computed from our approach (for components and statistical errors, see Figures S11 and S12) quickly converges to the reference value within the first 2.5 ns. The statistical uncertainty of our approach is much smaller than other methods, illustrating that the results are much more reliable even with a modest amount of sampling. These encouraging observations demonstrate that our approach is able to efficiently and accurately estimate the free energy difference between the low and high levels in the solvated case. We recall that only solute-derived features are used for model training. The fact that an accurate free energy difference can still be efficiently calculated suggests that our computational framework has the potential to be scaled up to systems with a large number of solvent molecules.

4. Concluding Remarks

In this work, we have developed a computational framework to compute the free energy difference between an inexpensive, low level of theory and a more costly, high level of theory, $\Delta F^{L\to H}$, which is a key quantity of interest in multi-level FES. The major challenge we aim to tackle is to automatically identify collective coordinates that dictate the distinct energy difference distributions at the two levels of theory. The solution we propose is a neural network model based on the adversarial autoencoder (AAE), which matches the distribution of the latent space to a prior distribution, thus allowing a better-controlled mapping from the input data to the latent space than the traditional AE. We implemented two modifications to the basic AAE architecture to make our approach particularly effective in the context of multi-level FES. Since the neural network model is aware of both configurational distributions and the energy difference distributions at the two levels of theory, the latent space effectively captures the nature of the DOFs that dictate the poor distribution overlaps between the two levels, thus overcoming a major remaining bottleneck of previous studies. Due to the multi-layer nature of the neural network used in our work, it remains difficult to explicitly define the contributions of individual features. Nevertheless, examination of the weight matrix (Figures S3 and S10) suggests that no single feature dominates and the latent space is spanned by multiple conformational degrees of freedom, highlighting the power and scalability of the AAE framework.

This optimal latent space is then used in US simulations to efficiently construct an intermediate model that has significantly improved distribution overlap with the high level. In turn, the staged transformation strategy proposed in our previous study²⁵ is adopted to calculate the total free energy difference between the low and high levels (eq. 2). This data-driven method has been applied to both gas phase and condensed phase systems. Compared to a straightforward BAR calculations that require extensive sampling at both low and high levels, more than an order-of-magnitude increase in the computational efficiency is observed in both cases with at least comparable accuracy. The success supports our expectation that a modest amount of sampling at the high level is both necessary and sufficient to guide the discovery of latent space variables for substantially improving the distribution overlap with the high level of theory by judiciously biasing low-level simulations in such a latent space.

The unique aspect of our approach compared to existing multi-level FES methods^{21,23,24,26,41} is that we use a neural network to automatically learn a compact latent space that represents the multiple degrees of freedom that encode the distinct distributions at the two levels. This feature avoids the need of identifying and/or correcting individual problematic DOFs as required by methods that explicitly modify the low-level potential function,^21,26 making our approach more scalable to large systems where many DOFs are expected to contribute to the poor overlap between the low and high levels of theory. For example, the $\Delta$-learning approach^39–42 needs to learn the energy (and force) differences between two levels of theory, which are complex functions of high dimensionality; an accurate and transferable training thus is expected to require a large number of data points for condensed phase systems, although many configurations may not make a significant contribution to the free energy difference. By contrast, by focusing only on the DOFs that greatly impact the $\Delta U$ distribution overlap between the two levels of theory, our approach explicitly drives the sampling of configurations that contribute predominantly to the free energy difference, which explains the high efficiency and accuracy observed in the example applications.

The performance of our neural network-assisted approach, which first learns an optimal latent space and then constructs an intermediate distribution to facilitate the convergence of free energy correction, can be further enhanced with more sophisticated molecular features, such as those representing the nearby solvent distributions.^39,40,85 In fact, we note that the two levels of theory can differ significantly in not only structural distributions but also electronic properties, such as charge distributions,^25,46 which lead to different interactions with the environment. For such cases, including the environmental features in the training of the latent space is expected to be particularly important. Overall, our study provides a uniquely effective avenue to compute the low-to-high free energy correction by integrating state-of-the-art machine learning models with enhanced sampling simulations.

Figure 2: The proposed workflow for multi-level free energy simulations.

A modified AAE model is trained with the input simulation data from both the low and high levels. The input data are the pairwise distances between selected atoms (for example, the heavy atoms of the solute molecule). The output of the decoder is the reconstructed $\Delta \hat{U}$ associated with the input data. The prior distribution is a two-dimensional Gaussian mixture with two components. As indicated in the figure, by design, the low level and high level data are supposed to be mapped to their respective components. The learned encoder and latent space is then used for the on-the-fly calculation of latent space variables during enhanced sampling simulations for computing the free energy correction between the low and high levels of theory (eq. 2).