Contrastive Representation Learning for 3D Protein Structures

1 INTRODUCTION . In recent years , learning on 3D protein structures has gained a lot of attention in the fields of protein modeling and structural bioinformatics . These neural network architectures process the 3D position of the atoms and/or amino acids in 3D space in order to make predictions of unprecedented performance , in tasks ranging from protein design ( Ingraham et al. , 2019 ; Strokach et al. , 2020 ; Jing et al. , 2021 ) , over protein structure classification ( Hermosilla et al. , 2021 ) , protein quality assessment ( Baldassarre et al. , 2020 ; Derevyanko et al. , 2018 ) , and protein function prediction ( Amidi et al. , 2017 ; Gligorijevic et al. , 2021 ) – just to name a few . Unfortunately , learning on the structure of proteins leads to a reduced amount of training data , as compared for example to sequence learning , since 3D structures are harder to obtain and thus less prevalent . While the Protein Data Bank ( PDB ) ( Berman et al. , 2000 ) today contains only around 182K macromolecular structures , the Pfam database ( Mistry et al. , 2020 ) contains 47M protein sequences . Naturally , the number of available structures decreases even further when only the structures labeled with a specific property are considered . We refer to these as annotated protein structures . The SIFTS database , for example , contains around 220K annotated enzymes from 96K different PDB entries , and the SCOPe database contains 226 K annotated structures . These numbers are orders of magnitude lower than the data set sizes which led to the major breakthroughs in the field of deep learning . ImageNet ( Russakovsky et al. , 2015 ) , for instance , contains more than 10M annotated images . As learning on 3D protein structures can not benefit from these large amounts of data , model sizes are limited or overfitting might occur . In order , to take advantage of unlabeled data , researchers have , over the years , designed different algorithms , that are able to learn meaningful representations from such data without labels ( Hadsell et al. , 2006 ; Ye et al. , 2019 ; Chen et al. , 2020a ) . In natural language processing , next token prediction or random token masking are commonly used unsupervised training objectives , that are able to learn meaningful word representations useful for different downstream tasks ( Peters et al. , 2018 ; Devlin et al. , 2019 ) . Recently , such algorithms have been used to learn meaningful protein representations from unlabeled sequences ( Alley et al. , 2019 ) , or as a pre-trained method for later fine-tuning models on different downstream tasks ( Rao et al. , 2019 ) . In computer vision recently , contrastive learning has shown great performance on image classification when used to pre-train deep convolutional neural network ( CNN ) architectures ( Chen et al. , 2020a ; b ) . This pre-training objective has also been used in the context of protein sequence representation learning by dividing sequences in amino acid ’ patches ’ ( Lu et al. , 2020b ) , or by using data augmentation techniques based on protein evolutionary information ( Lu et al. , 2020a ) . Most recently , the contrastive learning framework has been applied to graph convolutional neural networks ( You et al. , 2020 ) . These techniques were tested on protein spatial neighboring graphs ( graphs where edges connect neighbor amino acids in 3D space ) for the binary task of classifying a protein as enzyme or not . However , these algorithms were designed for arbitrary graphs and did not take into account the underlying structure of proteins . In this work , we introduce a contrastive learning framework for representation learning of 3D protein structures . For each unlabeled protein chain , we select random molecular sub-structures during training . We then minimize the cosine distance between the learned representations of the sub-structures sampled from the same protein , while maximizing the cosine distance between representations from different protein chains . This training objective enables us , to pre-train models on all available annotated , but more importantly also unlabeled , protein structures from the PDB . The obtained representation can later be used as a weight initialization strategy to improve performance on different downstream tasks , such as protein structure classification and function prediction . Moreover , we show how the learned protein representation is able to capture protein structural similarity and functionality , by embedding proteins from the same fold or with similar functions close together in this space . The remainder of this paper is structured as follows . First , we provide a summary of the state-of-theart in Section 2 . Then , we introduce our framework in Section 3.Later , in Section 4 , we describe the experiments conducted to evaluate our framework and the representations learned , and lastly , we provide a summary of our findings and possible lines of future research in Section 5 . 2 RELATED WORK . 3D protein structure learning . Early work on learning from 3D protein structures used graph kernels and support vector machines to classify enzymes ( Borgwardt et al. , 2005 ) . Later , the advances in the fields of machine learning and computer vision brought a new set of techniques to the field . Several authors represent the protein tertiary structure as a 3D density map , and process it with a 3D convolutional neural network ( 3DCNN ) . Among the problems addressed with this technique , are protein quality assessment ( Derevyanko et al. , 2018 ) , protein enzyme classification ( Amidi et al. , 2017 ) , protein-ligand binding affinity ( Ragoza et al. , 2017 ) , protein binding site prediction ( Jiménez et al. , 2017 ) and protein-protein interaction interface prediction ( Townshend et al. , 2019 ) . Other authors have used graph convolutional neural networks ( GCNN ) to learn directly from the protein spatial neighboring graph . Some of the tasks solved with these techniques , are protein interface prediction ( Fout et al. , 2017 ) , function prediction ( Gligorijevic et al. , 2021 ) , protein quality assessment ( Baldassarre et al. , 2020 ) , and protein design ( Strokach et al. , 2020 ) . Recently , several neural network architectures , specifically designed for protein structures , have been proposed to tackle protein design challenges ( Ingraham et al. , 2019 ; Jing et al. , 2021 ) , or protein fold and function prediction ( Hermosilla et al. , 2021 ) . Protein representation learning . Protein representation learning based on protein sequences is an active area of research . Early works used similar techniques as the ones used in natural language processing to compute embeddings of groups of neighboring amino acids in a sequence ( Asgari & Mofrad , 2015 ) . Recently , other works have used unsupervised learning algorithms from natural language processing such as token masking or next token prediction ( Peters et al. , 2018 ; Devlin et al. , 2019 ) to learn representations from protein sequences ( Alley et al. , 2019 ; Rao et al. , 2019 ; Min et al. , 2020 ; Strodthoff et al. , 2020 ) . This year , Lu et al . ( 2020b ; a ) have suggested using contrastive learning on protein sequences , to obtain a meaningful protein representation . Despite the advances in representation learning for protein sequences , representation learning for 3D protein structures mostly has relied on hand-crafted features . La et al . ( 2009 ) proposed a method to compute a vector of 3D Zernike descriptors to represent protein surfaces , which later can be used for shape retrieval . Recently , Guzenko et al . ( 2020 ) used a similar approach , to compute a vector of 3D Zernike descriptors directly from the 3D density volume , which can be used later for protein shape comparison . The annual shape retrieval contest ( SHREC ) usually contains a protein shape retrieval track , in which methods are required to determine protein similarity from different protein surfaces ( Langenfeld et al. , 2019 ; 2020 ) . Some of the works presented here , make use of 3DCNNs or GCNNs to achieve this goal . However , they operate on protein surfaces , and are either trained in a supervised fashion on the binary shape similarity task , or pre-trained on a classification task . Contrastive learning . In 1992 , Becker & Hinton ( 1992 ) suggested training neural networks through the agreement between representations of the same image under different transformations . Later , Hadsell et al . ( 2006 ) proposed to learn image representations by minimizing the distance between positive pairs and maximizing the distance between negative pairs ( see Figure 1 ) . This idea was used in other works by sampling negative pairs from the mini-batches used during training ( Ji et al. , 2019 ; Ye et al. , 2019 ) . Recently , Chen et al . ( 2020a ; b ) have shown how these methods can improve image classification performance . You et al . ( 2020 ) have transferred these ideas to graphs , by proposing four different data transformations to be used during training : node dropping , edge perturbation , attribute masking , and subgraph sampling . These ideas were tested on the commonly used graph benchmark PROTEINS ( Borgwardt et al. , 2005 ) , composed of only 1 , 113 proteins . However , since this data set is composed of spatial neighboring graphs of secondary structures , the proposed data augmentation techniques can generate graphs of unconnected chain sections . In this paper instead , we suggest using a domain-specific transformation strategy , that preserves the local information of protein sequences . 3 3D PROTEIN CONTRASTIVE LEARNING . 3.1 PROTEIN GRAPH . In this work , the protein chain is defined as a graph G = ( N , R , F , A , B ) , where each node represents the alpha carbon of an amino acid with its 3D coordinates , N ∈ Rn×3 , being n the number of amino acids in the protein . Moreover , for each node , we store a local frame composed of three orthonormal vectors describing the orientation of the amino acid wrt . the protein backbone , R ∈ Rn×3×3 . Lastly , each node has also t different associated features with it , F ∈ Rn×t . The connectivity of the graph is stored in two different adjacency matrices , A ∈ Rn×n and B ∈ Rn×n . Matrix A is defined as Aij = 1 if amino acids i and j are connected by a peptide bond and Aij = 0 otherwise . Matrix B is defined as Bij = 1 if amino acids i and j are at a distance smaller than d in 3D space and Bij = 0 otherwise . 3.2 CONTRASTIVE LEARNING FRAMEWORK Inspired by recent works in the computer vision domain ( Ye et al. , 2019 ; Ji et al. , 2019 ; Chen et al. , 2020a ) , our framework is trained by maximizing the similarity between representations from sub-structures of the same protein , and minimizing the similarity between sub-structures from different proteins . More formally , given a protein graph G , we sample two sub-structures Gi and Gj from it . We then compute the latent representations of these substructures , hi and hj , using a protein graph encoder , hi = E ( Gi ) . Based on the findings of Chen et al . ( 2020a ) , we further project these latent representations into smaller latent representation , zi and zj , using a multilayer perceptron ( MLP ) with a single hidden layer , zi = P ( hi ) . Lastly , the similarity between these representa- tions is computed using the cosine distance , s ( zi , zj ) . In order to minimize the similarity between these representations , we use the following loss function for the sub-structure Gi : li = −log exp ( s ( zi , zj ) /τ ) ∑2N k=1 1 [ k 6=i , k 6=j ] exp ( s ( zi , zk ) /τ ) ( 1 ) where τ is a temperature parameter used to improve learning from ’ hard ’ examples , 1 [ k 6=i , k 6=j ] ∈ [ 0 , 1 ] is a function that evaluates to 1 if k 6= i and k 6= j , and N is the number of protein structures in the current mini-batch . To compute lj we use again Equation 1 , but exchange the role of i and j . This loss has been used before in the context of representation learning ( Chen et al. , 2020a ; van den Oord et al. , 2018 ) , and as in previous work , our framework does not explicitly sample negative examples but uses instead the rest of sub-structures sampled from different proteins in the mini-batch as negative examples . In the following subsections , we will describe the different components specific to our framework designed to process protein structures .

Summary: This paper presents an unsupervised deep learning approach for learning representations of 3D protein structures. They use an objective function motivated by the recent contrastive learning approaches (from computer vision). They show the utility of their model in two downstream applications: protein fold classification, enzyme classification and protein similarity prediction. The main contributions are: 1. they show how the contrastive learning framework can be used in the context of protein structures 2. they analyze the learned representations 3. they show the utility of the model to downstream applications 4. they show how fine-tuning their model leads to an improved performance