Retrieval-Augmented Language Model for Knowledge-aware Protein Encoding
Authors: Jiasheng Zhang
Delvin Ce Zhang
Shuang Liang
Zhengpin Li
Rex Ying
Jie Shao
Venue: ICML 2025 poster (2025)
Keywords: Knowledge Graphs
Protein Science
Representation Learning
Primary area: applications->chemistry_physics_and_earth_sciences
TL;DR: We propose a knowledge-aware retrieval-augmented protein language model, achieving the first unified and direct integration of protein knowledge graphs and protein language models. Performance on 6 downstream tasks verify its superiority.

Abstract: Protein language models often struggle to capture biological functions due to their lack of factual knowledge (e.g., gene descriptions). Existing solutions leverage protein knowledge graphs (PKGs) as auxiliary pre-training objectives, but lack explicit integration of task-oriented knowledge, making them suffer from limited knowledge exploitation and catastrophic forgetting. The root cause is that they fail to align PKGs with task-specific data, forcing their knowledge modeling to adapt to the knowledge-isolated nature of downstream tasks. In this paper, we propose Knowledge-aware retrieval augmented protein language model (Kara), achieving the first task-oriented and explicit integration of PKGs and protein language models. With a knowledge retriever learning to predict linkages between PKG and task proteins, Kara unifies the knowledge integration of the pre-training and fine-tuning stages with a structure-based regularization, mitigating catastrophic forgetting. To ensure task-oriented integration, Kara uses contextualized virtual tokens to extract graph context as task-specific knowledge for new proteins. Experiments show that Kara outperforms existing knowledge-enhanced models in 6 representative tasks, achieving on average 5.1% improvements.

Retrieval-Augmented
Language Model for Knowledge-Aware Protein Encoding
Jiasheng Zhang' Delvin Ce Zhang²
Shuang Liang'
Zhengpin Li Rex Ying*
Jie Shao
Abstract
application tasks such drug discovery (Hoang et al., 2024) and function prediction (Xu et al., 2024; Shaw et al., 2024). However, pointed out by Kalifa et al. (2024); Zhou et al. (2023); Zhang et al. (2022a), lacking factual knowledge (e.g., gene descriptions) makes them struggle to capture intricate biological function encoded within protein sequences.
Protein language models often struggle to capture biological functions due to their lack of factual knowledge (e.g., gene descriptions). Existing solutions leverage protein knowledge graphs (PKGs) as auxiliary pre-training objectives, but lack explicit integration of task-oriented knowledge, making them suffer from limited knowledge exploitation and catastrophic forget- ting. The root cause is that they fail to align PKGs with task-specific data, forcing their knowledge modeling to adapt to the knowledge-isolated nature of downstream tasks. In this paper, we propose Knowledge-aweel aware retrieval augmented protein language model (Kara), achieving the first task-oriented and explicit integration of PKGs and protein language models. With knowledge retriever learning to predict linkages between PKG and task proteins, Kara unifies the knowledge integration of the pre-training and fine-tuning stages with a structure-bases regularization, mitigating catastrophic forgetting. To ensure task-oriented integration, Kara uses contextualized virtual tokens to extract graph con- text as task-specific knowledge for new proteins. Experiments show that Kara outperforms existing showledge-enhaneee models in 6 representative tasks, achieving on average 5.1 1% improvements.
Existing solutions leverage protein knowledge graphs (PKGs) that describe the relationships between proteins and gene ontology (GO) entities with biological relations (Chen et al., 2023b). These models use protein sequences and associated GO annotations as complementary encoding objectives during pre-training. For example, OntoProtein (Zhang et al., 2022a) uses the TransE objective (Bordes et al., 2013) to optimize the alignment between protein representa- tions and associated GO entity representations. KeAP (Zhou et al., 2023) uses GO entity representations to guide masked token prediction of protein sequences via attention mechanism. Despite their effectiveness, unfortunately, they lack explicit integration of taik-trie knowledge.
Limitations. 1) Implicitly embed knowledge information. Existing methods use knowledge only as encoding objectives to supervise the pre-training of the model, assuming that knowledge information can be well embedded within model parameters. However, as highlighted by Kandpal et al. (2023), LMs often struggle to precisely embed knowledge, particularly long-tail knowledge. Storing knowledge within model parameters also makes them unable to adapt to knowl- edge graph updates (e.g.. adding new knowledge), which further diminishes their usability. 2) Overlook the structure information. Existing methods treat each knowledge triplet (i.e., (protein, perootinnreeation www.tian,aiinnnno (O)) independently. However, the neighboring GO entities of protein are often correlated, and the high-order connections between proteins (e.g., proteins linked to GO entity through similar relations) can provide additional insights into their functional similarities. Ignoring the structural relevance makes existing methods fail fully exploit knowledge information within PKGs. 3) Lack of lask-triented knowledge modeling. Existing methods are unable to incorporate knowledge modeling during task fine-tuning, leading to inconsistent optimization objectives between the pre-training and fine-tuning stages. This inconsistency can cause the knowledge learned during pre-training to be catastrophically forgotten when applied downstream tasks (Lee al., 2020), while also making the
1. Introduction
Proteins are essential for understanding biological processes and recent advances in artificial intelligence led to growing interest learning generalized vector representations pro- teins (Hu et al.. 2024). By viewing amino acids as language tokens, protein language models (PLMs) such as ESM (Lin al., 2023), ProteinBert (Brandes et al., 2022), and ProtBert (Ahmed et al., 2022) have proven highly valuable in various
'University of Electronic Science and Technology of China 'The Pennsylvania State University Fudan University Yaale University. Correspondence to: Jie Shao <shaojic @ ueste.etecc ed.......
Proceedings of the nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Retrieval-Augmented Language Model for Knowledge-Aware Protcin Encoding
S-Contuct
M-Costact
2. Preliminaries
0.65
PPI (SHS27K)
Protein Knowledge Graph. A protein knowledge graph (PKG) G g. where V is the protein set and Vgo the gene ontology (GO) entity set. is the set of relations among proteins and GO entities. The knowledge set F consists of two kinds of triplets: (p.r.g) which describes the properties of proteins, and (T.2) which describes the relationships between GO entities, Each protein p V,, has an amino acid sequence 8. Each GO entity gE go includes a text description bg explaining the gene's function. Similarly, each relation E R comes with a text description ty. We first generate pre- trained embeddings of items in PKG and store them in vector databases for further usage. Specifically, relation and GO entity g are encoded based on their text descriptions using frozen PubMedBERT model (Gu et al., 2021), resulting in relation embedding r and GO embedding g. Protein p is encoded based on its amino acid sequence via a frozen ProtBert model (Ahmed et al., 2022), resulting in protein embedding p. These stored embeddings will be further used to construct virtual tokens in Kara. As in previous works, we use the ProteinK knowledge graph (Zhang et al., 2022a). Detailed introduction of ProteinKG25 can be found in Appendix B. Problem Formulation. Given PKG G, we aim to pre-train knowledge-awer aware protein language model f so that for each protein with amino acid sequence 8, we generate its knowledge-integrated representation f consists of protein encoder, knowledge projector, a protein projector, and a knowledge retriever. We use ProtBert (Ahmed et al., 2022) as the backbone of the protein encoder, the same as previous works (Zhou et al., 2023) for a fair comparison. We will test different backbones in Section 4.6, By fine-tuning on task-specific data, we further verify its capabilities to generalize pre-trained knowledge to downstream tasks (e.g., protein homology detection).
CoCoteet
8.56
IL.2
0.5L
17.30
9.43
PP1 (SH5148K)
12.82
M1.26
Semantic Similarity
A5,10
0,02
$7.59
PPI (STRING)
Dinding Affility
9.50
0.32
+ Karn (ours) KeAP
ESMI-2
Stability
Homology
Figure J. Performaacce downstream tasks. S-, M-, and L-Contact are the short-range. medium-range, and long-range contact prediction. PPI the protein proeein interaction prediction.
knowledge extraction process lack task-oriented adjustment.
The root cause is that proteins in downstream tasks often fall outside the PKG, restraining the use of knowledge during fine-tuning. Existing methods fail to align data from down- stream tasks with PKGs, foreing their knowledge modeling to adapt the knowledge-isolates nature of these tasks.
Proposed Work. To tackle these limitations, we propose Kara, a Knowledge-awee aware retrieval -aggeeted protein lan- guage model, achieving the first unified and direct integration of PKGs and protein language models. As the core of Kara, we propose a knowledge retriever that can accurately predict potential gene descriptions for new proteins and thus align them with PKGs. This alignment allows the pre-training and fine-tuning stages of Kara to be enhanced through a unified knowledge modeling process, and seamlessly adapt to knowledge updates. By employing contextualized virtual tokens, we achieve token-level information fusion between protein sequence and knowledge. Specifically, we categorized the virtual tokens into knowledge tokens and structure tokens, enabling the explicit injection of high-order graph context as trik-trie knowledge. To unify the optimization objectives, we incorporate structure-based regularization into both two stages, bringing function similar- ities into protein representations and helping the pre-trained knowledge to be effectively transferred to downstream tasks.
regreeennteareeeeeeeeennnttttt as p= f(G,s). In Kara,
3. Methodologies
As shown in Figure 2, with a knowledge retriever to align new proteins with the protein knowledge graph, Kara can uniformly integrate knowledge information during both the pre-training and fine-tuning stages. Specifically, the contextualized virtual tokens allow Kara to explicitly inject task-oriented knowledge and high-order structure infor- mation into protein representations. During pre-training masked language modeling (MLM) helps the protein encoder learn to fuse the information of protein sequences and structured knowledge the token level. During fine-tuning, downstream task modeling helps the protein encoder learn to extract task-specific useful knowledge from PKGs via virtual tokens. Additionally, based on the high-order connectivity between proteins, structure- regularization is incor- porated during the two stages to unify their optimization
As shown in Figure 1, experiments in representative tasks show the effectiveness of Kara. It outperforms powerful baselines (i.e,, KeAP and ESM-2) across all the tasks. For instance, Kara exceeds the ofate-of-the-art knowledge- enhanced model KeAP by 11.6% in the long-range contact prediction and by 10.3% in the protein homology detection, highlighting Kara as a better paradigm for integrating protein knowledge graphs into protein language models.
2

---

