OntoProtein: Protein Pretraining With Gene Ontology Embedding

1 INTRODUCTION . Protein science , the fundamental macromolecules governing biology and life itself , has led to remarkable advances in understanding the disease therapies and human health ( Vig et al . ( 2021 ) ) . As a sequence of amino acids , protein can be viewed precisely as a language , indicating that they may be modeled using neural networks that have been developed for natural language processing ( NLP ) . Recent self-supervised pre-trained protein language models ( PLMs ) such as ESM ( Rao et al . ( 2021b ) ) , ProteinBERT ( Brandes et al . ( 2021 ) ) , ProtTrans ( Elnaggar et al . ( 2020 ) ) which can learn powerful protein representations , have achieved promising results in understanding the structure and functionality of the protein . Yet existing PLMs for protein representation learning generally can not sufficiently capture the biology factual knowledge , which is crucial for many protein tasks but is usually sparse and has diverse and complex forms in sequence . By contrast , knowledge graphs ( KGs ) from gene ontology2 contain extensive biology structural facts , and knowledge embedding ( KE ) approaches ( Bordes et al . ( 2013 ) , Zheng et al . ( 2021 ) ) can efficiently embed them into continuous vectors of entities and relations . For example , as shown in Figure 1 , without knowing PEX5 has specific biological processes and cellular components , it ∗Equal contribution and shared co-first authorship . †Corresponding author . 1Code and datasets are available in https : //github.com/zjunlp/OntoProtein . 2http : //geneontology.org/ is challenging to recognize its interaction with other proteins . Furthermore , since protein ’ s shape determines its function , it is more convenient for models to identify protein ’ s functions with the prior knowledge of protein functions having similar shapes . Hence , considering rich knowledge can lead to better protein representation and benefits various biology applications , e.g. , protein contact prediction , protein function prediction , and protein-protein interaction prediction . However , different from knowledge-enhanced approaches in NLP ( Zhang et al . ( 2019b ) , Wang et al . ( 2021b ) , Wang et al . ( 2021a ) ) , protein sequence and gene ontology are two different types of data . Note that protein sequence is composed of amino acids while gene ontology is a knowledge graph with text description ; thus , severe issues of structured knowledge encoding and heterogeneous information fusion remain . In this paper , we take the first to propose protein pre-training with gene ontology embedding ( OntoProtein ) , which is the first general framework to integrate external knowledge graphs into protein pre-training . We propose a hybrid encoder to represent language text and protein sequence and introduce contrastive learning with knowledge-aware negative sampling to jointly optimize the knowledge graph and the protein sequence embedding during pre-training . For the KE objective , we encode the node descriptions ( go annotations ) as their corresponding entity embeddings and then optimize them following vanilla KE approaches ( Bordes et al . ( 2013 ) ) . We further leverage gene ontology of molecular function , cellular component , and biological process and introduce a knowledge-aware negative sampling method for the KE objective . For the MLM ( Mask Language Modeling ) objective , we follow the approach of existing protein pre-training approaches ( Rao et al . ( 2021b ) ) . OntoProtein has the following strengths : ( 1 ) OntoProtein inherits the strong ability of protein understanding from PLMs with the MLM object . ( 2 ) OntoProtein can integrate biology knowledge into protein representation with the supervision from KG by the KE object . ( 3 ) OntoProtein constitutes a model-agnostic method and is readily pluggable into a wide range of protein tasks without additional inference overhead since we do not modify model architecture but add new training objectives . For pre-training and evaluating OntoProtein , we need a knowledge graph with large-scale biology knowledge facts aligned with protein sequences . Therefore , we construct ProteinKG25 , which contains about 612,483 entities , 4,990,097 triples , and aligned node descriptions from GO annotations . To the best of our knowledge , it is the first large-scale KG dataset to facilitate protein pre-training . We deliver data splits for both the inductive and the transductive settings to promote future research . To summarize , our contribution is three-fold : ( 1 ) We propose OntoProtein , the first knowledgeenhanced protein pre-training approach that brings promising improvements to a wide range of protein tasks . ( 2 ) By contrastive learning with knowledge-aware sampling to jointly optimize knowledge and protein embedding , OntoProtein shows its effectiveness in widespread downstream tasks , including protein function prediction , protein-protein interaction prediction , contact prediction , and so on . ( 3 ) We construct and release the ProteinKG25 , a novel large-scale KG dataset , promoting the research on protein language pre-training . ( 4 ) We conduct extensive experiments in widespread protein tasks , including TAPE benchmark , protein-protein interaction prediction , and protein function prediction , which demonstrate the effectiveness of our proposed approach . 2 METHODOLOGIES . We begin to introduce our approach of protein pre-training with ontology embedding ( OntoProtein ) , as shown in Figure 2 . OntoProtein incorporates external knowledge from Gene Ontology ( Go ) into language representations by jointly optimizing two objectives . We will first introduce the hybrid encoder , masked protein modeling , and knowledge encoder , and then we will present the details of contrastive learning with knowledge-aware negative sampling . Finally , we will illustrate the overall pre-training objects . 2.1 HYBRID ENCODER . We first introduce the hybrid encoder to represent protein and GO knowledge . For the protein encoder , we use the pre-trained ProtBert from Elnaggar et al . ( 2020 ) . ProtBert is pre-trained using the BERT architecture with UniRef100 datasets . Compared to BERT Devlin et al . ( 2019 ) , ProtBert encodes amino acid sequences into token level or sentence level representations , which can be used for downstream protein tasks such as contacts prediction tasks . The encoder takes a protein sequence ofN tokens ( x1 , ... , xN ) as inputs , and computes contextualized amnio acid representationHiProtein and sequence representationHProtein via mean pooling . To bridge the gap between text and protein , we utilize affine transformation ( an extra linear layer ) to project those representation to the same space . We will discuss details of learning protein representation in Section Mask Protein Modeling . For the Go encoder , we leverage BERT ( Devlin et al . ( 2019 ) ) , a Transformer ( Vaswani et al . ( 2017 ) ) based text encoder for biological descriptions in Gene Ontology entities . The encoder takes a sequence of N tokens ( x1 , ... , xN ) as inputs , and computes contextualized representations HiGO ∈ RN×d through L layers of Transformer encoders . Finally , we obtain contextualized representations HGO for each sequence by pooling output using [ CLS ] token representation . Since the relations in Gene Ontology are important for representing the knowledge of biology features , thus , we utilize a relation encoder with the random initialization , and those embeddings of relations will be optimized and updated during pre-training . 2.2 KNOWLEDGE EMBEDDING . We leverage the knowledge embedding ( KE ) objective to obtain representations in the pre-training process since Gene Ontology is actually a factual knowledge graph . Similar to Bordes et al . ( 2013 ) , we use distributed representations to encode entities and relations . The knowledge graph here consists of lots of triples to describe relational facts . We define a triplet as ( h , r , t ) , where h and t are head and tail entities , r is the relation whose type usually is pre-defined in the schema3 . Note that there are two different types of nodes eGO and eprotein in our knowledge graph . eGO is denoted as nodes that exist in the gene ontology , such as molecular function or cellular component nodes , and eGO can be described by annotation texts . eprotein is the protein node that links to the gene ontology , and we also represent eprotein with amnio acids sequences . Concretely , the triplets in this knowledge graph can be divided into two groups , tripleGO2GO and tripleProtein2GO . To integrate multi-modal descriptions into the same semantic space and address the heterogeneous information fusion issue , we utilize hybrid encoders introduced in the previous Section . Note that protein encoder and GO encoder represent protein sequence and GO annotations separately . 2.3 MASKED PROTEIN MODELING . We use masked protein modeling to optimize protein representations . The masked protein modeling is similar to masked language modeling ( MLM ) . During model pre-training , we use a 15 % probability to mask each token ( amino acid ) and leverage a cross-entropy loss ` MLM to estimate these masked tokens . We initialize our model with the pre-trained model of ProtBert and regard ` MLM as one of the overall objectives of OntoProtein by jointly training KE ( knowledge embedding ) and MLM . Our approach is model-agnostic , and other pre-trained models can also be leveraged . 2.4 CONTRASTIVE LEARNING WITH KNOWLEDGE-AWARE NEGATIVE SAMPLING . Knowledge embedding ( KE ) is to learn low-dimensional representations for entities and relations , and contrastive estimation represents a scalable and effective method for inferring connectivity patterns . Note that a crucial aspect of contrastive learning approaches is the choice of corruption distribution that generates hard negative samples , which force the embedding model to learn discriminative representations and find critical characteristics of observed data . However , previous approaches either employ too simple corruption distributions , i.e. , uniform , yielding easy uninformative negatives , or sophisticated adversarial distributions with challenging optimization schemes . Thus , in this paper , we propose contrastive learning with knowledge-aware negative sampling , an inexpensive negative sampling strategy that utilizes the rich GO knowledge to sample negative samples . Formally , the KE objective can be defined as : ` KE = − log σ ( γ − d ( h , t ) ) − n∑ i=1 1 n log σ ( d ( h′i , t ′ i ) − γ ) ( 1 ) ( h′i , t ′ i ) is the negative sample , in which head or tail entities are random sampled to construct the corrupt triples . n is the number of negative samples , σ is the sigmoid function , and γ means the margin . d is the scoring function , and we use TransE ( Bordes et al . ( 2013 ) ) for simplicity , where dr ( h , t ) = ‖h+ r − t‖ ( 2 ) Specifically , we define triple sets and entity sets as T and E , all triplets are divided into two groups . If the head entity is protein node and the tail entity is GO node , we denote the triple as Tprotein−GO . Similarly , if head and tail entities are both GO nodes , we denote them as TGO−GO . As Gene Ontology describes the knowledge of the biological domain concerning three aspects , all entities in Gene Ontology belong to MFO ( Molecular Function ) , CCO ( Cellular Component ) , or BPO ( Biological Process ) . To avoid plain negative samples , for those TGO−GO triples , we sample triples by replacing entities with the same aspect ( MFO , CCO , BPO ) 4 . Finally , we define the negative triple sets T ′ and positive 3The schema of the knowledge graph can be found in Appendix A.1 4For Tprotein−GO triples , it is also intuitive to replace the proteins with their homologous proteins to gen- erate hard negative triples , and we leave this for future works . triple as ( h , r , t ) , the negative sampling process can be described as follows : T ′ GO−GO ( h , r , t ) = { ( h ′ , r , t ) | h′ ∈ E′ , h ∈ E′ } ∪ { ( h , r , t′ ) | t′ ∈ E′ , t ∈ E′ } T ′ Protein−GO ( h , r , t ) = { ( h , r , t ′ ) | t′ ∈ E′ } ( 3 ) where E′ ∈ { EMFO , ECCO , EBPO } , and we only replace the tail entities for TProtein−GO triples .

Summary: This paper argues that informative biology knowledge in KGs can enhance protein representation with external knowledge. To show this, the author proposes a computational framework named OntoProtein that makes use of structure in Gene Ontology into protein pre-training models. This paper conducts various experiments to show the superiority of the proposed method and the benefits of the information of Gene Ontology.