Machine Learned Coarse-Grained Protein Force-Fields: Are We There Yet?

Abstract

The successful recent application of machine learning methods to scientific problems includes the learning of flexible and accurate atomic-level force-fields for materials and biomolecules from quantum chemical data. In parallel, the machine learning of force-fields at coarser resolutions is rapidly gaining relevance, as an efficient way to represent the higher-body interactions needed in coarse-grained force-fields to compensate for the omitted degrees of freedom. Coarse-grained models are important for the study of systems at time and length scales exceeding those of atomistic simulations. However, the development of transferable coarse-grained models via machine learning still presents significant challenges. Here we discuss recent developments in this field and current efforts to address the remaining challenges.

Graphical Abstract

1. Introduction

The definition of simplified models is central to physical sciences; proteins are no exception [1, 2]. Statistical mechanical approaches to describe protein folding and dynamics [3, 4, 5], as well as the analysis of long molecular dynamics (MD) trajectories [6, 7, 8], have demonstrated that slow processes in large biomolecular systems can be described by a reduced number of variables despite hundreds of thousands of atoms comprising the full system. In this spirit, many coarse-grained (CG) models have been proposed to study proteins through MD and energy minimization. These CG models have been used to investigate the principles underlying protein folding [9, 10, 11, 12], intermolecular binding/interactions [13, 14], protein-mediated membrane phenomena [15, 16], and to make predictions about novel biological systems of immediate medical interest [17, 18].

Despite their successes, CG models of proteins have not yet achieved the predictive performance of their atomistic counterparts. CG models are primarily designed by specifying their resolution, which defines the coarse degrees of freedom (referred to as “sites” or “beads”, see Fig. 1), and by their effective energy function, which dictates how these beads interact. Traditionally, the resolution is first chosen using either chemical intuition or through optimization designed to reproduce chosen properties (e.g., [19]). The model’s effective energy function is then parameterized to reproduce experimental or simulation data. The fundamental goal of the transferable CG models discussed in this article is to predict the conformational landscape of proteins not used for their parameterization, ideally using only the primary structure of the proteins of interest. Atomistic models have been able to explore the relevant landscape of small globular proteins [20, 21, 22]; however, it is still an open question as to whether there exists a resolution at which a chemically transferable CG model can quantitatively describe the configurational landscape of arbitrary proteins.

Figure 1:

Sequential reduction in resolution of a variant of the miniprotein Chignolin (CLN025) from a solvated all-atom representation containing many thousands of atoms, to an implicit solvent representation, to a heavy-backbone representation with C_β beads, and finally to a C_α CG representation containing 10 beads.

A transferable CG protein model would have significant consequences. By employing special-purpose supercomputers [23] or distributed simulation combined with Markov State Models (MSMs)[24, 21, 25], the dynamics of small solvated proteins can be simulated over millisecond timescales [20]. However, biological phenomena routinely involve larger complexes and span longer timescales (seconds or more). CG models promise to reach such scales by reducing the computational cost via decreasing the number of degrees of freedom and increasing the effective simulation timestep. This increased efficiency would vastly improve the use of MD for both fundamental research and applications, e.g., in protein design.

There has been a surge of interest in using machine learning (ML) methods for molecular simulation [26], including learning of CG models from large amounts of data. In a sense, the development of ML CG models can be seen as an extension of ongoing research on the design of accurate atomistic force-fields from quantum mechanical calculations. In this area, ML has already produced highly accurate force-fields which have facilitated groundbreaking computational studies [27, 28, 29]. When combined with the field of bottom-up CG [30, 31], these approaches provide a seemingly clear strategy to lever-age ML to learn a CG force-field from existing atomistic MD trajectories. Indeed, thanks to the flexibility of ML algorithms, some frameworks developed for the atomistic resolution [32, 33] have been transferred to the CG resolution [34, 35, 36, 37].

Despite these advances, a completely bottom-up transferable CG model still does not exist for proteins or other biopolymers. This limited progress is due to multiple outstanding challenges, which together firmly differentiate the creation of ML bottom-up CG force-fields from their atomistic counterparts. We here discuss these difficulties and current efforts to overcome them.