Retrieval-Augmentc Language Model for Knowledge-Aware Protcin Encoding
Protein Knowledge Graph
A.L. Contextualized Tokens
A.Z. Knowledge-gulde Pre tra-irainng
Sortiiinng wi of miti
Karadola Projense
Mean Prefing
Evceriedyy lnflrmation
Proteis Encedkr
kamneedg Frnbedding Danahase
Musked LarorLaggen Lanpuage Madeliyg
Amiia Acid Segarasi
Nabrity
Proide Projectr
w.twayy (swi))
Preling
Meen
leformation
AfANMK22
A33144
AR3I44
Wnneedding Datahaue
A. Pre-training Stuge
Stnwstrrchardd Regularbaatinn
Olidgo
R.T. Knowledge Retriever
B.2. Tack-arinted Fino-tuning
014910
Projeter
Visi
F Query Erabwoding
Danasoron Fasi
GO Katly
Menlarit Searee
AMR32
Eacoler
Projecse-45
New probis Intreduced dariag
Praton:
Caerear Erabeddinn
Candittas
Token Sequence of New Pronca
Tabebliig
Proie
B. Fine-tuning Stage
Structone harre Regolarbatioe
Figure 2. Overall architecture. During pre-training. Kara directly integrates knowledge information via contextualized virtual tokens and structure-basse based regularization. During fine-tuning, the knowledge information can be similarly integrated into protein representations through knowledge retriever, which can align new proteins in downstream tasks with the protein knowledge graph.
objectives and inject funetion similarities into protein representations. We summarize main notations in Table 11 and provide a detailed analysis of Kara in Appendix C.
concatenation operation. MLP knowledge is a trainable multi- layer perceptron used to project text-modal information into uniform semantic space. Similarly, incorporate structure information of ps. we construct its structure virtual token as
3.1. Prettrainii Stage
MLP strustrctureeD (Py):
(2)
[2((P))
3. CONTEXTUALIZED VIRTUAL TOKENS
P.EN((P))
where Pj is the pre-trained embedding of protein Pj MLP structure is another trainable multi-layer perceptron used to project the amino acid sequence-modal information. We then construct the input embedding sequence for the pro- tein encoder by concatenating virtual tokens with amino acid tokens. Given the amino acid sequence s, = sf ..... [+1]] of protein p;. where s." represents an amino acid, we lookup the embedding vocabulary of protein encoder to initialize the input embedding sequence as S, = (s] s? ...... s!el| E RRlelxd, then concatenate it as
Existing protein language models struggle to encode knowledge information, since 1) knowledge in PKG is interconnected, providing the context of proteins based the graph structure, but language models are only designed to encode sequential data, limiting their ability to capture graph information; and 2) PKGs contain multi-modal information (e.g., amino acid sequences and GO text descriptions), and protein language models can only encode amino acid sequences, failing to achieve effective multi-modal information fusion. As shown in Figure 2 A. tackle the above challenges by introducing contextualized virtual tokens. By summarizing the associated knowledge of a protein as knowledge virtual tokens and summarizing its high-order structure as structure virtual tokens, Kara can directly inject the knowledge and graph information into protein representations. These virtual tokens are then concatenated with the amino acid token sequences the knowledge context, so that each amino acid can query them to integrate helpful knowledge information, enabling effective teken-level multi-modal information fusion. Specifically, for each protein P. E Vp. we extract its one-hop GO entities with ) as its knowledge, and use its two-hop proteins N2(p.) {p, (py,rr rs. g.) E Fs(ri,9)) EM (px) as its structure context. The knowledge virtual token of protein p, is then constructed as Σ
S [v.. v { +x))x
|s is the length of amino acid sequence s, and d is embed- ding dimension. During inference, any related knowledge updates can be perceived by constructing virtual tokens.
1.2. KNOWLEDGE-GUIDER PRE-TRAINING
The pre-training of Kara has two purposes: 1) achieving effec- tive information fusion of the contextualized virtual tokens (i.e., knowledge and structure information) and the amino acid tokens (iee. protein information); and 2) integrating the knowledge-based relevance (i.e., function similarities) among proteins into their representations. For the first purpose, we introduce knowledge-guide masked language modeling, allowing each amino acid to query the virtual tokens to extract helpful knowledge information for restoring masked tokens, which achieves token-level information fusion at each layer of the protein encoder. Specifically, given the input embedding sequence S,, we use 15% probability to mask each amino acid token (i.e., replace the amino acid
relations Mg (p.) = (rr...) ( (p.. F)
MLP knowledge gg)) (1)
[p)(
(T.. where r; and g, are respectively the pre-trained embeddings of relation r; and GO entity Is (see Section 2). [:] is the
3

---

Retrieval-Augmented Language Model for Knowledge-Aware Protcin Encoding
embedding with the embedding of special token '[MASKJ'. The masked embedding sequence is encoded by the
fail to explicitly extract helpful knowledge for downstream tasks, leading to unsatisfactory performance. 3) Knowledge graphs are consistently updated (e.g., correcting obsolete knowledge). Existing models cannot adapt to these updates without undergoing retraining. To tackle these challenges, we propose a knowledge retriever that can accurately predict potential knowledge for new proteins, and thus align them with PKGs. This allows the pre-training and finettniing stages directly integrate with knowledge through unified modeling process, thus unifying the optimization objectives and seamlessly adapting to knowledge updates.
Transformer component (Vaswani et al., 2017) as follows:
S: = L +
(4)
s(tt)= = LN( + ML))
(5)
where S is initiated by S, LN is layer-norm unit and MHA denotes the multi-head attention unit. After modeling the correlations among virtual tokens and amino acid tokens layer by layer, we leverage cross-entropy loss Cals on the last-layer token embeddings (i.e., SY. where I. is the number of layers in protein encoder) to estimate the masked tokens.
While the affrrmmentioned masked language modeling achieves token-level multi-modal knowledge infusion, we further introduce a sequence- level regularization based on graph connectivity between proteins, integrating biological function similarities into their representations. As we men- tioned before, each protein p5 E (P:) is two-hop connected with p. in graph structure. This high-order connectivity indi- cates that y3 and share the same knowledge (ri,9:) and thus should be similar in their biological functions. Therefore, each pair Pp..P (P:)) can be regarded as positive pair that we hope their embeddings are closer in semantic space (e.g., A9JR22 and A9JR44 in Figure 2), and (Pr.P% (p.)) can be regarded as negative pair (e.g., A9JR22 and 014910). Specifically, in Kara, we generate the sequence-level embed- ding of protein p; as P. = D0:MMAAN(SS [2:]), where MEAN is the mrea--polingg operation, and S [2:] is the last layer token em- beddings except the virtual tokens. Then, we apply the mar- gin loss on sequence-level protein embeddings to ensure high- order connected protein Pj is closer to pt than other proteins.
Generating Candidate Embeddings. We regard the GO entities in protein knowledge graphs as retrieval candidates. To achieve more accurate and stable retrieval, we integrate the neighboring structure information of each GO entity gm and generate its candidate embedding Cm MLP aggreegation [MLLe((g.) :MLP (g!".) :MLP P (gm
(7)
where gm is the stored embedding of 9m- We use g to incor- porate the information of neighboring GO entities of Im. de- fined as g [m Σ E. (Ng. Bk.. Similarly, giro is used to incorporate the information of gm 's neighboring proteins, defined as gpprot [eror Eps (pp Pk: (g)( and Ngrot(gm) are respectively the I-hop neighboring GO entities and 1-hop neighboring proteins of Im- All of MLP aggregation MI. G, and MLP p are trainable multi-layer perceptrons.
Retrieval Process. For each new protein Pa. we use a frozen ProtBert to generate its query embedding as q. -M-MPP y (MEAN Prott Bert(s)))) where sn is the amino acid sequence of pn. Intuitively, we can traverse the relation set R and the GO entity set Vgo to find potential knowledge for pa. However, the complexity of this strategy is unaccept- able due to the large size of V (i.e., 47K in ProteinKG25), Fortunately, we observe that each relation only connects with several GO entities in PKGs, inspiring us to reduce the retrieval complexity by finding relation- GO combinations. Specifically, for relation I'm E ER, we construct its candidate GO entity set E(...) = {gml(Pa,) (Pa.s"mmmm E F). During retrieval, we traverse each r.. E R and use each of its cor- responding candidate GO entity 9m E to construct the candidate knowledge (Fr..Tmm
Creg
Σ
Max(0.sim(p..D))
[a(pi)
Pj (p
(6)
-sim(p..px) +7),
where sim indicates the similarity function (e.g., cosine similarity). We finally pre-train the parameters within the protein encoder, knowledge projector, and structure projector by jointly optimizing Luxx and Creg. These three components are then used to handle downstream tasks.
3.2. Fine-tuning Stage
3.2. KNOWLEDGE RETRIEVER
Then we use the TransE objective (Bordes et al., 2013) to score (m.)) as S(pn, (..mm)mm.mm) = ||q - Fim -Cmll1. F. MLP rea (rm). Finally, we rank all candidate knowledge based on their scores, and then add top-K candidate knowl- edge into G to align new protein pr. with knowledge graph.
Proteins in downstream tasks often fail outside the PKGs (Zhou et al., 2023), restraining the use of knowledge during fine-tuning. Existing methods thus incorporate knowledge modeling solely during pre-training, leaving the fine-tuning process only guided by task-specific objectives. However, this strategy has several limitations. 1) The optimization objectives of the pre-training and fine-tuning stages are inconsistent (i.e., one is mhowledge-guidee while the other is knowledge-isolated), causing the pre-training knowledge to be catastrophically forgotten during fine-tume tuning (Lee et al., 2020). 2) Without PKGs during fine-tuning, these models
(8)
Training Strategy. We use triplets (ps,ri,9)) E as valid knowledge and by minimizing a margin-based ranking crite- rion, hope that valid knowledge can receive lower scores than invalid knowledge. The training objective is defined as Emargin = (0 -S(ptrr.g)) +7). (9)
MAX the maximum operation and ry is a syyper-parametes
4

---

