AMARO: All Heavy-Atom Transferable Neural Network Potentials of Protein Thermodynamics

Abstract

All-atom molecular simulations offer detailed insights into macromolecular phenomena, but their substantial computational cost hinders the exploration of complex biological processes. We introduce Advanced Machine-learning Atomic Representation Omni-force-field (AMARO), a new neural network potential (NNP) that combines an O(3)-equivariant message-passing neural network architecture, TensorNet, with a coarse-graining map that excludes hydrogen atoms. AMARO demonstrates the feasibility of training coarser NNP, without prior energy terms, to run stable protein dynamics with scalability and generalization capabilities.

layout=twocolumnachemso Institucìo Catalana de Recerca i Estudis Avançats (ICREA), Passeig Lluis Companys 23, Barcelona, 08010, Spain. Acellera Labs, Doctor Trueta 183, Barcelona, 08005, Spain.

1 Introduction

Molecular events at the individual macromolecule level reveal emergent and collective macroscopic behaviors, emphasizing the need for a comprehensive and hierarchical approach to deeply investigate the complexities of biophysical complexes critical to all cellular functions. ¹. Over the past decades, the integration of advanced computational methods, with the latest hardware and software, and the increasing availability of experimental molecular structure data, has deeply transformed molecular biology and drug discovery ^{2, 3}. These enhancements, supported by significant advancements in theoretical frameworks, have evolved molecular simulations from simple proofs of concept to detailed in silico studies of protein folding ^{4, 5} and dynamics ^{6, 7}.

Figure 1: The pipeline for developing all-heavy atom NNPs is reported here. A CG map is applied to the mdCATH dataset ⁸, and an embedding $z$ for each domain is created. TensorNet⁹ is then trained using the CG data. Generalization and scale-up properties are evaluated on a set of four fast-folding proteins and larger domains in the final stage.

Molecular dynamics (MD) simulations, in particular, serve as a powerful tool for capturing the behaviors of proteins and biomolecules at the atomic level, providing fine temporal resolution. However, classical all-atom MD simulations present multiple limitations due to the substantial increase in computational resources required for larger systems and extended time scales; furthermore, post-process and data analysis demands considerable effort ^{10, 11}. To investigate larger systems over an extended time scale, a leading approach involves reducing computational demands via coarse-grained (CG) simulations, where molecular systems are simulated using fewer degrees of freedom than those associated with the atomic positions ^{12, 13}.

Several CG models have been developed, each tailored to optimize molecular simulations and capture critical biophysical features. The MARTINI model ^{14, 15}, for instance, excels due to its adaptability across a range of biomolecular systems, including membrane structure formation and protein interactions. Models such as AWSEM ¹⁶ and UNRES ¹⁷ have been successful in simulating intramolecular protein dynamics, though they occasionally struggle to capture alternative metastable states. Similarly, Primo ^{18, 19} and Rosetta ²⁰ focus on specific molecular interactions and the design of protein structures and complexes, enhancing their accuracy in targeted applications. Currently, there is no universally accepted theory that precisely determines the most effective coarse-graining mapping for any given system, which in general is related to the intended application and computational constraints ^{21, 22}. Once a coarse-graining mapping has been established, different strategies exist to define the model energy function, either to reproduce the reference fine-grained statistic (bottom-up) ²³ or to match experimental observables (top-down) ²⁴. Neural network potentials (NNPs) demonstrate remarkable efficiency in rapidly learning accurate potential energy ²⁵ and effectively model many-body atomic interactions ^{26, 27}. These features are extremely beneficial for developing coarse-grained (CG) force-fields that need multi-body functional forms to accurately represent protein thermodynamics and incorporate implicit solvation effects. While previous CG-NNP models ^{28, 29, 30, 31, 32} rely on prior energies for stability and accuracy, here we present a first version of our Advanced Machine-learning Atomic Representation Omni-force-field (AMARO), a bottom-up CG-NNP without the need of prior terms to achieve stability and transferability.

2 Material and methods

