Predicting aggregation of sequence-deﬁned macromolecules with Recurrent Neural Networks

Self-assembly of dilute sequence-deﬁned macromolecules is a complex phenomenon in which the local arrangement of chemical moieties leads to the formation of a long-range structure. The dependence of this structure on the sequence necessarily implies that a mapping between the two exists, yet it has been diﬃcult to model so far. Predicting the aggregation behavior of these macromolecules is challenging due to the lack of eﬀective order parameters, a vast design space, inherent variability, and high computational costs associated with currently available simulation tech-niques. Here, we accurately predict the morphology of aggregates self-assembled from sequence-deﬁned macromolecules using supervised machine learning. We ﬁnd that regression models with implicit representation learning perform signiﬁcantly better than those based on engineered features such as k -mer counting, and a Recurrent-Neural-Network-based regressor performs the best. Further, we demonstrate the high-throughput screening of monomer sequences using the regression model to identify candidates for self-assembly into selected morphologies. Our strategy is shown to successfully identify multiple suitable sequences in every test we performed, so we hope the insights gained here can be extended to other increasingly complex design scenarios in the future, such as the design of sequences under polydispersity and at varying environmental conditions.

Due to the astounding number of patterns now possible to realize in synthetic copolymers, effectively leveraging the possibilities of polymer self-assembly requires a quantitative understanding of the design process. Fortunately, computational approaches are available across the many lengths and time scales relevant to this problem, including Density Func-for solving various problems in soft matter over the last decade. 43,[52][53][54][55][56][57] ML-based highthroughput screening is an emerging technique 58 used for tasks such as the computational discovery of transition-metal complexes, 59 exploration of nanomedicine design space, 60 and systematically identifying novel cancer drug targets, 61 among others. The premise for this strategy is that from a small number of training data points, the behavior of a much more extensive collection of possible systems can be predicted. Thus, MD simulations' relatively high computational cost can be amortized over many thousands of predictions from the ML model.
In this work, we apply supervised machine learning to predict the aggregation behavior of a model copolymer. Notably, even the formulation of this problem requires an order parameter for the aggregate morphology, which has only recently been made available. 62,63 Using these order parameters, we evaluate several classes of regression models and identify features that lead to the most successful ones. Since self-assembly is a stochastic process, we also quantify the intrinsic uncertainty as a benchmark for our model evaluation.
After identifying a suitable predictive model, we perform high-throughput screening of monomer sequences with a fixed composition. This amounts to proposing a small number of target morphologies and then identifying the best candidate sequences from the collection of all possible sequences based on the model prediction. We find that the model always finds sequences that yield qualitatively similar structures, and in most cases, the quantitative agreement is within the intrinsic process variability. The proposed method is also attractive because the inference is very fast after the model has been trained. Our results show that even highly complex behaviors of soft matter systems can be predicted using machine learning, given that sufficient attention is given to identifying suitable representations of the inputs and outputs.

A. Molecular Dynamics simulations
We use the same model as our other recent work, 64 adapted from disordered protein liquidliquid phase separation. 65 In the model, polymer chains contain a fixed number of coarsegrained beads that belong to one of two chemical groups: attractive A beads and purely repulsive B beads. The attractive beads interact via the standard Lennard-Jones potential (σ, ε, r cut = 3σ), 66 while the repulsive beads use the Weeks-Chandler-Andersen potential. 67 Bonds between adjacent monomers are handled by the standard finitely extensible nonlinear elastic (FENE) potential (R 0 = 1.5σ, K = 30ε/σ 2 ). 68 Additional details are available in Ref. 65.
In these simulations, only the monomer sequence was varied while all other parameters were fixed. We considered a fixed composition with 60% A beads and 40% B beads, with a chain length of 20 monomers for all polymers. The volume of the box was V ≈ (40σ) 3 with N = 500 chains. Equilibration was performed with HOOMD-blue 69,70 (version 2.9.4) in the NVT ensemble using a Langevin thermostat to enforce a temperature of T = 0.5ε.
Configurations were allowed to relax for 2 × 10 5 τ from an initially random state before the structure was evaluated.