Retrieval-Augmented Language Model for Knowledge-Aware Protcin Encoding
controlling the distance between valid and invalid knowledge. (pi)Tingg) & F is invalid knowledge constructed by perturb- ing gi in (pi.Tig9:) with random GO entity 9j- Since the retrieval process needs to match information from different modalities (i.e., text descriptions and amino acid sequences), we further propose a cross-modal matching loss to unity the semantic space of embeddings from different modalities.
Table 1. Performance in the amino acid contact prediction task, seq means the number of amino acids between two selected amino acids. P@I., P@ 1/2, and P@ L/5 denote the precision calculated upon top I. (i.e., I. most likely contacts), top L/2, and top L/5 predictions, respectively. The best results are bolded and The second best results are underlined.
6. 12
12.5
$2
24.5 acq
Models
Pisl.
PGLL7
P0515
Piel.
PG11/2
PGLL/
PIBL.
PG112
P0FL/S
1STM
0.26
0.36
0.49
0.30
0.26
0.34
0.20
0.23
0.27
ReNa
0.25
0.46
0.28
0.25
0.13
0.17
Lmatch = MMA |AX(0. -LPP p (=)"))))
Tansforne
0.28
0.46
0.39
0.25
0.33
0:17
0.20
ProtBent
0.30
0.40
0.52
0.27
0.35
0.47
0.20
0.26
0.34
(a.) prot (1 +7)
(10)
ESM ib
0.38
0.48
0.62
0.33
0.43
0.56
0.26
0.31
0.45
ESM-2
0.40
0.62
0.35
0.44
0.56
0.35
OutoProccin
0.37
0.57
0.32
0.80
0.31
KeAP
0.41
0.51
0.63
0.36
0.45
0.51
0.28
0.35
0.43
8) grrat is the neighboring protein embedding of a randomly sampled GO entity Jy. This loss forces the text modality information MLP (gg)) of to be closer to its corresponding neighboring protein MLP p(g,ro) (i.e., amino acid sequence modality) than other protein information MLP p (g))ro) Af- ter jointly optimizing Lmargin and Lmatch. the knowledge retriever can accurately predict the potential knowledge for new proteins, enabling its effective alignment with PKGs.
Kara
0.45
0.65
0.45
4. Experiments and Analyses
We evaluate Kara across six downstream tasks, such as amino acid contact prediction, homology detection, and stability prediction. Our analysis includes hyper-paraman sensitivity, component-wise ablations, detailed examinations of the generalization ability to unseen knowledge, and the analysis of model robustness to PKG incompleteness. Detailed task descriptions are in Appendix D. Experimental settings and implementation details are in Appendix E. Results are averaged over 3 independent runs.
3.2.2. TASK ORIENTRD FINE-TUME TUNING
After being aligned with PKGs, new proteins can be encoded with the enhancement of knowledge following Equations (1)-(5.. Any related knowledge updates will be perceived when constructing virtual tokens, as they can access the latest PKG to extract knowledge and structures. The downstream task objectives will be used to fine-tune Kara, enabling the protein encoder to extract task-specific knowledge from PKGs via virtual tokens. Note that for each new protein Pn. we exclude other new proteins from N1 (Px) when constructing structure virtual token vig to avoid noises.
4.1. Amino Acid Contact Prediction
Overview. This task aims to predict whether two amino acids within a protein are in contact, which is token-level classification task (Rao et al., 2019). Following Zhou et al. (2023), we use variants of LSTM, ResNet, and Transformer proposed by the TAPE benchmark (Rao et al., 2019), pre-trained language models ProtBert (Ahmed et al., 2022). ESM- (Rives et al., 2021), and knowledge-enhanee model OntoProtein (Zhang et al., 2022a) as baselines. The state-of-the knowledge-enhance enhanced model KeAP (Zhou et al., 2023) and the recent powerful protein language model ESM-2-30t (Lin et al., 2023) are also used for comparison.
Moreover, the structure- based regularization can also be seamlessly adapted to the fine-tuning stage. This brings two advantages. 1) Downstream tasks usually lack sufficient training data (Rao et al., 2019). The regularization term can introduce biological function similarities among new pro- teins as an auxiliary optimization objective, thus effectively avoiding over-fitting. 2) By using this regularization as a unified optimization objective of pre-training and fine-tuning, pre-trained knowledge can avoid being catastrophically forgotten and thus effectively transfer to downstream tasks.
Results. As shown in Table 1, Kara outperforms baselines by large margins in short- (6 < seg S 12), medium- (12 < < 24), and long-range (24 S seq) contact predic- tions, achieving on average 9.5% and 1.0% improvements in P@L and P@L/2 metries. Compared with the state-of-the- art PKG-enhanced model KeAP, Kara consistently surpasses it, especially in challenging long-range predictions. This is due to Kara's contextualized virtual tokens, allowing each amino acid token to explicitly extract task-oriented knowledge from PKG. However, KeAP fails to incorporate knowledge during the fine-tuning stage.
Complexity. Compared with vanilla protein language models, the extra complexity of Kara only stems from virtual tokens and retrieval process. Two virtual tokens let the com- plexity become O([[S] + 2) from (([S]2)) where [S] the length of amino acid sequences. Due the proposed strategy of relationn GOOiiooo-----o- combinations, the complexity of knowledge retriever is ((IR)kmaaz) where R2 is the size of the relation set, and kmaa is the maximum size of the candidate GO entity sets for relations. kmaz is much smaller than the size of the GO entity set (e.g., In proteinKG25, kmaz is 22,, while and the size of GO entity set is 47K).
4.2. Protein-Protein Interaction Identification
Overview. Protein- interaction (PPI) identification aims to predict the interaction types of protein pairs and is sequence-level classification. Experiments are done on three
5

---

Retrieval-Angmented Language Model for Knowledge-Aware Protcin Encoding
Table Performance in the protein-proteio interaction identifi- cation task. BFS (breadth- first search (BFS) and DFS (depth- first search) indicate the strategies to generate. We use Fl score as the
Table 3. Protein
homology
Table 4. Result of protein-
detection and stability prediction.
protein binding affinity
Models
Homology
Stability
prediction.
metrie.
LSTM
0.26
0.69
Models
Affinity
SHS27K
SHS148K
STRING
ResNet
0.17
0.73
PIPR
0.63
Modes
Avg
Avg.
Avg
Transformer
0.21
0.73
ProtBert
0.58
DNN-PPI
48.09
54.34
51.22
57.40
58.42
57.91
53.05
64.94
59.00
ProtBent
0.29
0.78
DPPI
41,43
46.12
43.77
52.12
52.03
52.08
56.68
66.82
61.75
ESM-1b
0.11
0.77
ESM-Ib
0.50
PIPR
44.48
57.80
51.14
61.83
63.98
62.91
55.65
67.45
61.55
ESM-2
0.13
0.80
ESM-2
0.50
GNN-PPI
ProtBert
63.81
70.94
73.36
71.72
69.27
72.15
71.37
70.32
82.67
78.86
77.02
74.59
78.37
67.61
91.07
87.44
84.72
77.55
OntoProtcin
0.24
0.75
OntoProtein
0.59
PB
74.92
78.83
76.88
17.49
82.13
79.31
78.54
88.59
83.57
KaAP
0.29
0.80
KeAP
0.52
ESM-2
75.05
79.55
77.30
77.19
83.34
80.26
81.32
89.19
Kara
0.32
0.83
Kara
0.50
OntoProtein
72.26
76.89
75.58
75:23
77.52
76.36
76,71
91.45
$1.08
KeAP
78.58
77.54
77.22
84.74
Kara
81,18
78.85
79.62
86.02
4.4. Protein-Protein Binding Affinity Prediction
Overview. This task aims to map each pair of proteins to real value to denote their binding affinity changes, i.e., a sequence-level regression. As in Unsal et al. (2022), we use Bayesian ridge regression to the element-wise multiplication of protein embeddings for prediction. The SKEMPI dataset (Moal & Fernández-Rees based mean square error of 10-fold cross- validation. We use the same baselines as (Zhou et al., 2023), additionally with KeAP and ESM-2.
widely-used datasets SHS27K (Chen al., 2019), SHS148K (Chen et al., 2019), and STRING (Lv et al., 2021). 7 types of interactions are included. As in Zhang et al. (2022a), we use DPPI (Hashemifar et al., 2018), DNNPPI (Li et al., 2018), PIPR (Chen et al., 2019). GNN-PPI (Lv et al., 2021) as baselines. The LM baselines are ProtBert, ESM-Ib, ESM-2. Knowledge-enhancee baselines are KeAP and OntoProtein.
2012) is used. Result is reported
Results. From Table 2, see that Kara outperforms baselines on nearly all datasets, showing its effectiveness in accurately understanding the relationships between proteins. An interesting observation is that the performance gains of KeAP compared with OntoProtein are very small on STRING dataset. As suggested in Zhou et al. (2023), this is because the large number of fine-tuning data in the STRING dataset reduces the impact of knowledge modeling in pre-training. In contrast, Kara incorporates knowledge modeling in both pre-training and fine-tuning stages, thus avoiding catastrophically forgetting pre- pre-trained knowledge.
Results. As shown in Table 4, both knowledge-enhance models fail to outperform ESM-1b. This because protein structure plays a vital role this task (Unsal et al., 2022), and existing models overlook the modeling of protein structures, while ESM- 1b achieves it via its network architecture. Kara achieves competitive performance with ESM-1b, since the protein knowledge graph contains the description of the protein structure, and Kara can directly inject such knowledge into protein embeddings via the virtual tokens.
4.5. Semantic Similarity Inference
4.3. Homology Detection and Stability Prediction
Overview. This task evaluates models' ability to extract biomolecular functional similarity among proteins. As in Un- sal et al. (2022), we use biological process (BP) and cellular component (CC) to divide protein attributes into two groups and calculate Lin similarity in each group as the ground- similarity. We then calculate Manhattan similarity between protein embeddings for prediction. The Spearman's 's rank correlation between these similarities is calculated as the metric. We include another powerful protein language model MSA Transformer (Rao et al., 2021) as baseline.
Overview. Homology detection aims to predict the remote homology of protein, which is sequence- classification task. We follow the datasets and experimental settings of Hou et al. (2018), and ask the model to predict the right fold type of protein from 1,195 different types. We report average accuracy on the fold-level heldout set. Stability prediction aims to predict the intrinsic stability of protein, which is a sequence-level regression task. As in Rocklin et al. (2017), we use Spearman's 's rank correlation scores for evaluation. The same baselines are used as in Table 1.
Results. Table 5 shows that Kara outperforms knowledge- enhanced models on both BP and CC. This can be attributed to the explicit incorporation of GO entity information in Kara, which describes the functionality of proteins. Kara is unable to outperform ESM- 1b on BP, potentially because of the larger number of parameters of ESM-Ib. However, it still outperforms the larger model ESM- 1b on CC, indicating its effectiveness in explicitly incorporating GO entities.
Results. As illustrated in Table 3, existing knowledge- enhanced models (i.e., OntoProtein and KeAP) cannot outperform traditional language models. Previous works (Zhang et al., 2022a) attributed this failure to the lack of sequence-level objectives during pre-training. Instead, using structure- based regularization, Kara incorporates knowledge- based relevance (i.e., function similarity) among proteins as unified sequence-level objective in both pre-training and fine-tuning stages, thus achieving better results.
6

---

Retrieval-Augmented
Language Model for Knowledge-Aware Protcin Encoding
Short-term Contact Prediction (P@)L)
Homology Defection (Accuracy)
Table 5. Performance the semantic similarity inference task.
0.32-
0.45
Models
BP
CC
MSA Transformer
0.31
0.30
0.31-
ProtBert
0.35
0.36
0.44
ESM-Ib
0.42
0.37
0.30-
ESM-2
0.41
0.39
0.43-
0.29
OntoPrutein
0.36
0.36
KeAP
0.41
0.40
0.42
0.28
Kara
0.41
0.41
5
50
100
0
50
100
K
K
(a)
(b)
4.6. Analysis of Kara
Figure 3. Kara with different numbers of knowledge K.
Ablation and Variants. In Table 6 we study the effectiveness of each component. We see that virtual tokens, structure- based regularization, and knowledge retriever are essential to achieve good performance. Specifically, removing con- textualized virtual tokens makes Kara unable to incorporate knowledge explicitly, thus degrading its performance in protein-protein binding affinity prediction task, which requires the property understanding of proteins. After moving structure- based regularization, Kara fails to integrate function similarities into sequence-level protein embeddings, resulting in performance degradation in sequence- tasks, e.g., homology detection and stability prediction.
knowledge (i.e., K=0), demonstrating Kara's ability to effectively extract useful knowledge for downstream tasks.
Backbone Model Analysis. We follow Zhou et al. (2023) and use ProtBert as the backbone model for a fair comparison to baselines. Here we further adapt different backbones to present the ficxibility of Kara in Table 7. Results show that Kara can flexibly adapt different backbones and produce competitive performance in different tasks. Specifically, by using larger protein models as backbone, Kara further improve the results, demonstrating the adaptability of Kara.
Protein Knowledge Graph Analysis. We follow Zhang et al. (2022a); Zhou et al. (2023) and use commonly adopted ProteinKG25 as PKG for a fair comparison to baselines. Here we analyze the effect of incomplete PKG on Kara's performance in Table 8. Results show that Kara with incomplete PKG still outperforms OntoProtein and KeAP with full PKG, showing Kara's ss robustness. We attribute the outperformance to knowledge retriever and virtual tokens, which well integrate knowledge updates into model learning.
To assess the effectiveness of our proposed knowledge retriever, compare it a variant that uses a protein similarity- based retriever. In this variant, use the frozen ProtBert model to calculate embedding similarities between new proteins and those in the PKG, selecting the top-K similar proteins and using their embeddings as virtual tokens. However, this approach does not outperform Kara. The reason is that similarity- -based retrievers struggle to accu- rately predict associated knowledge (i.e., gene descriptions) for proteins, but proteins with similar sequences can have different functions, so this approach may introduce irrelevant protein information as noise during encoding.
4.7. Analysis of Knowledge Retriever
Ablation Study. The accurate knowledge retriever is extremely important for Kara's performance in downstream tasks. Here we analyze how different components and hyper- parameters affect the retrieval performance of knowledge retriever. Knowledge retriever is trained ProteinKG25 knowledge graph, and we use the randomly sampled 2,000 proteins as test set to select the best model. During evaluation, for each test protein we first traverse each relation r R to construct query pairs (p:,r;?), and then use the knowledge retriever model to score the corresponding candidate knowl- edge (Pt,T,,(" where ge is the candidate GO entity from Er)). After traversing all the relations, we rank candidate knowledge based on their scores and calculate Precision @ (P@n) to evaluate the retrieval performance, indicating how much knowledge on the top-n ranked candidates is valid (i.e., exists in the protein knowledge graph).
Hyper-panee parametere Analysis. During pre-training, we use the ground-truth knowledge graph structure to construct the virtual tokens. However, in the fine-tur mee-tunng stage, because the new proteins not included in the protein knowledge graph, we need to use the knowledge retriever to predict its top-K K potential knowledge to construct the virtual tokens for fine-tuning and inference, where K is a hyper-parameter used to control the amount of predicted potential knowledge incorporated. Because the predicted potential knowledge can bring additional information but also inevitable noise, this part we study how K affects the performance of Kara. As shown in Figure 3, the performance improves across different tasks when K increases from 0 to 1, showcasing the value of incorporating knowledge into protein representations. As K continues to increase, performance fluctuates due to the introduction of noise from additional knowledge. Nevertheless, it still outperforms the variant without
Hyper-parameter Analysis. In Table 9, "without structure information" means that we remove neighbor information
7

---

Retrieval-Angmenteed Language Model for Knowledge-Aware Protcin Encoding
Table 6. Ablation study and performance of variants.
Tasks
Concate (6 S seg 12)
PPI (STRING)
Homology
Stability
Affinity +
w/o contextualized virtual tokens w/o structure- structure-based based regularizations Retrieval based on the protein sequence similarities
0,42
85.16
0.28
0.81
0.55
0.43
86.49
0.30
0.80
0.52
0.43
85,33
0.29
0.79
0.57
Kara
0.45
87.59
0.32
0.83
0.50
Table 7. Performance with different backbone models.
Table Performance comparison of different models after
Contact
pre-training and after task fine-tuning.
Models
Homology
Stability
Affinity
S seg 12)
After Pre-training
After Task Fine-tuning
OntoProtcin
0.460
0.240
0,750
0,590
Models
Precision
Similarity
Precision
Similarity
KeAP
0,510
0.290
0,800
0,520
OntoProtein
0.712
0.901
0.621
0.632
Kara (ProtRert)
0.553
0.323
0.830
0,501
KeAP
0.705
0.918
0.677
Kara (ProteinBert)
0.556
0.318
0.824
0,506
who structure- based regularizariin
0.722
0.906
0.624
0.749
Kara (BSM-1b)
0,563
0,327
0.833
0,510
w/o contextualized virtual tokcus
0.713
0.902
0.676
0.816
Kara
0.738
0.934
0.725
0,968
Table 8. Performance with incomplete protein knowledge graph
but also help filter out irrelevant GO candidates, thus im- proving retrieval accuracy. As shown in Figure 4, the higher neighbor sampling number helps achieve better retrieval result. We further analysis the generalization ability of our knowledge retriever to unseen knowledge in Appendix F.
Contact
Models
Homology
Stability
Affinity
(6.S sny S 12)
OntoProtein (full KG)
0.460
0,240
0.750
0.590
KeAP (Tull KG)
0.510
0.290
0.800
0.520
Kara (50% KG)
0.540
0,316
0.823
0.511
Kara (70% KG)
0.546
0.322
0.828
0.503
4.8. Case Study
Table 9. Ablation study results of the knowledge retriever. Metrics
Mitigating Catastrophic Forgetting. To evaluate Kara's effectivvness in mitigating catastrophic forgetting. 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 given protein. Higher accuracy suggests better embedding and preservation of knowledge information.
P@
P@5
Without structure information
0.681
0.669
Without cross cross-modal matching
0.733
0.721
Without relation GO combinations
0.649
0.538
Original
0.821
0.795
0.82
Precision
0.80-
0.78
As shown in Table 10, OntoProtein, KeAP, and Kara all per- form well after pre-training, confirming their ability to learn attribute knowledge. Kara achieves the highest performance, demonstrating its superior knowledge acquisition capability. A fter fine-tuning on downstream tasks, Kara's 's performance remains stable, whereas OntoProtein and KeAP show signif- icant drops, indicating that they lose some of the knowledge acquired during pre-trainne pre-training. Furthermore, removing the structure loss or virtual token leads to performance degrada- tion after fine-tuning, highlighting the importance of unified knowledge integration mitigating catastrophic forgetting.
0.76-
0.74
0.72
50
100
200
Neighbor Sampling Number
Figure 4. Performance of knowledge retriever with different neighbor sampling numbers.
in candidate GO embeddings (Equation (7)), and "without cross-modal matching" means that the knowledge retriever is only optimized based on Lmargin. Both components ben- eficial to the retrieval performance. "Without relation-GO combinations" means that for each relation, we use the whole GO entity set as candidates during retrieval. The worse per- formance of this variant shows that relation-GO combination strategy can not only reduce the retrieval time consumption,
Error Analysis. As shown in the first and second lines of Figure 5, Kara outperforms KeAP in predicting contacts for proteins with short sequences. However, as the sequence length increases, both Kara and KeAP struggle to accurately align with the ground -round-truth contact map. This limitation may stem from the lack of protein structural information mod- eling, which is crucial for handling long-sequene, proteins.
8