2. Thermodynamic Consistency: Why is it difficult?

Bottom-up coarse-graining typically models the following free-energy surface (U) [30, 26, 31] referred to as the effective CG (free) energy:

$$U(R)=-{\beta }^{-1}\text{ln}\int \delta [R-ℳ(r)]\text{exp}[-\beta u(r)]\text{d}r,$$ (1)

where ℳ maps all-atom configurations $r\in {ℝ}^{3n}$ to their CG counterparts $R=ℳ(r)\in {ℝ}^{3N}$, u is the reference all-atom energy, and β is the inverse temperature. Intuitively, ℳ defines the CG resolution and U defines how particles at this resolution interact; the design of a bottom-up CG model then entails defining ℳ and approximating U, and the two tasks are interdependent. A CG energy that, up to a constant, equals U is said to be thermodynamically consistent with the atomistic counterpart. Such a U produces free energy landscapes identical to the reference in any reaction coordinates that are a function of the CG coordinates. We note that the phrase thermodynamically consistent here does not refer to thermodynamic observables (e.g., pressure), but instead considers the configurational distributions of the CG and atomistic force-fields. For information on thermodynamic properties in CG models we defer to recent articles [38, 39, 40, 31].

Although the thermodynamically consistent CG energy is uniquely defined by Eq. (1), the integral cannot be solved for non-trivial systems [30]. As a result, multiple strategies [41, 30, 31] to approximate U have been proposed. Traditionally, the functional forms for CG (free) energy functions have been low body-order with physically motivated terms [30, 31]. However, recent studies have employed higher body-order terms parameterized by neural networks with success [42, 43, 44, 36, 35]. Kernel methods have also been proposed (e.g., [34]) but have not been applied to proteins. While exceptions exist [42, 45, 46, 47], for reasons of computational efficiency existing ML CG models [44, 36, 26] have been primarily based on the Multiscale Coarse-Graining [48] (“force-matching”) variational framework, primarily due to the fact that it does not require the CG model to be simulated during its parameterization.

In principle, (Eq. 1) suggests that once the CG resolution has been chosen the creation of the CG model should be straightforward. However, designing an accurate ML CG model is not trivial (see Fig. 2). The choices of reference atomistic system (u), of the resolution (ℳ), and of the different terms of U all entail challenges unique to models designed at a CG resolution. These challenges are compounded by difficulties with validation, which are also present in the development of ML atomistic force-fields from quantum chemical data. We address these challenges in detail: beginning with the data used for training, continuing by discussing the design of U and subsequently the resolution of the CG model, and finishing by discussing validation and robustness. For brevity we only discuss algorithms which are applicable to ML CG force-fields; for a more comprehensive introduction to bottom-up coarse-graining we refer readers to recent perspective articles [30, 31].

Figure 2:

A pipeline for creating and using ML CG models from atomistic simulation data and experimental measurements. A chosen CG mapping can reduce reference information into a CG dataset that can be used to train ML CG models. This training can rely on both simulation and experimental observables in order to reduce the complexity of the learning task and respect physical constraints. A trained ML CG model can then be validated through CG MD and used for general property predictions.

2.1. The difficulty in training CG force-fields

The principal challenge in bottom-up coarse-graining with machine-learned force-fields lies in finding a suitable ML formulation that directly or indirectly estimates the intractable integral described in Eq. (1). The situation is more difficult compared to learning atomistic potential energy surfaces from quantum mechanical data, where reference energies and forces are known: When learning a CG free energy, neither U nor its gradient for a given CG structure R are known because the integral Eq. (1) is intractable.

The most straightforward approach to Eq. (1) is to directly estimate the behavior of the probability density proportional to exp(−βU(R)) from simulation data. This requires an equilibrium sample of atomistic conformations r, e.g., obtained by MD simulations. After mapping them to the CG resolution, a ML model is then trained to approximate U(R), by minimizing the Kullback–Leibler divergence between the coarse-grained and atomistic probability densities. This is called relative entropy minimization in the coarse-graining literature [49, 50, 47] and maximum likelihood estimation in the ML energy-based model community [51]. Similar approaches [42, 45, 46] which estimate and reduce the difference between a CG force-field and U or optimize selected observables in turn expand on other approaches from the ML community (e.g., [52]).