A CG model typically comprises two main components: selecting the CG resolution (or mapping) and designing an effective energy function once the mapping has been established. Mapping schemes often draw from physical or chemical intuition. CG sites may represent functional groups, residues, or monomers, or be tailored to a specific resolution. In this study, we adopt a no-hydrogen coarse-graining map, and a machine learning model is employed to learn the potential energy function ³³. Traditional force fields have historically treated hydrogens primarily as charge carriers due to their negligible Lennard-Jones interactions with heavier atoms. However, significant advances in implicit atom models, including the united atoms representation and implicit water machine learning potentials, illustrate the feasibility of accounting for these interactions using a many-body approach ^{34, 30}.

2.1 Neural Network Model

In bottom-up CG approaches the interactions between CG beads are determined based on a more detailed model, and the many-body potential of mean force (PMF) is used to describe the free energy landscape of a system as a function of a collective coordinate or reaction coordinate, which is in other terms a configurational free energy in a reduced space. The term ”many-body” in the context of a CG model refers to the fact that the potential energy landscape is considered in terms of interactions between groups of atoms rather than individual atoms, and network models offer a straightforward approximation in this context ¹³. Our machine learning model of choice in this work is TensorNet⁹, a new neural network architecture that integrates O(3)-equivariance in message-passing and utilizes rank-2 Cartesian tensor representations.

2.2 No hydrogen CG map

The selection of an optimal mapping approach for transitioning from a fine- to coarse-grained representation is a critical aspect of a CG model definition. An effective CG map should significantly reduce the computational burden of the all-atom model while preserving sufficient information to prevent an excessively flat energy surface. In this study, a no-hydrogen (noh) and no-water coarse-grained map has been chosen to reduce the degrees of freedom by almost half compared to the all-atom counterpart and to align with the basic force aggregation method ³⁵.

Consider a dataset $\mathcal{D}$ consisting of $M$ coordinate-force pairs obtained using an all-atom MD force field. The coordinates are denoted by $\mathbf{r_{c}}\in\mathbb{R}^{3N_{c}}$ and the forces by $\mathbf{f_{c}}\in\mathbb{R}^{3N_{c}}$, with $c=1,...,\mathcal{M}$ and where $N_{c}$ is the number of atoms in the system depending on $c$ since different protein systems are considered. A linear map operator $\Xi:\mathbb{R}^{3N_{c}}\rightarrow\mathbb{R}^{3n}$ can be defined to map an all-atom conformation $\mathbf{r_{c}}$ to a coarse-grained conformation $\mathbf{R}$, where $n$ denotes the residual degrees of freedom. A paired coarse-grained force $\mathbf{F}$ is then considered for any conformation $\mathbf{R}$, and the force of the $i$-th noh-bead $\mathbf{F}_{i}$ is defined as:

$$\mathbf{F}_{i}=\mathbf{f}_{\textit{ih}}+\sum_{j\in\mathbb{B}}\mathbf{f}_{j},$$ (1)

where $\mathbf{f}_{\textit{ih}}$ represents the force of the $i$-th heavy atom, and $\mathbb{B}$ denotes the set of hydrogen atoms bonded to the $i$-th heavy atom in the fine-grained representation. Accounting for the noh coarse-graining map, this approach uniquely embeds each noh-bead by considering both its corresponding heavy atom and the number of bonded hydrogen atoms. A final set of 12 embedding values, as outlined in Table S1, was obtained enhancing the model’s ability to differentiate between various electronic hybridizations. This selection empowers the model to more effectively learn the atoms’ properties and molecular geometry.

2.3 Neural network training

