Co-evolution Transformer for Protein Contact Prediction

Proteins are the main machinery of life and protein functions are largely determined by their 3D structures . The measurement of the pairwise proximity between amino acids of a protein , known as inter-residue contact map , well characterizes the structural information of a protein . Protein contact prediction ( PCP ) is an essential building block of many protein structure related applications . The prevalent approach to contact prediction is based on estimating the inter-residue contacts using hand-crafted coevolutionary features derived from multiple sequence alignments ( MSAs ) . To mitigate the information loss caused by hand-crafted features , some recently proposed methods try to learn residue co-evolutions directly from MSAs . These methods generally derive coevolutionary features by aggregating the learned residue representations from individual sequences with equal weights , which is inconsistent with the premise that residue co-evolutions are a reflection of collective covariation patterns of numerous homologous proteins . Moreover , non-homologous residues and gaps commonly exist in MSAs . By aggregating features from all homologs equally , the non-homologous information may cause misestimation of the residue co-evolutions . To overcome these issues , we propose an attention-based architecture , Co-evolution Transformer ( CoT ) , for PCP . CoT jointly considers the information from all homologous sequences in the MSA to better capture global coevolutionary patterns . To mitigate the influence of the nonhomologous information , CoT selectively aggregates the features from different homologs by assigning smaller weights to non-homologous sequences or residue pairs . Extensive experiments on two rigorous benchmark datasets demonstrate the effectiveness of CoT . In particular , CoT achieves a 51.6 % top-L long-range precision score for the Free Modeling ( FM ) domains on the CASP14 benchmark , which outperforms the winner group of CASP14 contact prediction challenge by 9.8 % . ∗Work done during an internship at Microsoft Research . †Equal contribution . ‡Corresponding authors . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . 1 Introduction . Most in vivo biological processes are carried out by proteins , whose functions are mainly determined by their 3D structures [ 1 ] . Structural information is crucial for understanding the functions of a protein . The inter-residue contact map which measures the pairwise proximity between all amino acid pairs well characterizes the structural information of a protein . Protein contact prediction ( PCP ) is an important building block of many protein structure related applications , such as protein structure prediction [ 2 , 3 ] , protein complex assembly [ 4 ] , and protein design [ 5 ] . It is so important that it is one of the challenges of Critical Assessment of protein Structure Prediction ( CASP ) , which is the “ world championship ” of the computational structural biology field . The prevalent PCP approaches were built atop the co-evolution principle that the spatially proximate residues tend to co-evolve to maintain the functions of a protein [ 6–8 ] . The existing PCP methods generally infer pairwise coevolutionary patterns from MSAs . Among them , direct coupling analysis ( DCA ) [ 9 ] techniques are widely used to obtain coevolutionary features by fitting Potts models or calculating precision matrix [ 10 , 11 ] . Subsequently , a variety of deep-learning based methods were proposed to leverage the DCA-based features to estimate inter-residue contacts [ 3 , 12–15 ] . However , DCA techniques only consider single-residue and pairwise statistics of the sequences , ignoring high-order interactions among the residues within a sequence . To mitigate the information loss caused by hand-crafted features ( e.g. , DCA-based features ) , some recently proposed methods try to learn residue co-evolutions directly from MSAs [ 2 , 16 , 17 ] . For example , the SOTA of them , CopulaNet [ 2 ] learns residue representations from individual sequences and then aggregates these features with equal weights to derive residue co-evolutions . This causes two issues : 1 ) Inferring residue representations from individual sequences independently is inconsistent with the premise that residue co-evolutions are a reflection of collective covariation patterns of numerous homologs [ 18 ] ; 2 ) Assigning equal weights to the features coming from different homologs ignores the fact that , the quality of homologs varies a lot because of the existence of non-homologous residues and gaps [ 19 , 20 ] . In this paper , we propose an attention-based architecture , Co-evolution Transformer ( CoT ) , to address or mitigate the issues discussed above . The core component of CoT is the co-evolution attention ( CoA ) module . Different from the previous methods that focus on extracting information from individual sequences , CoA is capable of incorporating the residue co-evolution patterns derived from all homologous sequences into an attention function to learn residue representations . Moreover , the CoA learns to automatically weight the residue representations learned from different homologs , and then selectively aggregates the features to construct the co-evolution attention map . This design mitigates the influence of non-homologous information . CoT dramatically outperforms the baseline methods on two rigorous benchmarks CASP14 [ 21 ] and CAMEO ( Continuous Automated Model EvaluatiOn ) [ 22 ] . In particular , CoT achieves a 51.6 % top-L long-range precision score for the Free Modeling ( FM ) domains on the CASP14 benchmark , which outperforms the winner group of CASP14 contact prediction challenge by 9.8 % . Our code will be released at https : //github.com/microsoft/ProteinFolding/tree/main/ coevolution_transformer . 2 Related Work . Protein contact prediction is a binary classification task for amino acid residue pairs . A residue pair is called a contact if their distance is less than or equal to a distance threshold , typically 8Å , i.e. , 8× 10−10m . As widely acknowledged , co-evolution information is closely correlated to the contacts . In order to extract the co-evolution patterns for residue pairs , multiple sequence alignments ( MSAs ) are generated from raw protein sequences . For a target protein sequence , a generated alignment consists of multiple sequences with each corresponding position being an aligned residue or a gap ( annotated by a dash ) , and these aligned sequences as a whole are called a multiple sequence alignment . Many protein databases as well as various search schemes have been proposed to generate MSAs efficiently [ 19 , 20 ] . There are roughly three categories of methods for predicting contacts from MSAs , namely , unsupervised methods , supervised methods , and pre-training based methods . Unsupervised Methods To quantify the strength of the direct relation between the residue pairs of a protein sequence while excluding effects from other residues , many statistical modeling methods based on direct coupling analysis ( DCA ) have been proposed to fit Potts models [ 10 ] or precision matrix [ 11 ] to MSAs , e.g. , mean-field DCA [ 9 ] , sparse inverse covariance [ 11 ] and pseudo-likelihood maximization [ 10 , 23 , 24 ] , to name a few . These methods further exploit dedicated scores based on DCA for contact prediction . However , they only consider single-residue and pairwise statistics of the sequences , thus fail to capture high-order interactions among the multiple residues within a sequence . Supervised Methods By taking DCA-derived scores as features , deep neural networks based supervised methods significantly outperform the unsupervised methods [ 3 , 12–15 ] . However , the information lost by the DCA-based features from the sequences is still not recoverable . To mitigate this issue , several models are proposed to learn residue co-evolution information directly from the sequences in the MSAs [ 2 , 16 , 17 , 25 ] . Among them , CopulaNet [ 2 ] , the SOTA of the CASP13 benchmark , derives coevolutionary features differentially by aggregating the learned residue representations from the sequences . However , although they are capable of modeling the high-order interactions among the multiple residues within single sequence , the global information carried by the MSAs is ignored because they still model the protein sequences independently . AlphaFold2 [ 25 ] is claimed to be modeling the full MSAs , achieving an amazing performance on the CASP14 structure prediction task . Although its performance on the structure prediction task is remarkable , they did not participate in the CASP14 PCP task meanwhile no further details of their methods are publicly available . Pre-training Based Methods Following the pre-train and fine-tune paradigm , pre-trained language models are adapted to representation learning for single protein sequences from the unlabeled data [ 26–31 ] . Many of them take contact prediction as an important downstream task to validate their performance . While these methods show another solution to this task , they are still at an early stage thus can not achieve comparable performance to the SOTA approaches currently . To further improve the performance , a pre-trained language model named MSA Transformer is proposed to learn a better MSA representation directly [ 29 ] . MSA Transformer did solid work on learning coevolution information from unlabeled MSAs . However , non-homologous subsequences are inevitably introduced during the learning process . To exploit the co-evolution information from the full MSAs effectively , our proposed CoT model is built upon co-evolution attention , a novel attention mechanism dedicated to incorporate the coevolution information directly from MSAs in a supervised way . 3 Co-evolution Transformer . Co-evolution Transformer ( CoT ) is constructed by stacking several repeated CoT layers . Each CoT layer is composed of two attention modules , i.e. , a co-evolution attention ( CoA ) module and a self-attention module , as shown in Figure 1 . Given a prepared MSA , the stacked CoT layers are used to learn the residue representations . Residue co-evolutions are derived from the representations of the final layer and further employed to estimate the inter-residue contacts . In this section , we first start with a brief introduction to the vanilla Transformer Encoder , followed by the detailed descriptions of each CoT component . To better illustrate the co-evolution attention mechanism , the proposed co-evolution attention mechanism and the self-attention mechanism are discussed in the end . 3.1 Vanilla Transformer Encoder . Transformer is a network architecture built on the attention mechanism for machine translation . The Transformer encoder is widely adopted for machine learning tasks due to its excellent feature extraction capability when modeling long-distance interactions in sequences [ 32–34 ] . Each Transformer encoder layer consists of two modules , i.e. , a multi-head self-attention ( MHSA ) module and a position-wise fully connected feed-forward ( FFN ) module . To connect these two modules , residual connections [ 35 ] and layer normalization ( LAYERNORM ) [ 36 ] are applied as below : x = LAYERNORM ( x+ MHSA ( x ) ) , ( 1 ) x = LAYERNORM ( x+ FFN ( x ) ) , ( 2 ) The multi-head self-attention module is also adopted by the CoT layer as a component to learn the residue representations of different sequences in the MSA , however , iterating through the columns instead of rows/sequences . 3.2 Co-evolution Attention Module . For a target protein sequence and its MSA , the goal of the Co-evolution Attention Module is to leverage the whole MSA to derive pairwise inter-residue interaction features , namely residue co-evolutions . Intuitively , the residue co-evolutions are analogous to the covariance matrix in mathematics , depicting the correlations between any two residue positions ( two columns in the MSA ) . Then , these features are used as attention maps to guide representation learning of each sequence in the MSA . To achieve this goal , a CoA module employs two consequent submodules , i.e . co-evolution aggregation and co-evolution enhancement . The co-evolution aggregation submodule is designed to generate the coevolutionary features by aggregating pairwise interactions from all homologs in the MSA , while the co-evolution enhancement submodule further enhances the coevolutionary features and derives the final co-evolution attention . The overview of the CoA module is illustrated in Figure 2 . Given the target protein sequence ( r1 , r2 . . . rL ) , where L is the length of the sequence , the corresponding representation of its MSA is denoted as X ∈ RK×L×dmodel , where K is the number of homologous sequences and dmodel is the hidden dimension of the residue representation . A detailed representation schema of the MSA is described in Appendix . For the k-th sequence in the MSA , the overall procedure of the CoA module in each layer can be summarized as follows : Xk = LAYERNORM ( Xk + COATTN ( X ) ) , ( 3 ) Xk = LAYERNORM ( Xk + FFN ( Xk ) ) ( 4 ) where k ∈ [ 1 . . .K ] and COATTN is defined as : COATTN ( X ) = CONCAT ( head1 , . . . , headH ) , ( 5 ) headh = ATTNh ( X , A ) Xkh Wh , ( 6 ) where h ∈ [ 1 . . . H ] is the head index , A ∈ RL×L×dco is the co-evolution feature map generated with dco as its feature dimension , ATTNh ( X , A ) ∈ RL×L is the h-th co-evolution attention head , Xkh ∈ RL×dv is the h-th segment of Xk , and Wh ∈ Rdv×dv are learnable weights ( dv = dmodel/H ) . Note that the co-evolution attention ATTNh ( X , A ) is shared by all the K homologous sequences within the same head . Submodule 1 : Co-evolution Aggregation To depict inter-residue interactions , the co-evolution aggregation submodule exploits outer product on the residue representations . A selective pooling operation is then applied to the aggregated outer products from all homologous sequences , the disturbance of non-homologous information can be greatly reduced by this weighting mechanism . Two tensors P , Q ∈ RK×L×dmodel , are generated from X by separate linear projections . For the residue pair 〈ri , rj〉 in the k-th sequence of the MSA , the pair co-evolution features Aij are calculated by : Aij = PROJ ( K∑ k=1 Skij ( P ki ⊗ P kj ) ) , ( 7 ) and Skij = 1 Z exp ( Qki ⊗Qkj ) , ( 8 ) where ⊗ is the outer product operator , is the element-wise multiplication operator , PROJ is a flatten operation followed by a linear projection , converting the aggregated pair co-evolution features to A ∈ RL×L×dco , and Z stands for the normalization factor . The weight for the pair representation is denoted as Skij , where S k ij ∈ Rdmodel×dmodel whose elements fall in the range of [ 0 , 1 ] . The weights are normalized over the k sequences , thus we have ∑ k S k ij = 1 . Submodule 2 : Co-evolution Enhancement To model the high-order interactions among multiple residues , the co-evolution map A is concatenated with A′ from the previous CoT layer , the values are fed into a convolutional module ( CONV ) to generate an enhanced co-evolution map by : A = CONV ( [ A′ ; A ] ) . ( 9 ) In practice , we adopt a ResNet [ 35 ] as the convolution module . Note that we simply assign A′ an all-zero tensor for the first layer . The h-th co-evolution attention head is projected from A by ATTNh ( X , A ) = SOFTMAX ( AMh ) , ( 10 ) where Mh ∈ Rdco×1 are learnable weights .

Summary: The authors present a model that predicts contact maps from multiple sequence alignments. The paper proposes a weighted sum of k, LxL attention maps (k=#sequences, L=sequence length), where a per-residue-pair weight is learned across all sequences in the MSA. Afterwards a convolutional layer is applied to refine the prediction. The authors show that the co-evolution transformer outperforms existing supervised contact prediction methods on CASP14.