RLCG: When Reinforcement Learning Meets Coarse Graining
Abstract
Coarse graining (CG) algorithms are widely used to speed up molecular dynamics(MD) simulations.
Recent data-driven CG algorithms have demonstrated competitive performances to empirical CG methods.
However, these data-driven algorithmsoften rely heavily on labeled information (e.g., force), which is sometimes unavailable, and may not scale to large and complex molecular systems.
In this paper, wepropose Reinforcement Learning for Coarse Graining (RLCG), a reinforcementlearning-based framework for learning CG mappings.
Particularly, RLCG makes
CG assignments based on local information of each atom and is trained using anovel reward function.
This "atom-centric" approach may substantially improvecomputational scalability.
We showcase the power of RLCG by demonstrating itscompetitive performance against the state-of-the-arts on small (Alanine Dipeptideand Paracetamol) and medium-sized (Chignolin) molecules.
More broadly, RLCGhas great potential in accelerating the scientific discovery cycle, especially onlarge-scale problems.
Introduction
Molecular dynamics (MD) simulations [1, 2] are commonly used to probe and predict the atomicinteractions of molecules across a wide range of scientific fields, including drug discoveries [3, 4, 5],protein
folding predictions [6, 7] and battery material simulations [8, 9].
One major challenge in MDapplications is the high computational cost, often induced by the spatial and temporal scale of theproblem [10, 11].
Among the multiple approaches proposed to trade high resolution of the simulationfor a better speed, e.g., adopting a larger time step [12], the application of coarse graining (CG)algorithms has been the most popular one.
Coarse graining algorithms group atoms into different"beads" [13, 14], based on a mapping that can either be empirically defined or learned in a data-drivenfashion [15, 14].
By representing a group of atoms as a bead, coarse graining algorithms reduce thespatial resolution of the system and are thus easier to simulate.
Moreover, coarse-grained modelstend to have smoother energy landscapes because of the removal of small fluctuations.
Consequently,their simulations can be conducted with much larger time steps, leading to even faster simulationperformances [16, 15].
A critical component of a successful CG algorithm is the CG mapping that assigns atoms to beads.
To this end, several approaches have been proposed to define a CG mapping.
Empirical methods thatpredefine a grouping rule (e.g., a "four-to-one" mapping), are widely used, see, e.g., MARTINI [17,18]).
Other algorithms adopt more automated procedures to the mapping process, such as extensionsto MARTINI , as well as algorithms based on the graph structure of molecules [19, 20, 21].
More
[Figure]
orthogonalprojectionreward
[Figure]
recently, the rise of artificial intelligence (AI) has brought about data-driven algorithms that learn a CGmapping by training neural networks with MD trajectory data.
By learning the CG mapping purelyfrom data, these algorithms eliminate the presence of human-imposed biases, and have seen successon small and medium-sized molecules.
For example, Coarse-grained Autoencoder (CGAE) [16, 22]uses an autoencoder to learn the mapping by considering the CG space as its latent space.
Anothertype of data-driven algorithm adopts a supervised learning approach that uses manually identifiedmappings as labels for training the mapping algorithms [23, 24].
There are two major caveats with the aforementioned algorithms.
First, an appropriate loss functionneeds to be carefully designed in order to achieve a reasonable mapping.
Such a loss function (e.g.,in CGAE) may also require proper tuning and additional data beyond the coordinates of the atomsto compute, e.g., force labels, which can be hard to acquire.
At a higher level, there lacks of anintuitive metric to assess the quality of CG mappings.
As a consequence, measuring the quality of a CG mapping requires the presence of a downstream task.
Secondly, graph-based methods suffer fromlimited scalability because of the computational cost of performing matrix manipulations.
The sameissue also arises in data-driven methods because these algorithms (e.g. CGAE) learn a mapping onthe whole molecule, which typically has a computational complexity of at least O ( Nn ) with N, nbeing the number of atoms and CG beads, respectively.
In particular, these algorithms can take asubstantial amount of time to converge for large molecules with thousands of atoms.
To address these issues, we propose, to the best of our knowledge, the first reinforcement learningframework to learn CG mappings, i.e., Reinforcement Learning for Coarse graining (RLCG).
Ourframework uses an atom-centric approach by making bead assignments per atom, requiring learningonly a fixed number of parameters regardless of the size of the molecule.
More specifically, the beadassignment is facilitated by a novel policy network architecture based on a graph attention (GAT)neural network [25], together with a novel but simple design of the reward function.
These designsenable us to scale the algorithm to large molecules and provide us with an intuitive metric to measurethe quality of the learned CG mapping.
The rest of this paper is organized as follows: Section 2 details the RLCG framework; Section 3presents our experimental results on three molecules with different sizes, including Alanine Dipeptide,
Paracetamol, and Chignolin; Section 4 compares our method to several related works and brieflyconcludes the paper and discusses some future directions to further improve RLCG.
RLCG: Architecture and design
RLCG learns a CG mapping by sequentially assigning individual atoms to different CG beads.
Ateach iteration, it uses a graph attention policy network to featurize atom-specific information andmake CG assignments on individual atoms.
The updated CG mapping will receive an orthogonalprojection reward which quantifies the quality of the mapping.
RLCG learns the parameters of thepolicy network by routinely making assignments with different random initializations and updatingthe parameters under a reinforcement learning framework.
The overall procedure is depicted in Fig.
1,and we will introduce the components in the remainder of this section.
2.1 Preliminaries
Let T be the trajectory of the molecule and X t ∈ R n × 3 be the coordinates of the molecule attime t , where n is the number of atoms.
We want to learn a mapping M ( · ) :
R n × 3 (cid:55)→ R N × 3 thatmaps the atomic coordinates into the CG coordinates (geometric center of the corresponding atomiccoordinates), where N is the number of CG beads and N ≪ n.
Let M − 1 ( x t ; k ) ∈ R N k × 3 denotethe atomic coordinates of the k -th CG bead with N k atoms at the t -th time.
For simplicity, let
M ( i ) ∈
[ N ] represent the CG assignment of the i -th atom.
2.2 The orthogonal projection reward
We need a metric to evaluate the quality of the current CG mapping.
Intuitively, a good CG mappingshould preserve intra-bead rigidity.
In other words, the atoms mapped within one bead should havesmaller relative movements (Fig. 2.a).
Such within-bead rigidity may also benefit the training ofthe CG force field.
Following this intuition, we introduce the orthogonal projection loss.
Morespecifically, for two arbitrary frames of the atomic coordinates x and y , the orthogonal projectionloss of a CG mapping, M ( · ) , is defined as:
[Equation]
where w k is the mass of the k -th bead divided by the mass of the molecule, ∥ · ∥ F is the Frobeniusnorm, and Ω x,y,k,M is obtained via solving the orthogonal Procrustes analysis [26] by first performinga singular value decomposition:
M − 1 ( x ; k ) M − 1 ( y ; k ) ⊤ =
U Σ V ⊤ and letting Ω x,y,k,M = V U ⊤ .
Intuitively, the orthogonal projection loss is the residue of the projected atomic coordinates of each
CG bead between two frames, weighted by the mass of each bead.
A smaller orthogonal projectionloss corresponds to a better rigidity within a bead, as atoms are relatively static to each other whenthe loss is small.
The orthogonal projection reward of the CG mapping M is thus defined as the negative expectedorthogonal projection loss averaged over the entire trajectory T . :
[Equation]
In practice, we used 1000 randomly sampled pairs of frames of coordinates to approximate theexpectation above.
This reward (or loss) only depends on the coordinates of the atoms and accuratelyreflects the quality of the mapping as shown in Supplementary Fig.
S1.
2.3 The graph attention policy network
The policy network receives feature vectors of each atom in the molecule and generates two outputs:the selection output that determines which atom’s CG assignment to change, and the decision outputthat determines the selected atoms’ new CG assignment.
This design aims at explicitly decoupling from atom assignment, because not every atom needs to bere-assigned throughout the learning process.
Moreover, the assignment action is based on mergingthe CG id of one atom to the other, which lends itself to the use of a graph attention network thatexplicitly models the relative interaction between atoms.
Feature vector.
We adopt a similar featurization strategy as in literature using graph neural networksfor molecules and graphs[22, 27].
For the i -th atom, the node feature is defined as:
[Equation]
where (cid:76) denotes concatenation, one _ hot( M ( i )) is the one-hot encoding of the CG assignment, andload( i ) =
|{ k : M ( k ) =
M ( i ) , k = 1 , · · · , n }| n represents the proportion of CG i .
The node feature includes chemical information (atom type) aswell as the mapping related information.
The edge feature between a pair of atoms i, j connected with a bond is simply:
E ij = one _ hot(bond _ type( i, j )) .
Intuitively, the algorithm will make grouping assignments based on the type of the atom and how it isconnected to other atoms (e.g. a carbon atom may be more likely to be assigned to the same CG beadwith the oxygen atom through their C
=O bond than with another carbon atom through C-C bond,because the former combination has a stronger bond.).
Moreover, based on the loading of the atom,the algorithm can balance the loading of each CG bead by converting the assignment of an atom fromthe larger bead to that of a smaller bead.
Selection layer.
The selection layer consists of a graph attention network (GAT) layer:
[Equation]
where N ( i ) is the index set of the neighbors of atom i defined by bond connections, Θ is an MLPand α
i,j is the attention coefficient computed using Θ , F (1) i , F (1) j , E ij , and an attention network,a [28, 25].
The readout layer is an MLP combined with a masked softmax:
[Equation]
where bd is the boundary set such that i ∈ bd if there exists j ∈ N ( i ) such that CG ( i )
̸
= CG ( j ) .
In other words, the algorithm only seeks to update the assignment of atoms which has neighboringatoms with a different CG assignment.
This way the algorithm will avoid "incongruous" mappingwhere non-connected atoms are assigned to the same CG group.
Decision layer.
The decision layer outputs the CG assignment of each atom.
This is achieved withanother GAT layer:
[Equation]
The assignment of each atom is drawn from a multinomial distribution with class probabilitiesfollowing the attention coefficients: { α ij : j ∈ N ( i ) ∪ { i }} .
Hence, for each atom, the probability of being given the same assignment of its neighboring atomj is given by α ij .
Note that the decision is only limited to the neighboring atoms, which furtherreduces the undesirable scenarios where an atom is assigned with a CG bead different from any ofits neighbors.
Furthermore, the assignment is "relative": an atom is not directly assigned a CG id,but indirectly through the id of its neighboring atom.
This resolves the potential "label permutationinvariance" issue as the number that represents CG assignment may change with different episodes inthe RL cycle (i.e. CG bead 1 can be bead 2 in a different episode).
2.4 The Reinforcement learning framework
The framework follows the standard REINFORCE algorithm [29].
For each episode, we use a randombackbone initialization of the mapping and consecutively make CG assignments using the policynetwork until a predefined number of steps has been reached.
This is listed in Algorithm 1 andillustrated in Fig.
2.b.
Therefore, as an atom-centric approach, each individual atom makes themapping decision based on its surrounding environment.
Note that the actual reward we use to train the policy network is the difference between the reward ofthe current step and the previous one as we want to maximize "gain" in reward for each step.
[Figure]
GAT layer
GAT layerselection readout decision readout p r obab ili t y atom i
d bead id a tt en t i on w e i gh t s step T+ reward = -(ortho loss (T+1) - ortho loss (T)) update policy NN θ step T p r obab ili
t y a c t i o n a b
Input : policy network π θ ( · · · ) , initialization method, number of steps S .
M ← initialization
R = empty set for s ← 1 to S do r =
Reward( M )
F (0) i ← one _ hot(atom _ type( i ))
(cid:76) one _ hot( M ( i ))
(cid:76) load( i ) , ∀
i
F (1) i = α (1) i,i Θ F (0)
i + (cid:80) j ∈N ( i ) α (1) i,j Θ F (0) j η ← Softmax(MLP( F (1) i ) , i ∈ bd)
F (2) i = α (2)
i,i Θ F (1) i + (cid:80) j ∈N ( i ) α (2) i,j Θ F (1)
j
M ( η ) ← { α ij : j ∈ N ( i ) ∪ { i }}
F (0) η ← one _ hot(atom _ type( η ))
(cid:76) one _ hot( M ( η ))
(cid:76) load( η )
R = R ∪ { Reward( M ) − r }
Experiments
3.1 Experiment setup
We test the performance of RLCG on three molecules:
Alanine Dipeptide.
This molecule contains 10 heavy atoms ( 22 atoms in total).
The trajectory datawas obtained from the mdshare and downloaded from https
.io/
We molecule
Paracetamol.
This molecule contains 11 heavy atoms ( 20 atoms in total).
The trajectory data wasobtained from the MD17 dataset [30] and downloaded from http://www.sgdml.org/#datasets .
We randomly selected 10 , 000 frames for training and learned a CG mapping with 3 beads on thismolecule.
c
[Figure] [Figure]
Chignolin.
Chignolin is a 10-amino acid mini-protein with 99 heavy atoms ( 175 atoms in total).
The trajectory data was simulated using HTMD (High-Throughput Molecular Dynamics)
[31] anddownloaded from http://pub.htmd.org/Chignolin_trajectories.tar.gz .
We randomlyselected 10 , 000 frames for training and learned a CG mapping with 10 beads on this molecule.
The simulation details for Alanine Dipeptide and Chignolin can be found in [23].
These two datasetshave also been studied extensively as benchmark datasets in [22].
For all datasets, we only considerthe heavy atoms for simplicity.
To test the performance of RLCG, we trained the policy network for1500 episodes on Alanine Dipeptide and Paracetamol, and 2500 episodes on Chignolin.
Within eachepisode, we ran the policy network for 5 steps on Alanine Dipeptide and Paracetamol, and 50 stepson Chignolin with random initialization.
We also compared RLCG with CGAE.
The training loss for CGAE includes the reconstruction errorand the force regularization.
Note that the force label is only available in Paracetamol dataset.
Asa baseline, we compared RLCG with random partition where each CG bead has at least one atomassigned to it.
We performed a 5 -fold cross validation, where each algorithm was trained on 4 foldsand tested on the held-out fold.
See Fig.
S2 for the evolution of the testing orthogonal loss for
RLCG.
Even though the algorithms compared are not supervised approaches, 5 -fold cross validationmay account for potential variance during the learning and testing process and enable a more faircomparison.
3.2 Simulation results
As an illustrative example, we plotted the evolution of the CG mapping using RLCG on Alaninedipeptide within one episode (Fig. 3.a).
As can be seen (Fig. 3.a), RLCG automatically balances the
CG loading by stopping further updates after step 3 (step 4 is not shown because it is the same asstep 3 ).
On Paracetamal, RLCG also yields congruous mappings within 3 steps (Fig. 3.b).
Note that
RLCG will group the C atom with the one on the benzene ring instead of the N atom (step 2 to 3,red bead), indicating the algorithm has learned the correct relative bond information between atoms.
RLCG also learns congruent mapping with reasonably balanced CG loadings for Chignolin (Fig. 3.c).
RLCG achieves a competitive orthogonal loss across all the algorithms compared (Table 1).
Notethat for Alanine Dipeptide, the slight higher loss of RLCG than CGAE is due to the randomness ininitialization for RLCG, which can be further improved.
[Table]
Conclusions & Discussions4.1 Conclusions
In this paper, we present RLCG, a reinforcement learning framework that is designed to efficientlylearn the coarse-grained mapping of a molecule.
RLCG adopts an atom-centric decision mechanismand uses GAT to generate CG assignments, which is capable of fast simulations even on largemolecules.
Our algorithm has the potential to expedite the scientific discovery cycle if furtherextended to large molecules.
The RLCG framework has connections to other RL-based graph partitioning algorithms, but possesses distinct advantages.
For example, several vertex-centric methods are proposed for graphpartitioning [32, 33].
However, these algorithms do not use a standard RL framework and partition agraph "on the go", which may not generalize to different datasets.
GNN-based RL algorithms havealso been proposed [27].
However, such methods only consider graphs without chemical information,and only apply to bi-partitions.
There are also algorithms that fragment a molecule by breaking their
SMILES strings [34] or chemical bonds [35].
However, these methods are mainly designed fordecomposing a large molecule into motifs and hence may only result in larger fragments and may notbe directly applicable to coarse graining.
4.2 Discussions and future work
Multiple improvements can be made to speed up the RLCG framework.
First, the computation of theorthogonal projection reward is local: changing the assignment of one atom will only affect the CGpartition of two CG beads.
This property may enable a much faster computation of the orthogonalprojection reward.
Second, adding more layers to the policy network may enable the network tohandle more complicated structures.
For example, more GNN layers would allow the policy networkto see more hops away, and consequently handle rings of different numbers of atoms.
Third, the RLframework is flexible and can process chemical information, such as the rotation angle, as well asthe designer’s preference for the downstream task, e.g., using more CG beads for a benzene ring topreserve its planar property which will potentially improve the training of the CG force field.
Lastly,future work also includes experimenting on larger molecules and also testing the generality of thetrained policy network on new and unseen molecules.