As demonstrated in prior research, the acquisition of the many-body PMF involves minimizing the mean-squared deviation between a CG candidate force field and atomistic forces appropriately mapped. This method, known as variational force matching ³⁶, establishes a robust approach to effectively learning the PMF, laying the foundation for the exploration of intricate biomolecular interactions^{33, 37}. At CG resolution the force matching method becomes more complicated since the training data contain less information than their atomistic counterparts: energies are not available, and forces are noisy ³¹. To enhance the quality of input data, we employed the basic force aggregation method in dataset preparation. Coupled with the noh-CG map, this approach significantly improves PMF learning ³⁵. In essence, the model now learns the force of a noh-bead as the sum of the forces acting on its heavy atoms and the constrained (i.e., ’bonded’) hydrogen atoms. Notably, this model introduces a groundbreaking feature in the CG-NNP field moving from delta-learning to directly learning the forces acting on particles, no prior terms are considered. So, the effective energy of the system $\widetilde{U}(\mathbf{R};\theta)$ is represented by the network parameters $\theta$ which are optimized to minimize the mean-squared deviation between predicted and labeled forces via the loss function

$$\mathcal{L}_{\text{FM}}(\theta)=\frac{1}{M}\sum_{i}{\left\lVert-{\nabla}_{\mathbf{R_{i}}}\widetilde{U}(\mathbf{R};\{\theta\})-\mathbf{F_{i}}\right\rVert^{2}_{2}}.$$ (2)

The predicted forces are computed as the negative gradient of the potential energy $\widetilde{U}$ with respect to the noh-bead coordinates $\mathbf{R}$ and the result is averaged over the batch conformations, while $\mathbf{F}$ are the coarse-grained labeled forces.

2.4 Dataset

The mdCATH dataset ⁸ was the basis for applying the coarse-grained mapping approach. Specifically, the initial data were processed to retain only the heavy atoms’ coordinates and forces, with a basic force aggregation map applied to the latter. Additionally, the z dataset within each HDF5 file was modified to serve as the embedding for each system. A series of filters were implemented to exclude certain domains based on the following criteria: 1) domains containing more than 150 residues, 2) domains comprising more than 1000 noh-atoms, or 3) domains with less than 50% combined helix and sheet fractions. Various temperatures (320 K, 348 K, 379 K, 413 K, 450 K) were considered to facilitate the model’s learning of atomic proximity or separation dynamics. In total, 2,834 domains and more than 26 million conformations were selected from the mdCATH dataset after considering the filters and temperatures as described above.

The final dataset used to train the model was obtained by applying a stride of 25 to the 26M residual conformations, resulting in an approximate split of 900,000 conformations for training, 50,000 for validation, and 100,000 for testing.

The model’s scalability was then tested on larger mdCATH domains, specifically those with more than 150 residues and a combined helix and sheet content greater than 50%.

At the same time, transferability was assessed using four fast-folding proteins: Chignolin, Trp-cage, Villin, and $\alpha$3D; sequence similarity details are reported in Table S2.

2.5 AMARO Molecular Simulations

All MD simulations reported in this work, involving AMARO, were conducted with the mass of each CG-bead calculated as the combined mass of the heavy atoms and hydrogen atoms that constitute the bead. This will ensure: I) realistic dynamic properties of the system, such as diffusion coefficients and relaxation times, maintaining the mass-force relation; II) energy conservation and thermodynamic stability during the simulation by proper scaling of mass and force avoiding non-physical results and drifts in temperature and pressure.

2.6 Markov State Models

In this study, we analyze the dynamics of CG simulations using Markov State Models (MSMs) ^{38, 39, 40} available in HTMD ⁴¹ and compare them with those from corresponding all-atom simulations. MSMs partition the entire dynamics of a system into n discrete states and are particularly suited for systems that exhibit Markovian behavior—where future states depend only on the current state without memory of the past. We construct a transition probability matrix for these Markovian systems, characterized by n states and a lag time $\tau$, which records the system’s state. This matrix enables us to determine state populations and conditional pairwise transition probabilities, from which free energies are derived. We employ time-lagged independent component analysis (TICA) ⁴² to enhance our analysis, reducing the high-dimensional conformational space into a lower-dimensional, optimally reduced space. We then discretize this space using K-means clustering to construct the MSM. For each fast-folding protein, all-atom simulation data, characterized by pairwise C$\alpha$ distances, are projected onto the first four components using TICA. A reference free energy surface was then constructed for each system by binning the first two TICA dimensions into an 80×80 grid and averaging the weights of the equilibrium probability in each bin, as computed by the Markov state model. For the CG simulations, we adopt the approach outlined in ⁴³, utilizing covariance matrices from all-atom molecular dynamics (MD) to project the first three components. This method aligns with established methodologies, ensuring consistency and facilitating further analysis. HTMD provided the necessary computational tools and framework to perform these analyses ⁴¹. Finally, to avoid biasing the model with starting conformations, 10% of the initial frames of each trajectory were removed from the analysis except for $\alpha$3D, where 5% of the initial frames were discarded due to the longer simulation times considered.