B. Unsupervised representation learning
After equilibrating the polymer chains, we apply our recently developed unsupervised learning approach to obtain order parameters for the aggregates. 62,64 In short, this approach uses an extensive collection of local geometric features to create a low-dimensional representation of each bead's neighborhood via nonlinear manifold learning. These features are then aggregated in a global pooling scheme and reduced a second time to yield two snapshot-wide order parameters, Z 0 and Z 1 . This Z is a rotation-, translation-, and permutation-invariant measure of the aggregate morphology. The algorithm is described in detail in Refs. 64 and 62.
As introduced in Ref. 64, the Z manifold reveals the possible aggregate morphologies formed from the chosen 60%-A-type composition. The structure of the manifold is shown in Fig. 1 with some representative snapshots overlaid to illustrate the shape of the aggregates in each region of Z. It can be seen that aggregates around the periphery of the manifold correspond to "archetypes" while those closer to the middle is mixtures of multiple types. This is a result of the manifold learning scheme preserving the topology of the data during projection, such that the distance between archetypes is maximized while centrally located aggregates are equidistant to the locations of their respective pure archetypes. The 2 038 points shown in Fig. 1 comprise the training data for our supervised learning task described below.

C. Supervised learning
The main objective of this work is to approximate the function f : X → Z where X corresponds to the monomer sequence of a polymer chain and Z is a quantitative descrip- In order to apply regression models to the problem at hand, we must encode the monomer sequences in a machine-readable form (i.e., as a vector X). However, the sequences themselves are not suitable for regression as they are natively written as a string of characters (A or B). In other contexts, such as materials informatics and natural language processing, it has been shown that the quality of the encoding is critical to the performance of machine learning models. [74][75][76][77][78] Here, we focus on three methods of featurization: sequence vectors, token counting, and implicit feature learning.
One unusual feature of this problem is the symmetry invariance of the polymer sequences; each sequence represents a physical object (i.e., a polymer chain) without a concept of directionality. In fact, it would be nonsensical to have a discrepancy between the model prediction when feeding the sequence "forward" versus "reversed." As a result, we always consider the prediction of regression model to be: where X is a sequence and X is the mirror of that sequence. This choice ensures that all models yield exactly one result for each polymer sequence, which again represents a physical object without directionality.

D. Sequence vector encoding
The most straightforward featurization technique we investigated was encoding the sequence directly as a feature vector. In this scheme, the monomer chemical identity is represented as a vector of class labels (illustrated in Fig. 2(a)). Here, A is represented as 0, and B is represented as 1, yielding a vector of binary values corresponding to each monomer position in the chain. When provided with sequences encoded in this vector form, the regression models do not receive any physically meaningful information about the topology of the chain. That is, the performance would be identical if all the sequences were permuted in the same way (e.g., swapping beads 1 and 5). For instance, when a Random Forest Regressor was trained on ten random permutations of the sequence vectors, the coefficient of determination (R-squared) was 0.6596 ± 0.0011 compared to 0.6604 ± 0.0028 using the correct sequence ordering (i.e., statistically indistinguishable results). This presents a serious limitation in that the model is totally insensitive to the relative position of the monomers to each other, and can only consider the absolute position of A and B type monomers.

E. Token counting (k-mers)
In order to imbue the models with a sense of the relative positions of different monomers, we consider counting the frequency of k-mer tokens in the sequence (illustrated in Fig. 2(b)).
In this context, a k-mer is a substring of length k taken from the complete monomer sequence.
Because there is no sense of direction in the physical polymer chain, we count both "forward" and "reverse" copies of the k-mer substring. For instance, the substring AB is a k-mer of length k = 2 that appears twice in the string AAAAAABBBBBBBBAAAAAA -once on each side of the B block, even though the order of the k-mer is reversed in the two occurrences. Note that in this case the features X = X , sof = f (X) = f (X ) in Eq. 1. After counting the k-mers, a frequency vector is passed as input to the regression model.
Here we have considered k-mers with length k = 2 to k = 10. The lower limit of 2 ensures the model receives sufficient information to make a prediction, as k = 1 would reduce the count of A and B type beads in the chain, which is fixed in our study. The upper limit is a practical one to limit the number of tokens that must be counted; there are 1 085 possible 10-mers which makes the token frequency slow to compute.
This kind of engineered feature has been known to increase the predictive power of the machine learning models allowing the flexibility to use less complex models that are faster. 79,80 We do not need to train on mirrored copies of the sequences because the k-mers are combined with their mirrored versions in the frequency vector; the token frequency of a reversed sequence is identical to the original sequence.

