KGARevion: An AI Agent for Knowledge-Intensive Biomedical QA

Weakness 2.5 Details in the Fine-tuning Stage in Review

We apologize for any ambiguity. We have added the following to the last paragraph:

'Our KGARevion is compatible with any medical-related KGs (e.g., PrimeKG, OGB-biokg) and LLMs (e.g., LLaMA3-8B, LLaMA3.1-8B).'

Q1 Example in Section3

Thank you for your careful review, and we apologize for that.

We have revised it to “Which gene interacts with the Heat Shock Protein 70 family that acts as a molecular chaperone, and is implicated in Retinitis Pigmentosa 59 due to DHDDS mutation?”

Q2 Prediction Task for TransE

Thank you for your comment. TransE is pre-trained on PrimeKG on the node representation learning task in a self-supervised way.

Q3 Reason for TransE

TransE is a well-established and popular KG representation learning method and has advantages in capturing the relationship between entities in KG (Liu et al., 2015). For this reason, we chose TransE for our work.

References:
Lin Y, Liu Z, Sun M, et al. Learning entity and relation embeddings for knowledge graph completion $C$ //Proceedings of the AAAI conference on artificial intelligence. 2015, 29(1)

Q4 Sources for D(r)

The description is generated by the LLM based on the following instruction:
'Given a relationship between molecules, generate a description for the input relationship.'

The description dictionary can be found in Appendix Table 8, not Section 8. Apologies for the typo.

Q5 Expression on the Revise Action

We apologize for the confusion. The goal of the Revise action is to expand the number of triplets related to the input query, improving the coverage of its concepts.

The Revise action instructs a standalone LLM to revise an incorrect (‘False’) triplet by replacing either the head or tail entity. For example:

'The following triplet is incorrect: (HSPA1A, interactions, DHDDS). Please revise it to a correct triplet relevant to the input query.'

The LLM might then generate a new triplet, such as (HSPA1B, interactions, DHDDS), by replacing the head entity.

In this way, the Revise action increases the available triplets and ensures better coverage of the medical concepts in the input query.

Q6 Color in Figure 3

The difference between blue and orange lies in the backbone LLM: blue represents LLaMA3-8B, while orange represents LLaMA3.1-8B.

Thank you again for your thoughtful commentary. If you feel our responses are insufficient to motivate increasing your score, we would love to hear from you further about how we can better address your concerns. Thank you again!

Thank you for your valuable and thoughtful feedback. We sincerely appreciate your kind reviews on the figures and benchmarks and how reproducible and versatile our approach is. We also appreciate your constructive suggestions regarding our writing, datasets, and methodological details. To address your concerns, we have carefully responded to each point one-by-one. Most importantly, we have included a number of new experiments and added new datasets and benchmarks, where we can guarantee that there is i) no information leakage from the LLM into the newly added benchmark, as well as that there is ii) no information leakage from the KG into the newly added benchmark.

We respond to all your comments below. If you feel we have not sufficiently addressed your concerns to motivate increasing your score, we would love to hear from you further on what points of concern remain and how we can improve the work in your eyes. Thank you again!

Weakness 1 Data Leakage

Thank you for pointing out this area that we should clarify. Our response to this critical question has two parts. First, we present new experiments and benchmarks. Second, we offer a clarification on benchmarks and results included in the initial submission of this paper.

