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 studies unsupervised contrastive learning on protein structures, using sub-structure sampling as the data transformation strategy in contrastive learning. The protein structure representation from contrastive learning is then evaluated for three tasks: fold classification, enzyme classification, and protein similarity.  The main contributions are: 1) Demonstrating that contrastive learning with sub-structure sampling is a viable strategy for learning protein structure representations. 2) Improving empirical results for fold classification and enzyme classification when compared to existing methods.