Predicting protein variants with equivariant graph neural networks
Pre-trained models have been successful in many protein engineering tasks.
Most notably, sequence-based models have achieved state-of-the-art performance
on protein fitness prediction while structure-based models have been used
experimentally to develop proteins with enhanced functions. However, there is a
research gap in comparing structure- and sequence-based methods for predicting
protein variants that are better than the wildtype protein. This paper aims to
address this gap by conducting a comparative study between the abilities of
equivariant graph neural networks (EGNNs) and sequence-based approaches to
identify promising amino-acid mutations. The results show that our proposed
structural approach achieves a competitive performance to sequence-based
methods while being trained on significantly fewer molecules. Additionally, we
find that combining assay labelled data with structure pre-trained models
yields similar trends as with sequence pre-trained models.
  Our code and trained models can be found at:
https://github.com/semiluna/partIII-amino-acid-prediction.
Antonia Boca, Simon Mathis
# Predicting protein variants with

###### Abstract

Pre-trained models have been successful in many protein engineering tasks. Most notably, sequence-based models have achieved state-of-the-art performance on protein fitness prediction while structure-based models have been used experimentally to develop proteins with enhanced functions. However, there is a research gap in comparing structure- and sequence-based methods for predicting protein variants that are better than the wildtype protein. This paper aims to address this gap by conducting a comparative study between the abilities of equivariant graph neural networks (EGNNs) and sequence-based approaches to identify promising amino-acid mutations. The results show that our proposed structural approach achieves a competitive performance to sequence-based methods while being trained on significantly fewer molecules. Additionally, we find that combining assay labelled data with structure pre-trained models yields similar trends as with sequence pre-trained models.