---

Retrieval-Angmented Language Model for Knowledge-Aware Protcin Encoding
310
235
(a) Ground-truth
(b) KeAP
(c) Kara
1.0
579
200
ungzicoporr SSPCED KOUEOONE
Pocorieeetoooo coeu KEUSHONS
(d) Ground-truth
(e) KeAP
(i Kara
(g) Ground-truth
(h) KeAP
Kara
Figure 5. Ground-truths and predicted probability contact maps.
5. Related Work
relationships. Instead, Kara explicitly injects high-order knowledge during both pre-training and fine-tuning stages.
Protein representation learning has attracted much attention due to the rapid development of language models. Existing works treat amino acid sequences token sequences, and train the language model with either supervision (Bepler & Berger, 2019) or self-supervised objective (Alley et al., 2019; Rao et al., 2019; Xiao et al., 2021; Ahmed et al., 2022; Unsal et al., 2022: Lin et al., 2023; Brandes et al., 2022). However, they ignore factual knowledge (e.g., gene descriptions of proteins), resulting in inferior representations. Recently, OntoProtein (Zhang al., 2022a) incorporates PKG by proposing hybrid encoder. KeAP (Zhou et al., 2023) ex- tends it by performing teken-level knowledge exploration via cross-attention module. However, both are limited by ignor- ing graph structure and task-oriented knowledge modeling. Very recently, GOProteinGNN (Kalifa et al., 2024) explores the benefit of graph structure. However, it still suffers from inconsistent objectives and fails consider the high-order
Some models incorporate other modalities to improve protein representations (Chen et al., 2023a). For example, Otter- Knowledge (Lam et al., 2023) designs knowledge graphs for broadly biomedical concepts. ProtST (Xu et al., 2023) infers protein representations from biomedical texts, but with no graph structure. Kara captures text descriptions together with knowledge graphs for high-order knowledge incorporation.
6. Conclusion and Future Work
We develop a retrieval-augmented language model for knowledge-aware protein representation learning, which achieves direct integration of high-order knowledge graphs and protein language models. Experiments show Kara's 's superiority in 6 downstream tasks. A future direction is to integrate 3D structure for protein representation learning.

---

Retrieval-Angmented Language Model for Knowledge-Aware Protcin Encoding
Acknowledgements
Chen, J., Dong, H., Hastings, J., Jiménez-Ruiz, E.,, López, V., Monnin, P., Pesquita, C., Skoda, P., and Tamma, V. Knowledge graphs for the life sciences: Recent developments, challenges and opportunities. Transactions on Graph Data and Knowledge, 1(1):5:1-5:33, 25232..b
This work supported by the National Natural Science Foun- dation of China (No. 62276047) and Guangxi Key Research and Development Program (No. Guike AB24010112).
Impact Statement
Chen, M., Ju, C. J.-T., Zhou, G., Chen, X., Zhang, T, Chang, K.-W., Zaniolo, C., and Wang, W. Multifaceted protein protein-protein interaction prediction based on siamese residual remn. Bioinformatics, 35(14))3005-331 2019.
This paper presents work whose goal is to advance the field of Al for protein science. Protein representation learning aims to infer low-dimensional representations using amino acid sequences of proteins as well as their multi-modal data, such as protein knowledge graph and textual descriptions. The learned protein representations could fulfill diverse downstream tasks. One societal consequence of our work is to infer protein properties and functions, egg., stability, for a previously unseen protein, which facilitates scientific discovery process. We do not foresce any undesired implications stemming from our work.
Gu, Y., Tinn, R., Cheng, H., Lucas, M., Usuyama, N., Liu, X., Naumann, T., Gao, J., and Poon, H. Domain-specific language model pretraining for biomedical natural language processing. ACM Transactions on Computing for Healthcare (HEALTH), 3(1):1-23, 2021.
Hashemifar, S., Neyshabur, B., Khan, A. A., and Xu, J. Predicting protein-protein interactions through sequence- based deep learning. Bioinformatics, 34(17); i802-i810, 2018.
References
He, X., Tian, Y., Sun, Y., Chawla, N. V., Laurent, T., LeCun, Y., Bresson, X., and Hooi, B. G-retriever: Retrieval- augmented generation for textual graph understanding and question answering. In Advances in Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, 2024.
Ahmed, E., Heinzinger, M., Dallago, C., Rehawi, G., Wang, Y., Jones, L.., Gibbs, T., Fehér, T. B., Angerer, C., Steinegger, M., Bhowmik, D., and Rost, B. Prottrans: Toward understanding the language of life through self-supervised learning, IEEE Trans. Pattern Anal. Mach. Intell., (411444711777211111 2022.
Hoang, T. L., Sbodio, M. L., Galindo, M. M., Zayats, M., Bernandez-Diaz R., Valls, V., Picco, G., Berrospi, C., and Lopez, V. Knowledge enhanced representation learning for drug discovery. In Proceedings of the AAAI Conference on Artificial Intelligence, pp. 14544-10552. 2024.
Alley, E. C., Khimulya, G., Biswas, S., AlQuraishi, M., and Church, G. M. Unified rational protein engineering with sequence-basse based deep representation learning. Nature Methods, (6(12):1355-1322 2019.
AlQuraishi, M. Proteinnet: standardized data for ma- chine learning of protein structure. BMC bioinformatics, 20:1-10, 2019.
Hou, J., Adhikari, B., and Cheng. J. Deepsf: deep convo- lutional neural network for mapping protein sequences to folds. Bioinformatics, (4488):1995-111 2018.
Bepler, T. and Berger, B. Learning protein sequence embed- dings using information from structure. In International Conference on Learning Representations, 2019.
Hu, B., Tan, C., Wu, L., Zheng, J., Xia, J., Gao, Z., Liu, Z.,, Wu, F., Zhang, G., and Li, S. Z. Advances of deep learning in protein science: comprehensive survey. arXiv preprint arXiv:2403.053// 2024.
Bordes, A., Usunier, N., Garcia-Duran, A., Weston, J., and Yakhnenko, O. Translating embeddings for modeling multi-relational data. In Advances in neural information processing systems, pp. 2787-2795, 2013.
Hu, Y., Lei, Z., Zhang, Z., Pan, B., Ling, C., and Zhao, L. Grag: Graph retrieval-augmenned generation. In Findings of the Association for Computational Linguistics, pp. 1141,112222 2025.
Brandes, N., Ofer, D., Peleg. Y., Rappoport, N., and Linial, M. Proteinbert: universal deep-learning model of protein sequence and function. Bioinformatics, 38(8): 2102-2110, 2022.
Kalifa, D., Singer, U. and Radinsky, K. Goproteingnn: Leveraging protein knowledge graphs for protein represen- tation learning. arXiv preprint arXiv:2408.. 00057, 2024.
Chen, C., Zhang, Y., Liu, X., and Coates, M. Bidirectional learning for offline model-based biological sequence design. In International Conference on Machine Learning, pp. 5351-5366, 2023a.
Kandpal, N., Deng, H., Roberts, A., Wallace, E., and Raffel, C. Large language models struggle to learn long-tail knowledge. In International Conference on Machine Learning, pp. 15696-15707, 2023.
10

---

Retrieval-Angmented Language Model for Knowledge-Aware Protcin Encoding
Kim, D. E., DiMaio, F., Yu-Ruei Wang, R., Song, Y., and Baker, D. One contact for every twelve residues allows robust and accurate topology-level protein structure mod- eling. Proteins: Structure, Function, and Bioinformatics, 82(S2))208-22 2014.
Moal, H. and Fermández-Recio, J. Skempi: structural kinetic and energetic database of mutant protein interac- tions and its use in empirical models. Bioinformatics, 28 (20):2600 2607, 2012.
Moult, J., Fidelis, K., Kryshtafovych, A., Schwede, T., and Tramontano, A. Critical assessment of methods of protein structure prediction (casp) round xii. Proteins: Structure, Function, and Bioinformatics, (6651):7-15, 2018.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. In International Conference Learning Representations, 2015.
Lam, H. T., Sbodio, M. L., Galindo, M. M., Zayats, M., Fernandez-Diaz, R., Valls, V., Picco, G., Ramis, C. B., and Lopez, V. Otter-knowledge: benchmarks of multimodal knowledge graph representation learning from different sources for drug discovery. arXiv preprint arXiv:2306. 2023.
Pan, S., Luo, L., Wang, Y., Chen, C., Wang. J., and Wu, X. Unifying large language models and knowledge graphs: A roadmap. IEEE Trans. Knowl. Data Eng., 36 (7):3580-3599 2024.
Rao, R., Bhattacharya, N., Thomas, N., Duan, Y., Chen, P., Canny, J., Abbeel, P., and Song, Y. Evaluating protein transfer learning with tape. In Advances in neural information processing procemaiinnrooessinngg systems, pp. 9686-9698, 2019.
Lee, C., Cho. K., and Kang, W. Mixout: Effective regular- ization to finetune large-scale pretrained language models. In International Conference on Learning Representations, 2020.
Rao, R. M. Liu, J., Verkuil, R., Meier, J., Canny, J., Abbeel, P., Sercu, T., and Rives, A. Msa transformer. International Conference on Machine Learning, pp-
Li, H., Gong, X.-J., Yu, H., and Zhou, C. Deep neural network based predictions of protein interactions using primary sequences. Molecules, 23(8): 1923, 2018.
8844-8856, 2021.
Reidenbach, D. Evosbdd: Latent evolution for accurate and efficient structure-based drug design. In ICLR 2024 Workshop Machine Learning for Genomics Explorations, 2024.
Li. S., Li, X., Shang, L., Sun, C., Liu, B., Ji, Z., Jiang, X., and Liu, Q. Pre-training language models with deterministic factual knowledge. arXiv preprint arXiv:2210.11165, 2022.
Rives, A., Meier, J., Sercu, T., Goyal, S., Lin, Z., Liu, J., Guo, D., Ott, M., Zitnick, C. L., Ma, J., and Fergus, R. Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences. Proceedings of the National Academy of Sciences, 118 (15):e2016239133 18, 2021.
Lin, Z., Akin, H., Rao, R., Hie, B., Zhu, Z.,, Lu, W., Smetanin, N., Verkuil, R., Kabeli, O., Shmueli, Y., dos Santos Costa, A., Fazel-Zarandi, M., Sercu, T., Candido, S., and Rives, A. Exolutionary-scale prediction of atomic-level protein structure with a language model. Science, 379(6637); 1123-1130, 2023.
Rocklin, G. J., Chidyausiku, T. M., Goreshnik, I., Ford, A., Houliston, S., Lemak, A., Carter, L., Ravichandran, R., Mulligan, V. K., Chevalier, A., Arrowsmith, C. H., and Baker, D. Global analysis of protein folding using massively parallel design, synthesis, and testing. Science, 357(6347):168 175, 2017.
Liu, J.JJ., Orlova, N., Oakes, B. L., Ma, E., Spinner, H. B., Baney, K. L. M., Chuck, J., Tan, D., Knott, G. Harrington, I.. B., Al-Shayeb, B., Wagner, A., Brötzmann, J. Staahl, B. T., Taylor, K. L., Desmarais, J., Nogales, E., and Doudna, J. A. Casx enzymes comprise a distinct family of maa- na-gudded genome editors. nature,5666774) 218-223, 2019.
Shaw, P., Gurram, B., Belanger, D., Gane, A., Bileschi, M. L.,, Colwell, L. J., Toutanova, K., and Parikh, A. P. Protex: A retrieval- retrieval-uummentel augmented approach for protein function prediction. bioRxiv, pp. 2024-05, 2024.
Loshchilov, I. and Hutter, F. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019.
Sun, Y., Wang, S., Feng, S., Ding, S., Pang, C., Shang, J., Liu, J., Chen, X., Zhao, Y., Lu, Y., Liu, W., Wu, Z., Gong, W., Liang, J., Shang. Z., Sun, P., Liu, W., Ouyang, X., Yu, D., Tian, H., Wu, H., and Wang, H. Ernie 3.0: Large-scale knowledge enhanced pre-training for language understanding and generation. arXiv preprint arXiv:2107.021 2021.
Lv, Hu, 7., Bi, Y., and Zhang. correlations: Graph neural network for inver-novel-protein interaction prediction. In Proceedings of the Thirtieth International Joint Conference on Artificial Intelligence, pp. 3677-3633, 2021.
S. Learning unknown from

