Predicting protein stability changes under multiple amino acid
  substitutions using equivariant graph neural networks
The accurate prediction of changes in protein stability under multiple amino
acid substitutions is essential for realising true in-silico protein re-design.
To this purpose, we propose improvements to state-of-the-art Deep learning (DL)
protein stability prediction models, enabling first-of-a-kind predictions for
variable numbers of amino acid substitutions, on structural representations, by
decoupling the atomic and residue scales of protein representations. This was
achieved using E(3)-equivariant graph neural networks (EGNNs) for both atomic
environment (AE) embedding and residue-level scoring tasks. Our AE embedder was
used to featurise a residue-level graph, then trained to score mutant stability
($\Delta\Delta G$). To achieve effective training of this predictive EGNN we
have leveraged the unprecedented scale of a new high-throughput protein
stability experimental data-set, Mega-scale. Finally, we demonstrate the
immediately promising results of this procedure, discuss the current
shortcomings, and highlight potential future strategies.
Sebastien Boyer, Sam Money-Kyrle, Oliver Bent
Predicting protein stability changes under multiple amino acid substitutions using Equivariant Graph Neural Networks

###### Abstract

The accurate prediction of changes in protein stability under multiple amino acid substitutions is essential for realising true in-silico protein re-design. To this purpose, we propose improvements to state-of-the-art Deep learning (DL) protein stability prediction models, enabling first-of-a-kind predictions for variable numbers of amino acid substitutions, on structural representations, by decoupling the atomic and residue scales of protein representations. This was achieved using E(3)-equivariant graph neural networks (EGNNs) for both atomic environment (AE) embedding and residue-level scoring tasks. Our AE embedder was used to featurise a residue-level graph, then trained to score mutant stability (\(\Delta\Delta G\)). To achieve effective training of this predictive EGNN we have leveraged the unprecedented scale of a new high-throughput protein stability experimental dataset, Mega-scale. Finally, we demonstrate the immediately promising results of this procedure, discuss the current shortcomings, and highlight potential future strategies.

## 1 Introduction

Protein stability is a crucial component of protein evolution (Godoy-Ruiz et al., 2004), it lies at the root of our understanding of many human diseases (Peng & Alexov, 2016) and plays a major role in protein design and engineering (Qing et al., 2022). Protein stability is typically represented as the change in free energy, \(\Delta G\), between the unfolded and folded states (Matthews, 1993) and is a global feature of a protein. A negative \(\Delta G\) of folding indicates an energetically favourable protein conformation; the greater the magnitude of a negative \(\Delta G\), the more stable the conformation. Mutations can alter the favourability of a protein fold, with even single amino acid substitution events potentially disturbing the native conformation of a protein (Stefl et al., 2013). For example, a substitution from threonine to methionine in 12/15-Lipoxygenase is a cited potential cause of hereditary cardiovascular diseases (Schurmann et al., 2011); the mutation disrupts a chain of stabilising hydrogen bridges, causing structural instability and reducing catalytic activity. The mutational effect on protein stability is the difference in free energy of folding between the wild type (WT) and mutant proteins, \(\Delta\Delta G\)(Matthews, 1993). Mutagenic effects on protein stability can be determined experimentally using thermostability assays, with \(\Delta\Delta G\) being inferred from differences between WT and mutant denaturation curves (Bommarius et al., 2006). However, these assays are labourious and expensive; to adequately assess mutational effects at a higher throughput rate, researchers have turned to computational methods. The established precedent for computational modelling of mutant stability is empirical physics-informed energy functions, which rely on physical calculations to infer \(\Delta\Delta G\)(Marabotti et al., 2021). For example, Rosetta (Kellogg et al., 2011; Das & Baker, 2008) employs Monte Carlo runs to sample multiple protein conformations and predicts folding free energy from physical characteristics. These characteristics of Lennard-Jones interactions, inferred solvation potentials, hydrogen bonding and electrostatics are common to other packages such as FoldX (Schymkowitz et al., 2005). While Molecular Dynamics software, such as Amber (Case et al., 2005), utilises these characteristics in force fields to explore protein conformational landscapes and calculate potential energies by resolving classical physics calculations.
These physics-based models can provide scoring for both protein stability or mutation-induced change of protein stability, however, they are still not fully scalable to large data-sets given the computational expense necessary for each simulated prediction. For example, conformation sampling via Monte Carlo simulations in Rosetta requires extensive compute time. On the other hand, machine learning-based predictors and, more recently, Deep learning (DL) approaches have shown improved scalability and, in some cases, comparable accuracy with physics-based models (Iqbal et al., 2021). This work will continue to explore the advantages of an entirely data-driven DL approach for predicting protein stability changes under multiple amino acid substitutions.

