Coarse Graining Molecular Dynamics with Graph Neural Networks

I Introduction

Technologies facilitating molecular dynamics (MD) simulations, such as distributed computing Shirts and Pande 2000; Allen et al. 2001; Buch et al. 2010 and bespoke hardware Shaw et al. 2008, have made great strides in terms of the time- and length-scales accessible in silico. However, even the longest protein simulations still fail to reach total times exceeding milliseconds, and dedicated analysis methods are required to infer dynamics at longer timescales Lane et al. 2013; Plattner et al. 2017. In the context of such limitations at full atomistic resolution, coarse graining provides a crucial methodology to more efficiently simulate and analyze biomolecular systems. In addition to the practical advantages that arise from more efficient sampling, coarse graining can also elucidate the physical components that play key roles in molecular processes.

Coarse graining is especially useful for analyzing structures and processes that reach beyond the length- and time scales accessible to all-atom MD. Important examples include protein folding, protein structure prediction, and protein interactions Kmiecik et al. 2016. Some of the most-used coarse-grained models for such studies are structure-based models Clementi, Nymeyer, and Onuchic 2000, MARTINI Marrink et al. 2007; Monticelli et al. 2008, CABS Koliński et al. 2004, AWSEM Davtyan et al. 2012, and Rosetta Das and Baker 2008. These models differ with respect to their potential energy function, parameterization approaches, and resolution, which in combination determine their efficiency, accuracy, and transferability. In the past decade, coarse-grained models have become increasingly powerful due to an unprecedented wealth of experimental reference data and computational capabilities. In this context, the development of more realistic architectures and modeling approaches is of prime importance.

In the field of computer science, advances in hardware and autodifferentiation software have enabled enormous progress in machine learning algorithms, including at the intersection of computation and the physical sciences Noé et al. 2020; Gómez-Bombarelli and Aspuru-Guzik 2020. Crucial to the use of neural networks in the physical sciences is a consideration for the form the training data takes before it is input into the network. One strategy for representing molecules mathematically is through the use of graphs, whose nodes and edges intuitively correspond to atoms and bonds of (or interatomic distances within) a molecule, respectively. By performing multiple convolution operations on a graph, each node can influence other increasingly distant nodes. The use of graph neural networks Kipf and Welling 2017; Battaglia et al. 2018 in the molecular sciences is therefore a promising direction in a variety of applications, and graph convolutional architectures have been used to predict molecular Duvenaud et al. 2015; Kearnes et al. 2016; Gilmer et al. 2017; Feinberg et al. 2018 and material Xie and Grossman 2018 properties as well as atomic energies Schütt et al. 2017; Schütt et al. 2017 and forces Ruza et al. 2020.

In this work, we combine the use of graph representations of molecules with a supervised neural network architecture in the coarse graining context. We consider coarse graining to be the process of reducing structural degrees of freedom to facilitate more efficient simulation with specific goals in mind (e.g., reproducing system thermodynamics). Coarse graining can be implemented with a “top down” or “bottom up” approach, although other categories can be determined and strategies can be combined Noid 2013. In a “top down” scheme, coarse graining frameworks are explicitly designed to reproduce certain macroscale emergent properties Noid 2013. In a “bottom up” framework, which we consider here, implementations focus instead on reproducing specific features from a more detailed model.

The latter involves (i) a mapping from the entities in a fine-grained (e.g., atomistic) representation to a smaller set of interaction sites, often called “beads,” and (ii) a physical model (i.e., Hamiltonian function) for the coarse-grained system comprising those beads. Good choices for the mapping and model will lead to more efficient simulation while preserving the biophysical properties of interest to the researcher. Modern machine learning techniques have been recently employed to learn both the mapping Boninsegna, Banisch, and Clementi 2018; Wang and Gómez-Bombarelli 2019 and the model John and Csányi 2017; Zhang et al. 2018a; Wang et al. 2019; Wang et al. 2020; Ruza et al. 2020 components of bottom up coarse graining.

In the present contribution, we focus on the coarse graining model and employ a bottom up “force matching” scheme formulated as a supervised machine learning problem to reproduce the thermodynamics of small biomolecular systems. Particularly, we modify the architecture of the recently-introduced CGnet framework Wang et al. 2019 such that the molecular features it requires are learned via graph convolutional neural networks instead of hand-selected as in the original formulation. By leveraging the inherently transferable SchNet scheme Schütt et al. 2017; Schütt et al. 2018 to learn features, we render the entire CGnet framework transferable across molecular systems.

Our goal in this paper is to present the theory underlying CGSchNet—our new transferable coarse graining architecture—and to demonstrate its success on learning the thermodynamics of individual biomolecular systems. We find that our new protocol produces more accurate free energy surfaces in comparison with the use of hand-selected features, is more robust to hyperparameter choices, and requires less regularization. Presented alongside a machine learning software package that implements the methods introduced, the current contribution sets out a framework for the machine learning of transferable, coarse-grained molecular force fields and demonstrates its application to a small peptide system and the miniprotein chignolin Honda et al. 2008. The practical application of the methods described herein to larger protein systems, particularly those characterized by meaningful tertiary structure, remains an open challenge that will be explored in future work.

II Theory

Force matching was pioneered in the atomistic setting, in which forces obtained from an inexpensive calculation are matched to forces computed at a more computationally expensive level of theory (i.e., quantum) via an optimization scheme Ercolessi and Adams 1994. The method was later adapted by the coarse graining community; in that context, coarse-grained representations are sought such that the forces computed from the coarse-grained energy function for a given configuration match the average forces on corresponding atomistic representations Izvekov and Voth 2005a.

Because coarse graining away degrees of freedom entails that multiple atomistic structures will correspond to the same coarse-grained configuration, it is impossible to obtain zero error during force matching in the coarse graining context. However, it can be proved that the coarse graining model that matches the mean forces yields the correct thermodynamics, and that the objective is variationally bounded from below by a value that necessarily exceeds zero.

