Comparison of Three Autoencoder-Based Models Using Molecular Dynamics (MD) Simulations Data

1. Introduction

Molecular dynamics (MD) simulations have proven useful to study the dynamics of Biomolecules, for instance proteins. However, these simulations are computationally expensive and consume many core-hours of high performance computing (HPC) clusters to achieve enough sampling of the energy landscape. To enhance the sampled configurational space, several techniques have been used in the past but some of them make assumptions on the system for instance using predefined collective variables [1,2,3] or by modifications of the energy landscape [4] where empirical parameters are introduced.

More recently, machine learning (ML) models are being used to extract relevant features from the energy landscape which are then used to infer/interpolate new data. In particular, in the autoencoder (AE) machine learning model, one performs the training by first reducing the dimensionality of the data gradually to build a latent space and then use the compressed latent space to reconstruct the initial data through some metrics and regularization terms [5,6]. One advantage of the AE models is the reduced number of training parameters that are involved w.r.t. the fully connected neural networks. Once trained, these models can be used to generate new data. In the context of proteins, one can use existing data sets of protein structures collected through computational methods, classical (MD) or quantum mechanical (DFT, HF, among others) as the training data sets. These trained models can be used to extend the sampled ensemble of protein conformations using cheaper computational resources. In addition to the computational data sets, ML models open the window to include data from experiments, such as NMR, Cryo-EM, and X-ray crystallography.

In the present work, we used three different flavors of AE models to generate new protein conformations using MD trajectories of the Shikimate Kinase enzyme for the training data set. A key difference between these three models is the way in which the latent space is created. For the first AE model, we used the variational autoencoder (VAE) model of Zhu et al. [7] which was originally used to extend the sampled space of intrinsic disordered proteins (IDPs). In the second and third models, we studied the Wasserstein autoencoder (WAE) [8,9], and the Denoising autoencoder (DAE) [10] models, respectively. Details about the three models can be found in the Materials and Methods section. Because AE models are used as components of more complex workflows, understanding their performance can be beneficial when they are employed in standalone or as part of a larger workflow. The results of the three models are compared using standard analysis tools for MD.

2. Materials and Methods

2.1. Data Collection

The MD trajectory used in the present study was taken initially from [11], where the Shikimate Kinase (SK) enzyme was simulated in the absence of substrates, and further extended to achieve 5 μs. SK has three main domains called SB, β, and LID domains which are highly flexible and allow this enzyme to harbor substrates for catalysis. Details for the MD simulations can be found in this reference paper. Frames for data collection were saved every 0.05 ns resulting in 1$\times {10}^5$ frames.

2.2. Methods

A schematic view of the three AE models is presented in Figure 1 (a)-(c), VAE, WAE, and DAE, respectively.

2.2.1. Variational Autoencoder (VAE) Model

In the VAE model, the encoder network is modeled through the $Q_{VAE}\left((Ƶ,\aleph )\right)$ distribution that maps the input space (ℵ) to the latent space (Ƶ). The decoder or generative model is given by the $P_G\left((\aleph ,Ƶ)\right)$ distribution. The loss function is computed as in Ref. [7]:

$$L_{VAE}=L_R+L_{KL},$$

(1)

$L_R$ is the reconstruction loss computed as the logarithmic mean squared error between the input data and the reconstructed data. $L_{KL}$ is the Kullback-Leibler (KL) [12] divergence regularization term.

2.2.2. Wasserstein Autoencoder (WAE) Model

The WAE model uses a marginal distribution $Q\left(Ƶ\right)$ of the encoder distribution $Q_{WAE}\left((Ƶ,\aleph )\right)$. In this model the loss is given by,

$$L_{WAE}=L_R+\lambda L_{MM{D}^2}.$$

(2)

Here, the reconstruction error $L_R$ was computed as the mean squared error between the input and reconstructed data. The maximum mean discrepancy (MMD) [13,14] loss between two probability distributions $\pi$ and ${\pi }^{\prime }$, in the present case the Gaussian prior $P\left(Ƶ\right)$ and latent space marginal ($Q\left(Ƶ\right)$) distributions, was computed as follows:

$$L_{MM{D}^2}={\mathbb{E}}_{x,x^{\prime }\sim \pi }\left[k\left(x,x^{\prime }\right)\right]+{\mathbb{E}}_{y,y^{\prime }\sim {\pi }^{\prime }}\left[k\left(y,y^{\prime }\right)\right]-2{\mathbb{E}}_{x\sim \pi ,y\sim {\pi }^{\prime }}\left[k\left(x,y\right)\right],$$

(3)

where the kernel function is the radial basis function (RBF) defined by,

$$k\left(x,y\right)=\exp \left(-{\displaystyle \frac{x-y^2}{2{\sigma }^2}}\right)$$

(4)

with a standard deviation $\sigma$. The parameter $\lambda$ is a weighting factor for the MMD term which has a value of 10 in the present simulations.