F. Representation learning
Representation learning is an integral part of natural language processing to understand the representations of raw tokens like words or characters in a collection of texts. In this encoding scheme, we rely on learned embeddings of the raw monomer tokens to derive meaning from the sequence (illustrated in Fig. 2(c)). While this can also be thought of as using the sequence vector scheme together with a NN-based regression model, we discuss this class of encodings separately due to the ability of these regressors to apply nonlinear transformations on the input data to yield more meaningful representations inside the model. language to language translation, 82 natural language processing, 83 speech recognition 83 and so on. This should result in a greater ability to abstract sequence-based patterns from the input.
In the case of the MLP, we again train on the mirrored copies to ensure the model does not fit any spurious, direction-dependent trends.
For the RNN, we use bidirectional Long Short Term Memory (LSTM). 84 LSTM is a class of RNN that overcomes the problems of vanishing and exploding gradients. Bidirectional LSTMs 85 further train two networks, where one traverses the input sequence from left to right ("forward") while another traverses from right to left ("reverse"). This yields two lowdimensional embeddings of the sequence, which are merged via concatenation. This enables the network to have the sequence information in both directions and is known to perform better than unidirectional LSTM models in other sequence modeling tasks. [86][87][88] The specific RNN model employed here had three bidirectional LSTM layers with 4dimensional hidden states followed by a fully connected linear layer. The model was again trained on the sequence vectors. We trained with a batch size of 128 for 1 000 epochs using the Adam optimizer with a learning rate of 0.01. Furthermore, a custom loss function was used to ensure the model learned a symmetric representation of the polymer sequences: where L fwd indicates the "forward" pass, L rev indicates the "reverse" pass, and L sym is a penalty for violating symmetry. Each of these terms is simply a Mean Squared Error loss between respective observations and targets: This loss function is illustrated schematically in Fig. 3.

A. Intrinsic variability of aggregation
Our prior work 64 observed that repeated self-assembly simulations of the identical sequences at the same thermodynamic conditions resulted in considerable variability. However, we did not perform an exhaustive analysis of the intrinsic variability of the self-assembly process. Here, we simulate three replicas of 200 different sequences to systematically evaluate There may also be some spatial dependence on the uncertainty, such as reduced variance for Z 0 < 5 compared to Z 0 > 5. The results are colored by angle to show the spatial dependence on the direction of variance, such as those in the liquid-like Z 0 > 10 region being biased towards higher variance in the horizontal direction compared to those in the vesicle-like Z 1 > 5 region being biased towards higher variance in the vertical direction.

B. Regression with sequence vectors
We evaluated the performance of the following regression models using the sequence vector input scheme: Linear, 89 Lasso, 85 Ridge, 90 K-Nearest-Neighbors, 91 Random Forest, 92 and MLP. 92 Lasso and Ridge regressions are linear regression schemes that use L 1 and L 2 regularization, respectively, which results in Lasso preferring to shrink some weights to zero. Evaluating these three models gives insight into how effectively each input feature produces accurate output. The other models all have some nonlinear capability. K-Neighbors The root-mean-square error (RMSE) and coefficient of determination for different models on the label-encoded features are shown in Fig. 5 On the other hand, the nonlinear models appear to be much more susceptible to overfitting, with Random Forest being the worst offender. This is not surprising since Random Forest has the least grounding in physical reality. Even so, it still outperforms K-Neighbors in testing. The poor performance of K-Neighbors suggests that Euclidean distance in the "sequence vector space" does not approximate the structure of the aggregate morphology space well. We tried other metrics, such as Hamming distance and cosine similarity, but the results did not vary significantly. Finally, we see that MLP performs the best in testing.
This is a slightly unfair comparison since MLP is able to perform representation learning, as discussed above. Nevertheless, it technically does take the sequence vector as input, so we present the results here.
It is noteworthy that the nonlinear methods strictly outperform the linear methods in both training and testing. All evidence supports the notion that the sequence vector poorly represents the physics at play in the aggregation process. For instance, the overfitting of the Random Forest shows how unreliable it is to build decision trees based on the chemical identity of individual monomers. Likewise, regularization of the linear models reduces performance without reducing overfitting (because no overfitting is observed).