In Sec. II, we overview the major advances that enable the present contribution. The practically inclined reader may proceed directly to Sec. III, where we discuss the CGnet architecture and introduce this work’s methodological contribution: namely, the incorporation of learnable molecular features into CGnet via the use of continuous filter convolutions on a graph neural network (i.e., SchNet Schütt et al. 2017; Schütt et al. 2018). We will see in Sec. III that the scheme we introduce here enables, at least in principle, a coarse graining architecture that is transferable across system size and sequence. The practical use of this architecture to learn a force field in a transferable context will be addressed in future work.

where $\mathbf{R}$ is the set of all $M$ sampled atomistic configurations.

The objective (1) was introduced by Ercolessi and Adams to analyze ab initio simulations of elemental aluminum Ercolessi and Adams 1994. The authors highlight the method’s need to accommodate invariant properties of the system and discuss the requirement of a variety of geometries, physical conditions, and system identities in $\mathbf{R}$ if the learned potential is to be transferable across conformation, thermodynamic, or chemical space, respectively. Subsequent work has derived analytical approaches to this scheme in the context of liquids Izvekov et al. 2004; Guenza et al. 2018.

A decade later, Izvekov and Voth introduced the multiscale coarse graining (MS–CG) method, a groundbreaking advance that adapts force matching to the coarse graining context Izvekov and Voth 2005a; Izvekov and Voth 2005b. The MS–CG framework involves two steps: first, atoms are aggregated into “interaction sites” according to a linear mapping from $N$ atoms to $n$ interaction sites (henceforth “beads”),

Consider a coarse-grained energy function $U(\mathbf{x};\mathbf{\Theta})$. Let’s say we have a set of $M$ coarse-grained configurations that we have obtained by applying (2) to every configuration $\mathbf{r}_{i}\in\mathbf{R}$. To calculate the forces on the beads, we then take the negative derivative of $U$ with respect to the reduced coordinates; in other words, we evaluate,

for each configuration $i$. From here we have all the ingredients to write down the adaptation of (1) to the MS–CG method:

where $\mathbf{\Xi_{F}}\mathbf{F}$ is the instantaneous coarse-grained force (also called the local mean force); that is, the projection of the atomistic force into the coarse-grained space. A general expression for the force projection Ciccotti, Lelievre, and Vanden-Eijnden 2008 is $\mathbf{\Xi_{F}}=(\mathbf{\Xi}\mathbf{\Xi}^{\top})^{-1}\mathbf{\Xi}$. Other choices for the mapping $\mathbf{\Xi_{F}}$ are possible and used for coarse graining Noid et al. 2008a.

In principle, the coarse-grained energy $U(\mathbf{x})$ that is exactly thermodynamically consistent with the atomistic energy $V(\mathbf{r})$ can be expressed analytically as:

where $k_{B}$ is Boltzmann’s constant and $T$ is the absolute temperature. The function $p^{\text{CG}}$ is the marginal probability density,

The coarse-grained energy function (4) is called the potential of mean force (PMF) and is an analogue of the atomistic potential energy function. Via (5), it is a function of weighted averages of energies of atomistic configurations. For a given coarse-grained structure $\mathbf{x}^{\prime}$, in (5) we evaluate whether every possible $\mathbf{r}\in\mathcal{R}$ maps to $\mathbf{x}^{\prime}$. We expect multiple atomistic configurations $\mathbf{r}$ to map to $\mathbf{x^{\prime}}$ due to the reduction in degrees of freedom that results from structural coarse graining (n.b., this means the PMF is in fact a free energy, as it contains entropic information Noid 2013). In these cases, the Dirac delta function in (5) returns one, and the contribution of that atomistic configuration to the marginal probability distribution is a function of its Boltzmann factor. If $\mathbf{r}$ does not map to $\mathbf{x}^{\prime}$, then the evaluation of the delta function (and thus the contribution of that atomistic structure to the free energy of $\mathbf{x}^{\prime}$) is zero. The denominator of the right-hand side of (5) is the all-atom partition function, which serves as a normalization factor.

To calculate the forces on our coarse-grained beads, we must take the gradient of (4). However, since we cannot exhaustively sample $\mathcal{R}$, (5) is intractable, and we must approximate $U$ instead. One way to approximate $U$ is to employ force matching—that is, by minimizing 3—as we describe in Sec. II.2. Another method, which we do not discuss in this report, is through relative entropy Shell 2008, whose objective is related to that of force matching Rudzinski and Noid 2011; Noid 2013.

II.2 Coarse graining as a supervised machine learning problem

In 2008, Noid et al. 2008a formalized the notion of thermodynamic consistency and established the conditions under which it is guaranteed by the MS–CG approach: namely, that thermodynamic consistency is achieved when the coarse-grained coordinates are a linear combination of the all-atom coordinates (cf. (2)) and that the equilibrium distribution of the coarse-grained configurations is equal to the one implied by the equilibrium distribution of the atomic configurations (cf. (4)). Noid et al. then prove that, under certain restrictions of the coarse-grained mapping, the coarse-grained potential that achieves thermodynamic consistency at a given temperature is unique (up to an additive constant, cf. (5)) Noid et al. 2008a.

The authors define an error functional that is (uniquely) minimized for the thermodynamically consistent coarse-grained force field Noid et al. 2008a; Noid et al. 2008b. This framework provides the variational principle underlying the MS--CG method. It follows that a variational approach can be used to search for the consistent coarse-grained force field. In practice, such a search is limited by the basis of trial force fields as well as the finite simulation data used Noid et al. 2008b. The variational principle entails that we can refer to (1) and (3) as “loss functions” because they return a scalar that assumes a minimum value on the optimal model. In recent reports from both Wang, Clementi, et al. Wang et al. 2019 and Wang and Gómez-Bombarelli 2019, this fact is leveraged to formulate coarse graining via force matching as a supervised or semi-supervised machine learning problem, respectively. Here, we build upon on the supervised learning case introduced in Ref. Wang et al. 2019 as CGnet.