The difficulty with most of these approaches is that the CG model must be periodically re-simulated during training in order to evaluate the equilibrium density generated by the current model of U(R). While this is feasible for quickly equilibrating CG models, such as those of liquids [47], it is extremely challenging for models that exhibit rare events, such as realistic CG models of protein folding. This limitation is even more problematic for complex parameterizations of U (e.g., neural networks) and significantly impedes simultaneously training over multiple molecules when creating transferable models. Approaches have attempted to reduce this burden by, for example, reweighting the density of previous iterations [50, 45, 53] or by modifying the sampling of the atomistic system [46]. However, these approaches have not yet reached the simplicity and applicability of non-iterative approaches.

The most common bottom-up approach for approximating U is force-matching [48], which fits a CG free energy such that its negative gradient matches projected instaneous atomistic forces on average. Critically, this does not require simulations of the CG model during training, and was proposed for ML CG protein force-fields in [44]. As many atomistic coordinates r map to the same CG coordinate R, the instantaneous force is noisy, and the signal-to-noise ratio becomes smaller the more degrees of freedom are “CGed away”; thus CG force-matching requires more data as compared to atomistic force-fields. A second difficulty comes from the fact that U is obtained by implicitly integrating the mean force, and as a result, obtaining the free energy difference between minima depends on estimates of the forces along the transition path, where the uncertainties are the largest.

The recently proposed flow-matching method [54] combines relative entropy estimation and force-matching by employing generative deep learning: the CG density is estimated by a normalizing flow, a neural network that can generate one-shot samples of equilibrium CG conformations. This flow can then generate samples to train a downstream ML CG force-field by force-matching. The limitation of this approach still lies in finding flow architectures that can scale to large macromolecules.

The distribution of atomistic configurations is fundamental to the discussed algorithms. Rare events are important but infrequently sampled in the canonical distribution; directing atomistic sampling towards barriers and areas of “high uncertainty” may be beneficial. While ML models are expressive than, e.g., pair potentials, they require more data. For example, ML CG force-matching may use upwards of one million canonically distributed samples covering the configurational space for small proteins [36], in contrast to harmonic models parameterized using short trajectories in the folded state [55]. Modifying the distribution of samples may reduce data requirements [56], but it is unclear how such approaches scale with system complexity.

Concurrently, iterative methods may overcome their computational barriers if non-canonical sampling is used; expanding discriminative training may remove the need for repeated training simulations [46], and biasing potentials may promote diversity and produce more accurate parameters [57]. However, these approaches often require data to be drawn from a modified distribution, which impedes the use of preexisting atomistic trajectories. Nevertheless, these approaches will be critical to expanding current ML CG success to multi-domain proteins.

For a transferable ML CG model, more requirements for the training dataset arise. It is straightforward, and important, to simultaneously force-match a model using reference data from multiple proteins as evidence has shown that extended ensembles can act as regularization [58]. Previous pioneering work developing bottom-up transferable CG models used this approach, but fell short of unassisted folding and relied on artificially lowering simulation temperatures to stabilize states of interest [59]; we associate these inaccuracies to limitations of the force-field basis and training set. We anticipate that the proportion of structural motifs in the dataset plays an important role. In the ideal case, a general CG model of proteins would likely include globular, fibrous, and intrinsically disordered proteins in its training procedure. Such a transferable training setup naturally expands the amount of atomistic data available to train a given model; whether this will improve predictions on individual proteins remains to be seen.

2.2. Choice of the coarse-grained representation