---

Retrieval-Augmented Language Model for Knowledge-Aware Protcin Encoding
Unsal, S., Atas, H., Albayrak, M., Turhan, K., Acar, A. C., and Dogan, T. Learning functional properties of proteins with language models. Nature Machine Intelligence, 4 (3):227-245,200
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L.,, Gomez, A. N., Kaiser, I.., and Polosukhin, I. Attention is all you need. In Advances in Neural Information Processing Systems, pp. 1998-6008, 2017.
Wang. X., Gao, T. Zhu, Zhang, Z., Liu, Z., Li, J., and Tang, J. Kepler: A unified model for knowledge embedding and pre-trained language representation. Transactions of the Association for Computational Linguistics, 9: 9:176-194, 2021.
Xiao, Y., Qiu, Li. Z., Hsieh, C.-Y,, and Tang, Modeling protein using large-scale pretrain language model. arXiv preprint 2ww.v:2/08.007... 07435, 2021.
Xu, M., Yuan, X., Miret, S., and Tang, J. Protst: Multi- modality learning of protein sequences and biomedical texts. In International Conference Machine Learning, pp. 38749-38767, 2023.
Xu, Y., Chen, X., Wang. T., He, H., Chen, T., and Kellis, M. Demystify the secret function in protein sequence via conditional diffusion models. In ICLR 2024 Workshop Generative and Experimental Perspectives for Biomolecular Design, 2024.
Zhang, N., Bi, Z., Liang, X., Cheng. S., Hong. H., Deng, S., Zhang, Q., Lian, J., and Chen, H. Ontoprotein: Protein pre- training with gene ontology embedding. In International Conference on Learning Representations, 2022a.
Zhang, T., Wang, C., Hu, N., Qiu, M., Tang, C., He, X., and Huang, J. Dkplm: decomposable knowledge-enhancee pre-trained language model for natural language under- standing. In Proceedings of the AAAI Conference on Artificial Intelligence, pp. 1 1773-11771, 2022b.
Zhang, Y., Chen, Z., Guo, L.,, Xu, Y., Zhang, W., and Chen, H. Making large language models perform better in knowledge graph completion. In Proceedings of the 32nd ACM International Conference on Multimedia, pp. 233-242 2024.
Zhou, H.-Y., Fu, Y., Zhang, Z., Cheng, B., and Yu, Y. Protein representation learning via knowledge enhanced primary structure reasoning. In International Conference on Learning Representations, 2023.
12

---

