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 introduces a method to enrich the representations that are learnt by protein language models with knowledge encapsulated in gene ontologies. To do so, it curates a knowledge graph (ProteinKG25) and applies existing methods in multi-relational data embedding (Bordes et al.) to jointly train knowledge embeddings and protein embeddings, with the objective to enhance the latter and subsequently improve the performance on various downstream tasks (namely TAPE benchmark, protein-protein interaction and protein function prediction).