## 2 Related Work

In moving away from established molecular modelling approaches, machine learning methods EASE-MM (Folkman et al., 2016) and SAAFEC-SEQ (Li et al., 2021) leverage 1D sequences and protein evolutionary information to predict \(\Delta\Delta G\) with decision trees and Support Vector Machines, respectively. While ACDC-NN-Seq (Pancotti et al., 2021) explored utilising DL by applying Convolutional neural networks (CNNs) to protein sequences. As sequence data is more widely available than experimental structures, it is probable that the insight of these models into 3D structural characteristics, such as free energy of folding, is limited by their 1D representation. PON-tstab (Yang et al., 2018) implemented a combination of sequence and structure-based features in tabular format with random forests. DeepDDG (Cao et al., 2019) relies on tabular empirical features obtained from structure, such as solvent-accessible surface area, to predict stability with neural networks. However, tabular features engineered from structure are a restrictive depiction of protein geometry; graph-based approaches provide a promising alternative representation, with encouraging results when applied to protein structure prediction (Delaunay et al., 2022).

In particular, three DL models; ThermoNet (Li et al., 2020); RASP Blaabjerg et al. (2022); and ProS-GNN (Wang et al., 2021), have engaged in combining the two physico-scales involved in understanding protein geometry: the atomic scale and the residue scale of interactions. Both ThermoNet and RASP learn a representation of the atomic environment (AE) around the pertinent (mutated) residue using 3D CNNs before passing this representation through a Multi-layer perceptron (MLP) to score the mutational effect on protein stability. While obvious similarities exist between those two models, they are very different at their core. ThermoNet determines the AE representations on the fly, utilising both WT and simulated mutant structures as inputs for the MLP in the same loop. RASP initially trains a self-supervised AE embedder on a masked amino acid task, then uses this embedding as input features for a coupled WT and mutant amino acid encoding to feed a MLP trained on stability scoring. Moreover, ThermoNet is trained on a rather small experimental data-set (n \(\sim 3,500\)), while RASP is trained on a large data-set (n \(\sim 10^{6}\) for the AE embedder and n \(\sim 2.5\times 10^{5}\) for scoring) of Rosetta simulated scores, making it an emulator of the physics-based score. The third DL approach, ProS-GNN (Wang et al., 2021), replaced the CNN ThermoNet atomic environment embedding layer with GNNs. ProS-GNN also shares with ThermoNet and other DL models like ACDC-NN the constraint of being antisymmetric to reversed mutation. The aforementioned state-of-the-art stability prediction models in the literature share the following caveats:

1. Their underlying architecture allows only single amino acid substitutions.
2. Big experimental data-sets with the necessary structural data for these models are lacking.

Indeed, RASP is constrained to predicting on a fixed number of amino acid substitutions by the MLP scorer, which requires a fixed input shape; additional mutations increase the dimensions of the AE embedding to an incompatible size. In ThermoNet and ProS-GNN, the impossibility of decoupling the atomic and residue scales prohibits multiple amino acid substitutions; the required size of voxel or graph for multiple, even proximal, substitutions would be rapidly unmanageable.