In their study, Wang et al. present several crucial contributions Wang et al. 2019. First, they decompose the error term implied by (3) into three physically meaningful components; namely, bias, variance, and noise. Second, the authors introduce CGnet: a neural network architecture designed to minimize the loss in (3). Once a CGnet is trained, it can be used as a force field for new data points in the coarse-grained space while enforcing known properties of the system such as symmetries and equivariances (see Sec. III.3). Third, Wang et al. augment their initial framework to introduce regularized CGnets Wang et al. 2019. Regularized CGnets avoid catastrophically wrong predictions observed in their “unregularized” counterparts by introducing the calculation of prior energy terms before training. This adjustment means that, instead of learning the forces directly, the neural network learns a correction to the prior terms in order to match the atomistic forces.

Using regularized CGnets (henceforth, we assume all CGnets are regularized) on two peptide systems, the authors demonstrated effective learning of coarse-grained force fields that could not be obtained with a few-body model approach Wang et al. 2019. It is from this baseline that we present CGSchNet, an augmentation of the CGnet methodology.

III Methods

In the quantum community, supervised machine learning has been used to predict energies on small molecules through a variety of approaches Behler and Parrinello 2007; Bartók et al. 2010; Rupp et al. 2012; Bartók et al. 2013; Smith, Isayev, and Roitberg 2017; Chmiela et al. 2017; Bartók et al. 2017; Schütt et al. 2017; Smith et al. 2018; Schütt et al. 2018; Grisafi et al. 2018; Imbalzano et al. 2018; Nguyen et al. 2018; Zhang et al. 2018b; Zhang et al. 2018c; Bereau et al. 2018; Wang and Yang 2018. In particular, the SchNet architecture is based on the use of continuous filter convolutions and a graph neutral network Battaglia et al. 2018; Schütt et al. 2017; Schütt et al. 2018. SchNet is a scalable, transferable framework that employs representation learning to predict the properties and behavior of small organic molecules. In the vein of the original force matching procedure of Ercolessi and Adams 1994, SchNet has also been used to predict forces on atomic data from a quantum mechanical gold standard Schütt et al. 2018.

In Sec. III.1 we briefly overview the CGnet scheme upon which we base the method introduced in this work. Then, in Sec. III.2, we describe SchNet and introduce our adaptation of SchNet to the coarse graining problem by incorporating it into a CGnet to create a hybrid “CGSchNet” architecture. The original implementation of CGnet is not transferable across different systems due to its reliance on hand-selected structural features Wang et al. 2019. We recognized that SchNet could be leveraged as a subcomponent of CGnet in order to learn the features, thereby converting CGnet—i.e., force matching via supervised machine learning—to a transferable framework for the first time.

While the mapping is permitted to be more general, in our work we restrict it to the special case where the matrix $\mathbf{\Xi}$ contains zeroes and ones only. With this choice of mapping, the projection of the forces in (3) becomes simply $\mathbf{\Xi}_{\mathbf{F}}=\mathbf{\Xi}$. Our mapping thus “slices” the original atomic configuration such that the corresponding coarse-grained representation comprises a subset of the original atoms. For example, a useful mapping might retain only protein backbone atoms or $\alpha$-carbons.

To construct a CGnet, the structural data is preprocessed such that it is represented by features with the desired properties. Wang et al. use a set of distances, planar angles, and torsional angles Wang et al. 2019. In the present work, on the other hand, instead of using hand-selected structural features, we require only distances and bead identities from which features are learned; this is described in Sec. III.2.

For their regularized implementation, Wang et al. use up to two types of prior terms in CGnets Wang et al. 2019. The first is a harmonic prior on selected distances (i.e., bonds or pseudobonds) and angles. The second is a repulsion prior that can be used on nonbonded distances. Respectively, these priors are defined as follows for a given feature $f_{i}$ calculated from the data (e.g., a particular distance),

The constants in (6b) can be determined through cross-validated hyperparameter optimization as in Ref. Wang et al. 2019. The prior energy is the sum of each prior term for all relevant features $f_{i}$. In principle, any scalar function of protein coordinates can be used to construct a prior energy term.

III.2 Replacing structural features with graph neural networks

Wang et al. show that CGnets constructed upon hand-selected structural features produce machine-learned force fields yielding accurate free energy surfaces Wang et al. 2019. The model architecture is found to be somewhat sensitive to various hyperparameters and required individual tuning for each system (see e.g. Fig. 5 in Wang et al. 2019). Furthermore, a new system will in general require retraining because the feature size is fixed according to the system geometry.

In the present contribution, we replace the fixed structural features employed in the original CGnet formulation (i.e., distances, angles, and torsions) Wang et al. 2019 with learned features computed using continuous filter convolutions on a graph neural network (SchNet Schütt et al. 2017; Schütt et al. 2018). The SchNet architecture thereby becomes a subunit of CGnet with its own, separate neural network scheme; we refer to this hybrid architecture as CGSchNet.

The term graph neural network was introduced in Battaglia et al. 2018 as a generalization for networks operating on graph structures, including but not limited to graph convolutional networks Kipf and Welling 2017 and message passing networks Gilmer et al. 2017. These networks have in common that they have a notion of $n$ nodes $\mathcal{V}$ connected by edges $\mathcal{E}$ in a graph structure. In each neural network layer, information is passed between nodes and representations of the nodes and/or edges are updated. The various types of graph neural networks differ according to whether there are node updates, edge updates, or both; as well as by how functions are shared across the network and how the network output is generated. A fairly general formulation of a graph neural network with node updates is as follows: each node $i$ is associated with an initial node representation $\mathbf{h}^{(0)}_{i}$; in other words, $\mathbf{h}^{(0)}_{i}$ is a vector that represents the type or identity of the node. In each layer of the neural network, the node representations are updated according to,

where $i,j=1,\dots,n$, $\mathbf{e}_{ij}$ are edge features defined for all edges in $\mathcal{E}$, and $f$ is a trainable neural network function. After $T$ such layers, an output is generated,

In the present contribution, we make the following choices:

Graph nodes represent coarse-grained beads.

Because multi-body interactions are important for the coarse graining problem, edges are defined between all beads (or all beads within a specified cutoff radius).

