## GOProteinGNN: Leveraging Protein Knowledge Graphs for Protein Representation Learning

Dan Kalifa Technion

Kira Radinsky Technion kalifadan@cs.technion.ac.il

Uriel Singer Meta AI urielsinger@meta.com

## ABSTRACT

Proteins play a vital role in biological processes and are indispensable for living organisms. Accurate representation of proteins is crucial, especially in drug development. Recently, there has been a notable increase in interest in utilizing machine learning and deep learning techniques for unsupervised learning of protein representations. However, these approaches often focus solely on the amino acid sequence of proteins and lack factual knowledge about proteins and their interactions, thus limiting their performance. In this study, we present GOProteinGNN, a novel architecture that enhances protein language models by integrating protein knowledge graph information during the creation of amino acid level representations. Our approach allows for the integration of information at both the individual amino acid level and the entire protein level, enabling a comprehensive and effective learning process through graph-based learning. By doing so, we can capture complex relationships and dependencies between proteins and their functional annotations, resulting in more robust and contextually enriched protein representations. Unlike previous methods, GOProteinGNN uniquely learns the entire protein knowledge graph during training, which allows it to capture broader relational nuances and dependencies beyond mere triplets as done in previous work. We perform a comprehensive evaluation on several downstream tasks demonstrating that GOProteinGNN consistently outperforms previous methods, showcasing its effectiveness and establishing it as a stateof-the-art solution for protein representation learning.

## CCS CONCEPTS

· Applied computing → Bioinformatics ; · Computingmethodologies → Artificial intelligence .

## KEYWORDS

Knowledge Graphs, Protein Language Models, Bioinformatics

## 1 INTRODUCTION

Proteins play a vital role in biological processes and are central to understanding the mechanisms of living organisms. In the pursuit of unraveling their complexity, the challenge of accurately representing proteins has gained considerable attention. Accurate representation holds paramount significance, especially for drug development [20].

Amino acids serve as the fundamental building blocks of proteins. In recent years, there has been a surge of interest in employing deep learning techniques to create unsupervised amino acid representations that can be used to represent a protein [27, 28]. Prevailing approaches have predominantly focused on unraveling amino acid relationships using text-based methodologies, such as ProtBert [8].

kirar@cs.technion.ac.il

The amino acid representations are then averaged to create the protein representation.

Recently, several approaches have attempted to directly create a protein representation by integrating external information about proteins, rather than relying solely on the amino acid representations. Those methods, such as KeAP [45], harness vast protein knowledge graphs. These graphs contain diverse information about interactions the protein participated in, the processes in which it is involved, etc. The factual knowledge is conveyed through text descriptions. Figure 1 illustrates an example of such a protein knowledge graph, showcasing the intricate relations between proteins and Gene Ontology (GO) terms. However, current state-of-the-art (SOTA) methods still encounter challenges in effectively incorporating this extensive information.

Existing paradigms, at times, reduce the intricate protein knowledge graph to mere triplets (Protein, Relation, GO Term), neglecting the complex tapestry of the knowledge graph [45]. These methods have overlooked the importance of capturing a holistic graph representation. For instance, although they may encode annotations, crucial information such as broader relational nuances and contextual dependencies within the graph may be lost unless the entire graph structure is considered. The aforementioned models also incorporate knowledge post-encoder, either via loss objectives or pre-training enrichment [8, 45]. Language models often struggle to accurately embed knowledge, particularly long-tail knowledge [15]. This limitation hinders their ability to effectively incorporate external knowledge into protein representations during the encoder stage. Moreover, these models have predominantly concentrated on either amino acid level learning [8, 45] or protein level learning [17, 43], with none addressing both simultaneously.

In this work, we present GOProteinGNN, a novel architecture that integrates amino acid level and protein level representations within a single framework. Our method addresses existing limitations by comprehensively utilizing the entire protein knowledge graph structure during pre-training, ensuring the capture of intricate relationships. The main technical advancement of GOProteinGNN lies in its ability to capture broader relational nuances and dependencies between protein attributes, which existing models often overlook. In a specific context, when an attribute is connected to another attribute, determining how to embed these connections during protein representation learning is a complex challenge. Compared to GOProteinGNN, which addresses this limitation, KeAP [45] and previous studies have failed to account for a holistic view of the protein's interactions within a biological context, potentially leading to a fragmented understanding of its functional significance.

Akeyinnovation of our approach is the integration of knowledge graph information into protein language models (PLMs), using our novel Graph Neural Networks (GNN) Knowledge Injection (GKI)

mechanism, which leverages the [CLS] token representation during the encoder stage. This unique integration enables holistic learning of protein representations, as reflected in our SOTA results over extensive real-world bioinformatics tasks.

Specifically, GOProteinGNN interprets amino acid sequences as coherent 'sentences' that describe proteins, learning representations for individual amino acids in its initial encoder layers. The GKI layer then assigns representations to protein nodes in the protein knowledge graph and incorporates knowledge into these representations by leveraging neighboring representations through a Relational Graph Convolution Network (RGCN) [30]. These refined representations are processed by subsequent encoder layers, generating comprehensive amino acid representation. This process alternates between amino acid level representations and graph representation, enhancing the overall representation by allowing new insights to spread across the amino acid sequence and ensuring a thorough encoding that incorporates both sequential and graph-enhanced information.

Weperform an empirical evaluation over numerous protein tasks, spanning a broad spectrum of biological domains, comparing to superior methods and achieving SOTA performance across various benchmarks. We study the performance of our method through ablation tests.

The contributions of this work are threefold: (1) We present a novel algorithm for protein representation learning that allows the integration of knowledge graph information into PLMs during the encoder stage. This enables a comprehensive and effective learning process through graph-based learning. (2) We introduce a novel GNN Knowledge Injection (GKI) mechanism that enables protein knowledge information to propagate through the amino acid sequence representations. We leverage the [CLS] token representation to integrate the GKI with the PLM in a novel manner. (3) We present an extensive empirical evaluation of our work over numerous real-world bioinformatics tasks establishing SOTA results. Also, we contribute our code to the community 1 .

## 2 RELATED WORK

## 2.1 Proteins Representation Learning

Proteins are essential for living organisms and have vital roles in biological processes. Accurate representation of proteins is crucial, particularly in fields such as drug development. In recent years, there has been a surge of interest in applying machine learning and deep learning techniques to learn protein representations in an unsupervised manner. The popularity gain of Transformers [36] introduced many novel PLMs such as ESM-1b [28], ESM-2 [19], MSA Transformer [27], ProteinBERT [4], ProteinLM [40] and ProtTrans [8]. These models were trained on data from UniRef [34] which contains hundreds of billions of amino acids.

A protein's biological mechanism is determined by its threedimensional (3D) native structure [7]. AlphaFold [14, 31] is a recent model predicting the 3D structure of a protein. A common SOTA method that resembles the AlphaFold architecture but tuned for protein representation is ESM [19, 28]. Unlike AlphaFold, which primarily focuses on 3D structure prediction and relies on multiple

1 https://github.com/kalifadan/GOProteinGNN

2 https://www.zjukg.org/project/ProteinKG65

Figure 1: An example of a Protein Knowledge Graph. The upper figure illustrates a protein with associated biological knowledge, while the lower figure depicts the corresponding knowledge graph. The central node, Q14028, represents the protein connected to multiple GO terms through matching relations (edges), providing insights into Molecular Function, Cellular Component, and Biological Process . 2

