Protein Representation Learning by Geometric Structure Pretraining
Abstract
Learning effective protein representations is critical in a variety of tasks in biology such as predicting protein function or structure.
Existing approaches usually pretrain protein language models on a large number of unlabeled amino acid sequences and then finetune the models with some labeled data in downstream tasks.
Despite the effectiveness of sequence-based approaches, the power of pretraining on known protein structures, which are available in smaller numbers only, has not been explored for protein property prediction, though protein structures are known to be determinants of protein function.
In this paper, we propose to pretrain protein representations according to their 3D structures.
We first present a simple yet effective encoder to learn the geometric features of a protein.
We pretrain the protein graph encoder by leveraging multiview contrastive learning and different self-prediction tasks.
Experimental results on both function prediction and fold classification tasks show that our proposed pretraining methods outperform or are on par with the state-of-the-art sequence-based methods, while using much less data.
All codes and models will be published upon acceptance.
1. Introduction
Proteins are workhorses of the cell and are implicated in a broad range of applications ranging from therapeutics to material.
They consist of a linear chain of amino acids (residues) folding into a specific 3D conformation.
As the function annotation of a new protein sequence remains costly and time-consuming, accurate and efficient in silico protein function annotation methods are needed to bridge the existing sequence-function gap.
However, existing sequence-based methods neither directly capture nor leverage the available protein structural information that is known to be the determinants of protein functions (Bepler & Berger, 2019; 2021; Rives et al., 2021).
To better utilize this critical structural information, we first propose a simple yet effective structure-based encoder called
GeomEtry-Aware Relational Graph Neural Network (GearNet) , which encodes spatial information by relational and edge message passing on protein residue graphs.
We further introduce five geometric pretraining methods to learn the protein structure encoder.
These methods follow popular self-supervised learning frameworks such as contrastive learning (Chen et al., 2020) and self-prediction (Devlin et al., 2018).
For contrastive learning, we aim to maximize the similarity between representations of different augmented views from the same protein, while minimizing the similarity between those from different proteins.
For self-prediction, the model performs masked prediction of different geometric or biochemical properties, such as residue types, Euclidean distances, angles and dihedral angles.
Extensive experiments on several benchmarks verify our
GearNet augmented with edge message passing can consistently outperform existing protein encoders on most tasks in a supervised setting.
Further, by employing the proposed pretraining methods, our model trained on less than a million samples achieves comparable or even better results than the state-of-the-art sequence-based encoders pretrained on million- or billion-scale datasets.
2. Structure-Based Protein Encoder
Existing protein encoders are either designed for specific tasks or cumbersome for pretraining due to the dependency on computationally expensive convolutions.
Here we propose a simple yet effective protein structure encoder, named
GeomEtry-Aware Relational Graph Neural Network (GearNet) , which is further enhanced with edge message passing .
2.1. Geometry-Aware Relational Graph Neural Network
Given protein structures, our model aims to learn representations encoding their spatial and chemical information.
These representations should be invariant under translations and rotations in 3D space.
To achieve this requirement, we first construct our protein graph based on the spatial features invariant under these transformations.
Protein graph construction.
We represent the structure of a protein as a residue-level relational graph G =
( V , E , R ) , where V and E denotes the set of nodes and edges respectively, and R is the set of edge types.
We use ( i, j, r ) to denote the edge from node i to node j with type r .
We use n and m to denote the number of nodes and edges, respectively.
In this work, each node in the protein graph represents the alpha carbon of a residue with the 3D coordinates of all nodes x ∈ R n × 3 .
Then, we add three different types of edges into our graphs: sequential edges, radius edges and K-nearest neighbor edges.
Among these, sequential edges will be further divided into 5 types of edges based on the relative sequential distance d ∈ {− 2 , − 1 , 0 , 1 , 2 } between two end nodes, where we add sequential edges only between the nodes within the sequential distance of 2.
These edge types reflect different geometric properties, which all together yield a comprehensive featurization of proteins.
More details of graph construction can be found in Appendix C.1.
Relational graph convolutional layer.
To balance model capacity and memory cost, we use a relational graph convolutional neural network (Schlichtkrull et al., 2018) to learn graph representations, where a convolutional kernel matrix is shared within each edge type and there are |R| different kernel matrices in total.
Formally, the relational graph convolutional layer used in our model is defined as
[Equation]

 .
Specifically, we use node features
f i as initial representations h (0) i .
Then, given the node representation h ( l ) i for node i at the l -th layer, we compute updated node representation by aggregating features from neighboring nodes
N r ( i ) , where N r ( i ) =
{ j ∈ V| ( j, i, r ) ∈ E} denotes the neighborhood of node i with the edge type r , and W r denotes the learnable convolutional kernel matrix for edge type r .
Here BN denotes a batch normalization layer and we use a ReLU function as the activation σ ( · ) .
Finally, we update h ( l ) i with a residual connection from last layer.
2.2. Edge Message Passing Layer
Inspired by models on small molecules (Klicpera et al., 2020), we propose a variant of GearNet enhanced with an edge message passing layer, named as GearNet-Edge .
We first construct a relational graph G ′
=
( V ′ , E ′ , R ′ ) among edges, also known as line graph in the literature (Harary & Norman, 1960).
Each node in the graph
G ′ corresponds to an edge in the original graph.
G ′ links edge ( i, j, r 1 ) to ( w, k, r 2 )
if and only if j = w and i ̸
= k .
The type of this edge is determined by the angle between ( i, j, r 1 ) and ( w, k, r 2 ) .
We discretize the range [0 , π ] intobins and use the index of the bin as the edge type.
Then, we apply a similar relational graph convolutional network on the graph G ′ for the message function.
Formally, the edge message passing layer is defined as
[Equation]
 . Here we use m ( l ) ( i,j,r 1 ) to denote the message function for edge ( i, j, r 1 ) in the l -th layer with m (0)