A solution for both caveats exists. The self-supervised AE embedder of RASP already decouples the atomic and residue scales, and GNNs allow for some flexibility in graph topology, enabling consideration of multiple residues rather than only the embedding of the residue of interest. Integrating the RASP AE embedder with a graph-based approach would enable scoring of multiple substitution events. On the experimental data front, a new data-set, Mega-scale (Tsuboyama et al., 2022), based on high-throughput protein stability measurements, was published in late 2022. With over 600,000 data points of single and double mutants spanning over 300 WT structures, it provides a consistent
(in terms of experimental set-up) and large data-set, with the express purpose of training models to score the effects of single or double mutations on protein stability. In light of these observations, we contribute a JAX-implemented solution for resolving these constraints using two E(3) equivariant graph neural networks (EGNNs) (Garcia Satorras et al., 2021). The first EGNN is trained in a self-supervised way. The second is trained on the Mega-scale data set for scoring mutational effects on protein stability.

## 3 Method

### Atomic Environment (AE) Embedder

We followed the RASP protocol to design and train our AE embedder in a self-supervised masked amino acid manner, with two key differences:

1. We used an EGNN (Figure 2) with its own set of graph features describing the AE (Figure 1) instead of a CNN.
2. We used a macro averaged F1 score as our metric on the validation set to select model parameters from the highest-performing epoch.

The training and the validation sets are from the same data-set described in RASP (Blaabjerg et al., 2022). Our EGNN was built with layers described in Garcia Satorras et al. (2021), with an average message aggregation strategy (Equation 1). Recalling from Garcia Satorras et al. (2021) that \(\mathbf{h}^{\mathbf{l}}\) are node embeddings at layer \(l\) and \(\mathbf{x}^{\mathbf{l}}\) are coordinate embeddings at layer \(l\) (atoms coordinates), we defined the equivariant graph convolutional layer (EGCL), as they do, up to the use of this \(\frac{1}{N_{i}^{neighbors}}\) coefficient which allows the re-scaling of different messages according to the node of interest's number of neighbors (hence the average). As with their implementation, \(\phi_{e},\phi_{x},\phi_{h}\) are MLPs, \(a_{i;j}\) defines edge features between node \(i\) and \(j\), and finally, \(\mathfrak{N}(i)\) is the set of neighbors of node \(i\).

\[\mathbf{m}_{i,j} =\phi_{e}(\mathbf{h}^{\mathbf{l}}_{i},\mathbf{h}^{\mathbf{l}}_{j},\left\|\mathbf{x}^ {\mathbf{l}}_{i}-\mathbf{x}^{\mathbf{l}}_{j}\right\|^{2},a_{i,j}) \tag{1}\] \[\mathbf{x}^{\mathbf{l+1}}_{i} =\mathbf{x}^{\mathbf{l}}_{i}+\frac{1}{N_{i}^{neighbors}}\times\sum_{j\neq i,j\in\mathfrak{N}(i)}(\mathbf{x}^{\mathbf{l}}_{i}-\mathbf{x}^{\mathbf{l}}_{j})\phi_{x}(\mathbf{m} _{i,j})\] \[\mathbf{m}_{i} =\frac{1}{N_{i}^{neighbors}}\times\sum_{j\neq i,j\in\mathfrak{N }(i)}\mathbf{m}_{i,j}\] \[\mathbf{h}^{\mathbf{l+1}}_{i} =\phi_{h}(\mathbf{h}^{\mathbf{l}}_{i},\mathbf{m}_{i})\]

Node embeddings are passed sequentially through each \(N\) layer of the network. After each layer, node embeddings are copied, aggregated with an average graph level readout (global mean pooling) Equation 2, and saved. Finally, all the graph representations derived from the different layers are concatenated, Equation 2, to form the graph-level embedding \(\mathbf{h}_{\mathbf{G}}\) for the AE sub-graph of a residue \(G\), and processed through an MLP toward the desired prediction shape.