![Image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00057v2-figure-20fbd0b6-73d3-468c-aaf8-a937f764d429.jpg)

sequence alignments and template structures, ESM requires only a single input sequence (i.e., protein) and has demonstrated similar performance across various tasks. Our comparative analysis focuses on ESM, emphasizing its utilization of single-sequence input and acknowledging its performance similarity to AlphaFold.

## 2.2 Knowledge-Enhanced Language Models

Prior research has demonstrated the effectiveness of leveraging knowledge in language models to enhance their performance [33, 44]. Specifically, KEPLER [38] endeavored to integrate factual knowledge through knowledge embedding optimization, while KG-BERT [42] and K-BERT [22] treated knowledge graph triples as textual sequences or transformed the original input sentence into a knowledgerich sentence tree, respectively. In contrast, our algorithm differs significantly, as it considers graph topology and diverse node types present in protein knowledge graphs.

While PLMs traditionally learn protein representations from sequences, ProtST [41] was trained on biomedical text descriptions. Similarly, OntoProtein [43] integrated external knowledge graphs into protein pre-training, focusing on Gene Ontology (GO) term and protein entity alignment. KeAP [45] explored more granular knowledge graph interaction for improved protein decoding. OntoProtein's contrastive learning and masked modeling were enriched by KeAP's focus on the masked language modeling objective, achieving superior performance in various tasks. Otter-Knowledge [17] extended knowledge graph integration beyond proteins, incorporating diverse entities into a comprehensive knowledge graph.

(b) The GOProteinGNN architecture.

![Image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00057v2-figure-5bf91733-deb9-443a-9a65-b12f8fc6f996.jpg)

![Image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00057v2-figure-c9624019-a692-4e75-97c8-3e1fab9654c0.jpg)

Figure 2: The GOProteinGNN pre-training architecture. The model incorporates a protein's knowledge graph with relations and GO terms. It uses encoder layers to process the amino acid sequence, producing amino acid representations and a protein representation (which is represented by the [CLS] representation). Then, the GKI refines the protein's representation with knowledge graph information using graph learning techniques, yielding a knowledge-enhanced protein representation that captures essential biological contexts and interactions. Subsequent encoder layers further enhance the amino acid representation, while the model performs masked language modeling (MLM) to restore masked amino acids.

The aforementioned models also incorporate knowledge postencoder, either via loss objectives or pre-training enrichment [8, 45]. Language models often struggle to accurately embed knowledge, particularly long-tail knowledge [15]. This limitation hinders their ability to effectively incorporate external knowledge into protein representations during the encoder stage.

Both OntoProtein [43] and KeAP [45], together with other existing models, have predominantly focused on either amino acid level learning or protein level learning. Also, they reduce the intricate protein knowledge graph to mere triplets, neglecting the intricate tapestry of the knowledge graph. GOProteinGNN sets itself apart by uniquely encompassing both levels of representation within a single framework thus learning the entire protein knowledge graph and allowing it to capture broader relational nuances and dependencies beyond mere triplets.

We specifically target protein-related tasks, capturing critical biological contexts and interactions unique to proteins.

## 3 METHODS

We present the GOProteinGNN architecture (Fig. 2), which extends the capabilities of PLMs by incorporating knowledge graph information into the models, using our proposed GKI layer. We leveraged the [CLS] token representation, to integrate between the proposed GKI and the PLM in a novel manner. This integration of language model-based representation learning with knowledge graph processing allows GOProteinGNN to generate enriched protein representations that capture essential biological contexts and interactions. The detailed algorithm is presented in Supplementary Material A.

## 3.1 Protein Knowledge Graph

GOProteinGNN is trained on a knowledge graph consisting of two different types of nodes: Proteins and GO-Terms , with Relations that act as edges. Proteins are represented as a sequence of amino acids, while the relation and the GO term are terms of factual knowledge described in natural language. The GO terms cover three aspects of biology: Molecular Function, Cellular Component, and Biological Process.

Protein Knowledge Graph Definition. Let 𝑉 𝑝𝑟𝑜𝑡𝑒𝑖𝑛𝑠 be a set containing proteins (amino acid sequences). The GKI leverages the accompanying descriptions of GO terms extracted from the protein knowledge graph. These descriptions are processed through a language encoder 3 to generate comprehensive knowledge representations. To bridge the gap between text and protein, the GKI utilizes affine transformation to project those representations to the same space. These representations in the protein knowledge graph are denoted by 𝑉 𝑡𝑒𝑟𝑚𝑠 . We define the knowledge graph 𝐺 = ( 𝑉, 𝐸 ) where: 𝑉 = 𝑉 𝑝𝑟𝑜𝑡𝑒𝑖𝑛𝑠 ∪ 𝑉 𝑡𝑒𝑟𝑚𝑠 and 𝐸 = {( 𝑢, 𝑣 ) | 𝑢 ∈ 𝑉 𝑝𝑟𝑜𝑡𝑒𝑖𝑛𝑠 , 𝑣 ∈ 𝑉 𝑡𝑒𝑟𝑚𝑠 , ( 𝑢, 𝑣 ) ∈ 𝑅𝑒𝑙𝑎𝑡𝑖𝑜𝑛𝑠 } .

3 We use PubMedBERT [9]

## 3.2 GKI Layer

We introduce a method called GNN Knowledge Injection (GKI), inspired by GNN [29], a powerful neural network designed for graphstructured data [1, 21]. GKI is particularly suitable for effectively learning from protein knowledge graphs, analogous to applications in other domains [32, 37]. By utilizing a protein knowledge graph, denoted as 𝐺 , the GKI layer generates an updated protein representation that incorporates information from the knowledge graph into the input representation.

GNN Method. In response to the complex diversity of edge types in the protein knowledge graph, we select a GNN that accommodates the multifaceted edge information. Recognizing that conventional GCNs [16] lack the capacity to fully exploit this characteristic, we turn to the Relational Graph Convolution Network (RGCN) [30], a specialized variant of GNN. By doing so, we can effectively harness the rich information present in the knowledge graph to refine and enhance the protein representation. This approach allows us to leverage the graph structure of the protein knowledge graph and to capture complex relationships and dependencies between proteins and their variable number of functional annotations, leading to more powerful and contextually enriched protein representations, a capability not attainable through other methods, such as representation concatenation. Moreover, other graph algorithms could also be applied. We compare different graph algorithms in Section 6.1.

3.2.1 GKI Procedure. Given a protein representation (i.e., the [CLS] representation) 𝑧 𝑝𝑟𝑜𝑡𝑒𝑖𝑛 ∈ R 𝐷 corresponding to node 𝑖 in the knowledge graph 𝐺 , set: ℎ ( 0 ) 𝑖 = 𝑧 𝑝𝑟𝑜𝑡𝑒𝑖𝑛 .

The given protein is connected to multiple nodes in the knowledge graph 𝐺 , such as the GO terms nodes ( 𝑉 𝑡𝑒𝑟𝑚𝑠 ). The ( 𝑙 + 1 ) RGCN layer performs the following update to the node representation 𝑧 𝑝𝑟𝑜𝑡𝑒𝑖𝑛 = ℎ ( 𝑙 ) 𝑖 :

$$h _ { i } ^ { ( l + 1 ) } = \sigma \left ( \sum _ { r \in R } \sum _ { j \in N _ { r } ( i ) } \frac { 1 } { c _ { i, r } } W _ { r } ^ { ( l ) } h _ { j } ^ { ( l ) } + W _ { 0 } ^ { ( l ) } h _ { i } ^ { ( l ) } \right )$$