( i,j,r 1 ) =
f ( i,j,r 1 ) .
Similarly, the message function for edge ( i, j, r 1 ) will be updated by aggregating features from its neighbors N ′ r
(( i, j, r 1 )) , where N ′ r (( i, j, r 1 ))
= { ( w, k, r 2 ) ∈
V ′ | (( w, k, r 2 ) , ( i, j, r 1 ) , r ) ∈
E ′ } denotes the set of incoming edges of ( i, j, r 1 ) with relation type r in graph
G ′ .
Finally, we replace the update function in the original graph with the following one:
[Equation]
where FC ( · ) denotes a linear layer over the message.
Notably, it is novel to use relational message passing to model spatial interactions among residues.
In addition, to the best of our knowledge, this is the first work that explores edge message passing for protein representation learning.
3. Geometric Pretraining Methods
In this section, we study how to leverge massive unlabeled protein structures to boost protein representation learning.
We follow two popular self-supervised learning frameworks: contrastive learning and self-prediction, and propose five different pretraining methods for structure-based encoders.
3.1. Multiview Contrastive Learning
Inspired by recent contrastive learning methods (Oord et al., 2018; Chen et al., 2020), our framework learns representations by maximizing the similarity between representations of different augmented views of the same protein while minimizing the agreement between views of different proteins.
More specifically, given a protein graph G , we first sample two different views G x and G y via a stochastic augmentation module.
We then compute the graph representations h x and h y of two views using our structure-based encoder.
Following SimCLR (Chen et al., 2020), a two-layer MLP projection head is further applied to map the representations to a lower-dimensional space, denoted as z x and z y .
Finally, a contrastive loss function is defined by distinguishing views from the same or different proteins using their similarities.
For a positive pair x and y , we treat views from other proteins in the same mini-batch as negative pairs.
Mathematically, the loss function for a positive pair of views x and y can be written as:
[Equation]
where B denotes the batch size, τ denotes the temperature,[ k ̸ = x ] ∈ { 0 , 1 } is an indicator function that is equal to
if and only if k ̸ = x .
And the similarity function sim ( u , v ) is defined by the cosine similarity between u and v .
Data augmentation.
As suggested in (Chen et al., 2020), data augmentation schemes play an important role in contrastive learning.
In this work, we propose a principled way to generate diverse views of a protein.
We randomly choose a cropping function to make the protein graph smaller so that larger-size batches can be used for pretraining.
Here we consider two cropping functions:
• Subsequence: randomly sample a consecutive segment of protein sequences and take the corresponding subgraph from the protein residue graph.
• Subspace: randomly sample a residue as the center and select all residues within a Euclidean ball with a predefined radius.
Then, we randomly choose one of the following two transformations with equal probability to apply on the cropped protein graphs.
• Identity: no transformation.
• Random edge masking: randomly mask edges with a fixed mask rate equal to 0 .
3.2. Self-Prediction Methods
Another line of research is based on the recent progress of self-prediction methods in natural language processing (Devlin et al., 2018; Brown et al., 2020).
Given a protein, our objective is to predict one part of the protein given the remainder of the structure.
Here, we propose four self-supervised tasks based on geometric or biochemical properties.
Residue Type Prediction.
Our first method is based on the masked language modeling objective, which has been widely used in pretraining large-scale protein language models (Bepler & Berger, 2021).
For each protein, we randomly mask node features of some residues and then let structurebased encoders to predict these masked residue types.
For a masked node i , the learning objective is defined as:
[Equation]
where h ′ i denotes the representations of node i after masking, f residue ( · ) is an MLP classification head and CE ( · , · ) denotes the cross entropy loss.
This method is also known as Attribute Masking in the literature of pretraining on small molecules (Hu et al., 2019).
Distance Prediction.
In order to learn local spatial structures, we use our learned representations to predict the Euclidean distance between two nodes connected in the protein graph.
First, we randomly select a fixed number of edges from the original graph.
Then, these edges will be removed when feeding the graph into the encoder.
For a selected edge, the representations of its two end nodes will be concatenated to predict the distance between them.
More concretely, the loss function for an edge ( i, j, r ) will be defined as:
[Equation]
where f dist ( · ) is an MLP prediction head.
Angle Prediction.
Besides distances, angles between edges are also important features that reflect the relative position between residues.
Similarly, we can define a masked geometric loss by 1) randomly selecting two adjacent edges, 2) removing these edges from the graph, 3) using the three end nodes of the pair of edges to predict the angle between them.
Here we discretized the angles by cutting the range [0 , π ] into 8 bins.
The objective for a selected pair of edges ( i, j, r 1 ) and ( j, k, r 2 ) is to predict which bin the angle between them will belong to:
[Equation]
where f angle ( · ) is an MLP classification head and bin ( · ) is used to discretize the angle.
Dihedral Prediction.
As shown in (Klicpera et al., 2021), the dihedral angles between three edges can provide important clues about the relative directional information.
Therefore, we can also construct a masked geometric objective by predicting the dihedral angle between three consecutive edges ( i, j, r 1 ) , ( j, k, r 2 ) and ( k, t, r 3 ) :
[Equation]
where f ( · ) is an MLP classification head and bin ( · ) is used to discretize the dihedral angles. [Table] [Table]
w / p re t r a i n i n g
[Table]
4. Experiments
In this section
[Table]
we conduct experiments on four downstream tasks
[Table]
where the experimental setups are described in Append. E.1 and the results are reported in Table 1, including all models with and without pretraining.
The following conclusions can be drawn from the results:
Our structure-based encoders outperform all baselines without pretraining on 7 of 8 datasets.
By comparing the first three blocks, we find that the vanilla GearNet can already obtain competitive results against other baselines on three function prediction tasks (EC, GO, Reaction).
After adding the edge message passing mechanism, our method
GearNet-Edge significantly outperforms other baselines on
EC, GO-BP and GO-MF and achieves competitive results on GO-CC.
By further adding IEConv (Hermosilla & Ropinski, 2022) layers into our model, the GearNet-Edge-IEConv model can achieve the best results on the fold classification task.
These strong performance demonstrates the advantages of our proposed structure-based encoders.
Structure-based encoders benefit a lot from pretraining with unlabeled protein structures.
Comparing the results in the third and fifth blocks, it can be observed that models with all five proposed pretraining methods show large improvements over models trained from scratch.
Among these methods, Multiview Contrast is the best on 7 ofdatasets and achieve the state-of-the-art results on EC, GOBP, GO-MF, Fold and Reaction classification tasks.
This proves the effectiveness of our pretraining strategies.
with or even better than sequence-based encoders pretrained with much more data.
The last two blocks show the comparision between pretrained sequence-based and structure-based models.
It should be noted that our models are pretrained on a dataset with fewer than one million structures, whereas all sequence-based pretraining baselines are pretrained on million- or billion-scale sequence databases.
Though pretrained with an order of magnitude less data, our model can achieve comparable or better results against these sequence-based models.
Besides, our model is the only one that can achieve good performance on all four tasks, given that sequence-based models do not perform well on fold classification.
This again shows the potential of structurebased pretraining for learning protein representations.
Structure-based encoders perform well under different sequence identity cutoffs.
Besides the experiments above, where 95% is used as the sequence identity cutoff for EC and GO dataset splitting, we also test our models and several important baselines under four lower sequence identity
[Table]
cutoffs and show the experimental results in Table 6.
It can be observed that, at lower sequence identity cutoffs, our model can still achieve the best performance among models without pretraining and get comparable or better results against ESM-1b after pretraining
[Table] 5. Conclusions
In this work [Table] we propose a simple yet effective structurebased encoder for protein representation learning, which performs relational and edge message passing on protein residue graphs.
Moreover, five self-supervised pretraining methods are proposed following two standard frameworks: contrastive learning and self-prediction.
Comprehensive experiments over multiple tasks verify that our model outperforms previous encoders when trained from scratch and achieve comparable or even better results than the state-ofthe-art baselines while pretraining with much less data. Acknowledgments
The authors would like to thank Meng Qu, Zhaocheng Zhu,
Shengchao Liu, Chence Shi, Minkai Xu and Huiyu Cai for their helpful discussions and comments.
This project is supported by AIHN IBM-MILA partnership program, the Natural Sciences and Engineering Research Council (NSERC) Discovery Grant, the Canada CIFAR AI Chair Program, collaboration grants between Microsoft
Research and Mila, Samsung Electronics Co., Ltd., Amazon
Faculty Research Award, Tencent AI Lab Rhino-Bird Gift
Fund, a NRC Collaborative R&D Project (AI4D-CORE-06) as well as the IVADO Fundamental Research Project grant
PRF-2019-3583139727. A. Related Work Previous works seek to learn protein representations based on different modalities of proteins, including amino acid sequences (Rao et al., 2019; Elnaggar et al., 2021;
Rives et al., 2021), multiple sequence alignments (MSAs) (Rao et al., 2021; Biswas et al., 2021; Meier et al., 2021) and protein structures (Hermosilla et al., 2021; Gligorijevi´c et al., 2021; Somnath et al., 2021).
These works share the common goal of learning informative protein representations that can benefit various downstream applications, like predicting protein function (Gligorijevi´c et al., 2021; Rives et al., 2021) and protein-protein interaction (Wang et al., 2019), as well as designing protein sequences (Biswas et al., 2021).
A.1. Sequence-Based Methods
Sequence-based protein representation learning is mainly inspired by the methods of modeling natural language sequences.
Recent methods aim to capture the biochemical and co-evolutionary knowledge underlying a large-scale protein sequence corpus by self-supervised pretraining, and such knowledge is then transferred to specific downstream tasks by finetuning.
Typical pretraining objectives explored in existing methods include next amino acid prediction (Alley et al., 2019; Elnaggar et al., 2021), masked language modeling (MLM) (Rao et al., 2019; Elnaggar et al., 2021;
Rives et al., 2021), pairwise MLM (He et al., 2021) and contrastive predictive coding (CPC)
(Lu et al., 2020).
Compared to sequence-based approaches that learn in the whole sequence space, MSA-based methods (Rao et al., 2021;
Biswas et al., 2021; Meier et al., 2021) leverage the sequences within a protein family to capture the conserved and variable regions of homologous sequences, which imply specific structures and functions of the protein family.
A.2. Structure-Based Methods
Although sequence-based methods pretrained on large-scale databases perform well, structure-based methods should be, in principle, a better solution to learning an informative protein representation, as the function of a protein is determined by its structure.
This line of works seeks to encode spatial information in protein structures by 3D CNNs (Derevyanko et al., 2018) or graph neural networks (GNNs) (Gligorijevi´c et al., 2021; Baldassarre et al., 2021;
Jing et al., 2021).
Among these methods, IEConv (Hermosilla et al., 2021) tries to fit the inductive bias of protein structure modeling, which introduced a graph convolution layer incorporating intrinsic and extrinsic distances between nodes.
Another potential direction is to extract features from protein surfaces (Gainza et al., 2020; Sverrisson et al., 2021; Dai &
Bailey-Kellogg, 2021).
(Somnath et al., 2021) combined the advantages of both worlds and proposed a parameterefficient multi-scale model.
Besides, there are also works that enhance pretrained sequence-based models by incorporating structural information in the pretraining stage (Bepler & Berger, 2021) or finetuning stage (Wang et al., 2022).
Despite progress in the design of structure-based encoders, there are few works focusing on structure-based pretraining for proteins.
To the best of our knowledge, the only attempt is two concurrent works (Hermosilla & Ropinski, 2022) and (Chen et al., 2022), which apply contrastive learning and self-prediction methods on a small set of tasks, respectively.
Compared with these existing works, our proposed encoder is conceptually simpler and more effective on many different tasks, thanks to the proposed relational graph convolutional layer and edge message passing layer, which are able to efficiently capture both the sequential and structural information.
Furthermore, we introduce five structure-based pretraining methods within the contrastive learning and selfprediction frameworks, which can serve as a solid starting point for enabling self-supervised learning on protein structures.
A.3. Structure-based Encoders for Biological Molecules
Following the early efforts (Behler & Parrinello, 2007;
Bart ´ ok et al., 2010; 2013; Chmiela et al., 2017) of building machine learning systems for molecules by hand-crafted atomic features, recent works exploited end-to-end message passing neural networks (MPNNs)
(Gilmer et al., 2017) to encode the structures of small molecules and macromolecules like proteins.
Specifically, existing methods employed node/atom message passing (Gilmer et al., 2017;
Sch¨utt et al., 2017a;b), edge/bond message passing (Jørgensen et al., 2018;
Chen et al., 2019) and directional information (Klicpera et al., 2020; Liu et al., 2021;
Klicpera et al., 2021) to encode 2D or 3D molecular graphs.
Compared to small molecules, structural representations of proteins are more diverse, including residue-level graphs, atom-level graphs and protein surfaces.
There are some recent models designed to encode residue-level graphs (Hermosilla et al., 2021; Hermosilla & Ropinski, 2022) and protein surfaces (Gainza et al., 2020; Sverrisson et al., 2021), and they achieved impressive results on various tasks.
However, these models are either not expressive enough to capture edge interactions or too complicated for representation learning.
A.4. Pretraining Graph Neural Networks
Our work is also related to the recent efforts of pretraining graph neural networks (GNNs), which sought to learn graph representations in a self-supervised fashion.
In this domain, various self-supervised pretext tasks, like edge prediction (Kipf & Welling, 2016; Hamilton et al., 2017), context prediction (Hu et al., 2019; Rong et al., 2020), node/edge attribute reconstruction (Hu et al., 2019) and contrastive learning (Hassani & Khasahmadi, 2020; Qiu et al., 2020;
You et al., 2020; Xu et al., 2021), are designed to acquire knowledge from unlabeled graphs.
In this work, we focus on learning representations of residue-level graphs of proteins in a self-supervised way.
To attain this goal, we design several novel protein-specific pretraining methods to learn the proposed structure-based GNN encoder.
B. Broader Impact
This research project focuses on learning effective protein representations via pretraining with a large number of unlabeled protein structures.
Compared to the conventional sequence-based pretraining methods, our approach is able to leverage structural information and thus provide better representations.
This merit enables more in-depth analysis of protein research and can potentially benefit many real-world applications, like protein function prediction and sequence design.
However, it cannot be denied that some harmful activities could be augmented by powerful pretrained models, e.g. , designing harmful drugs.
We expect future studies will mitigate these issues.
C. More Details of GearNet
In this section, we describe more details about the implementation of our GearNet.
The whole pipeline of our structurebased encoder is depicted in Figure 2.
C.1. Protein Graph Construction
For graph construction, we use three different ways to add edges:
Sequential edges.
The i -th residue and the j -th residue will be linked by an edge if the sequential distance between them is below a predefined threshold d seq , i e., | j − i | < d seq .
The type of each sequential edge is determined by their relative position d = j − i in the sequence.
Hence, there are 2 d seq − 1 types of sequential edges.
Radius edges.
Following previous works, we also add edges between two nodes i and j when the Euclidean distance between them is smaller than a threshold d radius .
K-nearest neighbor edges.
Since the scales of spatial coordinates may vary among different proteins, a node will be also connected to its k-nearest neighbors based on the Euclidean distance.
In this way, the density of spatial edges are guaranteed to be comparable among different protein graphs.
Since we are not interested in spatial edges between residues close with each other in the sequence, we further add a filter to the latter two kinds of edges.
Specifically, for a radius or KNN edge connecting the i -th residue and j -th residue, it will be removed if the sequential distance between them is lower than a long range interaction cutoff
d long , i.e. , | i − j | < d long .
In this paper, we set the sequential distance threshold d seq = 3 , the radius d radius = 10 .
˚ A , the number of neighbors k
= 10 and the long range interaction cutoff d long = 5 .
By regarding radius edges and KNN edges as two separate edge types, there will be totally 2 d seq + 1 = 7 different types of edges.
Necessity of spatial edges.
Here we explain the necessity of radius and KNN edges by statistics and intuitions.
These two kinds of edges result in very different degree distributions.
In Figure 1, we plot the average degree distribution over all proteins in AlphaFold Database v1.
If we only consider KNN edges, the node degrees in protein graphs are close to a constant, which makes it difficult to capture those areas with dense interactions between residues.
If we only consider radius edges, then there will be about 45,000 proteins with average degrees lower than two.
In these sparse graphs, pretraining cannot capture structural information effectively, e.g. , Angle Prediction with limited edge pairs and Dihedral Prediction with limited edge triplets.
Such sparsity can hard be overcome by tuning radius cutoff, for the various scales of average distance on different proteins.
By simply combining two kinds of edges, we can overcome these issues.
C.2. Node and edge features.
Most previous structure-based encoders designed for biological molecules (Baldassarre et al., 2021; Hermosilla et al., 2021) used many chemical and spatial features, some of which are difficult to obtain or time-consuming to calculate.
In contrast, we only use the one-hot encoding of residue types with one additional dimension for unknown types as node features, denoted as f ∈ { 0 , 1 } n × 21 , which is enough to learn good representation as shown in our experiments.
The feature f ( i,j,r ) for an edge ( i, j, r ) is the concatenation of the node features of two end nodes, the one-hot encoding of the edge type, and the sequential and spatial distances between them:
[Equation]
where Cat ( · ) denotes the concatenation operation.
C.3. Enhance GearNet with IEConv Layers
In our experiments, we find that IEConv layers are very useful for predicting fold labels in spite of their relatively poor performance on function prediction tasks.
Therefore, we enhance our models by adding a simplified IEConv layer as an additional layer, which achieve better results than the original IEConv.
In this section, we describe how to simplify the IEConv layer and how to combine it with our model.
Simplify the IEConv layer.
The original IEConv layer relies on the computation of intrinsic and extrinsic distances between two nodes, which are computationally expensive.
Hence, we follow the modifications proposed in (Hermosilla & Ropinski, 2022), which show improvements as reported in their experiments.
Although these modifications are not proposed by us, we still briefly describe the model for completeness.
In the IEConv layer, we keep the edges in our graph G and use ˜ h ( l ) i to denote the hidden representation for node i in the l -th layer.
The update equation for node i is defined as:
[Equation]
where N ( i ) is the set of neighbors of i , f ( G , i, j ) is the edge feature between i and j and k o ( · ) is an MLP mapping the feature to a kernel matrix.
Instead of intrinsic and extrinsic distances in the original IEConv layer, we follow
New IEConv, which adopts three relative positional features proposed in (Ingraham et al., 2019) and further augments them with additional input functions.
We aim to apply this layer on our constructed protein residue graph instead of the radius graph in the original paper.
Therefore, we simply remove the dynamically changed receptive fields, pooling layer and smoothing tricks in our setting.
Combine IEConv with GearNet.
Our model is very flexible to incorporate other message passing layers.
To incorporate IEConv layers, we just use our graph and hidden representations as input and replace the update equation with
[Equation]
D. Demonstration of Geometric Preatraining
The high-level ideas of our proposed five geometric pretraining methods are illustrated in Figure 3.
For multiview contrastive learning methods, we propose two cropping functions and two transformation functions to generate different views of proteins.
Here we illustrate these augmentation functions in Figure 4.
E. Experimental Details E.1. Setup
Pretraining datasets.
We use the AlphaFold protein structure database (CC-BY 4.0 License)
(Jumper et al., 2021;
Varadi et al., 2021) for pretraining.
This database contains the protein structures predicted by the AlphaFold2 model, and we employ both 365K proteome-wide predictions and 440K Swiss-Prot (Consortium, 2021) predictions in our experiments.
In Appendix G, we further report the results of pretraining on different datasets.
Downstream tasks.
We adopt two tasks proposed in (Gligorijevi´c et al., 2021) and two tasks used in (Hermosilla et al., 2021) for downstream evaluation.
Enzyme Commission (EC) number prediction (BSD 3-Clause License) seeks to predict the EC numbers of different proteins, which describe their catalysis of biochemical reactions.
The EC numbers are selected from the third and fourth levels of the
EC tree (Webb et al., 1992), forming 538 binary classification tasks.
Gene Ontology (GO) term prediction (BSD 3-Clause License) aims to predict whether a protein belongs to some GO terms.
These terms classify proteins into hierarchically related functional classes organized into three ontologies: molecular function (MF), biological process (BP) and cellular component (CC).
Fold classification (CCBY 4.0 license) is first proposed in (Hou et al., 2018), with the goal to predict the fold class label given a protein.
Reaction classification (BSD 3-Clause License ) aims to predict the enzyme-catalyzed reaction class of a protein, in which all four levels of the EC number are employed to depict the reaction class.
Although this task is essentially the same with EC prediction, we include it to make a fair comparison with the baselines in (Hermosilla et al., 2021).
Dataset splits.
For EC and GO prediction, we follow the multi-cutoff split methods in (Gligorijevi ´ c et al., 2021) to ensure that the test set only contains PDB chains with sequence identity no more than 95% to the training set, which is also used in (Wang et al., 2022).
For fold classification, (Hou et al., 2018) provides three different test sets: Fold , in which proteins from the same superfamily are unseen during training;
Superfamily , in which proteins from the same family are not present during training; and Family , in which proteins from the same family are present during training.
For reaction classification, we adopt dataset splits proposed in (Hermosilla et al., 2021), where proteins have less than 50% sequence similarity in-between splits.
Baselines.
Following (Wang et al., 2022) and (Hermosilla et al., 2021), we compare our encoders with many existing protein representation learning methods, including four sequence-based encoders (CNN (Shanehsazzadeh et al., 2020), ResNet (Rao et al., 2019), LSTM (Rao et al., 2019) and Transformer (Rao et al., 2019)), six structure-based encoders (GCN (Kipf & Welling, 2017),
GAT (Veli ˇ ckovi ´ c et al., 2018), GVP (Jing et al., 2021), 3DCNN MQA (Derevyanko et al., 2018), GraphQA (Baldassarre et al., 2021) and New IEConv (Hermosilla & Ropinski, 2022)).
We also include two models pretrained on largescale sequence datasets (ProtBERT-BFD (Elnaggar et al., 2021), ESM-1b (Rives et al., 2021)) and two models combining pretrained sequence-based encoders with structural information (DeepFRI (Gligorijevi´c et al., 2021) and LMGVP (Wang et al., 2022)).
For LM-GVP and New IEConv, we only include results reported in the original paper due to the computational burden and the lack of codes.
Training.
On the four downstream tasks, we train GearNet and GearNet-Edge from scratch.
As we find that the
IEConv layer is important for predicting fold labels, we also enhance our model by incorporating this as an additional layer (details in Appendix C.3).
These models are referred as GearNet-IEConv and GearNet-IEConv-Edge, respectively.
Following previous works (Wang et al., 2022) and (Hermosilla & Ropinski, 2022), the models are trained for 200 epochs on EC and GO prediction and for 300 epochs on fold and reaction classification.
For pretraining, the models with the best performance when trained from scratch are selected, i.e. , GearNet-Edge for EC, GO, Reaction and
GearNet-Edge-IEConv for Fold Classification.
The models are pretrained on the AlphaFold protein database with our proposed five methods for 50 epochs.
All these models are trained on 4 Tesla A100 GPUs.
More details can be found in Appendix E.4.
Evaluation.
For EC and GO prediction, we evaluate the performance with the protein-centric maximum Fscore F max , which is commonly used in the CAFA challenges (Radivojac et al., 2013)
(See Appendix E.3 for details).
Experimental results with another popular metric, pair-centric area under precision-recall curve AUPR pair , are reported in Appendix.
F. For fold and reaction classification, the performance is measured with the mean accuracy.
Models with the best performance on validation sets are selected for evaluation on test sets.
E.2. Dataset Statistics
[Table]
Dataset statistics of our four downstream tasks are summarized in Table 3.
More details are introduced as follows.
Enzyme Commission and Gene Ontology.
Following
DeepFRI (Gligorijevi´c et al., 2021), the EC numbers are selected from the third and fourth levels of the EC tree, forming 538 binary classification tasks, while the GO terms with at least 50 and no more than 5000 training samples are selected.
The non-redundant sets are partitioned into training, validation and test sets according to the sequence identity.
We retrieve all protein chains from PDB used the code provided in their codebase and remove those with obsolete pdb ids, so the statistics will be slightly different from the number reported in the original paper.
Fold Classification.
We directly use the dataset processed in (Hermosilla et al., 2021), which consolidated 16, proteins with 1,195 different folds from the SCOPe 1.
database (Murzin et al., 1995).
Reaction Classification.
The dataset comprises 37, proteins categorized into 384 reaction classes.
The split methods are described in (Hermosilla et al., 2021), where they cluster protein chains via sequence similarities and ensure that protein chains belonging to the same cluster are in the same set.
E.3. Evaluation Metrics
Now we introduce the details of evaluation metrics for EC and GO prediction.
These two tasks aim to answer the question: whether a protein has a particular function, which can be seen as multiple binary classification tasks.
The first metric, protein-centric maximum F-score F max , is defined by first calculating the precision and recall for each protein and then taking the average score over all proteins.
More specifically, for a given target protein
i and some decision threshold t ∈ [0 , 1] , the precision and recall are computed as: precision
[Equation]
and recall
[-(cid:80) f 
-][Equation]
where f is a function term in the ontology, T i is a set of experimentally determined function terms for protein i , P i ( t ) denotes the set of predicted terms for protein i with scores greater than or equal to t and 1 [ · ] ∈ { 0 , 1 } is an indicator function that is equal to 1 iff the condition is true.
Then, the average precision and recall over all proteins at threshold t is defined as:
[Equation]
and
[Equation]
where we use N to denote the number of proteins and M ( t ) to denote the number of proteins on which at least one prediction was made above threshold t , i.e. , | P i ( t ) | > 0 .
Combining these two measures, the maximum F-score is defined as the maximum value of F-measure over all thresholds.
That is,
[Equation]
The second metric, pair-centric area under precision-recall curve AUPR pair , is defined as the average precision scores for all protein-function pairs, which is exactly the micro average precision score for multiple binary classification.
E.4. Implementation Details
In this subsection, we describe implementation details of all baselines and our methods.
For all models, the outputs will be fed into a three-layer MLP to make final prediction.
The dimension of hidden layers in the MLP is equal to the dimension of model outputs.
CNN (Shanehsazzadeh et al., 2020).
Following the finding in (Shanehsazzadeh et al., 2020), we directly employ a shallow convolutional neural network (CNN) to encode protein sequences.
Specifically, 2 convolutional layers with 1024 hidden dimensions and kernel size 5 constitute this baseline model.
ResNet (Rao et al., 2019).
We also adopt a deep CNN model, i.e. , the ResNet for protein sequences proposed by (Rao et al., 2019), in our benchmark.
This model is with 12 residual blocks and 512 hidden dimensions, and it uses the GELU (Hendrycks & Gimpel, 2016) activation function.
LSTM (Rao et al., 2019).
The bidirectional LSTM model proposed by (Rao et al., 2019) is another baseline for protein sequence encoding.
It is composed of three bidirectional
LSTM layers with 640 hidden dimensions.
Transformer (Rao et al., 2019).
The self-attention-based
Transformer encoder (Vaswani et al., 2017) is a strong model in natural language processing (NLP), (Rao et al., 2019) adapts this model into the field of protein sequence modeling.
We also adopt it as one of our baselines.
This model has a comparable size with BERT-Small (Devlin et al., 2018), which contains 4 Transformer blocks withhidden dimensions and 8 attention heads, and it is activated by GELU (Hendrycks & Gimpel, 2016).
GCN (Kipf & Welling, 2017).
We take GCN as a baseline to encode the residue graph derived by our graph construction scheme.
We adopt the implementation in TorchDrug (Zhu et al., 2022), where 6 GCN layers with the hidden dimension of 512 are used for encoding.
We run the results of GCN on EC and GO by ourselves and take its results on Fold and Reaction classification from (Hermosilla et al., 2021).
GAT (Veliˇckovi´c et al., 2018).
We adopt another popular graph neural network, GAT, as a structure-based baseline model.
We follow the implementation in TorchDrug and use 6 GAT layers with the hidden dimension of 512 andattention head per layer for encoding.
The results on EC,
Fold and Reaction classification are based on our runs, and the results on GO are taken from (Wang et al., 2022).
GVP (Jing et al., 2021).
The GVP model (Jing et al., 2021) is a decent protein structure encoder.
It iteratively updates the scalar and vector representations of a protein, and these representations possess the merit of invariance and equivariance.
In our benchmark, we evaluate this baseline method following the official source code.
In specific,
GVP layers with 32 feature dimensions (20 scalar andvector channels) constitute the GVP model.
3DCNN MQA (Derevyanko et al., 2018).
We implement the 3DCNN model from the paper (Derevyanko et al., 2018) with a box width of 40.0 and input resolution of 120 × 120 × 120 .
The model has 6 residual blocks and 128 hidden dimensions with ELU activation function.
GraphQA (Baldassarre et al., 2021).
Following the hyperparameters in the original paper, we construct the residue graphs based on bond and spatial information and re-implement the graph neural network in our codebase.
The best model has 4 layers with 128 node features, 32 edge features and 512 global features.
New IEConv (Hermosilla & Ropinski, 2022).
Since the code for New IEConv has not been made public when the paper is written, we reproduce the method according to the description in the paper and achieve similar results on Fold and Reaction classification tasks.
Then, we evaluate the method on EC and GO prediction tasks with the default hyperparameters reported in the original paper and follow the standard training procedure on these two tasks.
DeepFRI (Gligorijevi´c et al., 2021).
We also evaluate
DeepFRI (Gligorijevi ´ c et al., 2021) in our benchmark, which is a popular structure-based encoder for protein function prediction.
DeepFRI employs an LSTM model to extract residue features and further constructs a residue graph to propagate messages among residues, in which a 3-layer graph convolutional network (GCN) (Kipf & Welling, 2017) is used for message passing.
We directly utilize the official model checkpoint for baseline evaluation.
ESM-1b (Rives et al., 2021).
Besides the from-scratch sequence encoders above, we also compare with two state-ofthe-art pretrained protein language models.
ESM-1b (Rives et al., 2021) is a huge Transformer encoder model whose size is larger than BERT-Large (Devlin et al., 2018), and it is pretrained on 24 million protein sequences from
UniRef50 (Suzek et al., 2007) by masked language modeling (MLM) (Devlin et al., 2018).
In our evaluation, we finetune the ESM-1b model with the learning rate that is one-tenth of that of the MLP prediction head.
ProtBERT-BFD (Elnaggar et al., 2021).
The other protein language model evaluated in our benchmark is
ProtBERT-BFD (Elnaggar et al., 2021) whose size also excesses BERT-Large (Devlin et al., 2018).
This model is pretrained on 2.1 billion protein sequences from BFD (Steinegger & S ¨ oding, 2018) by MLM (Devlin et al., 2018).
The evaluation of ProtBERT-BFD uses the same learning rate configuration as ESM-1b.
LM-GVP (Wang et al., 2022).
To further enhance the effectiveness of GVP (Jing et al., 2021), (Wang et al., 2022) proposed to prepend a protein language model, i.e. ProtBERT (Elnaggar et al., 2021), before GVP to additionally utilize protein sequence representations.
We also adopt this hybrid model as one of our baselines, and its implementation follows the official source code.
Our methods.
For pretraining, we use Adam optimizer with learning rate 0.001 and train a model for 50 epochs.
Then, the pretrained model will be finetuned on downstream datasets.
For Multiview Contrast , we set the cropping length of subsequence operation as 50, the radius of subspace operation as 15 ˚ A , the mask rate of random edge masking operation as 0.15.
The temperature τ in the InfoNCE loss function is set as 0.07.
When pretraining GearNet-Edge and GearNetEdge-IEConv, we use 96 and 24 as batch sizes, respectively.
For Distance Prediction , we set the number of sampled residue pairs as 256.
And the batch size will be set asand 32 for GearNet-Edge and GearNet-Edge-IEConv, respectively.
For Residue Type , Angle and Dihedral Prediction , we set the number of sampled residues, residue triplets and residue quadrants as 512.
And the batch size will be set as 96 and 32 for GearNet-Edge and GearNet-Edge-IEConv, respectively.
For downstream evaluation, the hidden representations in each layer of GearNet will be concatenated for the final
[Table]
prediction.
Table 4 lists the hyperparameter configurations for different downstream tasks.
For the four tasks, we use the same optimizer and number of epochs as in the original papers to make fair comparison.
And for EC and GO prediction, we use ReduceLROnPlateau scheduler with factor 0 .and
patience 5 , while we use StepLR scheduler with step size 50 and gamma 0 .
5 for fold and reaction classification.
F. Additional Experimental Results F.1. AUPR on EC and GO Prediction
In Section 4, we have reported experimental results on EC and GO prediction with F max as the metric.
Here we report another popular metric AUPR in Table 5.
Note that we still use the best model selected by F max on validation sets.
It can be observed that our model can still achieve the best performance on EC prediction in both from scratch and pretrained settings.
However, there are still non-trivial gaps between our models with the state-of-the-art results.
This probably is because of the inconsistency between the two evaluation metrics.
It would be interesting to study the relationship between these two metrics and develop a model good at both.
We leave this problem as future works.
F.2. Experiments under Different Sequence Cutoffs
Besides the main experiments, where 95% is used as the sequence identity cutoff for EC and GO dataset splitting, we also test our models and several important baselines under four lower sequence identity cutoffs and show the experimental results in Table 6.
It can be observed that, at lower sequence identity cutoffs, our model can still achieve the best performance among models without pretraining and get comparable or better results against ESM-1b after pretraining.
F.3. Ablation Studies
To analyze the contribution of different components in our proposed methods, we perform ablation studies on the EC prediction task.
The results are shown in Table 7.
Relational graph convolutional layers.
To show the effects of relational convolutional layers, we replace it with graph convolutional layers that share a single kernel matrix among all edges.
As reported in the table, results can be significantly improved by using relational convolution, which suggests the importance of treating edges as different types.
Edge message passing layers.
We also compare the results of GearNet with and without edge message passing layers, the results of which are shown in Tables 1, 7.
It can be observed that the performance consistently increases after performing edge message passing.
This demonstrates the effectiveness of our proposed mechanism.
Different augmentations in Multiview Contrast.
We investigate the contribution of each augmentation operation proposed in the Multiview Contrast method.
Instead of randomly sampling cropping functions and transformations, we pretrain our model with four deterministic combinations of augmentations, respectively.
As shown in Table 7, all the four combinations can yield good results, which suggests that arbitrary combinations of the proposed cropping and transformation schemes can yield informative partial views of proteins.
Sampling schemes in Self-Prediction methods.
Different sampling schemes may lead to different results for self-prediction methods.
We study the effects of sampling schemes using Dihedral Prediction as an example.
Instead of sampling dihedral angles formed by three consecutive edges, we try to predict the dihedrals formed by four randomly sampled nodes.
We observe that this change of sampling schemes will make the self-prediction task more difficult to solve, which even brings negative effects after pretraining.
G. Pretraining on Different Datasets
We use the AlphaFold protein structure database as our pretraining database, as it contains the largest number of protein structures and is planned to cover over 100 million proteins in the future.
However, the structures in this database are not experimentally determined but predicted by AlphaFold2.
Therefore, it is interesting to see the performance of our methods when pretraining on different datasets.
To study the effects of the choice of pretraining dataset, we build another dataset using structures extracted from Protein
Data Bank (PDB) (Berman et al., 2000).
Specifically, we extract 123,505 experimentally-determined protein struc-
[Table]
[Table]
w / p re t r a i n i n g
[Table]
tures from PDB whose resolutions are between 0.0 and 2.
angstroms
[Table]
and we further extract 305,265 chains from these proteins to construct the final pretraining dataset
[Table]
Next
[Table]
we pretrain our five methods on AlphaFold Database v1 (proteome-wide structure predictions), AlphaFold
Database v2 (Swiss-Prot structure predictions) and Protein
Data Bank and then evaluate the pretrained models on the
EC prediction task.
The results are reported in Table 8.
As can be seen in the table, our methods can achieve comparable performance on different pretraining datasets.
Consequently, our methods are robust to the choice of pretraining datasets.
H. Results on Protein Engineering Tasks
Besides four standard tasks considered in Section 4, another important kind of downstream tasks is related to protein engineering, which is heavily relied on mutations on protein sequences.
These tasks aim to predict the ability of a protein to perform a desired function, termed protein fitness.
Good models are expected to have sufficient precision to distinguish between closely-related protein sequences upon mutations.
In this section, we further evaluate our model on four protein engineering related tasks.
H.1. Setup
We choose two protein engineering tasks from (Rao et al., 2019) and two landscape prediction tasks from (Dallago et al., 2021), all of which are standard benchmarks to evaluate protein language models.
The statistics of four datasets are shown in Table 9 and we describe each task as follows.
Fluoresence landscape prediction (Sarkisyan et al., 2016)
This task aims to predict the log-fluorescence intensity of mutants of the wild type green fluorescent protein (GFP).
The training set consists of single, double and triple mutants, while the test set includes variants with four or more mutations.
Stability landscape prediction (Rocklin et al., 2017)
This is a regression task that maps each input protein to a value measuring the most extreme circumstances in which the protein can maintain its fold above a concentration threshold.
To test the generalization ability from a broad set of relevant sequences with multiple mutations to local sequence neighbors with fewer mutations, the training set includes proteins from four rounds of experimental design, whereas the test set only contains 1-hop neighbors of top candidate proteins.
[Table]
GB1 (Wu et al., 2016)
This task uses the GB1 landscape to test the model’s ability to predict the effects of interactions between mutations, termed epistasis.
We adopt the low-vs-high split proposed in (Dallago et al., 2021), where sequences with fitness value equal to or below wild type are used to train, while sequences with fitness value above wild type are used to test.
Thermostability (Jarzab et al., 2020)
We use the screening landscape curated from (Dallago et al., 2021) to measure the model’s ability to predict thermostability of proteins.
This landscape includes both global and local variation instead of only mutants of a single protein.
Similarly, we use the low-vs-high split proposed in (Dallago et al., 2021).
All these four tasks are evaluated via Spearman’s ρ (rank correlation coefficient).
H.2. Implementation Details
Protein structure generation.
Since all these tasks are originally designed for evaluating sequence-based encoders, the experimentally-determined structures for proteins in these datasets are not available.
To solve this issue, we use
AlphaFold2 to generate structures for all datasets except
GB1.
For GB1 dataset, we only generate the structure of the wild type protein.
Because the differences between mutant structures are almost negligible on the residue level, we directly use the residue graph constructed from the wild type protein for all mutants and replace node and edge features with corresponding residue types after mutations.
We scale up the number of predictions that can be performed by AlphaFold2 by using the fast homology search of MMSeqs2 (Steinegger & S ¨ oding, 2017) and running AlphaFold in batch mode using ColabFold (Mirdita et al., 2021).
Five predictions were made for each sequence and prediction with the highest pLDDT score was chosen.
The number of recycles were set to 3 and relaxation using amber force fields was not used.
For those proteins the structures of which AlphaFold2 fails to generate, we only add sequential edges in the protein residue graph based on sequential information.
Training details.
We use the same pretraining and downstream setting as in Section E.4.
Hyperparameters for each dataset are described in Table 10.
H.3. Results
The results on four protein engineering tasks are reported in Table 11.
First, our model achieves the best result on all tasks among models without pretraining.
Moreover, it obtains comparable or even better performance than pretrained sequence encoders on Fluoresence and Stability tasks.
This can be understood because our model is a generalization of
CNN enhanced with structural information encoded in the protein graph and thus at least as good as CNN.
Surprisingly, the improvements of edge message passing aren’t significant on these four tasks.
One potential reason is that the structures of protein mutants are undistinguishable on the residue level and some structures are missing in the dataset.
As for pretraining, structure-based methods show significant improvements on Stability and GB1 datasets, which achieve the state-of-the-art performance.
However, no positive effects of pretraining are shown on the other two tasks.
Besides the reason mentioned above, it may also be attributed to the fact that sequential information plays a more important role in these tasks.
[Table]
[Table]
[Table]
Overall, our methods achieve the state-of-the-art performance on three of four datasets.
Nevertheless, it still needs further exploration and ablation studies on these mutationbased datasets, which should be considered in future work.
I. Latent Space Visualization
For qualitatively evaluating the quality of the protein embeddings learned by our pretraining method, we visualize the latent space of the GearNet-Edge model pretrained by
Multiview Contrast.
Specifically, we utilize the pretrained model to extract the embeddings of all the proteins in AlphaFold Database v1, and these embeddings are mapped to the two-dimensional space by UMAP (McInnes et al., 2018) for visualization.
Following (Akdel et al., 2021), we highlight the 20 most common superfamilies within the database by different colors.
The visualization results are shown in
Fig.
It can be observed that our pretrained model tends to group the proteins from the same superfamily together and divide the ones from different superfamilies apart.
In particular, it succeeds in clearly separating three superfamilies, i.e. , Protein kinase superfamily, Cytochrome P450 family and TRAFAC class myosin-kinesin ATPase superfamily.
Such a decent capability of discriminating protein superfamilies, to some degree, interprets our model’s superior performance on Fold Classification.
J. Residue-Level Explanation
Protein functions are often reflected by specific regions on the 3D protein structures.
For example, the binding ability of a protein to a ligand is highly related to the binding interface between them.
Hence, to better interpret our prediction, we apply Integrated Gradients (IG) (Sundararajan et al., 2017), a model-agnostic attribution method, on our model to obtain residue-level interpretation.
Specifically, we first select two molecular functions, ATP binding (GO:0005524) and Heme binding (GO:0020037), from GO terms that are related to ligand binding.
For each functional term, we pick one protein and feed it into the best model trained on the GO-MF dataset.
Then, we use IG to generate the feature attribution scores for each protein.
The method will integrate the gradient along a straight-line path between a baseline input and the original input.
Here the original input and baseline input are the node feature f and a zero vector, respectively.
The final attribution score for each protein will be obtained by summing over the feature dimension.
The normalized score distribution over all residues are visualized in Figure 6.
As can be seen, our model is able to identify the active sites around the ligand, which are likely to be responsible for binding.
Note that these attributions are directly generated from our model without any supervision, which suggests the decent interpretability of our model.
[Table]