\[\mathbf{h}_{\mathbf{G}}=\texttt{Concat}(\texttt{Average}(\{\mathbf{h}^{\mathbf{l}}_{i}|i\in G \})|l=0,...,N) \tag{2}\]

For building the AE graph, we followed part of the RASP protocol:

* We considered only atoms within a 9A radius of the C\({}_{\alpha}\) of interest.
* We removed the atoms that were part of the residue of interest.

Nodes are atoms featureised with a single number (atomic number). Edges are drawn between nodes if two nodes are within a 4A distance. Edges are featureised by a binary label distinguishing whether the edge is intra- or inter-residue, as well as 2 numbers encoding a notion of the typical distance between the two atoms linked by these edges:

* The sum over the two atoms involved in the edge, of their covalent radius.
* The sum over the two atoms involved in the edge, of their Van der Waals radius.
In this particular instance of the model, distances between atoms are not directly encoded as an edge feature, but given the use of an EGCL, this distance is present as a distance vector (Equation 1) (rather than the usual scalar distance, hence the necessity of E(3) equivariance) by design. Finally, we trained the model on a classification task consisting of retrieving the amino acid around which the AE has been built. Model parameters were selected from the epoch with the best macro F1 scores on the validation set. A detailed description of the model in terms of its hyper-parameters is provided in the Appendix (Table 6).

### Mutant Stability Scoring

We used the same model architecture as presented for the AE embedding (Figure 2), for the regression task of predicting \(\Delta\Delta G\). The set of hyper-parameters differs as described in the Appendix (Table 7). For this task, the graph is built at the residue-level with additional atomic-level features to

Figure 1: Definition of the atomic environment (AE) graph.

Figure 2: Backbone of the E(3) equivariant graph neural network (EGNN) used for both AE embedding and scoring tasks. The EGNN layer is the EGCL taken from (Garcia Satorras et al., 2021).
bridge the gap between the two fundamental scales. Indeed, in this representation nodes are residues represented in terms of their spatial positioning by the residue mean atomic position coordinates. Nodes are featureised with the vector output of the previously trained AE embedder. More node features are included with an 11-dimensional representation of the physico-chemical properties of the WT amino acid (Kawashima et al., 2007; Xu et al., 2020), concatenated to the same representation, except for the mutated amino acid nodes. When a particular node is not mutated this concatenation is just twice the WT physico-chemical 11-dimensional representation.

At the edge level, an edge is drawn between two nodes if the mean atomic position distance between the two nodes is within 9A. The graph is centered around the mutant residue and residues are added given the distance threshold up to n (here 1) edges away from the mutant nodes. In the case of multiple mutants, we allow the different graphs centered around their mutant node, to be disconnected from each other. Features for the edges follows a similar strategy to the atomic graph:

* A single number to stipulate if the two residues are linked by a backbone, bound or not.
* Two numbers to provide a specific scale for the distance between two WT residues: the sum of the residue side chain sizes, defined as (i) the maximum distance between the C\({}_{\alpha}\) and any atoms of the residue; (ii) the maximum distance between two atoms from the same residue.
* The same two numbers are produced for mutants involved in the edges. When there are no mutants involved then they are just duplicated from the WT numbers.
* The C\({}_{\alpha}\)/C\({}_{\alpha}\) distance and the mean atomic position/mean atomic position distance.

Finally, to help the training while homogenizing the ranges and variance of the target variable (here the experimental \(\Delta\Delta G\)) we used a Fermi transform, also as described in RASP (Blaabjerg et al., 2022). Our loss function is a simple Root Mean Squared Error, and the best model as well as the best epoch is chosen based on the spearman rank correlation coefficient on the validation set.

Figure 3: Definition of the residue sub-graph build around mutated residues.
## 4 Results

### Atomic Environment Embedder