The edge features $\mathbf{e}_{ij}$ are taken to be the distances between beads, implying translational and rotational invariance of the network output.

The update function $f$ in (7) is chosen to be a continuous convolution update as in SchNet Schütt et al. 2017.

The entire trainable part of a CGnet Wang et al. 2019—in this case a beadwise multilayer perceptron/dense neural network—becomes the output function $g$ in (8). Because this output is beadwise the learnable coarse grain energy is invariant with respect to permutations of identical beads.

The output $\mathbf{o}$ is a scalar; namely, the coarse-grained energy before the addition of the prior energy term.

Below we describe the SchNet updates in more detail and how to incorporate SchNet into CGnet to create CGSchNet.

One key motivating factor for the original development of SchNet is that, unlike the images and videos that comprise the datasets for much of modern machine learning, molecular structures are not restricted to a regular grid. Therefore, Schütt et al. introduced continuous-filter convolutions to analyze the structures of small molecules with the goal of predicting energies and forces according to a quantum mechanical gold standard Schütt et al. 2017. This development builds upon previous work in predicting atomic properties directly from structural coordinates Chmiela et al. 2017; Schütt et al. 2017; Gilmer et al. 2017.

SchNet is a graph neural network where the nodes correspond to particles embedded in three-dimensional space and the convolutional filters depend on interparticle distances, which preserves invariances expected in the system Schütt et al. 2017; Battaglia et al. 2018. While SchNet was originally used to predict quantum-chemical energies from atomistic representations of small molecules, here we employ it to learn a feature representation that replaces the hand-selected features in a CGnet for the purpose of predicting the coarse-grained energy on the coarse-grained bead coordinates $\mathbf{x}_{i}$.

As in other graph neural networks, SchNet learns feature vectors on the nodes (here, coarse-grained beads). The initial node features at the input are called node or bead embeddings $\mathbf{h}_{i}^{(0)}$, which are given by trainable, shared vectors with $d_{h}$ dimensions (“Embeddings” in Fig. 1),

Here, $k(i)$ is a lookup table that maps the bead index $i$ to its type $k$. In the present applications we use nuclear charges (capped alanine) or amino acid identities (chignolin) as bead types. The bead embeddings are shared among beads of the same type and are optimized during training. Crucially, this entails that SchNet learns a molecular representation, which avoids the common paradigm of fixed, heuristic feature representations.

Next, we describe how bead representations are updated (“Interaction block” and “cfconf” in Fig. 1). In our current architecture, these updates are implemented as in the original SchNet unless noted otherwise (see Refs. Schütt et al. 2017 and Schütt et al. 2018 for details).

In each interaction layer, we perform a continuous convolution between beads. For this, the inter-bead distances $|\mathbf{x}_{j}-\mathbf{x}_{i}|$ are featurized using radial basis functions $\mathbf{e}$, e.g., one-dimensional Gaussians centered at different distances. These featurized distances serve as the input to a filter-generating neural network $w$ that maps the featurized distance input $\mathbf{e}(|\mathbf{x}_{j}-\mathbf{x}_{i}|)$ to a $d_{h}$-dimensional filter. This filter is applied to the bead representations $\mathbf{h}_{i}$ as follows (“cfconf” in Fig. 1),

Here, $w$ and $b$ are trainable functions and $\cdot$ is element-wise multiplication. As in the original SchNet implementation Schütt et al. 2017, $w$ is a dense neural network and $b$ is a beadwise linear layer. The sum in (10) is taken over every bead $j$ within the neighborhood of bead $i$, which can be all other beads in the system or a subset thereof if a finite neighborhood is specified. Even when interactions are limited to particles within a cutoff radius, a sequence of multiple interaction layers will eventually allow all particles to be interacting, and therefore be able to express complex multi-body interactions.

In each layer, the bead representations are updated in interaction blocks, each of which comprises a residual update of the bead representation via a nonlinear function of the continuous convolution outputs $\mathbf{z}_{i}^{(t)}$ (“interaction block” in Fig. 1),

The residual update step is an “additive refinement” that prevents gradient annihilation in deep networks He et al. 2016. As described by Schütt et al. Schütt et al. 2017; Schütt et al. 2018, the trainable function $g$ involves beadwise linear layers and a nonlinearity. Instead of the softplus nonlinearity used in the original SchNet Schütt et al. 2017, here we use the hyperbolic tangent.

Following the last interaction layer we must choose an output function (8). As in the original SchNet implementation, the output of the final SchNet interaction block is input into an beadwise CGnet multilayer perceptron/dense network. An important feature of transferability is permutation invariance of beads with identical type. In the context of coarse graining, this means the contribution of a bead to the coarse-grained energy should depend on its location in the molecular graph, but not at which index this bead is positioned in the input. SchNet layers are permutation-equivariant, i.e., any exchange of the input representations $\mathbf{h}^{(0)}_{i}$ will correspond to the same exchange of the learned representations $\mathbf{h}^{(T)}_{i}$. In order to obtain permutation-invariant energies (as in the original SchNet publications Schütt et al. 2017; Schütt et al. 2018), the beadwise output CGnet network contracts down to a scalar energy prediction that is then summed over all beads to yield the total learnable part of the coarse-grained energy. It is important to note that the models used in this study employ priors that are not permutation invariant, and so the non-learnable part of the coarse-grained energy (i.e., the prior terms) breaks permutation invariance in the model overall. The development of permutation invariant priors is left for future work.

In the present paper we do not use CGSchNet in a transferable manner, but rather demonstrate its capabilities when trained on individual molecular systems as a foundation for future work. For this reason, here we use a (regularized) beadwise CGnet as the output function; i.e., a beadwise multilayer perceptron/dense neural network at whose output the learned part of the coarse-grained energy is predicted. In so doing, the SchNet interaction layers learn the input representation for a beadwise CGnet, and the beadwise CGnet “fine-tunes” the bead energies predicted by SchNet.

III.2.2 CGSchNet: a transferable architecture for coarse graining