C. Regression with token frequency
Next, we evaluated the performance of the same regression methods using the token frequency of k-mers. Critically, this provides a nonlinear featurization that captures interactions between monomers up to k − 1 beads away in the chain, which should improve the performance relative to using sequence vectors. The results for varying k are shown in Fig. 6, and the result from sequence vectors (i.e., Fig. 5a) is shown for comparison on the far right side. As before, lower RMSE is better, and the discrepancy between train and test performance indicates overfitting.
In general, RMSE decreases as k increases, except for K-Neighbors, which shows a local minimum around k = 4 before increasing again. This is likely because the distance metric becomes more meaningful as additional tokens are added to the vector up to a point. Still, at some point, the space becomes too high-dimensional, and the neighborhood is less meaningful. For the other models, having access to more features and features that convey more non-local information improves the predictive performance. At the same time, this also results in increased overfitting for higher k, as indicated by the growing gap between test and train curves for each model.
Again evaluating the linear model variants, we see that Lasso performs significantly worse than Linear or Ridge, which are nearly identical. It is interesting to observe this in the kmer featurization strategy in addition to the sequence vector one since the k-mer tokens are supposed to include nonlinearity. We assumed at the outset that certain patterns would indicate different aggregate morphology, such as alternating monomers corresponding to more liquid-like structures. In contrast, repeated monomers would correspond to micellelike structures. Based on Fig. 6, this is not the case, and instead, the regressors benefit from utilizing all available features, even at k = 10, which corresponds to 1 085 features.
MLP and Random Forest perform about the same across the k range, with MLP performing slightly better overall. Random Forest can be seen significantly overfitting even at low k, indicating again that k-mer token counts are not very informative. This is also reinforced by the improved performance of the nonlinear models on the label encoded vectors compared to any k.

D. Representation learning
Our most sophisticated model was a bidirectional LSTM regressor, which can only be trained on the label encoded sequence (and not the k-mer tokens since they are not a sequence). Compared to the other models, the RNN has the greatest ability to learn a meaningful internal representation of the sequence data due to its notion of the input sequence. As expected, the RNN performed the best out of all the tested models, with RMSE of 1.49 ± 0.12 compared to the next-best model, which was the MLP with RMSE of 1.79 ± 0.11. As previously noted, the MLP should also be considered to use (limited) representation learning, demonstrating the unsuitability of both the sequence vector and token frequency features for the regression task. Indeed, the next-best model that has no representation learning capability is Ridge regression for k = 10 with RMSE 2.48 ± 0.15, which is only slightly better than Random Forest and Linear regression.

E. Comparison of regression models
To summarize, we considered three types of features as input to the regression models in this study: label-encoded features, k-mer token counts, and implicit representation learning by NN. Six different regressors comprising three linear and three nonlinear models were trained on the label-encoded features and on the token counts with varying lengths of pattern from k = 2 to k = 10. Of the encoding schemes, representation learning by NN consistently performed the best. Supporting this, the RNN model (which had the greatest representation learning potential) performed the best of all the models, with MLP achieving a close second.
We summarize the results in Fig. 7 wherein the results are shown for the best version of each model overall representations considered. The plot shows that token counts perform better than sequence vectors for linear models, while for nonlinear models, the performance is of 1.49 ± 0.12 in testing. As the best-performing model, it is used in the high-throughput screening below. Note that this compares favorably to the RMSD of 0.67 reported in Fig. 4, which would be the minimum error we could theoretically expect.
In order to provide more context to these RMSE values, we also plot the 50th percentile isocurve in ∆Z space on the same scale as the learned manifold in Fig. 8. Note that these residuals all represent test data from different folds. An ideal model should produce a residual plot with a small area inside the isosurface and should be both centered on and symmetrical about the origin. The size of the isocurve can be compared to the ellipses in