3 Results

TensorNet was trained on the filtered mdCATH dataset, as described in Section 2.4, using TorchMD-Net ⁴⁴ for 100 epochs. Detailed information on model architecture and training hyper-parameters can be found in the supplementary material (see Tables S3 and S4).

The final L1 test loss for the model is reported as 5.07 kcal/mol/Å, while MSE loss for training and validation are reported in Figure 2.

Figure 2: Traning- and validation- MSE loss, in blue and orange respectively, for AMARO as a function of training epoch.

3.1 Generalization to larger domains

We explore the ability of the AMARO to scale up when larger protein domains than those considered in the training set. To achieve this, we selected a subset of 5,000 conformations from the mdCATH dataset, specifically targeting domains with between 150 and 250 residues and a combined helix and sheet fraction   50%. This selection criteria ensured that the domains were representative of complex protein structures while still maintaining a manageable size for computational analysis. The forces acting on the noh-bead within these larger domains were evaluated using the CG-NNP. AMARO exhibited a mean absolute error (MAE) of 4.98 kcal/mol/Å, assessed for each force component (x, y, z). The error value here recorded, compared also to the one obtained at the end of the training, proves that the learned potential can scale up without loss in accuracy. Moreover, Figure 3 presents a direct comparison between the expected and observed force values, with each dot color-coded by CG atom type, illustrating the model’s precision. The results underscore the robustness of the CG-NNP model and its potential applicability to larger biological systems, confirming its feature to maintain performance across an expanded range of system sizes and complexities.

Figure 3: Scale-up validation of AMARO on larger domains from mdCATH. Comparison between labeled (i.e. reference) and predicted force component values (x, y, z). Each data point in the scatter plot is color-coded according to the CG atom type.

Detailed errors per atom type are reported in Table S5, with NH3+ exhibiting the highest error, a mean absolute error (MAE) of 8.10 kcal/mol/Å. However, this is not due to NH3 being underrepresented in the training datasets, as shown in Figure S1, but rather due to the physical and chemical properties of this group. Five domains with the highest number of NH3+ groups (i.e. $\geq$20) were selected for further investigation: 1kvnA00, 1nu7D01, 1w9rA00, 2c5zA00, 2jzvA00, 2nc9A00 and 3qneA01. In all of these cases, the protonated amino group corresponds to the terminal amino group along the lysine side chain. These groups are oriented outward, suggesting solvent interactions, which are not accounted for in the current modeling. Moreover, in more compact structures, these groups might be involved in salt bridges with carboxyl groups. Such complex interactions contribute to the challenge of generalizing the model’s parameters for this specific atom type.

3.2 Validation on fast-folding proteins

To assess the model’s transferability, we selected four fast-folding proteins not included in the training set: Chignolin (175 atoms), Trp-Cage (210 atoms), Villin (573 atoms), and $\alpha$3D (1149 atoms); the number of atoms in their all-atom representations are indicated in parentheses. After applying the noh mapping, the sizes were reduced to 97, 112, 286, and 576 atoms, respectively. TorchMD ⁴⁵ was used to run 32 replica MD simulations for each protein and considering at least 320 ns of aggregated simulation time, see Table S6. The initial coordinates for these simulations were uniformly sampled from the respective TICA surfaces of each fast folder, as shown in Figure S2. For Trp-Cage, only residues from 2 to 16 have been considered to focus on capturing the overall folding and not the cis-trans Proline (residue 20) isomerization.