2.2.3. Denoising Autoencoder (DAE) Model

In this model, the initial data are corrupted with Gaussian noise. The idea of adding/subtracting noise from input data has been used in other ML models [15,16]. The loss function is computed in DAE by considering only the reconstruction error:

$$L_{DAE}=L_R$$

(5)

This error is calculated as the mean squared error between the input and reconstructed data.

2.3. Postprocessing and Analysis of Data

The MD trajectories were postprocessed with the Visual Molecular Dynamics (VMD) [17] (v. 1.9.4) software by extracting all atoms in the protein structure except for the hydrogen atoms. The data set for the autoencoder models was partitioned in 80% for the training set (TrS) and 20% for the testing set (TsS) using the preprocessing script given in Ref. [7]. Models’ training was done with Keras (v. 2.11) and Tensorflow (v. 2.11) libraries. To collect the performance metrics for the models (Table 1) we used 5 epochs as simulations displayed convergence already for this value.

The integrity of reconstructed structures from the three AE methods was analyzed through the Spearman correlation coefficient (SpCC) which measures the correlation in the monotonic growing between the ranked distributions for the original and reconstructed data. For instance, a SpCC value close to +1 indicates that both distributions grow in the same direction for each data sample. In addition to this, we also used the mean square error (MSE) between these two distributions.

The conformations generated (a total of 1$\times {10}^5$) were inspected with standard analysis tools for MD trajectories including the root mean square deviation (RMSD), root mean square fluctuation (RMSF), and principal component analysis (PCA). For the RMSD calculations, the reference structure was the first frame of the MD trajectory in all cases and the backbone atoms were selected for the computation. In the case of the RMSF results, only the C^α atoms were selected. The RMSF curves for the AE models were translated to match the value of the MD RMSF at 5 Å. Regarding PCA, all protein atoms except for hydrogen atoms were considered for the computation.

2.4. Code Availability

The codes for the three autoencoder models can be found here https://github.com/pojeda/AE-variants.

3. Results and Discussion

The SpCC for the SK protein showed values for both the training and testing sets larger than those reported in Ref. [7]. This can be due to the fact that SK has a more ordered structure in comparison to the IDPs discussed in that reference. The values for SpCC were close for the three different methods for this SK protein ~0.99. Regarding, the MSE values were similar for all methods in the training set, but we noticed that DAE had a lower value for the testing set w.r.t. to the other two methods. Notice that we have not conducted an extensive exploration of the hyperparameters involved in the model and we used values similar to those of Zhu et al. [7]. Results for SpCC and MSE can be seen in Table 1.

We noticed that even though the SpCC and MSE values showed that the reconstructed data was close to the original data, the generated protein structures presented distortions in the bond lengths especially in the flexible LID region. This issue has been previously faced by postprocessing the structures [7,18]. Other approaches can be used, for instance graph neural networks [19] and rigid body transformations [20].

By using the trained models, new conformations were generated. We investigated the capability of the models to generate an ensemble of conformations similar to the original MD ensemble. Thus, as a first approach we took the initial structure from the MD simulations and computed the RMSD values for the models. This can give us an idea of how the conformational space was explored. Using the RMSD values we constructed histograms of counts as it can be seen in Figure 2. The distribution for the MD case showed a peak ~2.4 Å while the AE models had a peak ~2.25 Å. The exploration range for the AE models was limited, especially for the DAE model where generated structures were close to the peak value. In all models, structures with RMSD values lower than 1.4 Å and close the initial one were not generated.

The variations in the displacement of the protein residues can be monitored through the RMSF values which are presented in Figure 3. Because the RMSF values can be correlated to the temperature of the environment, especially for the highly flexible regions, we argue that the three different AE models display a lower temperature than the MD data. However, the trend of the fluctuations follows a similar pattern as in the MD case for the flexible SB, β, and LID domains. Fluctuations are considerably reduced in the DAE model.

We also monitored the sampling of the conformational space for the generated structures in a reduced dimensional space through principal component analysis (PCA), see Figure 4. Here, we can see that the VAE model learns some structure from dynamical data which is reflected in the pattern displayed by the corresponding plot.

The PCA plots show that the different AE models can help to sample the conformational space in different ways. In the case of WAE the points corresponding to the generated structures are spread in the projected space while in DAE model the data points are more localized.

4. Conclusions

We have compared three different flavors of AE models, i.e. Variational, Wasserstein, and Denoising to generate protein structures not included in the original data from MD simulations. According to the RMSD and PCA analysis, the three models explore the conformational space in a different way, for instance the WAE model showed spread distribution of structures while the DAE displayed a more restricted distribution.

The RMSF analysis suggests that the generated conformations, considered as an ensemble, are not at the same temperature as the original MD ensemble. In particular, the flexible regions displayed attenuated fluctuations w.r.t. to the original MD data.

One shortcoming of the present simulations is that the bond lengths and angles of the generated structures are distorted, which opens the possibility for further improvement of the models.