Our code and trained models can be found at: [https://github.com/semiluna/partIII-amino-acid-prediction](https://github.com/semiluna/partIII-amino-acid-prediction).

Machine Learning, ICML

## 1 Introduction

In recent years, pre-trained models have garnered significant attention in the field of protein representation. Notably, models have been developed to deal with both the sequence and structure modalities of proteins (Rives et al., 2021; Elnaggar et al., 2022; Zhang et al., 2023). These models have demonstrated their potential in various applications such as protein fitness prediction (Meier et al., 2021; Notin et al., 2022) while being employed in a "zero-shot" manner, without the need for additional training data. Their success has also shown promising experimental results in protein engineering (Shroff et al., 2020; Lu et al., 2022). Additionally, Hsu et al. (2021) have observed that augmenting simple models for assay labelled data with fitness predictions extracted from pre-trained sequence models can enhance their performance.

Despite the experimental success of pre-trained structural methods for protein engineering, particularly those based on predicting residues given local atom environments (Torng and Altman, 2017; Lu et al., 2022), several crucial aspects remain unexplored. Firstly, these methods have not been systematically compared with sequence-based approaches using the same datasets. Secondly, their potential to augment assay labelled data, when available, has not been evaluated.

This paper aims to fill this research gap by conducting a study of the comparative performance of structure-based and sequence-based methods on predicting variants _that are better than the wildtype protein_. We compare representatives of the most successful equivariant graph neural networks (EGNNs) on the task of residue identity prediction, namely GVP (Jing et al., 2021) and EQGAT (Le et al., 2022), with representatives of the most successful sequence-based approaches: Tranception (Notin et al., 2022), ESM-1v (Meier et al., 2021) and the MSA Transformer (Rao et al., 2021).

By undertaking this comparative analysis, we aim to provide insights into the performance and suitability of geometric GNNs in protein engineering, specifically in the context of predictions based on the local atomic environment. Our contributions are as follows:

* We apply the most successful pre-training approach for structural methods (Shroff et al., 2020) to equivariant GNNs by using the ATOM3D RES dataset (Townsend et al., 2022) for residue identity prediction (Table 1);
* We benchmark the resulting structure-based pre-trained models with the most successful zero-shot sequence-based approaches (Table 2). We observe that structure does not trump sequence in downstream tasks when used in this way, although the amount of available structures used during pre-training is significantly lower than the number of sequences used in training large language models;
* We extend the simple combination approach for assay labelled data and pre-trained model outputs (Hsu et al., 2021) to the structure pre-trained domain. We find the same general trends as with sequence pre-trained models, as assay-labelled data quickly allows us to surpass zero-shot pre-trained sequence-based models with at few as 100 datapoints (Figure 2).

## 2 Methodology

We pre-train two equivariant graph neural networks on the task of residue identity prediction, also known as the RES task (Townshend et al., 2022). We choose the Geometric Vector Perceptron (Jing et al., 2021) and the Equivariant Graph Attention Network (Le et al., 2022). While Lu et al. (2022) used 3D-CNNs to engineer plastic enzymes, Jing et al. (2021) benchmark 3D-CNNs on the RES task and show that the GVP outperforms them, so we choose to focus on this structural method instead.

Table 1 shows a comparison between the reported accuracies of the two models and the accuracies achieved in this paper. We achieve a higher performance on the GVP model than originally reported in Jing et al. (2021). This jump in performance can be explained by the fact that Jing et al. (2021) only use a third of the original training dataset to train the GVP, possibly due to computational constraints. More details on our training parameters can be found in A.9.

### RES task formalism

We formalise the RES classification task as follows. For a given point-cloud atomic graph \(G=(V,E)\) with nodes \(i,j\in V\) and edges \((i\to j)\in E\). Given a node \(t\in V\) representing the \(\text{C}_{\alpha}\) of a residue in the atomic graph, we can define the _node classification function_\(\text{RES}:\mathcal{V}\times\mathcal{G}\rightarrow\mathbb{R}^{20}\) that takes as input node \(t\) and a _masked_ atomic graph \(G_{t}\) from which we have removed the side-chain atoms of node \(t\) and returns the likelihood scores of each of the 20 naturally occurring amino-acids to be part of the side-chain of node \(t\). A more extended version of this formalism can be found in A.1.

### The scoring function

We now formalise the function we use to score each amino-acid mutation in a sequence. For a wildtype protein sequence \(x_{1}\dots x_{n}\) with \(x_{i}\in\mathcal{A}=\{1,2,\dots,20\}\) we associate the point-cloud atomic graph \(G=(V,E)\) corresponding to the protein's structure. Edges are drawn between any two atoms that are less than 4.5A apart.

Then, using the formalism defined in 2.1, the score associated with the presence of amino-acid \(a\in\mathcal{A}\) at position \(i\) can be defined as:

\[S(i,a)=[\text{RES}(g(i),G)]_{a} \tag{1}\]

Where \(g:\{1,2,\dots,n\}\rightarrow\{1,2,\dots,|V|\}\) is a mapping function from positions to the index of the node representing the central \(\text{C}_{\alpha}\) of the amino-acid residue present at each position. Here, \(G_{g(i)}\) denotes the masked graph from which we removed the side-chain attached to node \(g(i)\).

Equation 1 essentially represents the score of amino-acid \(a\) for target position \(i\), associated with node \(g(i)\) in the atomic graph. Note that the true amino-acid at the same position is denoted by \(x_{i}\).

### Mutation generation

Once the equivariant models have been trained on the RES task, we use them to inform the generation of single-point mutations in monomers and homo-oligomers from the ProteinGym substitutions dataset (Notin et al., 2022). For each wildtype sequence we recover its structure, mask each amino-acid residue in turn, and retrieve the scores generated by the EGNN model for each of the 20 naturally occuring amino-acids. These scores are then ranked according to two strategies to determine the most promising single-point mutations. Figure 1 illustrates this approach visually.

Structure recovery.The ProteinGym substitutions dataset contains 87 molecular sequences; for each of these sequences, a number of experimentally tested mutations are scored according to their _fitness_. We evaluate our methods

\begin{table}
\begin{tabular}{l c c} \hline \hline Model & Reported & Our \\  & test accuracy & test accuracy \\ \hline EQGAT & 0.540 & 0.524 \\ GVP & 0.527 & **0.580** \\ \hline \hline \end{tabular}
\end{table}
Table 1: Classification accuracies on the ATOM3D RES dataset.

Figure 1: For every sequence, we recover the structure from the PDB and mask each amino-acid in turn. We pass the masked graph through a pre-trained EGNN model to recover the score associated with each amino-acid, which we then rank. The key idea is that this pre-training allows the model to identify amino acids which seem “unusual” given their local environment and propose better fitting candidates instead.
on a subset of the original dataset for which we could find either monomeric or homo-oligomeric structures. For each wildtype sequence, we recover the corresponding biological assembly from the Protein Data Bank (Berman et al., 2000). When multiple assemblies are available, we choose one at random. When assemblies are incomplete, we instead use the monomeric AlphaFold prediction (Jumper et al., 2021) if available. Otherwise, we discard the sequence.

### Mutation ranking

Our approach allows us to score every possible residue mutation for each position in a sequence. Our goal is to generate meaningful mutations that have a higher chance of being bio-physically relevant, so we discard positions where the equivariant model makes the wrong prediction. A more detailed analysis of this design choice can be found in Appendix A.4.

Global ranking.We rank the remaining mutations according to two strategies: _global_ and _positional_. When performing global ranking, we sort mutations in descending order of their score, regardless of their position. If we denote the single-point mutation to amino-acid \(a\) at position \(i\) by \(\mathbf{m}_{i}^{a}\), then \(\forall i,j\) and \(\forall a,b\in\mathcal{A}\) s.t. \(a\neq x_{i}\) and \(b\neq x_{j}\), we say that:

\[\mathbf{m}_{i}^{a}\text{ is better than }\mathbf{m}_{j}^{b}\iff S(i,a)>S(j,b) \tag{2}\]

Positional ranking.The second approach follows when we prioritise the positions we want to mutate instead of the amino-acids we mutate to. Formally, this can be quantified as:

\[\begin{split}&\mathbf{m}_{i}^{a}\text{ is better than }\mathbf{m}_{j}^{b}\\ \iff&\Big{(}S(i,x_{i})<S(j,x_{j})\Big{)}\vee\\ &\Big{(}S(i,x_{i})=S(j,x_{j})\wedge S(i,a)>S(j,b)\Big{)}\end{split} \tag{3}\]

Note that when we perform positional ranking, we only keep the 3 top mutations for each position.

### Protein fitness prediction

The GVP and EQGAT trained on the ATOM3D RES task can be thought of as unsupervised models that can suggest amino-acid mutations. We extend our original approach to perform fitness prediction using a ridge regression model augmented with the positional scores generated by equivariant GNNs, in a similar manner to that introduced by Hsu et al. (2021). For a given sequence of amino-acids \(x_{1}\dots a_{i}\dots x_{n}\) with a single-point mutation at position \(i\), we embed each amino-acid using either the one-hot encoding or _AAIndex_ embeddings (Kawashima et al., 1999) on which we perform PCA to render 19-dimensional features per amino-acid. We flatten and concatenate these encodings to render feature vectors \(\mathbf{h}_{\text{one-hot}}\in\mathbb{R}^{20\times n}\) and \(\mathbf{h}_{\text{aa-index}}\in\mathbb{R}^{19\times n}\). To this feature vector we concatenate the score predicted by the GNN model for amino-acid \(a_{i}\) at position \(i\):

\[\mathbf{x}_{\text{one-hot}} =[\mathbf{h}_{\text{one-hot}}\mid\mid S(i,a_{i})] \tag{4}\] \[\mathbf{x}_{\text{aa-index}} =[\mathbf{h}_{\text{aa-index}}\mid\mid S(i,a_{i})] \tag{5}\]

Here, \(S(i,a_{i})\) is the same scoring function defined in Equation 1. These features are then used to train a ridge regression model to predict protein fitness using subsets of single-point mutated sequences for each of the ProteinGym DMS assays we have model scores for.

## 3 Results

### Mutation generation

We generate single-point mutations for 49 out of the 87 DMS assays in the ProteinGym substitutions dataset (Notin et al., 2022). When we generate mutations, we discard any that we cannot find in the experimental dataset of the target sequence. We are interested in understanding how good our models are at suggesting mutations that are _better than the wildtype_ sequence, hence we propose three metrics through which to perform comparisons: (1) Spearman's rank correlation restricted to better than wildtype sequences, (2) the precision of the top 10 mutations, and (3) the recall of the top 10 mutations. To compute the last two metrics we only considered whether a mutation proposed by the model is better than the wildtype, disregarding its actual score. Table 2 shows the performance of our models, depending on the type of ranking used. We note that the equivariant models have a competitive performance to Tranception (Notin et al., 2022) when ranking mutations that are better than the wildtype, indicating that they represent a viable strategy for aiding the discovery process in protein engineering. Per-dataset performance metrics can be found in A.2.

EGNN models require a significantly smaller number of protein structures during training in order to reach a similar ranking correlation coefficient to Tranception for mutations that are better than the wildtype. While Tranception is trained on the UniRef100 database (Suzek et al., 2015), which contains over 4 million source sequences, our models are trained on the ATOM3D RES dataset (Townshend et al., 2022), which contains fewer than 22k molecules from which local environments are sampled.

Structure vs Sequence.We believe EGNNs may require less training because structure is more informative than sequence for fitness prediction. While sequence-based models attend the full protein and subsequently learn to focus on the important bits, EGNNs attend only local environments, thus learning to identify important atoms faster. Further experiments could be run to compare the power of sequence
and structure-based models when their level of training is comparable. However, we point out that training EGNNs to the same level as present state-of-the-art sequence models may be infeasible due to both data and computational constraints.

Correlation to sequence-based models.As part of our analysis, we also compute the correlation between the better than wildtype predictions made by our EGNN models and Tranception. Per-model and per-dataset statistics can be found in A.3; we note that the highest rank correlation we find is **0.212**, in the case of the EQGAT model. Since these approaches seem to be weakly correlated, we believe there are improvements to be gained from ensembling both structure- and sequence-based approaches.

The impact of design choices.As mentioned in Section 2.4, we discard mutations at positions where the EGNNs make the wrong prediction, as we find that incorporating these is detrimental to the overall performance (see A.4). This indicates that these structure-based models are still undertrained, with potential for improvement coming both from larger datasets and more data engineering.

### Protein fitness prediction

We train 4 types of ridge regression models on each of the 49 DMS datasets separately. The baseline non-augmented model uses only features \(\mathbf{h}_{\text{one-hot}}\) or \(\mathbf{h}_{\text{aa-index}}\) defined in Section 2.5; the remaining 3 models are augmented with single-point mutation scores from GVP, EQGAT, and Tranception, respectively. For each model type and each DMS array we first set aside 20% of the single-point mutated sequences for testing; We train the regression on increasingly larger training subsets. We repeat the process 20 times with different random subsets and report the average Spearman rank correlation on better than wildtype sequences, as seen in Figure 2. The performance on other metrics can be found in A.8.

Similar to the results reported by Hsu et al. (2021), the augmented linear models allow us to surpass the baseline zero-shot fitness prediction models with as few as 100 datapoints in the case of the model augmented with EQGAT scores. While the linear model augmented with Tranception scores performs best overall, we point out that Tranception is fine-tuned to predict protein fitness, while the scores retrieved from our models merely represent the confidence in a certain amino-acid for a target position.

\begin{table}
\begin{tabular}{c c c c c c c} \hline \hline Model & Ranking strategy & \begin{tabular}{c} Top 10 \\ precision \\ \end{tabular} & 
\begin{tabular}{c} Top 10 \\ recall \\ \end{tabular} & \multicolumn{4}{c}{Spearman’s rank correlation} \\ \cline{3-6}  & & & Average & Worse than WT & Better than WT \\ \hline EQGAT & Positional & 0.486 & 0.187 & 0.223 & 0.128 & 0.118 \\ EQGAT & Global & 0.491 & 0.072 & 0.262 & 0.154 & 0.157 \\ GVP & Positional & 0.462 & **0.419** & 0.106 & \(-0.009\) & **0.276** \\ GVP & Global & 0.426 & 0.100 & 0.202 & 0.128 & \(-0.011\) \\ \hline Tranception & & 0.619 & 0.012 & 0.429 & 0.299 & 0.143 \\ ESM-1v & & 0.618 & 0.018 & 0.407 & 0.288 & 0.135 \\ MSA Transformer & & **0.638** & 0.018 & **0.434** & **0.327** & 0.135 \\ \hline \hline \end{tabular}
\end{table}
Table 2: Ranking performance of the models across 49 DMS assays. Numbers in **bold** represent the highest score per column, while numbers with an underline represent the second highest score per column. We note that two equivariant GNNs have the highest rank correlation for better than wildtype mutations.

Figure 2: Performance on mutations that are better than the wildtype for four regression models using two types of embeddings. Statistics are aggregated across 49 DMS assays. We note that we can improve the fitness prediction performance above the Tranception baseline (in black) across all regression models by training on as few as 144 data points.
## 4 Limitations and future work

We apply pre-trained EGNNs to both mutation generation and protein fitness prediction, and find that structural approaches are a competitive approach to sequence-based language models for the prediction of mutations that are better than the wildtype, while also requiring **181x** fewer molecules to train.

While the results look promising, this comparison is limited in scope, as our approach does not deal with more complex (hetero-oligomeric) structures from the ProteinGym dataset.

Types of fitness.Additionally, the benchmarking dataset contains a wide range of sequences for which "fitness" can be interpreted in many different ways. DMS assays in the ProteinGym dataset come from humans, viruses, prokaryotes, and eukaryotes. In particular, in the subset of 49 DMS assays used in this paper, **5** come from eukaryotes, **21** from humans, **18** from prokaryotes, and **5** from viruses. Fitness, in the case of viruses, is interpreted as infectivity or the likelihood of mutation. In the rest of the cases, fitness can range from stress resistance to efficiency. For example, in their experimental paper, Lu et al. (2022) focused on improving thermal stability. Hence, the fitness score used by Notin et al. (2022) in the ProteinGym dataset represents a "fuzzy" concept that is context-dependent. Future work could focus more closely on identifying the types of fitness structure-based approaches excel at.

AcknowledgementsSVM was supported by the UKRI Centre for Doctoral Training in Application of Artificial Intelligence to the study of Environmental Risks (EP/S022961/1).