«

‹

where ℎ ( 𝑙 + ) 1 𝑖 ∈ R 𝐷 is the new node representation, 𝜎 (·) is the activation function, 𝑁 𝑟 ( 𝑖 ) is the set of neighbor nodes of node 𝑖 connected by relation type 𝑟 , 𝑅 is the set of all relation types in the graph, ℎ ( 𝑙 ) 𝑗 ∈ R 𝐷 is the node representation of the neighbor node 𝑗 at layer 𝑙 , 𝑐 𝑖,𝑟 is a normalization term for relation 𝑟 specific to node 𝑖 , and 𝑊 ( 𝑙 ) 𝑟 ∈ R 𝐷 × 𝐷 and 𝑊 ( 𝑙 ) 0 ∈ R 𝐷 × 𝐷 are learnable weight matrices associated with relation 𝑟 and the self-loop, respectively.

The new protein representation after the 𝐿 RGCN layers is given by: 𝑧 ′ 𝑝𝑟𝑜𝑡𝑒𝑖𝑛 = ℎ ( 𝐿 ) 𝑖 where 𝑧 ′ 𝑝𝑟𝑜𝑡𝑒𝑖𝑛 ∈ R 𝐷 represents the updated protein representation that captures essential biological contexts and interactions.

## 3.3 GOProteinGNN Algorithm

3.3.1 Training Phase. Following the approach of KeAP [45], we use a large pre-trained PLM (ProtBert [8]) as our initialization for the encoder layers' weights, leveraging its contextual understanding of proteins. The usage in ProtBert is a common practice [43, 45] and allows us to show the significance of the injection method compared to the baselines in a fair manner. Alternatively, other advanced encoders can be used in the training algorithm.

During the training stage, each amino acid sequence (i.e., protein) is passed through multiple encoder layers, yielding an intermediate representation of the amino acid sequence, denoted by 𝑧 𝑝 ∈ R 𝐿 𝑝 × 𝐷 , where 𝐿 𝑝 is the sequence length and 𝐷 is the feature dimension (line 5 in the algorithm).

In language models, including a specialized classification token ([CLS]) at the beginning of the sequence is a common practice. The final hidden state of this token serves as a comprehensive sequence representation for various classification tasks [6]. In our framework, we uniquely leverage the [CLS] token representation to integrate external knowledge into our representation framework and enhance the protein representation.

Within the context of a given encoder layer denoted by 𝑙 , a GKI layer takes the input representation [ 𝐶𝐿𝑆 ] 𝑙 along with the protein knowledge graph denoted by 𝐺 , resulting in a refined representation [ 𝐶𝐿𝑆 ] ′ 𝑙 (line 9):

$$[ C L S ] _ { l } ^ { \prime } = G K I ( [ C L S ] _ { l }, G ) \in \mathbb { R } ^ { D }$$

Where the GKI utilizes the entirety of the knowledge graph during each injection process.

Subsequently, this refined [CLS] token representation assumes the role of the new [CLS] token in the ensuing encoder layer. Then, using the self-attention [36] of the encoder layer, the refined representation propagates and refines the representations of all other amino acids within the sequence, resulting in a new representation of the amino acid sequence (line 7).

This mechanism enables GOProteinGNN to achieve a holistic learning process, harmonizing amino acid level representations with protein level representations. The outcome is an integration of contextually enriched representations that intricately capture both sequential insights and the graph-enhanced knowledge pertaining to proteins.

As the protein representation undergoes transformations across the encoder layers, our aim is to uniformly integrate knowledge a total of 𝑘 times across the 𝐿 encoder layers. Therefore, we add a GKI layer after every 𝐿 ÷ 𝑘 encoder layer (line 8). Unlike prior language models that utilize the [CLS] token only in the final output [6], our architecture iteratively integrates knowledge across the encoder layers, ensuring a more consistent and dynamic integration of knowledge throughout the representation process.

3.3.2 Masked Language Modeling Training Objective. For each input protein, 20% of the amino acids are randomly masked. Each masked amino acid 𝑎 ( 𝑗 ) has an 80% chance of being masked for prediction, 10% chance of replacement with a random amino acid, and 10% chance of remaining unchanged. Suppose the number of masked amino acids is 𝑀 , the training objective L 𝑀𝐿𝑀 is to minimize:

$$\mathcal { L } _ { M L M } = - \sum _ { j = 0 } ^ { M - 1 } \log P ( a _ { ( j ) } | z _ { p }, \Theta )$$

where 𝑧 𝑝 is the protein representation and Θ are the parameters of the GOProteinGNN model.

Table 1: Results on amino acid contact prediction task. Short-range, medium-range, and long-range contacts are contacts between positions that are separated by 6 to 11, 12 to 23, and 24 or more positions, respectively. P@L, P@L/2, and P@L/5 denote the precision scores calculated upon top L (i.e., L most likely contacts), top L/2, and top L/5 predictions, respectively. Statistically significant results with 𝑝 &lt; 0 05 . using a t-test are marked with an asterisk (*). The best result is highlighted in bold.

| Method       | Short Range   | Short Range   | Short Range   | Medium Range   | Medium Range   | Medium Range   | Long Range   | Long Range   | Long Range   |
|--------------|---------------|---------------|---------------|----------------|----------------|----------------|--------------|--------------|--------------|
| Method       | P@L           | P@L/2         | P@L/5         | P@L            | P@L/2          | P@L/5          | P@L          | P@L/2        | P@L/5        |
| LSTM         | 0.26          | 0.36          | 0.49          | 0.20           | 0.26           | 0.34           | 0.20         | 0.23         | 0.27         |
| ResNet       | 0.25          | 0.34          | 0.46          | 0.28           | 0.25           | 0.35           | 0.10         | 0.13         | 0.17         |
| Transformer  | 0.28          | 0.35          | 0.46          | 0.19           | 0.25           | 0.33           | 0.17         | 0.20         | 0.24         |
| ProtBert [8] | 0.30          | 0.40          | 0.52          | 0.27           | 0.35           | 0.47           | 0.20         | 0.26         | 0.34         |
| ESM-1b [28]  | 0.38          | 0.48          | 0.62          | 0.33           | 0.43           | 0.56           | 0.26         | 0.34         | 0.45         |
| ESM-2 [19]   | 0.40          | 0.50          | 0.62          | 0.35           | 0.44           | 0.56           | 0.27         | 0.35         | 0.45         |
| KeAP [45]    | 0.41          | 0.51          | 0.63          | 0.36           | 0.45           | 0.54           | 0.28         | 0.35         | 0.43         |
| GOProteinGNN | 0.44 ∗        | 0.54 ∗        | 0.66 ∗        | 0.39 ∗         | 0.50 ∗         | 0.60 ∗         | 0.30 ∗       | 0.39 ∗       | 0.48 ∗       |

3.3.3 Inference Phase. During inference, as factual knowledge is not available, and to eliminate its dependency, we introduce a weight-less injection mechanism. Specifically, GOProteinGNN omits the GKI components utilized solely for pre-training. Furthermore, task-specific classification heads are incorporated to enable predictions for these specific tasks. This separation between pre-training and inference allows GOProteinGNN to generalize effectively across various protein tasks, including those involving unseen proteins.

## 4.2 Baselines

## 4 EMPIRICAL EVALUATION