F. High-throughput screening
The obvious practical application of a well-trained regression model is to identify sequences that assemble into a target aggregate morphology from among the hundreds of thousands of possible permutations. Accordingly, we deploy the trained RNN model to predict the aggregate morphology (i.e., Z embedding) for all 63 090 unique sequences with 60% A monomers. From these predictions, we perform high-throughput screening to select the most promising candidates for exhibiting the desired aggregate structure.
We present the results from two demonstrations of this high-throughput screening approach. First, we seek sequences with morphology similar to the six sequences reported in Table 2 of Ref. 64, corresponding to peripheral "unique" archetypes in the latent space.
Second, we select eight morphologies from the 2 038 labeled data points using K-Means clustering, such that they are well dispersed across the latent space. In each case, we choose the five sequences with the smallest Euclidean distance in Z between the prediction of the RNN model Z p and the target location Z t . For each sequence, five replicas were simulated to quantify intrinsic variability in the self-assembly process.
The results of the high-throughput screening are shown in Fig. 9. At least one of the five candidate sequences includes the target structure within one standard deviation of its mean for all but a few of the targets. As indicated by the ellipses, the covariance among replicas is significantly smaller in the archetypes (top row) compared to the K-Means clusters (bottom row). This is consistent with Fig. 4, which shows reduced variance on the periphery of the latent space compared to the center; since UMAP assumes a uniform density, morphologies with the higher disorder (and thus more possible states) are distorted to occupy a larger area than more ordered ones (with fewer possible states).
Note that variance is shown for both the simulation results (circles) and RNN predictions Labels correspond to those in Fig. 9.
(triangles) in Fig. 9; for the regression model, the variance comes from cross-fold validation.
Large uncertainty around the targets -especially (h) -reveals low confidence in the ensemble to predict successful candidates. The ensemble prediction uncertainty generally seems to be commensurate with the observed error from simulations.
To further illustrate the successful outcome of this procedure, we show the best and worst snapshots for each target from the K-Means clustering in Fig. 10. Visual inspection reveals that tightly clustered points in latent space such as (f) appear very similar, while those with a higher variance like (h) exhibit different local environments (as indicated by colors) and global morphology. To be clear, we emphasize that each target was matched with five sequences with five replicas each, so the configuration shown to the right of each panel (a)-(h) is the worst of 25 possible snapshots targeting that morphology.

IV. CONCLUSIONS
In this work, we demonstrate the use of supervised machine learning to identify monomer sequences for a model copolymer that self-assembles into aggregates with prescribed morphology. The objective of the supervised regression task is to learn a mapping between the monomer sequence and the morphology of the self-assembled aggregates under fixed environmental conditions. Since this target morphology is used to fit a regression model, we require a quantitative order parameter, which was only recently developed in Ref. 63. We emphasize that this provides a continuous description of morphology throughout the entire structural space, rather than consolidating aggregates into discrete classes.
To obtain an accurate predictive model, we evaluated three different representation schemes for the monomer sequences: raw sequence vectors, k-mer token counts, and feature learning by NNs. The effectiveness of each scheme was determined by considering the accuracy (via R-squared and RMSE metrics) of regression models using that representation and trained to predict the morphology of 2 038 sequences simulated in our prior work. 63 The results clearly show that feature learning is the most effective strategy, whereas raw sequence vectors perform poorly for linear models (low accuracy) and nonlinear kernel-based models (overfitting). Surprisingly, k-mer token counting was only more effective than raw sequence vectors for linear models, presumably because these models have no other way of accounting for the nonlinearity of the problem. The raw sequence vectors were more effective in all other cases, and it was most effective when feature learning was available via NN. We also found that RNNs with LSTM layers performed better than NNs with only fully connected layers (i.e., MLPs).
The best regression model we trained was a bidirectional-LSTM-based RNN, which gave an RMSE of 1.49 ± 0.12 when deployed on unseen test data. This compares favorably with the intrinsic uncertainty observed in the self-assembly process, which had RMSD = 0.67.
Using this model, we performed high throughput screening to identify suitable sequences from the over 63 000 possible sequences at the chosen composition. For each of 14 target aggregate morphologies, represented as a point in the continuous morphology manifold, we selected the five sequences predicted to be closest to that point by the RNN model. MD simulations of these sequences showed that the model successfully identified sequences that matched the qualitative morphology of the target in the vast majority of cases, and at least one of the five sequences matched quantitatively in every case. Furthermore, the accuracy was much better than the RMSE when targeting archetypes than the disordered structures, indicating that the high variability in the center of the manifold may dominate estimates of model accuracy.
A major limitation of the present study is the restriction to fixed composition, density, temperature, and so on, with only the monomer sequence varying between simulations.
This was primarily driven by the availability of existing data from prior work, but we also anticipate additional challenges when training the model on a much larger domain. However, the practical usefulness of the proposed approach will depend on its ability to scale up to increasingly complex design scenarios. We also limited our study to monodisperse chains, which does not reflect the reality of most experimental copolymer systems. Thus, it will be crucial to include polydispersity in the future to capture the physical reality of these materials. Our work here clearly demonstrates that these open challenges are worth pursuing with an RNN-based regression model.

CONFLICTS OF INTEREST
There are no conflicts to declare.