KG -Rank: Enhancing Large Language Models for Medical QA with Knowledge Graphs and Ranking Techniques (2024)

Abstract: Large Language Models (LLMs) have significantly advanced healthcare innovation on generation capabilities. However, their application in real clinical settings is challenging due to potential deviations from medical facts and inherent biases. In this work, we develop an augmented LLM framework, KG-Rank, which leverages a medical knowledge graph (KG) with ranking and re-ranking techniques, aiming to improve free-text question-answering (QA) in the medical domain. Specifically, upon receiving a question, we initially retrieve triplets from a medical KG to gather factual information. Subsequently, we innovatively apply ranking methods to refine the ordering of these triplets, aiming to yield more precise answers. To the best of our knowledge, KG-Rank is the first application of ranking models combined with KG in medical QA specifically for generating long answers. Evaluation of four selected medical QA datasets shows that KG-Rank achieves an improvement of over 18% in the ROUGE-L score. Moreover, we extend KG-Rank to open domains, where it realizes a 14% improvement in ROUGE-L, showing the effectiveness and potential of KG-Rank.

## introduction large language models (llms), such as @xcite and llama2 @xcite , have demonstrated powerful generative capabilities @xcite . despite their considerable potential in various domains, including medicine @xcite @xcite , their limited training on medical data raises concerns about the consistency of the generated con-tent with established medical facts @xcite . to address this challenge without additional computational cost, previous research, such as almanac @xcite and chatent @xcite , leverages external medical knowledge to enhance the accuracy and reliability of llmgenerated content. however, merely retrieving external knowledge risks introducing irrelevant or unreliable information @xcite , which can compromise the effectiveness of llms, and raise issues of credibility, data consistency, privacy, security, and legality. while previous studies have emphasized the advantages of utilizing external knowledge, they have overlooked a crucial question: how to better integrate external knowledge? in this work, we propose kg-rank, an augmented framework that integrates a structured medical knowledge graph (kg) with ranking techniques into llms to achieve more accurate and reliable long-form medical question-answering (qa). we first retrieve one-hop relations of related medical entities from the medical kg (unified medical language system (umls)) @xcite . to retain relevant information from the kg, we then propose to apply ranking and re-ranking methods to optimize the ordering of triplets. specifically, we introduce three ranking techniques to improve the integration of llm with kg by filtering irrelevant data, highlighting key information, and ensuring diversity. these techniques also streamline the process by reducing the number of triplets required for llm inference. additionally, we apply re-ranking models to reassess and emphasize the most relevant triplets, enhancing the factuality of kg-rank in the long-form medical qa task. to summarize, our contributions are: (1) we propose kg-rank, a kg-augmented llm framework for the medical qa task. to the best of our knowledge, this is the first application of kg com-bined with ranking techniques to enhance llms for medical qa with long answers. (2) we incorporate different ranking and re-ranking techniques to eliminate noise and redundancy in the kg-retrieval stage. (3) we validate the effectiveness of kg-rank on both medical and various open-domain qa tasks. all the data and code can be found at https://github.com/yangrui525/kg-rank . ## methodology as shown in fig. 1 , we introduce the kg-rank (knowledge graph -rank) framework for the longform medical qa task. atrial fibrillation heart failure diabetes mellitus ... query : a 56 year old male patient with atrial fibrillation presents to the clinic. given their history of heart failure, diabetes and pad, what is their risk of stroke? should they be placed on anticoagulation? ## experiments we conduct experiments on four selected medical qa datasets, in which the answers are free-text, as shown in tab. 1. liveqa @xcite consists of health questions submitted by consumers to the national library of medicine. it includes a training set with 634 qa pairs and a test set comprising 104 qa pairs, which is used for evaluation. expertqa @xcite is a highquality long-form qa dataset with 2177 questions spanning 32 fields, along with answers verified by domain experts. among them, 504 medical questions (med) and 96 biology (bio) questions were used for evaluation. medicationqa @xcite includes 690 drug-related consumer questions along with information retrieved from reliable websites and scientific papers. we evaluate the generated answers using rouge @xcite , bertscore @xcite , mover-score @xcite and bleurt @xcite . ## ablation study and analysis medical llm to further investigate the capability of the medical llm, we compare the zero-shot performance of llama2-7b and baizehealthcare @xcite without kg-rank. baize-healthcare, which is fine-tuned on llama-7b using medical data, consistently outperforms llama2-7b across all four datasets, as shown in fig. 2 . more comparison results can be found in appendix b.1. re-ranking models we employ gpt-4 with similarity ranking as the final setting and compare two re-ranking models: the medcpt cross-encoder model, trained on the extensive pubmed articles, and the cohere ( https://cohere.com ) re-ranking model, designed for broader domain applications. as shown in tab. 3, medcpt steadily outperforms the cohere re-rank model on all datasets, highlighting the importance of specialized re-rank models ## conclusion in this work, we propose kg-rank, an enhanced llm framework that integrates a medical kg and ranking techniques to improve the factuality of medical qa. as far as we know, kg-rank is the first application of kg combined with ranking techniques for long-answer medical qa. across four medical qa datasets, kg-rank demonstrates over an 18% improvement in rouge-l score. its application to open domains yields a 14% rouge-l score enhancement, underscoring kg-rank's effectiveness and versatility.