CGnet as originally presented is incapable of learning a transferable coarse-grained force field due to its reliance upon system-specific structural features Wang et al. 2019. Since SchNet is inherently a transferable framework, learning CGnet features using SchNet enables the transferability of the entire CGnet architecture across molecular systems. Here, we present the advance of incorporating SchNet Schütt et al. 2017; Schütt et al. 2018 into CGnet to replace hand-selected features with machine-learned ones.

In CGSchNet, instead of predetermined structural features—i.e., distances, angles, and torsions—a SchNet is used instead, enabling the model to learn the representation itself (see Fig. 1). By replacing fixed-size geometric features with SchNet, we obtain a more flexible representation that both scales better with system size and is amenable to a transferable architecture Schütt et al. 2018. While angles and torsions may still be included in the prior energy terms, they are no longer propagated through any neural networks.

The use of SchNet requires us not only to provide structural coordinates but also a type for every bead. In the original (i.e., non-coarse graining) implementation for systems at atomic resolution, the types are atomic numbers Schütt et al. 2017; Schütt et al. 2017; Schütt et al. 2018. In the new context presented here (i.e., leveraging SchNet for coarse graining), we may specify coarse-grained bead types—effectively, chemical environments—however we deem appropriate for the system under study; for example, amino acid identities may be used.

One can view the typing requirement of the SchNet architecture as the outlet through which to incorporate physical or chemical intuition about the system into the model, as opposed to through fixed structural features. The benefit of the SchNet choice is that it enables an architecture that is transferable across size and sequence space because a set of embeddings can apply to multiple systems with the same components (e.g., atoms as in SchNet Schütt et al. 2017 or amino acid types in proteins), whereas hand-selected structural features are not transferable across different systems.

Finally, we note that we train CGSchNet with the coarse-grained force matching loss (3), which compares our predicted forces to the known forces from the training set. Unlike in the original SchNet formulation Schütt et al. 2017, we cannot straightforwardly incorporate an additional “energy matching” term into the coarse graining framework. This is because we do not have labels for the coarse-grained free energies: these energies are defined by an integral over all microscopic configurations associated with the same coarse-grained configuration (cf. (4)), and this integral cannot be solved exactly.

III.3 Coarse-grained simulations

A trained CGSchNet can be used as a force field to simulate the system in the coarse-grained space. Specifically, Langevin dynamics Schneider and Stoll 1978; Fass et al. 2018 are employed to propagate coarse-grained coordinates $\mathbf{x}_{t}$ forward in time according to,

where the diagonal matrix $\mathbf{M}$ contains the bead masses, $\gamma$ is a collision rate with units ps-1, and $\mathbf{W}(t)$ is a stationary Gaussian process with $\langle W(t)\rangle=0$ and $\langle W(t)W(t^{\prime})\rangle=\delta(t-t^{\prime})$, where $\langle\cdot\rangle$ is the mean. In practice, we integrate (12) using a “BAOAB” Langevin integrator Leimkuhler and Matthews 2013, and the integral of $W(t)$ is a Wiener process.

A special case of Langevin dynamics are so-called “overdamped” Langevin dynamics, also referred to as Brownian dynamics. Overdamped Langevin dynamics lack inertia. After setting the acceleration to zero, dividing both sides by $\gamma$, and rearranging terms, in the overdamped case, (12) becomes,

where the diffusion matrix $\mathbf{D}\equiv\mathbf{M}^{-1}k_{B}T/\gamma$. Although $\mathbf{D}$ contains a notion of mass, we note that propagating coarse-grained dynamics via (13) does not actually require bead masses, since the product $\mathbf{M}\gamma$ can be considered without separating its factors. Wang et al. 2019 use exclusively (13) with the Euler method to simulate dynamics from CGnets, with a constant diffusion matrix proportional to the identity matrix.

In both formulations, the noise term is intended to indirectly model collisions—e.g., from and among solvent particles—that are not present in the coarse-grained coordinate space. Since Langevin dynamics depend only on the coordinates (and, unless overdamped, velocities) of the previous time step, these simulations can easily be run in parallel from a set of initial coordinates. The resulting coarse-grained simulation dataset can then be used for further analysis as we will show in Sec. IV.

IV Results

Capped alanine—often referred to alanine dipeptide for its two peptide bonds—is a common benchmark for MD methods development because the heavy-atom dynamics of the central alanine are completely described by the dihedral (torsional) angles $\phi$ and $\psi$ (see Fig. 2). We performed a single 1-$\mu$s all-atom, explicit solvent MD simulation for capped alanine and saved the forces to use for CGSchNet training (see Ref. Wang et al. 2019 and supplementary material Sec. A). We can visualize the occupancies of backbone angle conformations by creating a histogram of the data on $\phi$ $\times$ $\psi$ space and visualizing the populations of the histogram bins. This is called a Ramachandran map and is depicted in Fig. 4a for the atomistic simulation using a $60\times 60$ regular spatial discretization.

As an initial benchmark of the CGSchNet method, we aim to learn a force field for a coarse-grained representation of capped alanine such that we can reproduce its heavy-atom dynamics using a trained CGSchNet instead of a more expensive explicit solvent all-atom MD simulation. For our coarse-grained mapping, we select the backbone heavy atoms C–[N–Cα–C]${}_{\text{Ala}}$–N as well as the alanine Cβ for a total of six beads. As in Ref. Wang et al. 2020, we require the $\beta$-carbon in order to break the symmetry of the system (i.e., to enforce chirality). In their demonstration of CGnet, Wang et al. 2019 used only the five backbone heavy atoms as beads because chirality is enforced through dihedral features, which we do not use here. We use atomic numbers for the bead embeddings as in the original SchNet formulation Schütt et al. 2017. A CGSchNet is trained on the coordinates and forces of the all-atom simulation depicted in Fig. 4a. The learning procedure involves a hyperparameter selection routine and the training of multiple models under five-fold cross-validation for each hyperparameter set (see supplementary material Sec. B).