In the design of an atomistic force-field, the Born-Oppenheimer approximation justifies the separation between electronic and nuclear degrees of freedom and provides the framework for effective nuclear potential energy surfaces. However, the separation of scales is less clear for CG models. Consequently, the selection of the CG resolution (ℳ from Eq. (1)) is non-trivial and influences the free energy surface that must be learned. The fundamental questions in this area are which resolutions are “easy to learn” and which are conducive to creating transferable models. These points highlight the challenges of validating ML CG models that are capable of extrapolating to unseen systems. For certain resolutions, it may be easier to learn an effective CG energy and extrapolate into unknown regions of phase space. On the other hand, certain resolutions may be conducive to accurate ML CG models but may be difficult to interpret.

Current successful CG ML protein applications [42, 44, 36, 46, 54] typically focus on a single site per resolution resolution (typically C_α); however, this appears to be mostly due to simplicity and not systematic validation. “Optimal” resolutions have been studied [60, 61], but it is unclear how they impact ML CG models. Back-mapping, i.e., reconstructing details from CG models, is a current area of investigation [62] and may alleviate interpretability constraints on the CG resolution. Views which link back-mapping with potential optimization can facilitate a joint optimization of the representation alongside the CG energy model [63]; however, these approaches do not yet in themselves search for transferable resolutions.

2.3. Functional Form of the Many-body Effective CG (Free) Energy

In practice, training ML CG models via force-matching from equilibrium data requires a baseline (or “prior”) potential to reduce catastrophically incorrect extrapolation in unphysical regions of phase space [44, 36, 64]. Ultimately, a good prior potential incorporates physical principles, reduces learning complexity, and allows for stable simulation. Similar to Δ-learning [65] for atomistic force-fields [66, 67], the CG energy is usually decomposed into:

$$U(R;\theta )=U_{prior}(R)+U_{net}(R;\theta )$$ (2)

where U_net(R; θ) is a trainable multibody potential expressed by an ML model with parameters θ and U_prior(R) is the prior energy.

Designing the priors is non-trivial as it depends on an interplay between the CG mapping, the ML architecture, and the training data. Poor choices of priors can significantly reduce the performance of an ML force-field [47, 68]. Currently, the prevailing strategy involves proposing a prior inspired by the low body-order terms from classical force-fields, and then iteratively developing an ML CG model over both the prior terms and hyperparameters [44, 36, 35, 37, 47]. Systematic strategies have yet to be developed to design prior energies that are transferable to different molecules.

While priors help enforce important physical asymptotic interactions, the ML model architecture itself should respect basic physical constraints. These include invariance with respect to permutations of particles of the same type, invariance to translations and rotations of the reference frame, and curl-free force predictions [69, 44]. A way to allow the learnable energy U_net from Eq. (2) to be transferable [36] is to decompose it bead-wise such that:

$$U_{net}(R,a)=\sum _i^N{u}_{net}\left(R-R_i,a_i\right)$$ (3)

where R_i is the i^th bead (with type a_i) in the configuration R so that R−R_i are the relative displacements of all beads around bead R_i, and u_net is the bead-wise contribution to the potential.

On top of these constraints, bottom-up coarse-graining involves additional architectural challenges. Coarse-graining a variety of different groups of atoms leads to a large number of CG bead types, e.g., at least 20 types (one for each amino acid) for proteins at the C_α resolution. As mentioned above, training with noisy forces requires a large number of training configurations. As a result, the ML approach must accommodate large training sets (indicating that neural networks may be preferable over kernel methods) and should not scale with the number of bead types so that evaluation times do not increase when considering transferable models. This constraint favors the use of deep learning architectures like SchNet [33] over models based on fixed representation, e.g., symmetry functions [70, 71].

2.4. Validation and Robustness

While atomistic ML force-field development has matured, there exists no appropriate set of best practices for probing stability and robustness. It is common to assess atomistic model accuracy with pointwise metrics, such as the mean force error, over fixed test datasets[27]. However, without an understanding of how models extrapolate into data-poor regions, these metrics can not be used as indicators of simulation stability or accuracy [72, 64], as simulations may explore uncovered configurations. For ML CG models, even with the use of prior energy terms, force error does not guarantee a stable model [68].