In this section, we outline our empirical setting. We performed extensive analysis on both intrinsic (e.g., semantic similarity inference of proteins) and extrinsic tasks (e.g., protein-protein interactions), demonstrating the value of the novel representation we provide in real-world bioinformatics tasks. Detailed experimental settings, such as implementation details and computational complexity, are provided in the Supplementary Material C. For pre-training, we employed the ProteinKG25 dataset [43], consistent with all knowledgeenhanced baseline methods (See Supplementary Material B for full details). Furthermore, the limitations of knowledge graphs and potential mitigation strategies are discussed in Supplementary Material B.0.1.

## 4.1 Tasks

We present an extensive empirical evaluation of our method across a wide range of protein-level and amino acid-level representation tasks. We use Tasks Assessing Protein Embeddings (TAPE) [26] as a common benchmark to evaluate protein representation learning. TAPE is a set of biologically relevant semi-supervised learning tasks, such as Contact Prediction and Remote Homology Detection, spread across different domains of protein biology. Furthermore, we tested our representations on the Semantic Similarity Inference task from PROBE [35] and on Protein-Protein Interactions (PPI) from GNN-PPI [23]. All tasks were divided into specific training, validation, and test splits to ensure that each task tested a biologically relevant generalization that transferred to real-life scenarios. See Supplementary Material D for comprehensive task details and Supplementary Material E for an analysis of additional tasks.

Incorporating SOTA models as our baselines, we used ProtBert [8] and ESM-1b [28] as benchmarks, following KeAP [45] for fair comparison. Furthermore, we added the SOTA ESM-2 [19], which improves the ESM-1b model. These models are considered the topperforming models in the absence of factual knowledge integration. Furthermore, we subjected our analysis to the scrutiny of the most recent and advanced knowledge-driven pre-training model, KeAP [45]. We exclude the results of OntoProtein [43] in this paper, as it was presented by KeAP [45] that it is significantly inferior on a dataset without leakage (See Supplementary Material B). For each specific task, we followed the convention of selecting SOTA methods, just as prior research has done (See Supplementary Material C.1 for full details on task-specific baselines).

## 5 EMPIRICAL RESULTS

In this section, we present our empirical results. We validate the statistical significance of performance differences, using a one-tailed paired t-test with a 95% confidence level ( 𝑝 &lt; 0 05), comparing . observations from the tested models. Significant results are marked with an asterisk (*) in the tables.

## 5.1 Contact Prediction

In Table 1, we present GOProteinGNN's performance on the contact prediction task compared to six baseline methods. GOProteinGNN significantly outperformed all baseline methods across all categories and metrics. In particular, GOProteinGNN showed superior performance compared to ProtBert [8] and ESM-2 [19], which lack knowledge graph information. This highlights the importance of incorporating knowledge during representation learning, especially for pairwise amino acid tasks (e.g., contact prediction task), where success depends on understanding the proteins' functions. Moreover, KeAP [45] incorporates a knowledge post-encoder, limiting its capacity to seamlessly integrate external knowledge into protein representations, as reflected in the performance gains of GOProteinGNN over KeAP. These findings highlight the significant impact of incorporating knowledge graph information in GOProteinGNN via the GKI layers, leading to superior performance in the contact prediction task.

## 5.2 Semantic Similarity Inference

Table 2 presents the results of the semantic similarity inference task. GOProteinGNN outperformed all baselines significantly in capturing semantic similarities, achieving an increase of 23.8% compared to the second-best model. Biological processes play a crucial role in many drug development processes, making this category particularly important as it pertains directly to drug development. The success in this task underscores GOProteinGNN's efficacy in acquiring enriched representations that encompass the complex biological processes in which proteins are involved, due to the whole protein knowledge graph learning, surpassing the capabilities of alternative approaches. Furthermore, KeAP [45] showed inferior performance compared to ESM-1b [28], which does not use factual knowledge. This result is surprising given that both GOProteinGNN and KeAP utilize the same factual knowledge repository, yet KeAP's method of knowledge injection failed to outperform the model that does not use knowledge at all.

Table 2: Results on semantic similarity inference task. The Spearman's rank is the evaluation metric. Statistically significant results with 𝑝 &lt; 0 05 . using a t-test are marked with an asterisk (*). The best result is highlighted in bold.

| Method          | Biological Process   |
|-----------------|----------------------|
| MSA Transformer | 0.31                 |
| ProtT5-XL       | 0.21                 |
| ProtBert        | 0.35                 |
| ESM-1b          | 0.42                 |
| ESM-2           | 0.41                 |
| KeAP            | 0.41                 |
| GOProteinGNN    | 0.52 ∗               |

## 5.3 Protein-Protein Interaction Identification

Table 3 presents the results of the protein-protein interaction (PPI) identification task. GOProteinGNN showed strong performance across the board, achieving the highest F1 scores across all datasets with statistical significance, which can be beneficial for drug discovery or personalized medicine. In particular, GOProteinGNN obtained SOTA results with an increase of approximately 2.1% on the SHS27K dataset, and 2.7% on the SHS148K dataset compared to KeAP [45], with statistical significance.

In conclusion, the results highlight the strengths of our GOProteinGNN, which enables a better biological understanding of the protein, leading to success in predicting interactions between proteins.

## 5.4 Remote Homology Detection

In Table 4, we present GOProteinGNN's performance compared to seven baseline methods on the remote homology detection task. Despite the challenging nature of the task, due to the large number of potential protein folds, GOProteinGNN significantly outperformed all methods, demonstrating its effectiveness in handling the complexities of remote homology detection. Specifically, GOProteinGNN achieved an advantage over KeAP [45], further affirming its ability to capture essential biological contexts and interactions

Table 3: Results on protein-protein interaction (PPI) identification task. The table displays the results in two datasets: SHS27K and SHS148K. The breadth-first search (BFS) method is used. The F1 score is the evaluation metric. Statistically significant results with 𝑝 &lt; 0 05 . using a t-test are marked with an asterisk (*). The best result is highlighted in bold.

| Method       | SHS27K   | SHS148K   |
|--------------|----------|-----------|
| DNN-PPI      | 48.09    | 57.40     |
| DPPI         | 41.43    | 52.12     |
| PIPR         | 44.48    | 61.83     |
| GNN-PPI      | 63.81    | 71.37     |
| ProtBert     | 70.94    | 70.32     |
| ESM-1b       | 74.92    | 77.49     |
| ESM-2        | 75.05    | 77.19     |
| KeAP         | 78.58    | 77.22     |
| GOProteinGNN | 80.24 ∗  | 79.27 ∗   |

which is crucial for predicting the protein fold structure, as it processed the whole protein knowledge graph. As remote homology detection is a sequence-level classification task, we can see that both KeAP [45] and ProtBert [8] achieved equal results, suggesting KeAP fails to leverage the factual knowledge in the protein level in this task. This task highlights the novelty of our algorithm which allows for the integration of information at both the individual amino acid level and the entire protein level.

Table 4: Results on remote homology detection task. The evaluation metric is accuracy. Statistically significant results with 𝑝 &lt; 0 05 . using a t-test are marked with an asterisk (*). The best result is highlighted in bold.

| Method       | Homology   |
|--------------|------------|
| LSTM         | 0.26       |
| ResNet       | 0.17       |
| Transformer  | 0.21       |
| ProtBert     | 0.29       |
| ProteinBert  | 0.22       |
| ESM-1b       | 0.11       |
| ESM-2        | 0.13       |
| KeAP         | 0.29       |
| GOProteinGNN | 0.31 ∗     |