Once a final architecture has been selected, the trained model can serve as a force field in the coarse-grained space; i.e., by predicting the forces on a set of input coarse-grained coordinates. Along with an integrator, predicted forces can be used to propagate coarse-grained coordinates forward in time (recall Sec. III.3). This procedure (i.e., force prediction with CGSchNet followed by propagation with an integrator) is iterated until a simulation dataset of the desired duration has been obtained. Since we employ five-fold cross-validation during the model training procedure, we have five trained CGSchNet models with a common architecture at hand. To perform our coarse-grained simulation, we simultaneously predict the forces on each set of input coordinates from all five trained networks, and the mean force vector is used to propagate Langevin dynamics according to (12).

To facilitate sampling, 100 coarse-grained simulations of length 200 ns each are performed in parallel from various starting positions in Ramachandran space (see supplementary material Sec. C and Fig. S3). The time series of the $\phi$ and $\psi$ values for two of the trajectories that feature transitions among the major basins are plotted in Fig. 3. The same trajectories are also overlaid on the two-dimensional energy surface in Fig. S4 in the supplementary material.

Free energy surfaces resulting from the coarse-grained simulation dataset are presented in Fig. 4b. We can see qualitatively that the two-dimensional free energy surface from the CGSchNet simulation captures the same basins as the surface calculated from the baseline all-atom simulation. In the one-dimensional free energy surfaces, we see that the barriers are well-approximated by the CGSchNet simulation data.

To calibrate our understanding of the CGSchNet simulation dataset’s relationship to the baseline atomistic simulation dataset, we create a set of new systems by perturbing the Cartesian coordinates of the latter with noise distributed as $\mathcal{N}(0,\sigma^{2})$ for $\sigma\in\{0,0.01,0.02,\dots,0.30\}$ Å. From the perturbed Cartesian coordinates, the new $\phi$ and $\psi$ dihedrals are calculated and assigned to the same $60\times 60$ regularly spaced bins in Ramachandran space. Examples of the perturbed free energy surfaces are shown in Fig. 4c, d, and e for $\sigma=$ 0.1 Å, 0.2 Å, and 0.3 Å, respectively. We see that the surfaces become smeared and the free energy barriers are reduced with increasing noise.

This ensemble of perturbed simulation datasets enables us to understand the CGSchNet-produced simulation in the context of the baseline atomistic simulation. To quantify the relationship between two distributions, we can use the Kullback-Leibler (KL) divergence Kullback and Leibler 1951 and a mean squared error (MSE) formulation. The KL divergence is defined for discrete distributions as,

where $p$ and $q$ are the “reference” and “trial” distributions, respectively, and $m$ is the number of bins in each discrete distribution. In this case, $p$ and $q$ represent the normalized bin counts. The index $i$ returns the normalized count from the $i$th bin of a 60 $\times$ 60 regular discretization of $\phi\,\times\,\psi$ space. The distribution obtained from the baseline atomistic simulation always serves as the reference. The mean squared error used here is,

We see in Fig. 5 that as the noise increases, both divergence metrics also increase. The dashed lines in Fig. 5 show us that the error on the CGSchNet simulation dataset is approximately comparable to the corresponding error on the perturbed dataset with noise scale $\sigma=0.1$ Å (Fig. 4c) Similar results were obtained for the two-dimensional Wasserstein distance.. Upon qualitative comparison of the free energy surfaces, however, the former has more visual fidelity to the baseline surface in Fig. 4a than to the broader spread seen (and expected) in the latter. We know that coarse graining can result in increased population in transition regions that are rarely visited in an all-atom model; this is what we observe in Fig. 4b. As a corollary, we do not expect coarse graining to result in the absence of states known to exist in the baseline system.

IV.2 Chignolin

The CLN025 variant of chignolin is a 10-amino acid miniprotein Honda et al. 2008 featuring a $\beta$-hairpin turn in its folded state (Fig. 6). Due to its fast folding, its kinetics have been investigated in several MD studies Lindorff-Larsen et al. 2011; Beauchamp et al. 2012; Husic et al. 2016; McKiernan, Husic, and Pande 2017; Sultan and Pande 2018; Scherer et al. 2019. Our training data is obtained from an atomistic simulation of chignolin in explicit solvent for which we stored the forces (see Ref. Wang et al. 2019 and supplementary material Sec. A). To build our CGSchNet model, we retain only the ten $\alpha$-carbons for our coarse-grained beads. For the SchNet embeddings, we assign each amino acid type its own environment with a separate designation for the two terminal tyrosines. After determining hyperparameters for our CGSchNet model, we simulate chignolin in the coarse-grained space using Langevin dynamics (12) as in the previous section. The procedures for CGSchNet training and simulation are similar to those used for capped alanine and are described in the supplementary material Secs. B and C.

Given our CGSchNet simulation data, we are interested not only in performing a similar analysis to the one in the previous section for alanine dipeptide (i.e., comparison to the baseline dataset with and without noise added) but also to simulation data obtained from a CGnet trained according to the protocol in Ref. Wang et al. 2019 for the same system. We thus also construct a CGnet according to the parameters selected in Ref. Wang et al. 2019 (i.e., using fixed geometric features as described in Sec. III.1; see also Sec. B in the supplementary material). Then, we create a simulation dataset using the protocol described in the previous section and supplementary material Sec. C. Finally, we employ a similar protocol to the previous section by perturbing the raw Cartesian coordinates of the all-atom chignolin simulation dataset with noise distributed as $\mathcal{N}(0,\sigma^{2})$ for $\sigma\in\{0,0.03,0.06,\dots,0.90\}$ Å.

