Retrieval-Augmented Language Model for Knowledge-aware Protein Encoding

Thanks for your kind comments, we place all tables and figures in this anonymous link(https://anonymous.4open.science/r/Rebuttal-F1C0/README.md)

W1. Claims needing stronger justification: (1) The advantage of direct knowledge injection over implicit embedding is not clearly isolated. A controlled experiment comparing models with virtual tokens but without retrieval is needed to confirm its impact. (2) This paper lacks visualization or biological case studies to demonstrate how Kara's representations align with real-world protein functions.

(1) Thanks for your kind comment. Here, we present the performance of different models for proteins within knowledge graphs. For Kara with virtual tokens but without retrieval, we construct virtual tokens using the ground-truth knowledge of each protein from the knowledge graph. In contrast, the original Kara employs a retriever to predict relevant knowledge for virtual token construction.

As shown in Table 10, the variant without retrieval outperforms KeAP and OntoProtein, which embed knowledge implicitly, demonstrating the advantage of direct knowledge injection. Furthermore, the original Kara achieves performance comparable to the variant without retrieval, which utilizes ground-truth knowledge. This result highlights the effectiveness of the proposed retriever in accurately predicting protein knowledge.

(2) We have provided a visualization case study comparing Kara and KeAP on the contact prediction task in "case_study_Figures.png". The results indicate that Kara outperforms KeAP in predicting contacts for proteins with short sequences (e.g., cases 1, 4, 5, and 7). However, as the sequence length increases, both Kara and KeAP struggle to accurately align with the ground truth contact map (e.g., cases 2, 3, and 6). This limitation may stem from the lack of protein structural information modeling, which is crucial for effectively handling long-sequence proteins.

W2. The study does not evaluate Kara on low-identity or novel proteins, making its generalization unclear.

We would like to clarify, as stated in Appendix B (Lines 637–639), that we removed all proteins appearing in the downstream task datasets from the protein knowledge graph. Consequently, during inference, all proteins in the downstream tasks were unseen during training, and no related knowledge existed in the knowledge graph. Therefore, all results presented in the original paper reflect the model's performance on unseen proteins. The significant advantage of Kara over existing models demonstrates its strong generalization ability to unseen proteins.

W3. Only ProtBert is tested as the encoder, without comparisons to models like ESM-1b, limiting generalizability.

We would like to clarify that Table 7 in the original paper already presents the performance of Kara using different encoders, including ProtBert, ProteinBert, and ESM-1b.

W4. The approach closely follows retrieval-based NLP models and does not introduce fundamentally new ML architectures.

Thanks for your kind comment. As we have discussed in Appendix C (Lines 643-668), there are several key differences between our model and the retrieval-based NLP models, highlighting our model as an enhanced paradigm for integrating protein knowledge graphs into protein language models.

(1) NLP approaches using virtual tokens assume that all encoding objectives exist in a knowledge graph, allowing direct extraction of relevant information. However, this assumption fails in protein encoding, where many proteins are absent from KGs. Our model addresses this by introducing a knowledge retriever that predicts gene descriptions for unseen proteins, enabling generalization beyond predefined KG entities.

(2) Existing retriever-based NLP models use general KGs but cannot account for the unique complexities of protein KGs. Protein KGs contain multi-modal entity types requiring specialized retrieval mechanisms, and they contain large and complex gene descriptions, making the retrieval time-consuming. Our model overcomes these challenges through multi-modal matching loss and relation-go combination strategies.

(3) Previous retriever-based NLP models primarily target document encoding, where KG entities are words within text corpora. These methods fail in protein encoding, since both the encoding objective and KG entities are protein sequences. Moreover, they only incorporate one-hop neighbors, overlooking higher-order structural relevance critical for protein functionality. Our model incorporates structure-based regularizations to address these limitations.

Thanks for your kind comments, we place all tables and figures in this anonymous link(https://anonymous.4open.science/r/Rebuttal-F1C0/README.md)

W1. The baselines used in the paper should include more recent papers.

Thanks for your kind comment. As shown in Table 1, we compare the performance of Kara with more recent and powerful models, SaProt [1] and ProSST [2], on the ProteinGYM benchmark. Kara, with the ESM2 backbone, performs competitively with SaProt but falls short of ProSST due to the lack of structural information. In the inductive setting, where test proteins are unseen and no structural or prior knowledge is available, Kara outperforms both models (Table 2). This is because Kara can automatically align unseen proteins with the protein knowledge graph through the retriever, thus enhancing generalization. In contrast, SaProt and ProSST perform poorly as they heavily rely on manually curated protein structures.

[1] ProSST: Protein Language Modeling with Quantized Structure and Disentangled Attention. NuerIPS 2024

[2] SaProt: Protein Language Modeling with Structure-aware Vocabulary. ICLR 2024

W2. Combining retrieval information with sequence encoding lacks innovation and does not explore deeper graph structures.

Thanks for your kind comment. Firstly, our model is not a simple combination of retrieval information with sequence encoding, it aims to address three key technical problems for knowledge-aware protein encoding:

(1) KGs are consistently updated in the real world, how to avoid the model using outdated knowledge?

(2) Many newly observed proteins are understudied and thus do not exist in KG. How can we generalize the model to these understudied proteins?

(3) Usually, we need to fine-tune the model to adapt to various downstream applications. How can we ensure the knowledge learned during pre-training is not to be catastrophically forgotten during fine-tuning?

To address these problems, many protein modeling-specific challenges arise, such as how to align multiple modalities of information in protein knowledge graphs, how to scale the retriever to large-scale protein knowledge graphs, and how to unify the knowledge pre-training and task-oriented fine-tuning optimization objectives. The proposed Kara is incorporated with new methodologies designed especially to address these challenges. We propose a knowledge retriever to solve the modality alignment and scaleability challenges through multi-modal matching loss and relation-go combination strategies, and the contextualized virtual tokens with structure-based regularization are proposed to unify knowledge modeling during pre-training and fine-tuning.

In summary, Kara contains several unique technical designs (e.g., knowledge retriever and structure-based regularizations) to solve special challenges in protein encoding scenarios, making it different from previous methods.

Second, our model has incorporated both first-order and second-order graph structures, as these local structures contain the most informative knowledge for a protein. We agree that exploring more complex graph structures is an interesting direction, but given the increased retrieval complexity, we tend to consider this in future research.

W3. Have you considered applying retrieval augmentation directly to the inference stage instead of using the retriever only in the fine-tuning phase?

We have to clarify that the proposed Kara does not use the retriever only in the fine-tuning phase, instead, it uses the retriever both during fine-tuning and inference. During fine-tuning, the retriever is used to align downstream proteins with the knowledge graph, thus unifying the pre-training and fine-tuning objectives. During inference, the retriever is used to predict potential knowledge descriptions for unseen proteins, and thus explicitly integrate knowledge for protein encoding. Table 4 presents the model’s performance when the retriever is removed during fine-tuning or inference. The results show a substantial performance drop of Kara in the absence of the retriever, highlighting the critical role of the retriever in the effective knowledge integration.

Thanks for your kind comments, we place all tables and figures in this anonymous link(https://anonymous.4open.science/r/Rebuttal-F1C0/README.md)

W1. The following claim requires further clarification: Kara effectively avoids catastrophic forgetting due to the unified integration of knowledge across training stages.

We would like to clarify that "catastrophic forgetting" in this context refers to the loss of protein attribute knowledge learned during pretraining on the protein knowledge graph due to parameter updates during downstream fine-tuning.

To evaluate Kara’s effectiveness in mitigating catastrophic forgetting, we designed two experiments. The first measures the similarity between the embeddings of two proteins with the same attribute knowledge—a higher cosine similarity indicates better retention of knowledge information. The second requires the model to identify, from a set of candidate proteins, the one sharing attribute knowledge with a given protein. Higher accuracy suggests better embedding and preservation of knowledge information.

As shown in Table 5, OntoProtein, Keap, and Kara all perform well after pretraining, confirming their ability to learn attribute knowledge. Kara achieves the highest performance, demonstrating its superior knowledge acquisition capability. After fine-tuning on downstream tasks, Kara’s performance remains stable, whereas OntoProtein and Keap show significant drops, indicating that they lose some of the knowledge acquired during pretraining. Furthermore, removing the structure loss or virtual token leads to performance degradation after fine-tuning, highlighting the importance of unified knowledge integration in mitigating catastrophic forgetting.

W2. There is a limited examination of how sensitivity varies based on the quality or completeness of the knowledge graph.

First, we have to clarify that Table 8 in the original paper already analyzes the performance of different models when dealing with incomplete knowledge graphs (Edge dropout experiments). For retrieval noises analysis (Perturbation experiments), please refer to Reviewer rtvB-Q2. In Table 9, we provide the performance of Kara with ProteinKG65, showing its generalization ability to different knowledge graphs.

W3. There is a lack of comprehensive error analysis, particularly in situations where Kara underperforms or fails.

Thank you for your thoughtful comment. We have provided a visualization case study comparing Kara and KeAP on the contact prediction task in "case_study_Figures.png". The results indicate that Kara outperforms KeAP in predicting contacts for proteins with short sequences (e.g., cases 1, 4, 5, and 7). However, as the sequence length increases, both Kara and KeAP struggle to accurately align with the ground truth contact map (e.g., cases 2, 3, and 6). This limitation may stem from the lack of protein structural information modeling, which is crucial for effectively handling long-sequence proteins.

W4. Computational overhead and the practical scalability of very large-scale protein knowledge graphs

As discussed in lines 263–274 on page 5, we designed the relation-GO combinations strategy to generalize to large-scale KGs. Table 6 presents the time cost of Kara during training and inference with and without the retriever. The results show that, after applying the relation-GO combinations strategy, Kara's training and inference time only slightly increases compared to knowledge-free baselines (even on large-scale knowledge graph ProteinKG65), demonstrating its scalability.

Q1. More insights on how Kara mitigates catastrophic forgetting.

Existing approaches like KeAP and OntoProtein employ knowledge graph-supervised pre-training to encode knowledge information into model parameters, followed by task-specific fine-tuning using proteins from downstream tasks. However, the absence of knowledge descriptions for proteins in downstream tasks creates a knowledge supervision gap during fine-tuning. This causes the model optimization to focus solely on task objectives, potentially overwriting the previously acquired knowledge representations in parameters and leading to catastrophic knowledge forgetting.

In contrast, our proposed Kara framework addresses these issues through two key innovations. First, rather than implicitly storing knowledge in model parameters, Kara explicitly incorporates knowledge through virtual tokens. This architectural design decouples knowledge storage from model parameters, making the acquired knowledge resilient to parameter updates during downstream fine-tuning. Second, Kara introduces a knowledge retriever that aligns downstream task proteins with the knowledge graph to predict potential knowledge descriptions. This alignment mechanism enables continuous knowledge supervision during fine-tuning, ensuring that parameter updates simultaneously optimize for both task performance and knowledge consistency.

Thanks for your kind comments, we place all tables in this anonymous link(https://anonymous.4open.science/r/Rebuttal-F1C0/README.md)

W1. Stronger baselines may be considered.

As shown in Table 1, we compare the performance of Kara, SaProt, and ProSST on the ProteinGYM benchmark. Kara, with the ESM2 backbone, performs competitively with SaProt but falls short of ProSST due to the lack of structural information. In the inductive setting, where test proteins are unseen and no structural or prior knowledge is available, Kara outperforms both models (Table 2). This is because Kara can automatically align unseen proteins with the protein knowledge graph through the retriever, thus enhancing generalization.

W2. ESM2 performance reported in SaProt for contact prediction is higher.

They used the 33-layer ESM2-33t model, while we opted for the 30-layer ESM2-30t model to ensure a fair comparison with previous knowledge graph-based models, which used a 30-layer ProtBert as the backbone. Despite having fewer parameters and layers than ESM2-33t, our model outperforms ESM2-33t on the short-term Contact task (Table 3), and achieves comparable performance on the ProteinGYM benchmark, as shown in Table 1.

W3. Claims require further evidence. (1) The knowledge retriever enhances generalization to unseen proteins. (2) Kara mitigates catastrophic forgetting through unified knowledge integration.

(1) As stated in Appendix B (Lines 637–639), we removed all proteins appearing in the downstream task datasets from the protein knowledge graph. Consequently, during inference, all proteins in the downstream tasks were unseen during training, and no related knowledge existed in the knowledge graph. Therefore, all results presented in the original paper reflect the model's performance on unseen proteins.

Additionally, Table 4 presents the model’s performance when the retriever is removed during fine-tuning and inference. The results show a substantial performance drop in the absence of the retriever, further highlighting its critical role in enhancing generalization.

(2) Due to space limitations, please refer to Reviewer pk6Q-W1.

Q1. How well does Kara generalize to large-scale PKGs.

As discussed in lines 263–274 on page 5, we designed the relation-GO combinations strategy to generalize to large-scale KGs. Table 6 presents the time cost with and without the retriever. The results show that, after applying this strategy, Kara's training and inference time only slightly increases compared to knowledge-free baselines (even on large-scale ProteinKG 65), demonstrating its scalability.

Q2. How does retrieval noise or incomplete PKG affect the model.

Table 8 in the original paper already analyzes the performance of different models when dealing with incomplete knowledge graphs.

Table 7 presents the performance when varying levels of noise are introduced into the retrieved results (by replacing retrieved knowledge with random knowledge). The results indicate that retrieval noise does not significantly impact performance. This is because the model's fine-tuning process does not rely on ground-truth knowledge and is inherently noisy, enhancing the model's robustness.

Q3. Can retrieval be extended to handle proteins with missing knowledge?

Our retriever is inherently designed to handle unseen proteins without relevant prior knowledge. It achieves this by predicting the most likely knowledge description for an unseen protein, linking it to the KG, and retrieving relevant structure as knowledge information. Therefore, regardless of whether a protein has associated knowledge or useful GO entities, the retriever can still process it. While prediction errors may occur, the impact of retrieval noise on model performance due to such errors has already been discussed in Q2.

Q4. Is there a trade-off between retrieval depth and performance? Could a confidence-based retrieval mechanism improve accuracy?

We propose a relation-GO combinations strategy, which balances the breadth and performance of retrieval by identifying the interacting entities specific to each relation. Moreover, the retrieval method does not require deep-first graph search, making it scalable for large KGs, as discussed in response to Q1. The performance of using the confidence threshold is provided in Table 8, showcasing a slight performance improvement.

Q5. Why explicitly using knowledge (as of Kara) is better than the one encoded in pretraining?

First, previous works have demonstrated that encoding knowledge directly into parameters can not accurately retain the knowledge information [1]. Second, KGs are subject to updates. Models based on parameter encoding cannot easily adapt to such updates and require retraining. In contrast, Kara explicitly integrates knowledge through virtual tokens, enabling it to seamlessly adapt to updates in the knowledge graph.

[1] Large language models struggle to learn long-tail knowledge. ICML 2023