## 6 ABLATION TESTS

We conducted ablation tests on all tasks to comprehensively assess the performance of our proposed algorithm. Consistent trends were observed across tasks; hence, we focus on the results for contact prediction and PPI tasks, which are widely studied benchmarks.

## 6.1 GKI Ablation

The ablation test for the GKI component is presented in Table 5. The full model results highlight the effectiveness of GOProteinGNN's model components, where the model demonstrated superior performance across all metrics, underscoring its pivotal role. Ignoring relations through the GKI component implies using GCN [16] instead of RGCN [30] and ignoring edge type during the knowledge

graph learning. Without relation learning, we observe a slight decrease in results, emphasizing its importance. Moreover, excluding the GKI component (i.e., without knowledge integration) substantially degrades performance, reinforcing the significance of the GNN mechanism for effective knowledge integration. These experiments highlight the contribution of the complete GKI component to GOProteinGNN's ability to capture essential biological contexts and interactions, leading to performance improvements.

| Method                   | Short Range Contact   | Short Range Contact   | Short Range Contact   | PPI     |
|--------------------------|-----------------------|-----------------------|-----------------------|---------|
|                          | P@L                   | P@L/2                 | P@L/5                 | F1      |
| GOProteinGNN             | 0.44 ∗                | 0.54 ∗                | 0.66 ∗                | 80.24 ∗ |
| - GKI Ignoring Relations | 0.41                  | 0.51                  | 0.64                  | 79.14   |
| - Without GKI            | 0.30                  | 0.40                  | 0.52                  | 70.94   |

Table 5: Ablation test for the GKI component on the contact prediction task, and the PPI task (SHS27K). The metrics are precision scores (P@L, P@L/2, and P@L/5) for short-range contact predictions and the F1 score for the PPI task. The best result is highlighted in bold.

## 6.2 Number of GKI Components Ablation

Wewishtoexperiment with the number of optimal GKI components. Each GKI component adds an additional knowledge graph injection layer to the model. In the case of GOProteinGNN, the number of GKI components (referred to as 𝑘 ) is set to 3. In this ablation test, we reduce the number of GKI components to observe their impact. As shown in Table 6, reducing knowledge injections in GOProteinGNN results in a slight performance decrease. This decrease highlights the significance and meaningfulness of the GKI component.

| Method                | Short Range Contact   | Short Range Contact   | Short Range Contact   | PPI F1   |
|-----------------------|-----------------------|-----------------------|-----------------------|----------|
|                       | P@L                   | P@L/2                 | P@L/5                 |          |
| GOProteinGNN ( 𝑘 = 3) | 0.44 ∗                | 0.54 ∗                | 0.66 ∗                | 80.24 ∗  |
| GOProteinGNN ( 𝑘 = 2) | 0.43                  | 0.53                  | 0.64                  | 79.28    |

Table 6: Ablation test for the number of GKI components ( 𝑘 ) on the contact prediction task, and the PPI task (SHS27K). The metrics are precision scores (P@L, P@L/2, and P@L/5) for short-range contact predictions and the F1 score for the PPI task. The best result is highlighted in bold.

## 7 CONCLUSIONS

In this paper, we introduced GOProteinGNN, a novel architecture that enhances PLMs by integrating protein knowledge graph information during the creation of amino acid-level representations. Specifically, GOProteinGNN proposed a GKI layer that leverages the [CLS] token representation of the PLM during the encoder stage in a novel manner. Our methodology enables the incorporation of information at both the individual amino acid and protein levels, fostering a comprehensive learning process through graph-based techniques. Furthermore, GOProteinGNN uniquely learns the entire protein knowledge graph during training, allowing it to capture broader relational nuances and dependencies beyond mere triplets as done in previous work. GOProteinGNN showcases its potential across diverse real-world bioinformatics tasks, such as amino acid contact prediction and protein-protein interaction identification. We posit that the observed performance improvements arise from a deeper comprehension of proteins themselves and the associated factual knowledge. This understanding allows generalization for never-before-seen proteins, as demonstrated on downstream tasks. By leveraging both the protein knowledge graph and sequential data modeling, GOProteinGNN effectively propagates information regarding proteins with biological connections. In drug development, it can facilitate virtual screening by prioritizing compounds for binding to specific protein targets, underscoring its potential to provide deeper insights into complex biological phenomena and expediting the drug discovery process.

## REFERENCES

