[Paper Review] OntoProtein: Protein Pretraining With Gene Ontology Embedding

Ningyu Zhang, Zhen Bi+7 more|arXiv (Cornell University)|Jan 23, 2022

Machine Learning in Bioinformatics49 citations

TL;DR

OntoProtein integrates Gene Ontology knowledge into protein pre-training by jointly optimizing a protein MLM objective with knowledge embedding, using a ProteinKG25 knowledge graph, to improve protein function prediction, PPI, and TAPE tasks.

ABSTRACT

Self-supervised protein language models have proved their effectiveness in learning the proteins representations. With the increasing computational power, current protein language models pre-trained with millions of diverse sequences can advance the parameter scale from million-level to billion-level and achieve remarkable improvement. However, those prevailing approaches rarely consider incorporating knowledge graphs (KGs), which can provide rich structured knowledge facts for better protein representations. We argue that informative biology knowledge in KGs can enhance protein representation with external knowledge. In this work, we propose OntoProtein, the first general framework that makes use of structure in GO (Gene Ontology) into protein pre-training models. We construct a novel large-scale knowledge graph that consists of GO and its related proteins, and gene annotation texts or protein sequences describe all nodes in the graph. We propose novel contrastive learning with knowledge-aware negative sampling to jointly optimize the knowledge graph and protein embedding during pre-training. Experimental results show that OntoProtein can surpass state-of-the-art methods with pre-trained protein language models in TAPE benchmark and yield better performance compared with baselines in protein-protein interaction and protein function prediction. Code and datasets are available in https://github.com/zjunlp/OntoProtein.

Motivation & Objective

• Incorporate external biology knowledge from Gene Ontology into protein pre-training to enrich protein representations.
• Develop a framework that jointly optimizes protein sequence modeling and knowledge graph embeddings.
• Create and release ProteinKG25, a large-scale KG aligned with protein sequences and GO descriptions for pre-training.
• Demonstrate improvements on downstream protein tasks such as function prediction, PPI, and TAPE benchmarks.

Proposed method

• Use a hybrid encoder: ProtBert for proteins and PubMedBERT-based GO descriptions, with an additional relation encoder for GO relations.
• Employ masked protein modeling (MLM) to predict masked amino acids, initialized from ProtBert.
• Represent the KG with entities (proteins and GO terms) and relations; apply a knowledge embedding objective (KE) using a TransE-style scoring function d(h,t)=||h+r−t||.
• Introduce knowledge-aware negative sampling to construct hard negatives for KE, including GO-GO negatives by replacing leaves within the same GO aspect and Protein-GO negatives by replacing proteins with homologs (future work).
• Jointly optimize KE and MLM with a weighted objective: ell = alpha * ell_KE + ell_MLM.

Experimental results

Research questions

• RQ1 Does incorporating GO-based knowledge graphs into protein pre-training improve downstream protein understanding tasks beyond standard PLMs?
• RQ2 Can knowledge embedding with GO relations and knowledge-aware negative sampling provide complementary signals to MLM in protein pre-training?
• RQ3 How does OntoProtein perform on TAPE, PPI, and protein function prediction under transductive and inductive settings?

Key findings

• OntoProtein achieves higher token-level performance on TAPE tasks than several baselines, including ProtBert and the TAPE Transformer.
• In protein-protein interaction prediction, OntoProtein outperforms several baselines and is competitive with GNN-PPI methods, especially on smaller datasets.
• For protein function prediction, OntoProtein shows improvements in transductive settings (notably in BPO) and maintains competitive results in other GO aspects, with some tail-constraint considerations discussed.
• OntoProtein provides competitive results without changing model architectures, by adding training objectives and knowledge incorporation.
• The authors release ProteinKG25, a large-scale KG with 612,483 entities and 4,990,097 triples, aligned with GO terms and protein sequences.