For each type of system (i.e., baseline, CGSchNet, CGnet, and baseline with noise perturbation), we build Markov state models (MSMs) Zwanzig 1983; Schütte et al. 1999; Swope, Pitera, and Suits 2004; Singhal, Snow, and Pande 2004; Chodera et al. 2007; Noé et al. 2007; Buchete and Hummer 2008; Prinz et al. 2011 (see Refs. Husic and Pande 2018 and Noé 2020 for recent overviews). First, the data is “featurized” from Cartesian coordinates into the set of 45 distances between pairs of $\alpha$-carbons. From these distances, time-lagged independent component analysis (TICA) Pérez-Hernández et al. 2013; Schwantes and Pande 2013 is performed to yield four slow reaction coordinates for the system. These four reaction coordinates are clustered into 150 discrete, disjoint states using the $k$-means algorithm. An MSM is then estimated from the cluster assignments. The MSM for the baseline simulation dataset is constructed first; then, the other simulation datasets are projected onto the space defined by the former. MSM essentials are presented in supplementary material Sec. D from a theoretical standpoint, and the specific protocols used for the MSM analysis in this section are given in supplementary material Sec. E.

The stationary distribution of each MSM is then used to reweight the TICA coordinates used for its own construction. Histograms of the first two TICA coordinates are presented in the top row of Fig. 7 for the baseline, CGSchNet, and CGnet simulation datasets as well as the baseline dataset for $\sigma=0.3$. The first two reweighted TICA coordinates are also individually binned into one-dimensional free energy surfaces, which are depicted in the second and third rows of Fig. 7. We see that the free energy barriers along these reaction coordinates are reasonably approximated by the CGSchNet simulation.

Figure 8 shows the same divergence metrics calculated in Fig. 5 in the previous section. Again, we see that both the KL divergence and the MSE increase monotonically with the magnitude of the noise. In this case, we can assess the equivalent noise value for both the CGSchNet and CGnet simulation datasets. For both divergences measured, we see that the CGSchNet simulation corresponds to a lesser value of added noise than the CGnet simulation.

We can also obtain free energy surfaces from the MSMs constructed for the systems; the surfaces for the baseline and CGSchNet simulation datasets of chignolin are presented in Fig. 9a on the left and right, respectively. We see that the three major basins observed in the atomistic data are captured by CGSchNet. These basins represent folded, unfolded, and misfolded ensembles and are indicated in Fig. 9a with blue, green, and yellow stars, respectively. Each star represents one of the 150 MSM states and was manually selected from the MSM states near the relevant basin (see the supplementary material Fig. S8 for a visualization of all 150 MSM states). To verify that the protein conformations are similar in each of the states, we sample ten structures from each starred state per simulation dataset. The structures are visualized in Fig. 9b-d, and the similarity of the structures on the left-hand side (baseline simulation) to those on the right-hand side (CGSchNet simulation) from corresponding MSM states is apparent.

The analysis of the CGSchNet and CGnet simulation datasets so far used TICA reaction coordinates that were obtained by projecting the simulation data onto coordinates defined by a TICA model built for the baseline atomistic data (see supplementary material Sec. E). This was done in order to compare simulation results using the same reaction coordinates. We can also construct TICA models from the simulation data without projection to determine the scaling factor for the coarse-grained timescale. For this analysis, we build two further (independent) TICA models for the CGSchNet and CGnet datasets at a lag time long enough for the TICA timescales to have leveled off (see Fig. S9 in the supplementary material). A 100-round bootstrapping analysis of the longest TICA timescale from the CGSchNet simulation data yields a time scaling factor of 2.2 with a standard deviation of 0.4. From this time rescaling we determine that the effective collision rate (i.e., friction) of the coarse-grained simulations is 180–260 ps-1. This value is four orders of magnitudes larger than the friction constant in the all-atom model (0.1 ps-1) Wang et al. 2019, which we expect because we have coarse-grained out the solvent dynamics. The same analysis for the CGnet simulation data yields a scaling factor of $2.2\pm 0.3$ and a corresponding effective collision rate of 190–250 ps-1.

Although we use MSMs and TICA models to obtain thermodynamics and effective friction constants, we do not attempt a kinetic analysis in the present work because the scope of force matching is limited to thermodynamic consistency Izvekov and Voth 2005a; Noid et al. 2008a. The matching of dynamics in addition to thermodynamics is an open challenge that has been the subject of recent work Nüske, Boninsegna, and Clementi 2019. Given coarse-grained dynamics, analytical methods have been derived that enable their rescaling to the dynamics of the system’s all-atom counterpart Lyubimov and Guenza 2011.

V Discussion

Coarse graining holds the promise of simulating larger systems at longer timescales than are currently possible at the atomistic level. However, mathematical frameworks must be developed in order to ensure that the results obtained from a coarse-grained model are faithful to those that would be obtained from an atomistic simulation or experimental measurement. Force matching Ercolessi and Adams 1994; Izvekov and Voth 2005a is one such framework that, when certain restrictions are applied, guarantees thermodynamic consistency with atomistic data in the variational limit Noid et al. 2008a. Such a variational framework enables the formulation of the force matching problem as a supervised machine learning task, which is presented in Ref. Wang et al. 2019 as CGnet.

A key limitation of the original CGnet is that it is not transferable across different systems: a new network must be trained for each individual molecular system under study because the molecular features from which it learns the force field must be chosen by hand. Here, we replace manually determined features with a learnable representation. This representation is enabled by the use of continuous filter convolutions on a graph neutral network (i.e., SchNet Schütt et al. 2017; Schütt et al. 2018). SchNet is an inherently transferable architecture originally designed to match energies and forces to quantum calculations for small organic molecules. By leveraging SchNet in the coarse graining context—i.e., to learn the molecular features input into a CGnet—we render the hybrid CGnet architecture (i.e., CGSchNet) transferable across molecular systems of different sizes and sequences.

Our aim in the present contribution is threefold: to summarize the variational framework enabling a supervised learning approach to force matching, to provide an accompanying software package implementing the methods discussed herein (see Appendix A), and to demonstrate that CGSchNet produces results on individual systems that are superior to those obtained from bespoke features. The advances presented in this work prepare us to address the ultimate challenge of machine learning a coarse-grained force field that is transferable across molecular systems.