- [1] Andrea Agiollo and Andrea Omicini. 2022. GNN2GNN: Graph neural networks to generate neural networks. In Conference on Uncertainty in Artificial Intelligence . https://api.semanticscholar.org/CorpusID:252898899
- [2] Mohammed AlQuraishi. 2019. ProteinNet: a standardized data set for machine learning of protein structure. BMC bioinformatics 20 (2019), 1-10.
- [3] Michael Ashburner, Catherine A. Ball, Judith A. Blake, David Botstein, Heather L. Butler, J. Michael Cherry, Allan Peter Davis, Kara Dolinski, Selina S. Dwight, Janan T. Eppig, Midori A. Harris, David P. Hill, Laurie Issel-Tarver, Andrew Kasarskis, Suzanna E. Lewis, John C. Matese, Joel E. Richardson, Martin Ringwald, Gerald M. Rubin, and Gavin Sherlock. 2000. Gene Ontology: tool for the unification of biology. Nature Genetics 25 (2000), 25-29.
- [4] Nadav Brandes, Dan Ofer, Yam Peleg, Nadav Rappoport, and Michal Linial. 2021. ProteinBERT: a universal deep-learning model of protein sequence and function. Bioinformatics 38 (2021), 2102 - 2110.
- [5] Muhao Chen, Chelsea J.-T. Ju, Guangyu Zhou, X. Chen, Tianran Zhang, KaiWei Chang, Carlo Zaniolo, and Wei Wang. 2019. Multifaceted protein-protein interaction prediction based on Siamese residual RCNN. Bioinformatics 35 (2019), i305 - i314. https://api.semanticscholar.org/CorpusID:196809757
- [6] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. ArXiv abs/1810.04805 (2019).
- [7] Ken A. Dill and Justin L. MacCallum. 2012. The Protein-Folding Problem, 50 Years On. Science 338 (2012), 1042 - 1046. https://api.semanticscholar.org/CorpusID: 5756068
- [8] Ahmed Elnaggar, Michael Heinzinger, Christian Dallago, Ghalia Rehawi, Yu Wang, Llion Jones, Tom Gibbs, Tamas B. Fehér, Christoph Angerer, Martin Steinegger, Debsindhu Bhowmik, and Burkhard Rost. 2022. ProtTrans: Toward Understanding the Language of Life Through Self-Supervised Learning. IEEE Transactions on Pattern Analysis and Machine Intelligence 44 (2022), 7112-7127.
- [9] Yu Gu, Robert Tinn, Hao Cheng, Michael R. Lucas, Naoto Usuyama, Xiaodong Liu, Tristan Naumann, Jianfeng Gao, and Hoifung Poon. 2020. Domain-Specific Language Model Pretraining for Biomedical Natural Language Processing. ACM Transactions on Computing for Healthcare (HEALTH) 3 (2020), 1 - 23. https: //api.semanticscholar.org/CorpusID:220919723
- [10] Somaye Hashemifar, Behnam Neyshabur, Aly A Khan, and Jinbo Xu. 2018. Predicting protein-protein interactions through sequence-based deep learning. Bioinformatics 34, 17 (09 2018), i802-i810. https://doi.org/10.1093/ bioinformatics/bty573 arXiv:https://academic.oup.com/bioinformatics/articlepdf/34/17/i802/50582268/bioinformatics\_34\_17\_i802.pdf
- [11] Kaiming He, X. Zhang, Shaoqing Ren, and Jian Sun. 2015. Deep Residual Learning for Image Recognition. 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR) (2015), 770-778. https://api.semanticscholar.org/CorpusID: 206594692
- [12] Sepp Hochreiter and Jürgen Schmidhuber. 1997. Long Short-Term Memory. Neural Computation 9 (1997), 1735-1780. https://api.semanticscholar.org/CorpusID: 1915014
- [13] Jie Hou, Badri Adhikari, and Jianlin Cheng. 2017. DeepSF: deep convolutional neural network for mapping protein sequences to folds. Bioinformatics 34 (2017), 1295 - 1303. https://api.semanticscholar.org/CorpusID:4891436
- [14] John M. Jumper, Richard Evans, Alexander Pritzel, and et. al. 2021. Highly accurate protein structure prediction with AlphaFold. Nature 596 (2021), 583 589.
- [15] Nikhil Kandpal, H. Deng, Adam Roberts, Eric Wallace, and Colin Raffel. 2022. Large Language Models Struggle to Learn Long-Tail Knowledge. In International Conference on Machine Learning . https://api.semanticscholar.org/CorpusID: 253522998
- [16] Thomas N. Kipf and Max Welling. 2017. Semi-Supervised Classification with Graph Convolutional Networks. In International Conference on Learning Representations . https://openreview.net/forum?id=SJU4ayYgl
- [17] Hoang Thanh Lam, Marco Luca Sbodio, Marcos Mart'inez Galindo, Mykhaylo Zayats, Ra'ul Fern'andez-D'iaz, V'ictor Valls, Gabriele Picco, Cesar Berrospi Ramis, and Vanessa L'opez. 2023. Otter-Knowledge: benchmarks of multimodal knowledge graph representation learning from different sources for drug discovery. ArXiv abs/2306.12802 (2023).
- [18] Hang Li, Xiujun Gong, Hua Yu, and Chang Zhou. 2018. Deep Neural Network Based Predictions of Protein Interactions Using Primary Sequences. Molecules : A Journal of Synthetic Chemistry and Natural Product Chemistry 23 (2018). https://api.semanticscholar.org/CorpusID:51906277
- [19] Zeming Lin, Halil Akin, Roshan Rao, and et al. 2023. Evolutionary-scale prediction of atomic-level protein structure with a language model. Science 379, 6637 (2023), 1123-1130. https://doi.org/10.1126/science.ade2574 arXiv:https://www.science.org/doi/pdf/10.1126/science.ade2574
- [20] Jiajia Liu, Mengyuan Yang, Yankai Yu, Haixia Xu, Kang Li, and Xiaobo Zhou. 2024. Large language models in bioinformatics: applications and perspectives. ArXiv (2024). https://api.semanticscholar.org/CorpusID:266899789
- [42] Liang Yao, Chengsheng Mao, and Yuan Luo. 2019. KG-BERT: BERT for Knowledge Graph Completion. CoRR abs/1909.03193 (2019). https://api.semanticscholar. org/CorpusID:202539519

| [21]   | S. Liu, Bernardo Cuenca Grau, Ian Horrocks, and Egor V. Kostylev. 2021. INDIGO: GNN-Based Inductive Knowledge Graph Completion Using Pair-Wise Encod- ing. In Neural Information Processing Systems . https://api.semanticscholar.org/ CorpusID:245119728                                                                                                                    |
|--------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| [22]   | Weijie Liu, Peng Zhou, Zhe Zhao, Zhiruo Wang, Qi Ju, Haotang Deng, and Ping Wang. 2019. K-BERT: Enabling Language Representation with Knowledge Graph. In AAAI Conference on Artificial Intelligence . https://api.semanticscholar.org/ CorpusID:202583325                                                                                                                   |
| [23]   | Guofeng Lv, Zhiqiang Hu, Yanguang Bi, and Shaoting Zhang. 2021. Learning Unknown from Correlations: Graph Neural Network for Inter-novel-protein Interaction Prediction. In Proceedings of the Thirtieth International Joint Conference on Artificial Intelligence .                                                                                                         |
| [24]   | Iain H. Moal and Juan Fernández-Recio. 2012. SKEMPI: a Structural Kinetic and Energetic database of Mutant Protein Interactions and its use in empirical models. Bioinformatics 28 20 (2012), 2600-7. https://api.semanticscholar.org/CorpusID: 39983995                                                                                                                     |
| [25]   | Fleur Mougin. 2015. Identifying Redundant and Missing Relations in the Gene Ontology. Studies in health technology and informatics 210 (2015), 195-9. https: //api.semanticscholar.org/CorpusID:17903457                                                                                                                                                                     |
| [26]   | Roshan Rao, Nicholas Bhattacharya, Neil Thomas, Yan Duan, Xi Chen, John Canny, Pieter Abbeel, and Yun S Song. 2019. Evaluating Protein Transfer Learning with TAPE. In Advances in Neural Information Processing Systems .                                                                                                                                                   |
| [27]   | Roshan M Rao, Jason Liu, Robert Verkuil, Joshua Meier, John Canny, Pieter Abbeel, Tom Sercu, and Alexander Rives. 2021. MSA Transformer. In Proceedings of the 38th International Conference on Machine Learning (Proceedings of Machine Learning Research, Vol. 139) , Marina Meila and Tong Zhang (Eds.). PMLR, 8844- 8856. https://proceedings.mlr.press/v139/rao21a.html |
| [28]   | Alexander Rives, Siddharth Goyal, Joshua Meier, Demi Guo, Myle Ott, C. Lawrence Zitnick, Jerry Ma, and Rob Fergus. 2019. Biological structure and function emerge from scaling unsupervised learning to 250 million protein se- quences. Proceedings of the National Academy of Sciences of the United States of                                                             |
| [29]   | America 118 (2019). Franco Scarselli, MarcoGori, AhChungTsoi,MarkusHagenbuchner, and Gabriele Monfardini. 2009. The Graph Neural Network Model. IEEE Transactions on Neural                                                                                                                                                                                                  |
| [30]   | Networks 20 (2009), 61-80. https://api.semanticscholar.org/CorpusID:206756462 M. Schlichtkrull, Thomas Kipf, Peter Bloem, Rianne van den Berg, Ivan Titov, and Max Welling. 2017. Modeling Relational Data with Graph Convolutional Networks. In Extended Semantic Web Conference . https://api.semanticscholar. org/CorpusID:5458500                                        |
| [31]   | Andrew W. Senior, Richard Evans, John M. Jumper, and et. al. 2020. Improved protein structure prediction using potentials from deep learning. Nature 577 (2020), 706-710.                                                                                                                                                                                                    |
| [32]   | Ruoxi Sun. 2022. Does GNN Pretraining Help Molecular Representation? ArXiv abs/2207.06010 (2022). https://api.semanticscholar.org/CorpusID:250493139                                                                                                                                                                                                                         |
| [33]   | Yu Sun, Shuohuan Wang, Yukun Li, Shikun Feng, Hao Tian, Hua Wu, and Haifeng Wang. 2019. ERNIE 2.0: A Continual Pre-training Framework for Language Understanding. ArXiv abs/1907.12412 (2019).                                                                                                                                                                               |
| [34]   | Baris E. Suzek, Hongzhan Huang, Peter B. McGarvey, Raja Mazumder, and Cathy H. Wu. 2007. UniRef: comprehensive and non-redundant UniProt ref- erence clusters. Bioinformatics 23 10 (2007), 1282-8.                                                                                                                                                                          |
| [35]   | Serbulent Unsal, Heval Atas, Muammer Albayrak, Kemal Turhan, Aybar Can Acar, and Tunca Dogan. 2022. Learning functional properties of proteins with language models. Nature Machine Intelligence 4 (2022), 227 - 245. https://api. semanticscholar.org/CorpusID:247642074                                                                                                    |
| [36]   | Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is All you Need. In NIPS .                                                                                                                                                                                                  |
| [37]   | Hongwei Wang, Weijian Li, Xiaomeng Jin, Kyunghyun Cho, Heng Ji, Jiawei Han, and Martin D. Burke. 2021. Chemical-Reaction-Aware Molecule Representa- tion Learning. ArXiv abs/2109.09888 (2021). https://api.semanticscholar.org/ CorpusID:237581512                                                                                                                          |
| [38]   | Xiaozhi Wang, Tianyu Gao, Zhaocheng Zhu, Zhiyuan Liu, Juan-Zi Li, and Jian Tang. 2019. KEPLER: AUnified Model for Knowledge Embedding and Pre-trained Language Representation. Transactions of the Association for Computational Linguistics 9 (2019), 176-194.                                                                                                              |
| [39]   | Zhili Wang, Shimin Di, and Lei Chen. 2023. A Message Passing Neural Network Space for Better Capturing Data-dependent Receptive Fields. Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining (2023). https://api.semanticscholar.org/CorpusID:260500077                                                                                      |
| [40]   | Yijia Xiao, Jiezhong Qiu, Ziang Li, Chang-Yu Hsieh, and Jie Tang. 2021. Modeling Protein Using Large-scale Pretrain Language Model. ArXiv abs/2108.07435 (2021).                                                                                                                                                                                                             |
| [41]   | Minghao Xu, Xinyu Yuan, Santiago Miret, and Jian Tang. 2023. ProtST: Multi-Modality Learning of Protein Sequences and Biomedical Texts. ArXiv abs/2301.12040 (2023).                                                                                                                                                                                                         |
| [42]   | Liang Yao, Chengsheng Mao, and Yuan Luo. 2019. KG-BERT: BERTfor Knowledge                                                                                                                                                                                                                                                                                                    |