3.3 Recovering the energetic landscape

For three of the four fast-folding proteins the TICA landscape has been successfully recovered, see Fig. 4. In contrast, for $\alpha$3D, the recovery is only partial, with most microstates populating a middle region between the global minimum on the right and a local minimum on the left. This behavior could be attributed to the particular secondary structure of $\alpha$3D, which is characterized by a high proportion of $\alpha$-helices. Additionally, the relative shape anisotropy (RSA) of the system, 0.3, may not be well-represented in the training dataset, which has an average RSA of 0.17 ± 0.16. This discrepancy, coupled with the larger system size and the presence of NH3+ groups in the unstructured region between the first and second helices of the folded state, could contribute to the observed differences.

While Fig. 4 displays the two TICA dimensions as the principal axes and uses free energy as the third dimension, Fig. S3 presents the free energy on the y-axis, providing a quantitative analysis. For the simplest target, Chignolin, the absolute minimum is perfectly captured by AMARO. For more complex structures like Trp-Cage and Villin, the overall shape of the profile is well approximated. Conversely, for $\alpha$3D, extensive sampling in the central region results in a shift, leading to the identification of an absolute minimum between the two minima.

Figure 4: Comparative analysis of the free energy landscape obtained from all-atom simulations (left) and NNP coarse-grained simulations (right) across the first two TICA dimensions for four fast-folding proteins: Chignolin, Trp-cage, Villin and $\alpha$3D.

3.4 Sampling the native structures of unseen training proteins

Analysis of the CG simulations using MSMs, detailed in Section 2.6 and Table S7, revealed that the model successfully reproduced the experimental structure of the corresponding fast-folding proteins, as illustrated in Figure 5a. Sampling originated from the native macrostate, defined as the macrostate containing the conformation with the minimum root-mean-square deviation (RMSD) with respect to the experimental crystal structure. The models accurately predicted secondary and tertiary structural elements, with loops and unstructured terminal regions showing minimal variation, except for $\alpha$3D. Table S8 provides detailed information on the equilibrium probability, mean, and minimum RMSDs of the native macrostates. The results indicate extensive sampling of the native conformation, as reflected by the high equilibrium probabilities for these macrostates. For Chignolin, Trp-Cage, and Villin, an average RMSD below 1 Å was observed when comparing the native macrostate to the folded structure, more in detail: 0.15 Å, 0.30 Å, and 0.6 Å, respectively. In contrast, the larger and more complex $\alpha$3D system reported an RMSD of 2.30 Å. These results underscore the model’s accuracy in structural prediction and its adaptability across different molecular systems, though further investigation is required to fully understand the particular case of $\alpha$3D.

Figure 5: CG trajectories of Chignolin, Trp-Cage, Villin, and $\alpha$3D, selected based on the inclusion of microstates from the lowest RMSD macrostate. (a) Minimum RMSD conformation (blue) aligned with the experimental structure (grey) for each protein, labeled with the protein name and PDB ID. (b) C$\alpha$ RMSD of each trajectory compared to the crystal structure. (c) CG free energy surface, projected over the first two TICs with the folded state (red star) and sampled states indicated by RMSD color-coded dots. The trajectory’s progression is illustrated with arrows connecting the starting (yellow point) and ending (orange point) conformations. The all-atom equilibrium density is shown by a red contour.

3.5 Computational efficiency

To assess the computational efficiency of AMARO, we conducted a comparative analysis against traditional all-atom simulations, focusing on sampling free energy (FE) landscapes within the TICA space. Both simulations were constrained to a 12-hour time-frame to ensure fair comparison, with the all-atom simulations using the CHARMM22* force field ^{46, 47}, and explicit solvent, while the AMARO simulation focused on heavy atoms using a noh mapping approach. Both models were run on openMM ⁴⁸ using an NVIDIA RTX 4090, and employed a Langevin thermostat at 350 K, with differing friction coefficients: 0.1 ps^-1 for CHARMM22* and 1^-1 for AMARO (standard setup, see Section 2.5). Additionally, the hydrogen mass repartitioning (HMR) scheme ⁴⁹ was set at 4 a.m.u. for the classical force field, and increased mass was applied for the NNP system as detailed in the methods section.