In our computational experiments performed on capped alanine and the miniprotein chignolin, we find that CGSchNet’s performance exceeds that of CGnet in three ways. First, the free energy surface obtained from CGSchNet simulations of chignolin is more accurate than the free energy surface presented for the same systems in Ref. Wang et al. 2019. Second, CGSchNet is more robust to network hyperparameters than its predecessor. In fact, for the CGSchNet hyperparameters varied during model training (see supplementary material Sec. B), the same selections are used for both systems presented in Sec. IV. Third, CGSchNet employs less regularization; particularly, it does not require the extra step of enforcing a Lipschitz constraint Gouk et al. 2018 on its network’s weight matrices as was found to be necessary for CGnet Wang et al. 2019.

While our current protocol has demonstrated success for a capped monopeptide and a 10-amino acid miniprotein, adapting the CGSchNet pipeline to produce accurate coarse-grained force fields for larger protein systems remains an open challenge. Addressing this challenge may require specific sampling strategies when obtaining training data, the incorporation of new priors that inform tertiary structure formation, or modifications to the CGSchNet architecture itself such as regularization. Successfully modeling the thermodynamics of protein folding or conformational change via a transferable, machine-learned force field would signify a major success for the union of artificial intelligence and the computational molecular sciences.

The method introduced herein enables us to reproduce the thermodynamics of small protein systems using an architecture that is transferable across system size and sequence. However, CGSchNet is not readily transferable across thermodynamic states. Related work leveraging the same variational principle in a semi-supervised learning context allows the learning of coarse-grained representations over multiple thermodynamic states, enabling transferability across different temperatures Wang and Gómez-Bombarelli 2019; Ruza et al. 2020. This method has been demonstrated for ionic liquids, for which nonequilibrium transport properties are of prime interest. Finally, the reproduction of kinetics is an open research problem, and methods for the so-called “spectral matching” problem have recently been introduced Nüske, Boninsegna, and Clementi 2019. Ideally, both force and spectral matching could be pursued in conjunction to match both thermodynamics and kinetics simultaneously.

Structural, bottom up coarse graining consists of two aspects: the model resolution and the force field. Here, we assume the resolution is set and focus on the force field, but the choice of an optimal model resolution is itself a significant challenge that is interconnected to the goal of force field optimization. How to choose a resolution for coarse graining—and the interplay of this choice with transferable force field architectures—remains an open question. Recent work has employed machine learning and data-driven approaches to pursue an optimal resolution using various objectives Boninsegna, Banisch, and Clementi 2018; Wang and Gómez-Bombarelli 2019.

Altogether, the methodology we introduce in the present contribution establishes a transferable architecture for the machine learning of coarse-grained force fields, and we expect our accompanying software to facilitate progress not only in that realm but also towards the outstanding challenges of learning coarse-grained dynamics and optimizing a model’s resolution.

Supplementary Material

See the supplementary material for further specifics on simulation, model training, and MSM construction.

Acknowledgements

B.E.H. is immeasurably grateful to Moritz Hoffmann for his wisdom and support. The authors are grateful to Dr. Ankit Patel for discussions on model debugging and machine learning techniques, to Iryna Zaporozhets for discussions concerning model priors, to Dr. Stefan Doerr for software help, to Dr. Jan Hermann for discussions regarding SchNet, and to the reviewers for their constructive comments.

We acknowledge funding from the European Commission (ERC CoG 772230 “ScaleCell”) to B.E.H., A.K, Y.C., and F.N; from MATH+ the Berlin Mathematics Center to B.E.H. (EF1-2), S.O. (AA1-6), and F.N (EF1-2 and AA1-6); from the Natural Science Foundation (CHE-1738990, CHE-1900374, and PHY-1427654) to N.E.C., J.W., and C.C.; from the Welch Foundation (C-1570) to N.E.C, J.W., and C.C.; from the NLM Training Program in Biomedical Informatics and Data Science (5T15LM007093-27) to N.E.C.; from the Einstein Foundation Berlin to C.C.; and from the Deutsche Forschungsgemeinschaft (SFB1114 projects A04 and C03) to F.N. D.L. was supported by a FPI fellowship from the Spanish Ministry of Science and Innovation (MICINN, PRE2018-085169). G.D.F. acknowledges support from MINECO (Unidad de Excelencia Mariá de Maeztu AEI [CEX2018-000782-M] and BIO2017-82628-P) and FEDER. This project has received funding from the European Union’s Horizon 2020 research and innovation programme under Grant Agreement 823712 (CompBioMed2 Project). We thank the GPUGRID donors for their compute time.

Simulations were performed on the computer clusters of the Center for Research Computing at Rice University, supported in part by the Big-Data Private-Cloud Research Cyberinfrastructure MRI-award (NSF grant CNS-1338099), and on the clusters of the Department of Mathematics and Computer Science at Freie Universität, Berlin.

Part of this research was performed while B.E.H., N.E.C., D.L., J.W., G.dF., C.C., and F.N. were visiting the Institute for Pure and Applied Mathematics (IPAM) at the University of California, Los Angeles for the Long Program “Machine Learning for Physics and the Physics of Learning.” IPAM is supported by the National Science Foundation (Grant No. DMS-1440415).

Data Availability Statement

The data that support the findings of this study are available from the corresponding authors upon reasonable request.

Appendix A Software

The cgnet software package is available at https://github.com/coarse-graining/cgnet under the BSD-3-Clause license. cgnet requires NumPy Harris et al. 2020, SciPy Jones et al. 01, and PyTorch Paszke et al. 2019, and optional functionalities further depend on pandas McKinney et al. 2010, MDTraj McGibbon et al. 2015, and Scikit-learn Pedregosa et al. 2011. The examples are provided in Jupyter notebooks Kluyver et al. 2016 which also require Matplotlib Hunter 2007. The SchNet part of the code is inspired by SchNetPack Schutt et al. 2018 and the Langevin dynamics simulation code is adapted from OpenMM Eastman et al. 2017. In addition to cgnet and the packages already mentioned, visualization was aided by Seaborn Waskom et al. 2017 and VMD Humphrey, Dalke, and Schulten 1996. Analysis was facilitated by PyEMMA Scherer et al. 2015; Wehmeyer et al. 2018.

References