- [43] Ningyu Zhang, Zhen Bi, Xiaozhuan Liang, Siyuan Cheng, Haosen Hong, Shumin Deng, Qiang Zhang, Jiazhang Lian, and Huajun Chen. 2022. OntoProtein: Protein Pretraining With Gene Ontology Embedding. In International Conference on Learning Representations . https://openreview.net/forum?id=yfe1VMYAXa4

[44] Zhengyan Zhang, Xu Han, Zhiyuan Liu, Xin Jiang, Maosong Sun, and Qun Liu. 2019. ERNIE: Enhanced Language Representation with Informative Entities. In Annual Meeting of the Association for Computational Linguistics .

- [45] Hong-Yu Zhou, Yunxiang Fu, Zhicheng Zhang, Bian Cheng, and Yizhou Yu. 2023. Protein Representation Learning via Knowledge Enhanced Primary Structure Reasoning. In The Eleventh International Conference on Learning Representations . https://openreview.net/forum?id=VbCMhg7MRmj

## A PRE-TRAINING ALGORITHM

Algorithm 1: GOProteinGNN training algorithm

- 1 Input: G 𝑡 training set, G 𝑣 validation set , E 𝑚𝑎𝑥 maximum epochs, 𝑘 number of GKI components, E pretrained encoder (with 𝐿 layers).
- 2 Output: E pretrained model.

1: for 𝑒𝑝𝑜𝑐ℎ = 1, 2, ... , E 𝑚𝑎𝑥 do

2: Sample protein (amino acid sequence) 𝑝 ∈ G 𝑡

3: Extract protein knowledge graph 𝐺

4: Mask protein: 𝑝 = [ 𝐶𝐿𝑆 ] 𝑎 1 𝑎 2 ... &lt; 𝑀𝐴𝑆𝐾 &gt; ... 𝑎 𝐿 𝑝

5: 𝑧 𝑝 = 𝐸 𝐸𝑚𝑏𝑒𝑑𝑑𝑖𝑛𝑔 ( 𝑝 ) ∈ R 𝐿 𝑝 × 𝐷

6: for 𝑙 = 1, 2, ... , 𝐿 do 7: 𝑧 𝑝 = E 𝑙 ( 𝑧 𝑝 ) ∈ R 𝐿 𝑝 × 𝐷 8: if 𝑙 ( mod 𝐿 ÷ 𝑘 ) ≡ 0 then 9: 𝑧 𝑝 [ 0 ] = [ 𝐶𝐿𝑆 ] ′ 𝑙 = 𝐺𝐾𝐼 ([ 𝐶𝐿𝑆 ] 𝑙 , 𝐺 ) ∈ R 𝐷 10: end if 11: end for 12: L 𝑀𝐿𝑀 = - ˝ 𝑀 - 1 𝑗 = 0 log 𝑃 𝑎 ( ( 𝑗 ) | 𝑧 𝑝 , Θ ) 13: Minimize L 𝑀𝐿𝑀 using Adam optimizer 14: Evaluate the model on G 𝑣 15: if evaluation criterion improves then 16: Save model 17: end if 18: end for

## B PRE-TRAINING DATASET DESCRIPTION

We employed the ProteinKG25 dataset [43] for pre-training, consistent with all knowledge-enhanced baseline methods, such as OntoProtein [43] and KeAP [45]. ProteinKG25 provides a knowledge graph with aligned descriptions and protein sequences, respectively, to GO terms [3] and protein entities. GO terms can be a molecule, a cellular component, or a biological process. The GO terms and their relations to proteins are described using natural language.

The dataset contains around 5 million triplets, including nearly 600k proteins, 50k attribute terms, and 31 relation terms, with an average of 8.64 relations per protein.

Consistent with KeAP [45], we address the data leakage issue in OntoProtein [43] by removing amino acid sequences present in downstream tasks from the knowledge graph, ensuring they are unseen during pre-training. This inductive setup allows for a more accurate assessment of the model's generalization to novel proteins. KeAP, after mitigating data leakage, is considered SOTA across all tasks. Given the robust performance demonstrated by KeAP, surpassing methods such as OntoProtein, we focus our comparative analysis on KeAP.

B.0.1 Mitigating Limitations in Knowledge Graphs. Knowledge graphs may contain inconsistencies or incomplete data due to their curation process. To address this, we conducted a thorough analysis to detect inconsistencies within GO relations in our dataset, primarily leveraging lexical approaches and synonym exploitation [25]. Our findings demonstrate that the ProteinKG25 dataset is robust, with inconsistencies accounting for less than 0.02% of the triplets.

For a fair comparison, these minor inconsistencies were disregarded. Our primary contribution lies in advancing representation learning through a novel knowledge injection approach, building on prior SOTA methods that have shown promising results with existing knowledge graphs. Future work can explore additional strategies to mitigate these limitations further.

## C EXPERIMENTAL DETAILS