Retrieval-Augmented Language Model for Knowledge-Aware Protcin Encoding
Table 11. Important notations and descriptions.
Notation
Description
G Vo, R F
A protein knowledge graph Protein set, GO entity set, and relation set in G
Set of triplets (i.e., knowledge) in G protein, GO entity, and relation The amino acid sequence of protein p., and each amino acid s;
Per Sc,8,"
Stored pre-traied embeddings of protein pt. relation rj, and GO entity (see Section 2)
v!
Knowledge virtual token and structure virtual token of protein p;
S.
Input embedding sequence of protein pi, embedding sequence at the L-th layer
P.
Encoded embedding of protein by Kara
8m 88m
prot
Neighboring GO entity embedding and neighboring protein embedding of GO entity 9m
qn
Query embedding corresponds to new protein pr.
Fm
Query embedding corresponds to relation I'm
Cio
Candidate embedding corresponds GO entity gm
S()
Score function
MLP())
Trainable multi-layer perceptron
NI (p:)
One-hop GO entities with relations of protein Ps Two-hop connected proteins of protein p; One-hop neighboring GO entities of GO entity gm. One-hop neighboring proteins of GO entity gm Candidate GO entity set corresponding to relation I'm
Na(p:)
Ngo(gm)
Nprot(gm
E(rm)
A. Mathematical Notations
Here we summarize main mathematical notations used in the paper in Table 11.
B. Dataset Description
We train the proposed Kara using the ProteinKG25 knowledge graph (Zhang et al., 2022a), consistent with previous knowledge-enhanced models to achieve a fair comparison. ProteinKG25 includes about 4.5 million triplets describing relationships between protein and gene ontology (GO) entities, and 100K triplets describing relationships between GO entities. There are 31 kinds of relations, 600K proteins, and 50K GO entities ProteinKG25. Each GO entity in ProteinKG25 can be a molecule, a cellular component, or a biological process, and each protein in ProteinKG25 has an average of 8.64 relations. Fol- lowing the strategy provided by Zhou al. (2023), we removed proteins appearing in the datasets of downstream tasks avoid data leakage. The raw data of ProteinKG25 can be found in https: / /www. jukg. wtggprooeeeeeeooooooe wtwp:/eeeeeetteetteeeeeee
C. Detailed Analysis of Kara
Differences Compared with Retrieval-augmented LMs in Other Fields.
First, some previous works in the question- answering field also use virtual tokens to incorporate knowledge information (Zhang et al., 2022b; Sun et al., 2021: Zhang al., 2024). They typically assume every encoding objective has existed in KG and the knowledge information can be directly extracted after matching corresponding entities. However, in the protein-encoding scenario, many under-studied proteins do not exist in KG, making the previous "matching and extracting" strategy not work. To tackle this challenge, we propose a novel knowledge retriever to predict potential gene descriptions for new proteins, which enables our model to generalize to unseen encoding objectives (where previous work fell short). Second, some recent methods also propose using a retriever to find related entities from KG to enhance LLM generation (He et al., 2024: Hu al., 2025). However, they are designed for general KGs and cannot handle unique challenges for protein knowledge graphs. Specifically, protein KGs contain two types of entities with different modality information, requiring the retrieval process to consider multi-modal information alignment. Additionally, it contains large amount of different textual gene descriptions, bringing large candidate space with complex semantics. Our knowledge retriever
is specially designed to solve these challenges with multi-modal matching loss and relation-go combination strategies. Third, some previous works also integrate knowledge and structural information within KGs into training objectives (Wang et al., 2021; Li et al., 2022; Pan et al., 2024). They are typically designed for document encoding where the entities in KG are words that appeared in documents. They use structural information to assign mask possibilities for different
13

---

Retrieval-Augmented Language Model for Knowledge-Aware Protcin Encoding
words during masked language modeling or train the model predict graph neighbors. However, in the protein-encoding scenario, both the encoding objective and entities in KG are protein sequences, making previous training strategies not work. Moreover, they only predict one-hop neighbors, which overlooked high-order structural relevance in their objective functions. However, high-order relevance is important for protein encoding since it indicates the functional similarity between proteins.
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 surreval-augmeentt LMs and being an enhanced paradigm for integrating protein knowledge graphs into protein language models.
Detailed Comparison to KeAP.
From how to integrate knowledge into language models. KeAP implicitly embeds knowledge within the parameters of the language model. Specifically, during pre-training, it uses the protein language model to encode a protein's 's amino acid sequence into an embedding. A Transformer-base based decoder then takes this embedding along with related knowledge to predict masked amino acid tokens. KeAP proposes that this noowledg--guided pre-training approach helps retain knowledge within the model parameters. However, language models often struggle to retain knowledge precisely, Additionally, KeAP processes each piece of knowledge independently, failing to integrate the complete knowledge context of proteins. Kara directly uses knowledge of each protein as part of the language model's input. As described in Section 3, Kara summarizes 1-hop neighbors of a protein (gene descriptions) as "knowledge virtual tokens" and 2-hop neighbors (functionally similar proteins) as "structure virtual tokens". These virtual tokens are then concatenated with the amino acid sequence to form the model input. This approach not only can input precise knowledge information to the language model, but also provides a broader knowledge context by leveraging neighboring information.
From how to pre-train the language model. KeAP employs a decoder to predict masked amino acid tokens using knowledge input and embeddings encoded by protein language model. However, the Transformer-basee decoder introduces significant training complexity and large number of parameters. Additionally, KeAP's pre- pre-training overlooks the protein relevance provided by the KG structure, leading to insufficient knowledge exploitation. Kara predicts masked amino acid tokens directly using the protein language model with the prompt of virtual tokens. This eliminates the need for a decoder, reducing both training complexity and parameter size. Furthermore, Kara is also trained to embed the functionally similar proteins closer together in embedding space, integrating high-order graph structural relevance (i.e., functional similarity) into protein representations.
From how encode new proteins. Since KeAP assumes that knowledge has been embedded within parameters of the language model, they directly input the amino acid sequence of new protein into the pre-trained language model to get its embedding, which suffers from imprecise knowledge information, and fails to adapt to knowledge updates. Kara proposes a novel knowledge retriever, retrieving related knowledge for each new protein and summarizing the retrieved knowledge as virtual tokens to input into the language model, which can integrate precise knowledge into the protein language model. Moreover, any updates of the related knowledge of a protein can be perceived by the knowledge retriever during retrieving, and then integrated during encoding via the virtual tokens, ensuring that Kara can always use the most current knowledge for encoding.
From model complexity. Due to the incorporation of a Transformer-based decoder, the additional time complexity of KeAP compared with vanilla protein language models is O(1S xd). where |S| is the length of protein amino acid sequence (typically 2 500). and is the embedding hidden size (usually 768 or 1024). Kara's additional time complexity, compared with vanilla protein language models, arises only from the virtual tokens (increasing from 0(15} x d) to D(([S]+2) x d)) and the retrieval process (O(Rxk) ). where |R k|| is much smaller than | x d. Therefore, the time complexity of Kara is much smaller than that of KeAP.
From parameter number. For KeAP, the incorporation of Transformer-basee decoder brings a large number of parameters, including Q, K, V. and O weight matrices for n heads, the MLP for the multi-head mechanism, layer normalization, etc. The additional parameter of Kara only comes from four projection matrices: MLP knowledger MLP struture, MLPG, and MLP p. which is much smaller than KeAP.
D. Downstream Task Definitions
This is a pairwise token-level matching task, where each pair of input amino acids (s' ( from protein sequence s is mapped to a label yi.j E (0, 1}. indicating whether they are in contact or not (< 8Á apart). Accurate contact maps can facilitate robust modeling of full 3D protein structure (Kim et al., 2014). Following previous
Amino Acid Contact Prediction.
14

---

Retrieval-Augmentee Language Model for Knowledge-Aware Protcin Encoding
works (Zhou al., 2023), we use data that comes from ProteinNet (AIQuraishi, 2019) and report precision on the ProteinNet CASP12 sett which is standard metric reported in CASP (Moult et 2018).
Protein-protein Interaction Identification. This is pairwise sequence-level classification task. Given pair of proteins (pr.p)), the model aims to predict the interaction types Jk between them. Similar to previous works (Zhou et al., 2023). 7 types of interactions are included in our experiments, which are reaction, binding, post-translational modifications, activation, inhibition, catalysis, and expression. Each protein pair may belong to several types simultaneously so this is a multi-label clas- sification problem. We use three widely- used datasets SHS27K (Chen et al., 2019), SHS148K (Chen et al.. 2019), and STRING (Lv et al., 2021) in our experiments, where SHS27K and S48K can be regarded as two subsets of STRING, which remove proteins with no more than 50 amino acids or 2 40% sequence identity. The F1 score is used as the evaluation metric for this task.
Homology Detection. This is a sequence-level classification task where each input protein is mapped to a label y (1,2,...,1119 based on its representation generated by protein language models, which indicates its possible protein fold. This task requires the evolutionary understanding of proteins and thus is valuable in microbiology and medicine (e.g. discover new CAS enzymes (Liu et al., 2019)). We follow the previous works and use data from Hou et al. (2018). By holding out entire evolutionary groups from the training set, the model required to generalize across evolutionary gaps. Same as Hou et al. (2018), we report accuracy on the fold-level heldout set.
Stability Prediction. This is a sequence-level regression task. Each input protein p is mapped as a number y IR, which repre- sents the most extreme conditions under which the protein maintains its structure above concentration threshold, serving as a proxy for its intrinsic stability. Measuring the stability of proteins is important for finding top candidates from expensive protein engineering experiments (Rao et al., 2019). We use the data provided by Rocklin et al. (2017), where the training set includes proteins from four rounds of experimental design, while the test set contains proteins that are Hamming distance-1 I neighbors of the top candidates. We report the Spearman's 's rank correlation scores on the test set evaluate the model performance.
Protein- Protein-protein Binding Affinity Prediction. This is a pairwise vequence-leved regression task that maps each pair of proteins (Pt.P)) as a real value y E R. indicating the binding affinity changes between them. This task evaluates how well a protein representation can predict changes in binding affinity resulting from protein mutations, thus being valuable for many down- stream applications such as drug design (Reidenbach, 2024). Following Unsal et al. (2022), we use Bayesian ridge regression to the element-wis multiplication of protein embeddings for predicting the binding affinity. The SKEMPI dataset (Moal & Lennándz--Recio 2012) is used and the performance is reported based on the mean square error of 10-fold cross- validaaii,t
Semantic Similarity Inference. This is a pairwise sequence-leve regression task, which evaluates how well protein language models can capture information about biomolecular functional similarity between proteins. In this task, we emphasize the biological process (BP) and cellular component (CC) categories similar to previous works (Unsal et al., 2022). We first use BP and CC to divide protein attributes into two groups and calculate the Lin similarity in each group as the ground-truth similarity. We then calculate the Manhattan similarity between protein embeddings as the prediction. The Spearman's 's rank correlation between these similarities is calculated as the metric.
E. Experimental Details
Experimental Settings. Same as previous knowledge-enhanced protein language models such as KeAP and OntoProtein, we use the ProtBert model' as the backbone of the protein encoder within Kara for a fair comparison. The text descriptions of GO entities and relations are encoded by the PubMedBert model?, which is also consistent with previous works. While generating the pre-trained embeddings of items in the protein knowledge graph (see Section 2), we represent each item as averaging the embeddings of its amino acid or word tokens. Our model is implemented with Python and we refer to the official code released by Zhou al. (2023) to implement the downstream task experiments. All tasks use standard datasets and metrics, consistent with previous works, to ensure a fair comparison. Note that since the train/valid/test set splittings of SHS27K, SHS148K, and STRING datasets are not provided, we use the official code released by Lv et al. (2021) to split each dataset with three different random seeds, and the average performance of each dataset is reported. All the experiments are conducted on NVIDIA A40 with 48 GB memory.
Pre-training Implementation Details. In the pre-training stage, we set the protein encoder within Kara (i.e., a PortBert model) as full-parameter trainable similar to previous works (Zhang et al., 2022a). We only use proteins and knowledge preserved
hnttpps//hugggggggac.... co http:///a.....t.ooooaaaaa//ttttt "https://hugggnggtset
-ncasse-aabstratttttt abstract- -oseett-ulllll fulltext
15

---

Retrieval-Angmented Language Model for Knowledge-Aware Protcin Encoding
Table 12. Hyper-parameter settings for different downstream tasks.
Tasks
Train Steps
Batch Size
K
Lreg
Learning Rate Gradient Accumulation Step
Contact
30,000
1
5
False
3e-5
Homology
2,200
2
1
True
4e-5
16
Stability
4,800
5
5
True
le-5
16
Table 13. Performance of the knowledge retriever unseen knowledge. Hits @ 10
Models
@its@1
@its@3
without PubMedBert fine-tuning
0.430
0.608
0.796
with PubMedBert fine-tuning
0.495
0.683
0.859
in the ProteinKG25 knowledge graph to pre-train Kara, where the maximum token length is set as 1024 for proteins and 512 for text descriptions. For each protein, we randomly select 10 knowledge and 10 high-order connected proteins respectively from W and N to construct its virtual tokens. The margin 7 is set as 5 and the number of negative samples is set as 2. We set the batch size to 4 with the maximum number of update steps to 10,000, and the gradient accumulation step to 16. The learning rate is set as 1e-6 and we use Adam W (Loshchilov & Hutter, 2019) for optimization. The weight decay is set as 1e-2.
Knowledge Retriever Implementation Details. In the knowledge retriever, we set the sampling number of neighbors during the candidate embedding generation as 100. Similar to the pre-training stage, the maximum token length is 1024 for proteins and 512 for text descriptions. To train the knowledge retriever, randomly sample 2,000 proteins well as their associated knowledge from the ProteinKG25 knowledge graph as the test set, and the remaining proteins are used as training data. The best knowledge retriever model is selected based on the Precision @ 5 metric on the test set. We train the knowledge retriever with the Adam optimizer (Kingma & Ba, 2015). The number of training epochs is set as 500 with the batch size as 100, and we use the early stopping strategy with patience of 5. The learning rate is set as 1e-3 and the negative sampling number set as 20, The margin 7 is also set as 5. Note that we only train the parameters within MLPs and the embeddings items in the protein knowledge graph are frozen, thus making our knowledge retriever seamlessly generalize to knowledge updates. During inference, we rank all the candidate knowledge for a new protein based on their scores S (lower is better), and then select top-K K knowledge to add to the protein knowledge graph, where K (105,50,100).
Fine-tuning Implementation Details. In the fine-tuning stage, we freeze the knowledge projector MLP MT.PPnoodcddg knowlldge and the structure projector M struucarr and only optimize the parameters within the protein encoder for downstream tasks. Note that the protein- protein-protein interaction identification, the protein-preee protein binding affinity prediction, and the semantic similarity inference tasks do not need fine-tuning and we directly use the pre-trained Kara to encode proteins for these tasks. For the structure-based regularization term, we still set the margin y as and the number of negative samples and we summarize them in Table 12. Additionally, we follow
Different
downstream tasks require various fine-tuning the implementations in GNN-PPI (Lv et al., 2021) for PPI prediction, where the number of epochs is 600 and batch size is 2048. The learning rate is set as le-3 for the SHS27K dataset and le-4 for the SHS 148K and STRING datasets. We follow the implementations in PROBE (Unsal et al., 2022) for binding affinity prediction and semantic similarity inference.
F. Generalization Ability of Knowledge Retriever
To evaluate the generalization ability of the knowledge retriever on unseen knowledge, we employ new data-splitting strategy. First, randomly divide the triples (i.e., (protein, relation, go)) into training and testing sets in an 8:2 ratio. Next, we remove any triple (pi,Tigoo)) from the training set if go; appears in any test triples. This ensures that the knowledge (e.g.. gene descriptions) associated with test proteins is entirely absent from the training set, and thus unlearnable during training. This splitting method simulates under-studied knowledge that has gene descriptions not been observed before. The results presented in Table 13 demonstrate that our knowledge retriever can generalize to these proteins. Additionally, fine-tuning the last three layers of the PubMedBert encoder during training further improves its performance, highlighting its potential to generalize to unseen gene descriptions through -spccffcc fine-tuning.
16


