Lune

ICLR2025

KGARevion: An AI Agent for Knowledge-Intensive Biomedical QA

Xiaorui Su, Yibo Wang, Shanghua Gao, Xiaolong Liu, Valentina Giunchiglia, Djork-Arné Clevert, Marinka Zitnik

2025Year
4Citations

Abstract

Biomedical reasoning integrates structured, codified knowledge with tacit, experience-driven insights. Depending on the context, quantity, and nature of available evidence, researchers and clinicians use diverse strategies, including rule-based, prototype-based, and case-based reasoning. Effective medical AI models must handle this complexity while ensuring reliability and adaptability. We introduce KGAREVION, a knowledge graph-based agent that answers knowledge-intensive questions. Upon receiving a query, KGAREVION generates relevant triplets by leveraging the latent knowledge embedded in a large language model. It then verifies these triplets against a grounded knowledge graph, filtering out errors and retaining only accurate, contextually relevant information for the final answer. This multi-step process strengthens reasoning, adapts to different models of medical inference, and outperforms retrieval-augmented generationbased approaches that lack effective verification mechanisms. Evaluations on medical QA benchmarks show that KGAREVION improves accuracy by over 5.2% over 15 models in handling complex medical queries. To further assess its effectiveness, we curated three new medical QA datasets with varying levels of semantic complexity, where KGAREVION improved accuracy by 10.4%. The agent integrates with different LLMs and biomedical knowledge graphs for broad applicability across knowledge-intensive tasks. We evaluated KGAREVION on AfriMed-QA, a newly introduced dataset focused on African healthcare, demonstrating its strong zero-shot generalization to underrepresented medical contexts. Present work. We introduce KGAREVION, a KG-based LLM agent designed for complex biomedical QA that integrates the non-codified knowledge of LLMs with the structured, codified knowledge found in KGs. As illustrated in Fig. 2 , KGAREVION executes four key actions to ensure accurate and context-aware biomedical reasoning. First, it prompts the LLM to generate relevant triplets based on the input question. To effectively leverage structured KG data, KGAREVION fine-tunes the LLM on a KG completion task, incorporating pre-trained structural embeddings of triplets as prefix tokens. The fine-tuned model then evaluates the correctness of the generated triplets. Next, KGAREVION performs a 'Revise' action to correct erroneous triplets, refining the knowledge base before selecting the final answer. Given the complexity of medical reasoning, KGAREVION adaptively chooses the most appropriate reasoning approach for each question, allowing for nuanced and context-aware QA. This flexibility enables KGAREVION to handle both multi-choice and openended QA tasks. Our key contributions include: 1 ⃝ Developing KGAREVION, a versatile KG agent that dynamically adjusts reasoning strategies, achieving a 6.75% improvement over 15 baseline models in seven datasets, including three challenging newly curated benchmarks. 2 ⃝ Demonstrating that grounding through generated triplets significantly enhances KGAREVION's capabilities across multiple KGs. 3 ⃝ Showing that KGAREVION effectively answers complex, knowledge-intensive medical queries in both multi-choice and open-ended QA formats. 4 ⃝ Evaluating KGAREVION on African healthcare datasets: We benchmark KGAREVION on AfriMed-QA, a newly introduced Published as a conference paper at ICLR 2025 dataset focused on African healthcare. The results highlight KGAREVION's strong zero-shot generalization, demonstrating its ability to reason effectively in underrepresented medical contexts. 5 ⃝ Analyzing robustness to input variations: We analyze KGAREVION's sensitivity to changes in question structure, answer ordering, and answer relabeling. Unlike LLMs, which exhibit high variance when answer choices are reordered, KGAREVION maintains stable performance, demonstrating its stronger robustness in real-world settings. KGAREVION is available at https://github.com/mims-harvard/KGARevion .