With our AE implementation, we reached a macro averaged F1 score of 0.63 on the training set and of 0.61 (accuracy = 0.65) on the validation set, which is comparable to the RASP 0.63 accuracy also on a similar but different validation set (we shuffled the structures), full results in Appendix Table 4. The confusion matrix on the validation set (Figure 9) is also provided in the Appendix and shows a variable but strong ability of the model to match ground truth.

### Mutant Stability Scorer

Evaluation metrics for the different splits are available in Figure 10, as well as a description of the Mega-scale data-set in the Appendix: A. Given the unique qualities of the Mega-scale data-set, we decided to evaluate the model in what we believe is a more stringent way than simply looking at the Mega-Scale test split (metrics are provided for the split too). Indeed, the Mega-scale data-set only contains domains and not full proteins, and structures were resolved computationally using AlphaFold (Jumper et al., 2021). The Mega-scale data-set also only contains up to double mutations. Hence, we decided to evaluate our model on a more standard data-set with experimentally resolved entire protein structures: ThermoMutDB (Xavier et al., 2021) (a description of the ThermoMutDB data-set is also provided in the Appendix: A.3).

Over the pooled ThermoMutDB data-set our scorer achieved RMSE = 2.288; Spearmanr = 0.311; Pearsonr = 0.251 Table 1. Interestingly the model seems to generalise well to structures with more than two mutations (Figure 4), for which it has not been directly trained. Spearman correlation, for example, spans a range between 0.159 and 0.381, for a number of mutations going from 1 to 3.

At the level of individual structures (Figure 5), model performances can also vary quite drastically from WTs with at least 100 mutants(for an overview per pdbse see Figure 11).

Finally, we also compared, for a subset of one-point mutations, our work to RASP (Figure 12). On pooled single mutations our proposed approach performs significantly worse (Pearsonr RASP = 0.53; Pearsonr for this work = 0.42 ). Yet our approach outperforms RASP for some pdbs, and suffers from the same drawbacks as RASP on some proteins; structures for which RASP poorly predicts mutational effect are also, with a few exceptions, poorly predicted by our method. Yet, overall our performance is still significant, even more so when put in perspective with the fact that RASP regression is an ensemble model.

Figure 4: Evaluation of our scorer for differing numbers of mutations. Purple markers for Spearman rank correlation p-value\(<\)0.05 else orange. Marker size is proportional to the number of mutations.
## 5 Discussion

These preliminary results show that the combination of decoupling of the atomic and residue scales, with the usage of an EGNN architecture, to allow flexibility on the number of mutations accessible to score, is promising. In realising this exploratory work we faced two main challenges:

1. The scorer had a tendency to over-fit the Mega-scale data-set.
2. The current choice of a threshold for the residue graph is constrained. We ended up choosing 9A, where depending on the residues, a typical length for such interactions could go to 16A or more (for two tryptophhans, given the max distance between their own atoms). But such a threshold would lead to a hyper-connected graph that would hinder the training. Generally speaking the graph building hyper-parameters, for example, the number of hops around the recovered nodes of interest (here one: neighbours one hop from the mutant nodes), would influence hyper-connectivity and our ability to not over-fit.

We believe both of those problems, particularly the latter point, could be partially alleviated by finding a better encoding of meaningful distances of interaction, as well as including a more appropriate way to sum messages within the message passing loop (Ying et al., 2021).