Part 1
To address your concerns with absolute certainty that there is no leakage, we have included another newly published QA dataset, AfriMed-QA (https://huggingface.co/datasets/intronhealth/afrimedqa\_v2), as a benchmark.

• AfriMed-QA creates a novel multispecialty open-source dataset of 15,000 pan-African clinically diverse QA to rigorously evaluate LLMs in African healthcare.
• The dataset is sourced from over 500 clinical and non-clinical contributors across 16 countries and covers 32 clinical specialties.
• First, because the dataset is sourced from healthcare systems whose data are not online, it means that none of the multiple-choice and open-ended questions and answers can be in the knowledgebase of the LLM. Because of that, this benchmark has not been leaked into KGARevion’s LLM. Furthermore, the knowledge cut-off date for KGARevion’s LLM was before AfriMed-QA was released.
• Second, AfriMed-QA is sourced from clinical, private healthcare systems whereas PrimeKG, STaRK-Prime, and OGB-bioKG are sourced from non-clinical, public, biological data repositories. Furthermore, PrimeKG, STaRK-Prime, and OGB-bioKG were all released several years prior to when the AfriMed-QA benchmark was released. As such, AfriMed-QA has certainly not been leaked into KGARevion’s KG.

In summary, we can now guarantee that there is no data leakage from the newly added AfriMed-QA benchmark into KGARevion because:

• The queries in AfriMed-QA are entirely independent of the underlying KG, such as PrimeKG.
• The answers to AfriMed-QA queries are verifiably absent from the KG or the LLM's knowledge base.
• AfriMed-QA was released on September 12, 2024, after the publication of LLaMA3.1-8B on July 18, 2024.

Please let us know if that is the case, or if you have any other concerns that remain on this point. We fully agree that the issue of data leakage is incredibly important and recognize that many recently published papers suffer from the issue of potential leakage. We are grateful for a chance to clarify this point and include additional results that are scientifically robust and for which we can verify that there is no leakage from training into test sets.

From the 15,000 queries in the AfriMed-QA, we first selected 3,910 multiple-choice queries labeled ‘test’ and ‘Expert’ in the dataset. To save time, we tested GPT-4-Turbo on all 3,910 queries and identified 910 queries that could not be answered correctly. We then tested KGARevion and the baseline models on these 910 queries. The results are shown below:

These results further indicate the effectiveness of KGARevion, providing additional evidence that KGARevion is a well-performing KG-based agent and that KGARevion is scientifically robust. We will make sure to clarify these points in the revised manuscript and integrate these new experiments in the Results section of the paper.

Part 2
We think there may be a misunderstanding here, for which we apologize. While we will clarify the text on this point, we would like to highlight that in the case of benchmarks included in the initial submission, there is no leakage in evaluating the correctness of generated triplets in KGARevion for the following reasons:

• Both ‘Generate’ and ‘Revise’ actions are performed by a standalone LLM, e.g., LLaMA-3.1-8B, by instructing the LLM in generating triplets related to the input query or revising triplets, respectively.
• The fine-tuning of the LLM on the knowledge graph completion task is designed to enable the LLM to leverage structural embeddings learned from the KG, e.g., PrimeKG, to assess the correctness of triplets. This fine-tuned LLM carries out the Review action. It should be noted that the triplets used in fine-tuning the LLM are from PrimeKG.
• Crucially, the Review action is specifically intended to evaluate the correctness of triplets generated by the ‘Generate’ action, NOT those from PrimeKG or OGB-biokg.

We apologize if any of these points were unclear in the paper's original presentation, and will work hard to clarify so this message is better conveyed.

We hope the new experiments in Part 1 and additional explanations in Part 2 provide sufficient evidence that KGARevion and evaluations are scientifically robust. Thank you again for raising this critical point. We will work hard to clarify the text and include new experiments in the paper.

Weakness 2.1 Redundancy of Paragraphs 1&2 in Introduction

We are sorry that the paper wasn’t clear, and we will work to clarify it, specifically focusing on flow, removing unnecessary abbreviations, and removing any redundant parts. We have removed it in the latest version and started directly with Paragraph 3.

Weakness 2.2 Figure 1

Thank you for your suggestion. We have removed Figure 1a to the section of results.

Weakness 2.3 RAG vs. KG

Thank you for your comment. To clarify, our claim refers only to RAG systems potentially suffering from incorrect or incomplete information, not KGs. We aim to use the structured and accurate knowledge in KGs to improve LLM performance in medical QA and enhance its reasoning ability. In KGARevion, the KG is not used for knowledge retrieval like that in RAG systems but as a verifier to review the triplets generated by a standalone LLM.

Weakness 2.4 Writing on Review

Thank you for pointing out this area that we should clarify. We have refined the architecture of the Review action, as shown below:

“

3.2 Review Action

To enable LLMs to judge the correctness of generated triplets, beyond relying solely on semantic dependencies inferred by LLMs (Shinn et al., 2023), the Review action leverages relationships among various medical concepts in KGs. This is achieved by fine-tuning the LLM on a KG completion task and by explicitly integrating entity embeddings learned from KGs into the LLM. Then, the Review action is performed by the fine-tuned LLM to assess the correctness of triplets generated by the Generate action.

3.2.1 Fine-tuning Stage

Generating KG Embeddings and Triplet Descriptions: We use TransE (Bordes et al., 2013), a well-known KG representation learning method, to learn embeddings for entities and relations in $G$. For a triplet $(h, r, t) \\in G$, the learned pre-trained KG embeddings are denoted as $\**e**\_h \\in \\mathbb{R}^{d}$, $\**e**\_r \\in \\mathbb{R}^{d}$, and $\**e**\_t \\in \\mathbb{R}^{d}$, where $d$ represents the embedding dimension. These embeddings are kept fixed.

Additionally, we instruct the LLM to generate a description template for each relation $r \\in G$. We store these descriptions in a dictionary $D(r)$, where the key is the relation and the value is its description. The description of the dictionary can be found in Appendix Table 8.

Aligning Embeddings: Since LLM embeddings are based on token vocabularies (Radford et al., 2019), LLMs cannot directly use pre-trained KG embeddings. To use the pre-trained KG embeddings, we align them with the corresponding descriptions to generate new embeddings for each input triplet.

Specifically, given the description $D(r)$ for the input triplet $(h, r, t)$, we denote the embedding of $D(r)$ obtained from the LLM as $\\mathbf{X} \\in \\mathbb{R}^{\\vert l \\vert \\times d\_L}$, where $\\vert l \\vert$ is the maximum number of tokens and $d\_L$ is the embedding dimension in the LLM.

Next, we concatenate the embeddings of the head entity, relation, and tail entity, denoted as \\mathbf{V} \= g(\mathbf{e}_h); g(\mathbf{e}_r); g(\mathbf{e}_t) $\\in \\mathbb{R}^{3 \\times d\_L}$, where $g(\\cdot): \\mathbb{R}^d \\rightarrow \\mathbb{R}^{d\_L}$.We then apply an attention block (Vaswani, 2017), followed by a two-layer feedforward neural network (FFN, as shown in Fig. 2b) to obtain the aligned triplet embedding matrix $\\mathbf{Z} \\in \\mathbb{R}^{3 \\times d\_L}$ as follows:
$

 \\widehat{\\mathbf{V}} \= \\mathbf{V} \+ \\sigma(\\mathbf{V}\\mathbf{X}^{T})\\mathbf{X} 

\

\\mathbf{Z} \= \\widehat{\\mathbf{V}} \+ ((\\varphi(\\widehat{\\mathbf{V}})\\mathbf{W}\_{1}))\\mathbf{W}\_{2}  

$*where$\sigma(\cdot)$is the Softmax function,$\varphi(\cdot)$represents layer normalization,$\mathbf{W}_1 \in \mathbb{R}^{d_L \times d_h}$and$\mathbf{W}_2 \in \mathbb{R}^{d_h \times d_L}$are trainable parameters in the two-layer FFN, and$d_h$ is the dimension of the hidden layer in the FFN.*

Fine-tuning LLM: After obtaining the aligned embedding $\\mathbf{Z}$, we add it to the beginning of the instruction and fine-tune the LLM using LoRA (Hu et al., 2022) with the next-token prediction loss (Radford, 2018). The instruction is: 'Given a triplet from a knowledge graph, where each triplet consists of a head entity, a relation, and a tail entity, please determine if the triplet is correct and respond with True or False.'

3.2.2 Inference Stage

The fine-tuned LLM is then used in the Review action to check the accuracy of each triplet in $T$, which was generated by the Generate action (3.1). Specifically, we first use UMLS codes (Bodenreider, 2004) to map the entities in the KG and obtain pre-trained KG embeddings for the head entity, relation, and tail entity, respectively. These embeddings, along with the relation description and the instruction, are fed into the fine-tuned LLM to determines whether each triplet is correct or not.

However, not all entities in the generated triplet $(h, r, t) \\in T$ can be mapped to entities in KGs. To address this, the Review action applies a soft constraint rule to distinguish whether the generated triplet is factually wrong or the result of incomplete knowledge in KGs, as follows:

• Factually Wrong: if we can map $h$ and $t$ to entities in KGs and the output of fine-tuned LLM is False, then the triplet $(h, r, t)$ is factually wrong and is removed from $T$.
• Incomplete Knowledge: if we cannot map either $h$ or $t$ to entities in KGs, then the triplet $(h, r, t)$ is considered incomplete knowledge and is kept. ”

Dear Reviewer jrR1,

We sincerely appreciate your insightful question and valuable suggestions during the author-reviewer phase. Those are extremely useful, and I believe this submission experience is the best one we've had so far. Please never hesitate to leave us a comment if you have further ideas during our interaction with other reviewers. Thanks again!

Authors

We sincerely thank the reviewer for their thoughtful and encouraging feedback on the originality, quality, clarity, and significance of KGARevion. We appreciate your positive evaluation of our work and the opportunity to address your insightful comments. In response, we have provided detailed explanations to address your questions. Please do not hesitate to let us know if additional clarifications would be helpful. Thank you again for your valuable input.

Weakness 1 Implementation Details of Revise and Answer

Thank you for your kind suggestion. We have revised the Methods section and the Appendix to include complete information on data processing of the Revise and Answer actions. The updated manuscript now includes an overall description of the method as well as a step-by-step, detailed explanation as follows:

“The Revise and Answer actions are performed by prompting an LLM, such as Llama3.1-8B, instead of training them. This design ensures that our agent framework remains adaptable and generalizable to various new LLMs without fine-tuning.

The Revise action corrects the "False" triplets identified during the Review action. In this step, we prompt the LLM to modify the head or tail entity of each "False" triplet to generate a corrected triplet relevant to the input query. Once the revised triplets are created, the Review action verifies their accuracy. This Revise-Review cycle is repeated until all triplets are verified as correct or a maximum number of iterations is reached.

The Answer action uses the verified triplets to generate a response to the query. Specifically, we provide the verified triplets and the query to the LLM, prompting it to produce a final answer based on the verified triplets. Additional details on the prompts and processes for the Revise and Answer actions are included in the Appendix.”

In the Appendix, we give a detailed description of the methodology for Revise and Answer actions:

Implementation of the Revise Action:

The Revise action is performed by instructing an LLM, such as LLaMA3.1-8B, to correct "False" triplets identified by the Review action. The correction involves modifying either the head or tail entity to generate a valid triplet relevant to the input query.

• Input of False Triplets: The process begins with the Revise action receiving "False" triplets flagged by the Review action. For example, a triplet like (HSPA1A, interacts, DHDDS) may be marked as incorrect.
• Instruction to the LLM: The system provides clear instructions to the LLM, informing it that the given triplet is incorrect and needs revision. The instruction directs the LLM to replace either the head or tail entity and generate a corrected triplet. For instance: "The following triplet is incorrect: (HSPA1A, interactions, DHDDS). Please revise it to a correct triplet relevant to the input query." Additional details about the instructions are available in the Appendix.
• Revised Triplet Generation: The LLM generates a corrected triplet in a specified format, such as a JSON structure. For example: {'Revised_Triplets': (HSPA1B, interactions, DHDDS)}.
• Review of Revised Triplet: The revised triplet is then sent to the Review action for validation. The Review action determines whether the newly generated triplet is correct.
• Iterative Revision Process: 1. If the revised triplet is validated as correct by the Review action, it proceeds to the Answer action for inclusion in the final output. 2. If the revised triplet is still deemed incorrect, it is sent back to the Revise action for further modification. 3. This iterative process continues until a valid triplet is identified or the maximum number of revision rounds is reached.

This approach ensures that KGARevion generates accurate triplets while minimizing errors through iterative refinement.

Implementation of the Answer Action:

The Answer action is executed by prompting an LLM, such as LLaMA3.1-8B. This approach aligns with the Revise action, enabling the system to leverage pre-trained models without additional training.

• Input of Verified Triplets: The Answer action begins with all triplets that have been verified as True by the Review action. These verified triplets serve as the input for generating the final response.
• Instruction to the LLM: KGARevion sends instructions to the LLM, guiding it in generating the final answer using the provided verified triplets.
• Final Answer Generation: The LLM processes the input of verified triplets and generates a final answer relevant to the input query. This response is based entirely on the correctness and relevance of the verified triplets.

Weakness 1 Writing in Fine-tuning Stage

We apologize if these points were unclear in the paper's original presentation. We worked to clarify this section so it is easier to follow. We have refined this section as follows:

“

3.2 Review Action

To enable LLMs to judge the correctness of generated triplets, beyond relying solely on semantic dependencies inferred by LLMs (Shinn et al., 2023), the Review action leverages relationships among various medical concepts in KGs. This is achieved by fine-tuning the LLM on a KG completion task and by explicitly integrating entity embeddings learned from KGs into the LLM. Then, the Review action is performed by the fine-tuned LLM to assess the correctness of triplets generated by the Generate action.

3.2.1 Fine-tuning Stage

Generating KG Embeddings and Triplet Descriptions: We use TransE (Bordes et al., 2013), a well-known KG representation learning method, to learn embeddings for entities and relations in $G$. For a triplet $(h, r, t) \\in G$, the learned pre-trained KG embeddings are denoted as $\**e**\_h \\in \\mathbb{R}^{d}$, $\**e**\_r \\in \\mathbb{R}^{d}$, and $\**e**\_t \\in \\mathbb{R}^{d}$, where $d$ represents the embedding dimension. These embeddings are kept fixed.

Additionally, we instruct the LLM to generate a description template for each relation $r \\in G$. We store these descriptions in a dictionary $D(r)$, where the key is the relation and the value is its description. The description of the dictionary can be found in Appendix Table 8.

Aligning Embeddings: Since LLM embeddings are based on token vocabularies (Radford et al., 2019), LLMs cannot directly use pre-trained KG embeddings. To use the pre-trained KG embeddings, we align them with the corresponding descriptions to generate new embeddings for each input triplet.

Specifically, given the description $D(r)$ for the input triplet $(h, r, t)$, we denote the embedding of $D(r)$ obtained from the LLM as $\\mathbf{X} \\in \\mathbb{R}^{\\vert l \\vert \\times d\_L}$, where $\\vert l \\vert$ is the maximum number of tokens and $d\_L$ is the embedding dimension in the LLM.

Next, we concatenate the embeddings of the head entity, relation, and tail entity, denoted as \\mathbf{V} \= g(\mathbf{e}_h); g(\mathbf{e}_r); g(\mathbf{e}_t) $\\in \\mathbb{R}^{3 \\times d\_L}$, where $g(\\cdot): \\mathbb{R}^d \\rightarrow \\mathbb{R}^{d\_L}$.We then apply an attention block (Vaswani, 2017), followed by a two-layer feedforward neural network (FFN, as shown in Fig. 2b) to obtain the aligned triplet embedding matrix $\\mathbf{Z} \\in \\mathbb{R}^{3 \\times d\_L}$ as follows:
$

 \\widehat{\\mathbf{V}} \= \\mathbf{V} \+ \\sigma(\\mathbf{V}\\mathbf{X}^{T})\\mathbf{X} 

\

\\mathbf{Z} \= \\widehat{\\mathbf{V}} \+ ((\\varphi(\\widehat{\\mathbf{V}})\\mathbf{W}\_{1}))\\mathbf{W}\_{2}  

$*where$\sigma(\cdot)$is the Softmax function,$\varphi(\cdot)$represents layer normalization,$\mathbf{W}_1 \in \mathbb{R}^{d_L \times d_h}$and$\mathbf{W}_2 \in \mathbb{R}^{d_h \times d_L}$are trainable parameters in the two-layer FFN, and$d_h$ is the dimension of the hidden layer in the FFN.*

Fine-tuning LLM: After obtaining the aligned embedding $\\mathbf{Z}$, we add it to the beginning of the instruction and fine-tune the LLM using LoRA (Hu et al., 2022) with the next-token prediction loss (Radford, 2018). The instruction is: 'Given a triplet from a knowledge graph, where each triplet consists of a head entity, a relation, and a tail entity, please determine if the triplet is correct and respond with True or False.'

3.2.2 Inference Stage

The fine-tuned LLM is then used in the Review action to check the accuracy of each triplet in $T$, which was generated by the Generate action (3.1). Specifically, we first use UMLS codes (Bodenreider, 2004) to map the entities in the KG and obtain pre-trained KG embeddings for the head entity, relation, and tail entity, respectively. These embeddings, along with the relation description and the instruction, are fed into the fine-tuned LLM to determines whether each triplet is correct or not.

However, not all entities in the generated triplet $(h, r, t) \\in T$ can be mapped to entities in KGs. To address this, the Review action applies a soft constraint rule to distinguish whether the generated triplet is factually wrong or the result of incomplete knowledge in KGs, as follows:

• Factually Wrong: if we can map $h$ and $t$ to entities in KGs and the output of fine-tuned LLM is False, then the triplet $(h, r, t)$ is factually wrong and is removed from $T$.
• Incomplete Knowledge: if we cannot map either $h$ or $t$ to entities in KGs, then the triplet $(h, r, t)$ is considered incomplete knowledge and is kept. ”

Q1 Example for Fine-tuning

Taking the input “(HSPA8, interacts, DHDDS)” as an example:

• Get the embedding of “HSPA8”, “interacts”, “DHDDS”, respectively. These embeddings are learned from the KG;
• Get the description for the triplet. The description is “HSPA8 interactions with DHDDS, indicating a general interaction that may involve binding, modulation, or other forms of molecular communication”;
• Get aligned embeddings. We align the embeddings generated by KGs with the embeddings of description obtained from LLMs to produce aligned embeddings for (“HSPA8”, “interacts”, “DHDDS”);
• Finally, we send the aligned embeddings for (“HSPA8”, “interacts”, “DHDDS”) along with an instruction into LLM and fine-tune it with the next token prediction objective.

Q2 Removed vs. Kept Triples

Thank you for your suggestion. To address your comment, we have conducted new experiments to show the average number of removed vs. kept triplets. The results are shown in the following Table.

The numbers in this table show how often each case (removed vs. kept triplets) occurs on average across all queries in each dataset.

Q3 Experiments with Irrelevant Knowledge Graph

Thank you for your great point and suggestion. Please refer to our response in Weakness 2 Irrelevant Knowledge Graph, where we report on new experiments and discussion addressing this question.

Q4 Reasoning Paths

Thank you for your valuable comment. KGARevion supports paths indirectly through KG embeddings. Additionally, it supports multi-hop reasoning in the Answer action by considering all verified triplets (which can form a multi-hop reasoning path) when generating the final answer.

Q5 Review vs. Fact-checking

Thank you for your valuable suggestion and for sharing fact-checking approaches with us. We agree with the concept behind existing fact-checking models, which ground generated text on structured, accurate knowledge graphs. We apply this idea when designing the review action in KGARevion.

Our focus is specifically on improving the LLMs’ ability to use medical reasoning by enabling them to verify the correctness of generated triplets, not only using the general knowledge within LLMs but also the structural knowledge from knowledge graphs.

Q6 Applications

Thank you for your kind words. One key aspect of KGARevion is its flexibility. In this work, we have shown that KGARevion performs well in expert-level, domain-specific medical reasoning. Based on our results, KGARevion could also be considered for other fields where domain-specialized and expert-level reasoning is needed.

Thank you for your helpful feedback. If you feel our responses are insufficient to motivate increasing your score, we would love to hear from you further about how we can better address your questions. Thank you again!

Weakness 3 Related Work

Thank you for your valuable suggestion. We have incorporated the following related work into our latest version.

“Fact-checking. The vast amount of generated text has created an increasing need for fact-checking (Graves & Glaisyer, 2012). Early methods addressed this challenge by relying on manually defined rules to verify generated statements (Hassan et al., 2014; Wu et al., 2012; Jiang et al., 2011). KGs have also been leveraged for fact-checking through graph connectivity analysis. These methods validate statements by identifying paths between the entities mentioned (Tunstall-Pedoe, 2010; Shi & Weninger, 2016; Ciampaglia et al., 2015). Furthermore, some approaches use KG representation learning algorithms for KG completion tasks to assess the accuracy of statements (Li et al., 2011). With the advent of LLMs, novel fact-checking techniques have emerged. For example, retrieval-augmented generation (RAG)-based models integrate external evidence retrieval to ensure the consistency and reliability of LLM-generated outputs by verifying and supporting their content (Asai et al., 2024).

References:

• Graves L, Glaisyer T. The fact-checking universe in Spring 2012 $J$ . New America, 2012.
• Naeemul Hassan, Afroza Sultana, You Wu, Gensheng Zhang, Chengkai Li, Jun Yang, and Cong Yu. Data in, fact out: Automated monitoring of facts by factwatcher. Proceedings of the VLDB Endowment, 7(13):1557–1560, 2014.
• You Wu, Pankaj K Agarwal, Chengkai Li, Jun Yang, and Cong Yu. On” one of the few” objects. In Proceedings of the 18th ACM SIGKDD international conference on Knowledge discovery and data mining, pp. 1487–1495, 2012.
• Xiao Jiang, Chengkai Li, Ping Luo, Min Wang, and Yong Yu. Prominent streak discovery in sequence data. In Proceedings of the 17th ACM SIGKDD international conference on Knowledge discovery and data mining, pp. 1280–1288, 2011.
• William Tunstall-Pedoe. True knowledge: Open-domain question answering using structured knowledge and inference. Ai Magazine, 31(3):80–92, 2010.
• Baoxu Shi and Tim Weninger. Discriminative predicate path mining for fact checking in knowledge graphs. Knowledge-based systems, 104:123–133, 2016.
• Giovanni Luca Ciampaglia, Prashant Shiralkar, Luis M Rocha, Johan Bollen, Filippo Menczer, and Alessandro Flammini. Computational fact checking from knowledge networks. PloS one, 10(6):e0128193, 2015.
• Xian Li, Weiyi Meng, and Clement Yu. T-verifier: Verifying truthfulness of fact statements. In 2011IEEE 27th International conference on data engineering, pp. 63–74. IEEE, 2011.
• Akari Asai, Zeqiu Wu, Yizhong Wang, Avirup Sil, and Hannaneh Hajishirzi. Self-RAG: Learning to retrieve, generate, and critique through self-reflection. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview.net/forum?id=hSyW5go0v8.”

Weakness 2 Utility of Knowledge Graph

Thank you for bringing up this point, which we should clarify. We consider PrimeKG and OGB-biokg in our experiments because these two knowledge graphs are widely adopted in the field and have a broad representation of biomedicine. Results in the paper suggest that KGARevion is flexible in that the agent can be used with various KGs and that the agent is not apriori tied to any specific KG. Also, model performance remains strong with the agent using either PrimeKG or OGB-bioKG (see Section 4.4 on the versatility of KGARevion).

However, your comment points to another important and exciting question, which is the coverage of the knowledge domain (i.e., medical vs non-medical KGs). To address your question, we explored running experiments using non-medical Freebase15k-237 KG. However, upon closer inspection, we found that the Freebase KG does not cover at all medical domains, let alone specialized medical knowledge of molecules or medical concepts. Based on KGARevion’s design, if we cannot map entities in the generated triplets to entities in the knowledge graph, those triplets are still retained as "Incomplete Knowledge" and used to generate the final answers. As a result, using an irrelevant KG would lead to outcomes similar to removing the Review action entirely. To test this claim, we performed a new experiment where we excluded the Review action from KGARevion. Results of this new experiment are shown in the Table below:

Based on these results, we conclude that grounding triplets in biomedical KGs is valuable. However, even if KGARevion were used with an irrelevant or out-of-domain KG, it would not experience catastrophic degradation. In such cases, the agent would simply skip the Review action. In the extreme scenario where the KG provides no useful information, the system would rely entirely on the LLM as a fallback.

Thank you for your helpful feedback. We worked hard to improve our paper, and we sincerely hope the reviewers find our responses informative and helpful. If you feel the responses have not addressed your concerns to motivate increasing your score, we would love to hear what points of concern remain and how we can improve our work. Thank you again!

Thank you for your response! If you have any more questions, I’d be happy to discuss them. If not, could you please consider improving your score to reflect the improvements and clarifications in the rebuttal?

Weakness 1 Complex Approach

Thank you for your suggestions. KGARevion uses four main actions to support medical reasoning: Generate, Review, Revise, and Answer.

1. Generate: KGARevion uses the general knowledge in LLM to generate relevant triplets (subject-predicate-object) that could potentially answer the input query.
2. Review: KGARevion then grounds these generated triplets by cross-referencing them with a KG to verify their correctness.
3. Revise: Any triplets identified as "False" during the Review action are corrected by the LLM. The revised triplets are then sent back for further review, where they are re-evaluated until they are deemed correct or the maximum number of review rounds is reached.
4. Answer: KGARevion uses the grounded triplets to infer and provide the final answer to the input query.

By combining the power of LLMs with a knowledge graph, KGARevion ensures that the generated answers are both contextually relevant and factually grounded.

Entity recognizer

We use UMLS to map entities in KG to improve KGARevion's flexibility. Different KGs may use different molecular vocabularies, so adopting a unified recognizer allows KGARevion to be applied across various KGs.

The number of LLMs used in KGARevion

KGARevion uses the same LLM as the backbone for all actions. The Generate, Revise, and Answer actions are handled by the base LLM (e.g., LLaMA3.1-8B), while the Review action uses a fine-tuned version of the same model.

In implementation, the model is loaded just once. Only when the fine-tuned LLM is needed are the fine-tuned weights loaded. Additionally, the fine-tuned model can be used as a standalone module by other projects and contributors.

Q1 New Knowledge Graph

Thank you for your suggestion. Since KGARevion works with any medical-related KGs, we can explore using facts and relationships from UMLS for the KG completion task. To answer your question, we performed a new experiment using the following setup:

1. Download the latest version of UMLS from this link. Note that UMLS contains key terminology, classification, and coding standards, as well as many vocabularies, including CPT, ICD-10-CM, LOINC, MeSH, RxNorm, and SNOMED CT.
2. Fine-tune a large language model (e.g., LLaMA3.1-8B) using the UMLS data for KG completion task. The LLM takes the KG embeddings from UMLS for each triplet and the instruction as input to determine whether each tripe is True or False. The LLM is fine-tuned by LoRA with next token prediction loss.
3. Integrate the fine-tuned LLM with KGARevion and test KGARevion (UMLS) on all datasets.

The results we obtained are as follows:

We are very grateful for your constructive feedback about KGARevion. We appreciate that you recognized the core contributions of our work, as well as the novelty of multiple components of this agent. Please reach out if you have any additional questions or need clarification about this work. Thank you!

Thank you for your helpful feedback. We worked hard to improve our paper, and we sincerely hope the reviewers find our responses informative and helpful. If you feel the responses have not addressed your concerns to motivate increasing your score, we would love to hear what points of concern remain and how we can improve our work. Thank you again!

Readability. As also mentioned by all four reviewers, I found this manuscript hard to follow (though I spent hours on it). Specifically, the motivation in the Introduction section is not logically clear (see, e.g., the "reliability" point above). More importantly, the methods section seems to be over complex (but not in the positive way). This makes it hard for me to know what exactly KGARevion performs at each step. This in turns hampers the reproducibility of the manuscript, as one cannot get enough technical details for fully replicating the approach after reading the manuscript (note that here by reproducibility I am more referring to the writing, not the source code). Therefore, it seems that the manuscript requires a thorough revision before it can be published in top venues like ICLR.

Computational cost and accessibility. A final question I have is regarding the computational cost of the fine-tuning stage. Normally, this can be as least importance for a methodology paper in academia. However, since KGARevion uses KG to deal with medical QA---both are very sensitive to knowledge recency---one may then desire that the approach could be easily updated to leverage the newest KG at the moment (imagine that if KGARevion uses a KG proposed before 2019, then how could it accurately reason about questions regarding COVID-19?). However, I noticed that the authors mentioned in their Appendix, 4 H100 GPUs were used for fine-tuning the LLMs. This then seems to make the update of KGARevion very expensive. I am wondering about the exact cost of the fine-tuning stage (e.g., total GPU hours, memory consumption), and would this be accessible by labs with only one or two consumer-level GPU(s) (e.g., RTX 4090).

I would appreciate it if the authors could kindly address my concerns (though I know they would be busy with responding to the reviewers).

[1] H. Nori et al., Capabilities of GPT-4 on Medical Challenge Problems, arXiv (2023).

[2] D. Veen et al., Adapted large language models can outperform medical experts in clinical text summarization, Nat. Med. (2024).

[3] J. Yang et al., Poisoning medical knowledge using large language models, Nat. Mac. Intell. (2024).

[4] M. Polak and D. Morgan, Extracting accurate materials data from research papers with conversational language models and prompt engineering, Nat. Comm. (2024).

Evaluation (data leakage)

Thank you for your comment. These are great points, which we addressed thoroughly through new experiments. Kindly see our response to Review jrR1 on Weakness 1.

Reliability 1 (KG quality)

Thank you for your comment.

First, to clarify, we simply wanted to point out that existing RAG-based approaches are sensitive to the quality of retrieved information. There was no intention to criticize these methods. Second, thank you for sharing the recent work. However, the use of triplets to model and predict biomedical KGs is well-established, thoroughly validated, and widely used in research (for example, see following references), who provides a principled grounding for triplet-based representation of relational facts in KGARevion. Please see our response to Review yzKB on Weakness 2 for additional information and discussion of these questions.

$1$ Zifeng Wang, Zichen Wang, Balasubramaniam Srinivasan, Vassilis N. Ioannidis, Huzefa Rangwala and RISHITA ANUBHAI. BioBridge: Bridging Biomedical Foundation Models via Knowledge Graphs. In the Twelfth International Conference on Learning Representations, 2024.*
* $2$ Cattaneo A, Martynec T, Bonner S, et al. Towards Linking Graph Topology to Model Performance for Biomedical Knowledge Graph Completion $C$ //ICML'24 Workshop ML for Life and Material Science: From Theory to Industry Applications.*
* $3$ Pengcheng Jiang, Lang Cao, Cao Xiao, Parminder Bhatia, Jimeng Sun, and Jiawei Han. KG-FIT: Knowledge Graph Fine-Tuning Upon Open-World Knowledge. The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.*
* $4$ Yang, Junwei, et al. "Poisoning medical knowledge using large language models." Nature Machine Intelligence (2024): 1-13.*
* $5$ Teneva, Nedelina, and Estevam Hruschka. "Exploring the Properties and Structure of Real Knowledge Graphs across Scientific Disciplines." NeurIPS 2023 AI for Science Workshop. 2023.*

Reliability 2 (LLM for entity extraction)

Thank you for your points. KGARevion is proposed to improve the accuracy of LLMs by incorporating structural knowledge from biomedical KGs, rather than addressing hallucinations in LLMs. The multi-round strategy is designed to ensure completeness. Existing LLMs (e.g., LLaMA3.1-8B) follows instructions well, effectively extracting entities contained in the input query. Regarding triplets, our goal is not to ensure that the LLM extracts every triplet related to the input question. Instead, we focus on analyzing the triplets generated by the LLM to understand its interpretation of the query. By verifying these triplets and removing any factually incorrect ones, we improve the LLM’s accuracy in responding to the query.

Explainability & trustworthiness

Thank you for your comment. Explainability and trustworthiness are indeed important topics for LLMs and represent separate standalone research areas in the LLM field. However, they are not the primary focus of KGARevion, as the KGARevion falls in the area of KG-based AI agents and thus we don’t claim any algorithmic/methodological advances that would relate to the areas of explainability. However, to address your question, KGARevion derives its final answer by considering all verified triplets, which collectively provide a reasoning path to support the outcome. Regarding fact-checking methods, we have also included related work in response to Reviewer yzKB’s question Q5.

Readability

Thank you for your suggestion. We have rewrite the section of the ‘Review’ action. Please refer to the response to Reviewer jrR1 on Weakness 2.4.

Computational cost and accessibility

We fine-tuned the LLM backbone (e.g., LLaMA3.1-8B) using 4 H100 GPUs, which took approximately 4-5 hours. The computational cost of KGARevion depends on the choice of the backbone LLM model, not its design or architecture. KGARevion is compatible with any LLMs and KGs, meaning that one can use KGARevion with smaller or larger language models as well as knowledge graphs of varying sizes, both of which we demonstrate in the paper (see Section 3.2.2 and 4.4 in the paper). We have implemented it with both LLaMA3-8B and LLaMA3.1-8B. The fine-tuning process is efficient and can also be performed using an RTX 4090 GPU.