Reviews:
[{"review_id": "PjpU7joF4l", "reviewer": "Reviewer_rtvB", "summary": "The paper presents Kara, a knowledge-aware retrieval-augmented protein language model that explicitly integrates protein knowledge graphs (PKGs) with protein language models (PLMs), enhancing protein representation learning with task-specific knowledge and graph structure information.\nKara predicts potential gene descriptions for proteins (via knowledge retriever), aligning them with PKGs before injecting the task-relevant high-order graph structure information into protein representations with contextualized virtual tokens.\nThey propose the use of structure-based regularization to maintain consistency between pre-training and fine-tuning objectives.\nKara outperforms existing models like knowledge-enhanced PLMs KeAP** and ESM-2 across six downstream tasks.\nThe paper proposes different ablation studies, e.g., the effectiveness of each component, numbers of knowledge, and knowledge retriever.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "__Thanks for your kind comments, we place all tables in this anonymous link(https://anonymous.4open.science/r/Rebuttal-F1C0/README.md)__\n\n``W1. Stronger baselines may be considered.``\n\n 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.\n\n``W2. ESM2 performance reported in SaProt for contact prediction is higher.``\n\nThey 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.\n\n``W3.  Claims require further evidence. (1) The knowledge retriever enhances generalization to unseen proteins. (2) Kara mitigates catastrophic forgetting through unified knowledge integration.``\n\n(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.\n\nAdditionally, 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.\n\n(2) Due to space limitations, please refer to Reviewer pk6Q-W1.\n\n``Q1. How well does Kara generalize to large-scale PKGs.``\n\nAs 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.\n\n``Q2. How does retrieval noise or incomplete PKG affect the model.``\n\nTable 8 in the original paper already analyzes the performance of different models when dealing with incomplete knowledge graphs.\n\nTable 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.\n\n``Q3. Can retrieval be extended to handle proteins with missing knowledge?``\n\nOur 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.\n\n``Q4. Is there a trade-off between retrieval depth and performance? Could a confidence-based retrieval mechanism improve accuracy?``\n\nWe 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.\n\n``Q5. Why explicitly using knowledge (as of Kara) is better than the one encoded in pretraining?``\n\nFirst, 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.\n\n[1] Large language models struggle to learn long-tail knowledge. ICML 2023", "claims_and_evidence": "1. The main claim of the paper (*Kara improves protein function modeling by explicitly incorporating task-oriented knowledge*) is clear and supported by empirical evidence.\n* Kara outperforms KeAP and ESM-2 across six benchmark tasks, e.g., Kara achieves 11.6$\\%$ improvement in long-range contact prediction.\n* Stronger baselines may be considered (though they utilize further protein structural information) like SaProt (https://openreview.net/forum?id=6MRm3G4NiU) or ProSST (https://www.biorxiv.org/content/10.1101/2024.04.15.589672v3), which shows to achieve higher performance, e.g., on contact prediction. Furthermore, it seems the ESM2 performance reported in SaProt (Table 3) for contact prediction is higher than the scores in this paper.\n\n2. The following claims are clear and intuitive but require further evidence to be strongly supported.\n- (1) The knowledge retriever enhances generalization to unseen proteins.\n- (2) Kara mitigates catastrophic forgetting through unified knowledge integration.\n\n3. The ablation studies confirm the contributions of virtual tokens, structure-based regularization, and retrieval mechanisms.\nHowever, an error analysis on the effect of the retriever on the performance when it introduces  incorrect knowledge", "methods_and_evaluation_criteria": "**Methods**\n- The proposed framework is sound and intuitive. \n- However, it should indicate the assumptions in which the method works: (i) Relevant knowledge for a protein exists in the PKG, (ii) Graph structure provides meaningful relationships.\n\n**Evaluation**\n- The benchmark is comprehensive (amino acid contact prediction, PPI identification, homology detection, stability prediction). However, the fitness prediction (zero-shot) for PLMs should be considered together with perplexity to validate that the model prevents forgetting.\n- Ablation studies validate each component’s contribution and provide insights into retrieval behavior.\n- Baselines: As mentioned above, the paper may consider the latest baselines that show higher performance on these tasks.", "theoretical_claims": "No theoretical claim.", "experimental_designs_or_analyses": "Yes, I checked all experiments for 6 downstream tasks and ablation studies.", "relation_to_broader_scientific_literature": "- The problem of knowledge-aware protein representation learning is highly relevant to current works of protein language model and computational biology.\n- The paper clearly articulates the limitations of previous methods and proposes a novel retrieval-based augmentation.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "**Strengths**:\n  - The problem of interest is interesting and timely.\n  - The proposed method is sound and useful for further research on PLMs.\n  - The paper is well-structured \n\n**Weaknesses**:\n  - Kara depends on the retriever. If the retriever selects incorrect gene descriptions, it may degrade performance rather than improve it.\n  - The retrieval complexity scales better than brute-force retrieval, but it is unclear how well it generalizes to large-scale PKGs, which serves the purpose of the paper. Furthermore, how does the retrieval overhead compare to knowledge-free baselines?", "other_comments_or_suggestions": "- The \"limitations\" in the introduction may cause a misunderstanding as to the limitations of this work. The writing of this part can be polished for better understanding with cited evidence.\n- Some sections are very dense, e.g., 3.2 and 3.3, with some repetitive/unnecessary information/notations that can be trimmed for better readability.", "questions_for_authors": "1. What happens when the retrieval mechanism introduces incorrect or irrelevant gene descriptions? How does noisy or incomplete PKG knowledge affect the model?\n\n2. Can retrieval be extended to handle proteins with missing or sparse knowledge? What happens when a protein has no useful GO entities?\n\n3. Instead of fixed K retrieved entities, could a confidence-based retrieval mechanism improve accuracy? Is there a trade-off between retrieval depth and performance?\n\n4. Given the proteins in the PKGs for inference, can you elaborate why explicitly using knowledge (as of Kara) is better then the one encoded in pretraining?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4B8uzVFsOr", "reviewer": "Reviewer_pk6Q", "summary": "This paper proposes Kara, a knowledge-aware retrieval-augmented language model for protein representation learning, explicitly integrating protein knowledge graphs (PKGs) with protein language models (PLMs). The key innovation lies in using contextualized virtual tokens and a knowledge retriever, allowing explicit integration of structured and task-specific knowledge during both pre-training and fine-tuning phases. The model demonstrates superior performance across multiple downstream protein tasks, including amino acid contact prediction, protein-protein interaction (PPI) prediction, homology detection, and protein stability prediction, surpassing existing baselines (e.g., ProtBert, OntoProtein, KeAP) with significant margins.\n\n## Update after rebuttal\n\nI appreciate the authors' comprehensive and thoughtful rebuttal. The additional experiments, including the case study on contact prediction, robustness evaluation on alternate knowledge graphs, and further clarification on catastrophic forgetting, were informative and addressed most of my earlier concerns.\n\nThe strategy to mitigate catastrophic forgetting—via explicit virtual token design and continual alignment through the knowledge retriever—is clearly explained and well-supported by the updated results. I also found the scalability discussion reassuring in terms of Kara’s practicality for large-scale knowledge graphs.\n\nThat said, there is still room for **deeper empirical error analysis** and **broader generalizability validation on different KGs**. These are directions that could further improve the work, but do not critically undermine the current contribution.\n\nOverall, I keep my score of 3.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "__Thanks for your kind comments, we place all tables and figures in this anonymous link(https://anonymous.4open.science/r/Rebuttal-F1C0/README.md)__\n\n``W1. The following claim requires further clarification: Kara effectively avoids catastrophic forgetting due to the unified integration of knowledge across training stages.``\n\nWe 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.\n\nTo 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.\n\nAs 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.\n\n``W2. There is a limited examination of how sensitivity varies based on the quality or completeness of the knowledge graph.``\n\nFirst, 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.\n\n``W3. There is a lack of comprehensive error analysis, particularly in situations where Kara underperforms or fails.``\n\nThank 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.\n\n\n``W4. Computational overhead and the practical scalability of very large-scale protein knowledge graphs``\n\nAs 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.\n\n``Q1. More insights on how Kara mitigates catastrophic forgetting.``\n\nExisting 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.\n\nIn 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.", "claims_and_evidence": "Most claims are well-supported by comprehensive evidence from the experiments:\n\n1. The explicit integration of knowledge graphs into PLMs significantly enhances downstream protein representation tasks (see Tables 1, 2, and 3).\n\n2. Contextualized virtual tokens effectively incorporate graph structures and gene ontology information, as demonstrated by the ablation studies (refer to Table 6).\n\nHowever, the following claim requires further clarification:\n\n- Claim: Kara effectively avoids catastrophic forgetting due to the unified integration of knowledge across training stages.\n\nThis claim is partially supported, as shown in Table 2 and Table 8. However, more insights and case studies are needed to illustrate scenarios of catastrophic forgetting. What is the formal definition of catastrophic forgetting? Does it refer to the decline in performance of a Protein LM when fine-tuned on a large number of data?", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria align well with standard practices in protein representation learning.\n\n1. **Methodology**: \n   - Clearly detailed, integrating knowledge via contextualized virtual tokens.\n   - Incorporates structure-based regularization (Section 3) is appropriate.\n\n2. **Evaluation Criteria**: \n   - Includes different metrics for each downstream tasks, which are comprehensive.\n   - Choice of downstream evaluation tasks is suitable and diverse.", "theoretical_claims": "The paper makes no explicit theoretical claims, thus no theoretical evaluation is required.", "experimental_designs_or_analyses": "**Strengths:**\n\n1. A robust experimental design that includes multiple relevant downstream tasks and thorough comparisons with various competitive baselines. \n\n2. It features ablation studies that clearly illustrate the contributions of different components in the model.\n\n**Weaknesses:**\n1. There is a limited examination of how sensitivity varies based on the quality or completeness of the knowledge graph. Empirical studies that specifically analyze the effects of knowledge graph incompleteness or noise are lacking. It would be beneficial to test a different protein knowledge graph to determine if the proposed method is generalizable.\n\n2. There is a lack of comprehensive error analysis, particularly in situations where Kara underperforms or fails.", "relation_to_broader_scientific_literature": "The proposed method builds upon existing literature, situating itself within knowledge-enhanced protein language modeling (e.g., OntoProtein, KeAP). It explicitly extends prior work by incorporating structured, task-oriented knowledge into the fine-tuning and inference phases of PLM.", "essential_references_not_discussed": "I am not aware of any essential references being missing. I think the author did a great job in Appendix C discussing the differences compared to retrieval-augmented LMs in other fields.", "other_strengths_and_weaknesses": "**Strengths:**\n- Methodologically innovative, with a clear integration of knowledge graphs and PLMs.\n- Strong empirical validation showing consistent improvements across a variety of downstream tasks.\n- Clearly communicates the methodological advantages and potential real-world biological implications.\n\n**Weakness:**\n- The discussion does not address computational overhead and the practical scalability of very large-scale protein knowledge graphs.", "other_comments_or_suggestions": "1. Propose a deeper analysis of robustness against incomplete or noisy PKGs, including:\n   - Edge dropout experiments\n   - Perturbation experiments\n\n2. Recommend the following for real-world deployment:\n   - Exploration of computational overhead\n   - Scalability analysis", "questions_for_authors": "If the authors could provide insights into the questions raised in the previous sections and the below questions, I would be happy to consider raising the score.\n1. Could you provide explicit evidence or case studies on scenarios of catastrophic forgetting, as well as more insights on how Kara mitigates these issues?\n2. How sensitive is Kara to errors or noise in the predictions made by the knowledge retriever in the finetuning stage?\n3. Have you tested Kara's robustness with varying qualities of textual descriptions associated with GO entities? How does performance change based on the quality of these descriptions?\n4. Can you discuss the practical scalability of very large-scale protein knowledge graphs?\n5. Could you provide a case study for error analysis, particularly when Kara underperforms or encounters failures?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "N9YXWuMK2h", "reviewer": "Reviewer_VkW4", "summary": "This article proposes a knowledge-aware retrieval-augmented protein language model named Kara. During the pre-training phase, it extracts structural and knowledge information from protein KGs through contextualized virtual tokens, which are jointly embedded into the protein sequence encoding. The optimization objectives of both the pre-training and fine-tuning stages are unified through structure-based regularization. In the fine-tuning stage, a knowledge retriever predicts potential GO representations for new proteins within the PKG, thereby alleviating the issue of catastrophic forgetting of knowledge that has plagued previous models. Across multiple downstream tasks in protein prediction, Kara surpasses the current state-of-the-art models.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "__Thanks for your kind comments, we place all tables and figures in this anonymous link(https://anonymous.4open.science/r/Rebuttal-F1C0/README.md)__\n\n``W1. The baselines used in the paper should include more recent papers.``\n\nThanks 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.\n\n[1] ProSST: Protein Language Modeling with Quantized Structure and Disentangled Attention. NuerIPS 2024\n\n[2] SaProt: Protein Language Modeling with Structure-aware Vocabulary. ICLR 2024\n\n``W2. Combining retrieval information with sequence encoding lacks innovation and does not explore deeper graph structures.``\n\nThanks 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: \n\n(1) KGs are consistently updated in the real world, how to avoid the model using outdated knowledge?\n\n(2) Many newly observed proteins are understudied and thus do not exist in KG. How can we generalize the model to these understudied proteins? \n\n(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?\n\nTo 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. \n\nIn 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.\n\nSecond, 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.\n\n``W3. Have you considered applying retrieval augmentation directly to the inference stage instead of using the retriever only in the fine-tuning phase?``\n\nWe 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.", "claims_and_evidence": "Yes. The experimental design and ablation study validated the effectiveness of the model and the necessity of each structure.", "methods_and_evaluation_criteria": "The six experimental datasets and Tape benchmark mentioned in the article are commonly used prediction datasets in the field of protein language modeling. Protein language models have practical significance for predicting the properties and functions of newly discovered proteins. However, the baselines used in the paper should include more recent papers.", "theoretical_claims": "There is no theoretical claims.", "experimental_designs_or_analyses": "Yes. The article conducts comprehensive experiments on multiple tasks of protein prediction. It performs three independent experiments and takes the average to avoid the impact of randomness. Data partitioning also adheres to previous works, such as removing overlapping proteins between the training and test sets, ensuring the accuracy of the experiments. The results reported in the tables and figures can be more cafully compared to state-of-the-art.", "relation_to_broader_scientific_literature": "The article extensively cites relevant protein models (such as ProtBert, ESM) and knowledge-enhanced methods (including OntoProtein, KeAP), and provides detailed experimental comparisons.", "essential_references_not_discussed": "The technical elements in this paper are easy to follow, and no more references are in need to make the paper more readable.", "other_strengths_and_weaknesses": "Firstly, it introduces a dynamic retrieval mechanism from general language models into protein language models, supporting the dynamic updates of knowledge graphs, which is suitable for the rapid iteration needs in the biomedical field. The main issue is that combining retrieval information with sequence encoding lacks innovation and does not explore deeper graph structures.", "other_comments_or_suggestions": "No.", "questions_for_authors": "Have you considered applying retrieval augmentation directly to the inference stage instead of using the retriever only in the fine-tuning phase?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "D1byKSg7rK", "reviewer": "Reviewer_muG4", "summary": "The paper presents Kara, a knowledge-aware retrieval-augmented protein language model, designed to explicitly integrate knowledge from protein knowledge graphs (PKGs) into protein language models (PLMs). Unlike previous methods that implicitly embed knowledge, Kara directly injects structured knowledge through contextualized virtual tokens, allowing seamless knowledge integration during both pre-training and fine-tuning. A knowledge retriever dynamically retrieves gene descriptions for new proteins, ensuring the model continuously adapts to knowledge updates. Extensive experiments across six protein-related tasks demonstrate that Kara consistently outperforms existing knowledge-enhanced models by effectively capturing high-order biological relationships.\n\n## update after rebuttal\n\nThe author responses address most of my concerns. I will keep my positive rating.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "__Thanks for your kind comments, we place all tables and figures in this anonymous link(https://anonymous.4open.science/r/Rebuttal-F1C0/README.md)__\n\n``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.\n(2) This paper lacks visualization or biological case studies to demonstrate how Kara's representations align with real-world protein functions.``\n\n(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.\n\nAs 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.\n\n(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.\n\n``W2. The study does not evaluate Kara on low-identity or novel proteins, making its generalization unclear.``\n\nWe 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.\n\n``W3. Only ProtBert is tested as the encoder, without comparisons to models like ESM-1b, limiting generalizability.``\n\nWe 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.\n\n``W4. The approach closely follows retrieval-based NLP models and does not introduce fundamentally new ML architectures.``\n\nThanks 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.\n\n(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.\n\n(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.\n\n(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.", "claims_and_evidence": "Overall, the paper's main claims are mostly supported by experimental evidence, but some key assumptions require further verification.\nWell-supported claims:\n1. Kara outperforms existing knowledge-enhanced PLMs (e.g., KeAP, OntoProtein) across six protein-related tasks. The paper provides comprehensive evaluations and ablation studies showing consistent improvements.\n2. Kara effectively mitigates catastrophic forgetting by integrating knowledge into both pre-training and fine-tuning. The use of a knowledge retriever ensures continuous knowledge updates, which static models lack.\n3. Explicit knowledge injection and high-order structure modeling enhance PLM performance. The introduction of contextualized virtual tokens improves knowledge integration and retrieval efficiency.\nClaims needing stronger justification:\n1. 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.\n2. This paper lacks visualization or biological case studies to demonstrate how Kara's representations align with real-world protein functions.", "methods_and_evaluation_criteria": "The proposed methods are well-designed for integrating knowledge into protein language models and offer several unique advantages.\n1. Unlike previous models that only inject knowledge during pre-training, Kara ensures continuous knowledge updates via retrieval, preventing catastrophic forgetting.\n2. Contextualized virtual tokens enable direct knowledge and structure fusion. By representing gene ontology (GO) annotations and functionally similar proteins as learnable tokens, Kara effectively integrates biological insights at the sequence level.\n3. Knowledge retriever dynamically aligns new proteins with the knowledge graph. Instead of relying on static embeddings, Kara retrieves relevant gene descriptions for unseen proteins, improving generalization.\n4. Six tasks are well-chosen to assess Kara’s ability to model sequence-function relationships, demonstrating its applicability to real-world biological challenges.", "theoretical_claims": "The paper primarily focuses on methodological innovations rather than formal theoretical derivations. The key claims rely on empirical results rather than mathematical proofs. The loss functions and optimization strategies, such as structure-based regularization, are well-defined and align with standard machine learning principles.", "experimental_designs_or_analyses": "Strengths of experimental designs/analysis:\n1. The six protein-related tasks cover both structural and functional modeling, ensuring practical relevance.\n2. Ablation studies show that removing virtual tokens, retrieval, or structure-based regularization leads to performance drops, proving their importance.\nLimitations:\n1. The study does not evaluate Kara on low-identity or novel proteins, making its generalization unclear.\n2. Only ProtBert is tested as the encoder, without comparisons to models like ESM-1b, limiting generalizability.", "relation_to_broader_scientific_literature": "This paper builds on prior work in protein language modeling and knowledge-enhanced machine learning. It follows models like OntoProtein and KeAP, which integrate protein knowledge graphs (PKGs) into language models, but differs by introducing explicit knowledge injection through contextualized virtual tokens and a knowledge retriever. This aligns with trends in retrieval-augmented language models, similar to approaches in retrieval-based NLP models that dynamically fetch external knowledge. The use of structure-based regularization is also inspired by contrastive learning methods, commonly used in representation learning to enforce semantic consistency. While Kara advances knowledge integration in PLMs, it does not explore multimodal protein representations (e.g., 3D structures), which have gained traction in recent biological AI research.", "essential_references_not_discussed": "I have already mentioned the missing related research in the previous question “Relation To Broader Scientific Literature”.", "other_strengths_and_weaknesses": "Kara improves protein language models with structured knowledge, making it highly relevant to biological AI. It creatively integrates retrieval, virtual tokens, and structure-based regularization to enhance knowledge utilization. The paper is well-structured and clearly explains the methods, experiments, and results.\nHowever, the approach closely follows retrieval-based NLP models and does not introduce fundamentally new ML architectures. The paper also does not discuss how Kara performs on novel proteins with missing or incomplete knowledge.", "other_comments_or_suggestions": "See above.", "questions_for_authors": "Please refer to the previous sections for key points regarding Claims and Evidence, Methods and Evaluation Criteria, Experimental Designs or Analyses, etc.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper introduces Kara, a knowledge-aware retrieval-augmented protein language model that explicitly integrates protein knowledge graphs (PKGs) with protein language models (PLMs). This integration enhances protein representation learning by incorporating task-specific knowledge and leveraging the structural information inherent in protein knowledge graphs. Experimental validations are conducted across multiple tasks to verify the effectiveness of this approach. 

During the author-reviewer discussion, KeAP was thoroughly discussed, and most reviewers were convinced of its merit, although the novelty is not particularly significant. The paper is well-executed overall. Therefore, I recommend a weak accept
Decision: Accept (poster)