## C.1 Task Specific Baselines

Building upon the KeAP [45] baselines, we extend our evaluation to encompass additional SOTA models that are known to perform well on specific tasks within the protein domain. For the contact and homology tasks, we include variants of LSTM [12], ResNet [11], and Transformer [36], as well as ProteinBert [4] for the homology. Following Zhang et al., we examine DPPI [10], DNNPPI [18], PIPR [5], and GNN-PPI [23] for the protein-protein interactions (PPI) task. Also, we incorporate PIPR [5] for affinity estimation and MSA Transformer [27] and ProtT5-XL [8] for semantic similarity inference.

## C.2 Implementation Details

During pre-training, GOProteinGNN is trained for 50k steps using a learning rate of 1e-4, weight decay of 0.02, and a batch size of 512 (proteins) over four GPUs (NVIDIA A40, 48GB Memory each). RGCN layers 𝐿 are set to 1 and GNN Knowledge Injection (GKI) components (referred to as 𝑘 ) are set to 3. Our experiments demonstrated that the single-layer RGCN effectively captures most of the critical information, similar to prior studies. We found that adding additional layers did not yield significant gains and, in some cases, introduced noise into the information. For downstream tasks, we follow the SOTA KeAP [45] to ensure a fair comparison. Specifically, we use hyperparameters from GNN-PPI [23] for PPI prediction and PROBE [35] for binding affinity and similarity inference. For TAPE [26], we performed a grid search to select the optimal hyperparameters. For full implementation details, refer to the provided code.

## C.3 GOProteinGNN Computational Complexity

Efficient handling of computational complexity is crucial for graphbased models such as GOProteinGNN, particularly when applied to large-scale protein knowledge graphs.

In terms of computational complexity, GOProteinGNN, like other GCNs, operates at approximately 𝑂 (| 𝑉 | ∗ | 𝐸 |) . Several optimizations, including limiting the number of hops per protein, reduce the complexity to 𝑂 (| 𝑉 |) , achieving performance akin to KeAP. During application, the complexity of both algorithms is 𝑂 ( 1 . )

In our current setup, the full receptive graph [39] for each protein is loaded, enabling effective modeling. For scalability with increasing graph size, studies in graph algorithms offer promising solutions, such as dynamically caching receptive graphs [39] for individual nodes. These approaches ensure that our method remains practical for larger datasets.

## D DOWNSTREAM TASKS DEFINITIONS

D.0.1 Contact Prediction. Contact prediction is a pairwise amino acid task, where each pair ⟨ 𝑥 , 𝑥 𝑖 𝑗 ⟩ of input amino acids from a sequence (protein) 𝑥 is assigned to a label 𝑦 𝑖 𝑗 ∈ { 0 1 , } , where the label denotes whether the amino acids are in contact ( &lt; 8Å) [26]. The data comes from ProteinNet [2], and the evaluation is based on the precision metric.

D.0.2 Semantic Similarity Inference. This task aims to evaluate how much information representation models capture about biomolecular functional similarity, with a specific emphasis on the biological process (BP) category within the Gene Ontology (GO) dataset. We adopt the semantic similarity measures outlined by Unsal et al. and compute Spearman's rank correlation between representation vector similarities and the actual GO-based semantic similarities of corresponding protein pairs. Higher correlation values indicate greater success in representing functional similarity.

D.0.3 Protein-Protein Interaction Identification. Protein-protein interactions (PPI) refer to the physical contacts between two or more amino acid sequences, or two proteins, in our study. The goal is to predict the interaction type(s) of each protein pair, out of 7 possible interactions (i.e., multilabel classification task). The evaluation is conducted on the datasets: SHS27K [5] and SHS148K [5]. Breadthfirst search (BFS) is used to generate test sets from the two datasets. F1 score is used as the evaluation metric.

D.0.4 Remote Homology Detection. Remote Homology Detection is a sequence-level classification task. Each input sequence (protein) 𝑥 is assigned a label 𝑦 ∈ { 1 , ..., 1195 } representing different possible protein folds. The data is from DeepSF [13] and the evaluation is based on average accuracy on a fold-level holdout set.

## E PROTEIN-PROTEIN BINDING AFFINITY ESTIMATION

In this section, we present a task where GOProteinGNN provides limited advantages over existing methods. The Protein-Protein Binding Affinity Estimation task involves evaluating the capability of a protein representation to predict changes in binding affinity resulting from protein mutations. In this context, GOProteinGNN does not incorporate protein structural information and thus offers no additional benefits compared to other methods. The results indicate that while performance is not degraded, it also does not show significant improvement.

E.0.1 Task Definition. Each protein pair is associated with a real value in this task, making it a regression task that reflects binding affinity changes. We follow the approach of Unsal et al. and use the SKEMPI dataset [24] with an evaluation based on the mean squared error obtained by 10-fold cross-validation.

E.0.2 Task Results. Table 7 presents the binding affinity estimation task results. GOProteinGNN outperformed PIPR [5] and ProtBert [8] with a significant margin of improvement. Furthermore, it is noteworthy that GOProteinGNN exhibited a slight improvement over KeAP [45]. The performances of ESM-1b [28], ESM-2 [19], and GOProteinGNN are comparable, with no statistically significant differences. This similarity in performance may be attributed to the task's sensitivity to subtle structural changes, which ESM-1b and ESM-2 capture effectively through their extensive pre-training, which are critical for this task [35]. In this context, the external factual knowledge leveraged by GOProteinGNN does not provide an advantage. However, the protein knowledge graph includes descriptions of structural properties, allowing GOProteinGNN to inject such knowledge directly into protein representations. This capability contributes to its improved understanding of the protein knowledge graph, leading to modest improvements over KeAP [45].

Table 7: Results on protein-protein binding affinity prediction task. Affinity ( ↓ ) means the lower the better.

| Method       |   Affinity ( ↓ ) |
|--------------|------------------|
| PIPR         |             0.63 |
| ProtBert     |             0.58 |
| ESM-1b       |             0.5  |
| ESM-2        |             0.5  |
| KeAP         |             0.52 |
| GOProteinGNN |             0.51 |
Dan Kalifa Uriel Singer Kira Radinsky
GOProteinGNN: Leveraging Protein Knowledge Graphs for Protein Representation Learning
Proteins play a vital role in biological processes and are indispensable for
living organisms. Accurate representation of proteins is crucial, especially in
drug development. Recently, there has been a notable increase in interest in
utilizing machine learning and deep learning techniques for unsupervised
learning of protein representations. However, these approaches often focus
solely on the amino acid sequence of proteins and lack factual knowledge about
proteins and their interactions, thus limiting their performance. In this
study, we present GOProteinGNN, a novel architecture that enhances protein
language models by integrating protein knowledge graph information during the
creation of amino acid level representations. Our approach allows for the
integration of information at both the individual amino acid level and the
entire protein level, enabling a comprehensive and effective learning process
through graph-based learning. By doing so, we can capture complex relationships
and dependencies between proteins and their functional annotations, resulting
in more robust and contextually enriched protein representations. Unlike
previous methods, GOProteinGNN uniquely learns the entire protein knowledge
graph during training, which allows it to capture broader relational nuances
and dependencies beyond mere triplets as done in previous work. We perform a
comprehensive evaluation on several downstream tasks demonstrating that
GOProteinGNN consistently outperforms previous methods, showcasing its
effectiveness and establishing it as a state-of-the-art solution for protein
representation learning.