Due to the difficulty in constructing comprehensive test sets, the robustness and accuracy of a trained ML model can only be ascertained through extensive sampling, e.g., by using the model to run long MD simulations. Recent investigations into ML architectures have revealed the need for such metrics for both atomistic and CG ML models [64, 68]. Unfortunately, obtaining a converged CG MD simulation can require several million force evaluations; for large systems and complex architectures this may present a computational bottleneck [73]. Validation difficulties impede hyperparameter optimization (e.g., regularization strength, cutoff, or prior potential), as searches may become prohibitively expensive. We note, however, that existing applications provide suitable initial choices of hyperparameters for select architectures and resolutions (see [36]), but that the introduction of novel architectures naturally requires substantial effort for the initial hyperparameter search.

Even once MD has been used to characterize a ML CG model, validation still poses difficulties. When characterizing atomistic force-fields on selected configurations it is typically possible to compare the model’s energy and force predictions to reference values; unfortunately, these are not available at the CG resolution ((Eq. 1)). Instead, analysis typically projects CG configurations onto low-dimensional collective variables (e.g., Fig. 3). However, as ML CG models are now able to reproduce such collective variable surfaces, the need for more rigorous validation is emerging. Recent work [75] has proposed classification as an approach to generate energy-like errors for CG models and may provide an avenue for connecting atomistic and CG force-field validation.

Figure 3:

State-of-the-art performance for a C_α CG ML model on the benchmark protein CLN025. A) Comparison of the CG free energy landscape of CLN025 (produced using MD) for a learned CG ML model with the corresponding free energy for the reference all-atom dataset projected onto slow degrees of freedom (TICA) [74]. B) Ensembles of structures sampled from the CG ML model MD simulation (in red) are superimposed onto all-atom reference structure counterparts (in blue). Basin 1 represents the unfolded state, basin 2 the misfolded state, and basin 3 the folded state.

A related challenge is presented by model uncertainty: How robust is a ML model to different training seeds or data partitioning strategies? For neural networks, these can be expensive questions to answer. However, recent advances have started to enable estimates of uncertainty [76, 77]. A promising strategy involves estimating the uncertainty of predictions and minimizing it either before or during model deployment, either through iterative training or through “on-the-fly” frameworks [78] where data is added to the training set based on such estimates.

3. Conclusion

At the moment of writing, state-of-the-art ML CG models can quantitatively reproduce the behavior of small proteins, as shown in Fig. 3 for Chignolin and in Ref. [54] for Trpcage, BBA, and Villin. Currently, the largest barrier to describe larger proteins is gathering sufficient training data. To what extent such an approach can be extended to define transferable CG models remains an open question. It may be possible only for a class of proteins, or at particular resolutions. Before the advent of ML methods, these questions remained challenging to answer, as thermodynamic consistency between an atomistic and a CG model ((Eq. 1)) could only be approximately enforced; it was not clear whether the limitations of transferable models [59, 79, 80, 81, 39] were due to the limited expressivity of the CG energy and limited reference data or to more fundamental problems with transferability. Now, as ML CG models can quantitatively enforce thermodynamic consistency for single proteins (as shown in Fig. 3), we have the tools to address these questions and explore the trade-off between accuracy and transferability. Here, we have discussed the practical challenges towards this goal, but we remain optimistic that such a line of research can be pursued.

Even if a transferable bottom-up ML CG model can be defined, eventually, the success of a computational model relies on its comparison to experiments. Bottom-up CG models rely on the reference atomistic models and necessarily inherit their inaccuracies and flaws. With the improvements in atomistic force-fields, we expect CG models to also become more accurate. However, even small inconsistencies between the CG and atomistic models may compound into a significant deviation from experimental data. We believe that ultimately bottom-up ML CG will need to be merged with top-down models for their useful and predictive applications.

Highlights.

• Coarse-graining is a powerful tool for modeling complex macromolecular systems

• Machine learning is now enabling the definition of accurate bottom-up coarse-grained protein models

• Bottom-up protein coarse-grained models that are fully transferable in sequence space still do not exist

• We discuss the outstanding challenges towards the design of transferable coarse-grained protein models and possible ways forward