The analysis revealed that the all-atom simulations achieved 265.5 ns of molecular simulation while AMARO completed 8.4 ns within the same operational time window. However, to account for differences in temporal resolution between the coarse-grained and all-atom models, we compared the areas of the TICA landscapes recovered by both models, see Figure 6. The starting point (blue dot), the same for both cases, was selected randomly and resulted in a conformation near the global minimum, while the red star represents the crystal structure (PDB ID: 2A3D).

The all-atom simulation (left panel) explores a more localized region of the TICA space, with the trajectory (black arrows) predominantly staying within a single basin. This indicates limited sampling within the time frame due to the high dimensionality and energy barriers typical of fully atomistic representations. The final conformation is marked by a yellow dot, indicating that the simulation does not move significantly away from the initial state.

In contrast, the CG-NNP simulation (right panel), despite the shorter trajectory length, explores a broader region of the conformational space (red arrows), as seen by the wider spread of sampled conformations. The trajectory covers a larger portion of the landscape, crossing energy barriers more easily and reaching areas of higher free energy. This behavior, characteristic of coarse-grained models, is due to the reduced degrees of freedom and the smoothing of energy landscapes which produces an effective faster kinetics. Despite this broader sampling, the path still shows consistency with the overall free energy landscape, see Figure 4 for reference, suggesting thermodynamic consistency.

Notably, while the all-atom simulation reaches a minimum RMSD of 5.03 Åwith respect to the crystal structure, the CG-NNP model achieves a lower minimum RMSD of 4.38 Å. Moreover, the average RMSD values of 13.31 ± 4.82 Åfor CG-NNP and 6.41 ± 0.5 Åfor CHARMM22* suggest that the coarse-grained model explores a larger conformational space, as expected.

Figure 6: Sampling efficiency comparison in the TICA landscape for the $\alpha$3D system under two force fields: CHARMM22* (left) and AMARO (right). Both simulations, starting from a randomly selected initial conformation (blue dot) and ending at a final conformation (orange dot), were run within a fixed 12-hour time window for direct comparison.

The CG-NNP’s efficiency was further quantified relative to system size, measured in CG beads (Figure 7, providing a quantitative benchmark of AMARO’s performance. Both millions of simulation steps (left y-axis) and ns/day (right y-axis) are considered variables. The term ”simulation step” refers to a forward and backward step of the model. The efficiency demonstrated by the CG-NNP in small-to-medium systems (up to  1000 CG beads) makes it a valuable tool for exploring larger conformational landscapes or long-timescale events.

Figure 7: Log-log plot of AMARO’s computational speed relative to system size (number of CG Beads). The left y-axis shows the number of simulation steps (in millions) that can be executed within a single day. The right y-axis displays the performance in nanoseconds per day (ns/day), assuming a time-step of 4 fs.

4 Conclusions

This paper introduces the first version of AMARO a new fully machine-learning coarse-grained force field offering a new framework for molecular dynamics simulations. AMARO uses an all-heavy-atoms coarse-graining strategy paired with variational force matching, which simplifies protein representation while retaining essential dynamical information. This approach addresses previous challenges in balancing model simplicity with the retention of critical dynamics, which often require energy priors for stabilization. Notably, our model demonstrates remarkable transferability and scaling-up ability. However, further enhancements in computational efficiency and memory usage must be achieved in the future. The current study serves as a proof of concept, challenging the reliance on prior energies in developing stable NNPs for studying protein thermodynamics.

5 Code Availability

Code and relevant data to reproduce this work are available at https://github.com/compsciencelab/amaro.

AM is financially supported by Generalitat de Catalunya’s Agency for Management of University and Research Grants (AGAUR) PhD grant 2024 FI-1-00278; the project PID2023-151620OB-I00 has been funded by MCIN / AEI / 10.13039/501100011033.