In terms of evaluating scoring performance, when exploring the ThermoMutDB data-set as a potential out-of-distribution test set, we realised that Mega-scale has a significant advantage compared to all of the other available data-sets; it is experimentally consistent, both for \(\Delta\Delta G\) measurement and the use of AlphaFold for structure prediction. This is not the case for ThermoMutDB which is an aggregate of results obtained with a variety of methods for both stability measurement and structure determination, making it a challenge to understand why and how the model is failing to give accurate predictions. Training on such a data-set which will, for example, not include certain types of interactions, such as inter-domain interactions, as well as not containing the inherent real noise of protein structure prediction, is advantageous for its consistency and an inconvenience for its representational inaccuracies when compared to more "realistic" data-sets. In terms of computational performance, as we are using GNNs, we recognize that we lose an important aspect of the RASP model which is Rapid by naming. Yet, as the most time-consuming part is the construction of the residue sub-graph (roughly 5 seconds for subgraphs of less than 96 nodes with 8 CPUs, an A100 GPU and vectorization/jit features within JAX) saving it once and slightly modifying it to include the specific mutations later on, makes the model very efficient at assessing scores for multiple combinations of mutants within a pre-defined set of positions.

Finally, since we decoupled the atomic and the residue scales, it is now possible to swap elements from other successful models: for example ThermoNet. This exposes a new bottleneck, or rather a new further challenge, as it implies the creation of a new data-set including structures for each mutant present in the Mega-scale data-set. That would also have been the case if one wanted to include anti-symmetry properties within the model.

\begin{table}
\begin{tabular}{|c|c|c|c|c|c|c|} \cline{2-7} \multicolumn{1}{c|}{} & **Training Set** & **Validation Set** & **Test Set** & **ThermoMutDB** \\ \hline
**Spearman r** & 0.754 & 0.518 & 0.442 & 0.311 \\ \hline
**Pearson r** & 0.758 & 0.562 & 0.412 & 0.251 \\ \hline
**RMSE** & 0.794 & 0.740 & 0.935 & 2.288 \\ \hline \end{tabular}
\end{table}
Table 1: Evaluation metrics for the \(\Delta\Delta G\) scorer

\begin{table}
\begin{tabular}{|c|c|c|c|c|c|c|c|} \hline
**Number of mutations** & **1** & **2** & **3** & **4** & **5** & **6** & **7** & **9** \\ \hline
**Spearman r** & 0.349 & 0.159 & 0.381 & 0.012 & 0.271 & 0.714 & -0.500 & 1.000 \\ \hline
**Pearson r** & 0.342 & 0.077 & 0.346 & 0.079 & 0.202 & 0.613 & -0.242 & 1.000 \\ \hline
**RMSE** & 2.109 & 2.571 & 2.378 & 1.972 & 2.448 & 1.965 & 1.588 & 0.810 \\ \hline \end{tabular}
\end{table}
Table 2: Metrics for our scorer across different numbers of mutations.
## 6 Conclusion

In this work, we explored the possibility of using graph neural network models for scoring multiple substitution effects on protein stability. Our approach, based on the decoupling of atomic and residue scales by successively training two different scale-specific E-GNN models on massive experimental data-sets, shows promising results. Indeed, the model demonstrates an ability to predict effects of a variable number of mutations, even beyond what it has been trained on. Yet some key parameters of this modelling still need to be better understood; for example, a biologically reasonable edge distance threshold and an overall more appropriate way to handle connectivity in the created residue sub-graph.

Figure 5: Evaluation of our scorer on individual structures, PDBs Burley et al. (2017), chosen with at least 100 occurrences in the ThermoMutDB test data-set. All four structures have a significant prediction correlation (p-value\(<\)0.05) and the marker size is proportional to the number of mutations in the experiment. Further results breakdown in Table 3.

\begin{table}
\begin{tabular}{|c|c|c|c|c|} \hline
**PDB ID** & **1BNI** & **ISTN** & **1VQB** & **1RX4** \\ \hline
**Spearman r** & 0.503 & 0.462 & 0.519 & 0.479 \\ \hline
**Pearson r** & 0.456 & 0.456 & 0.523 & 0.453 \\ \hline
**RMSE** & 1.651 & 1.632 & 2.526 & 1.193 \\ \hline \end{tabular}
\end{table}
Table 3: Scorer performance metrics on proteins with over 100 data points, as shown in Figure 5.