PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Rui Yang, Haoran Liu, Edison Marrese-Taylor, Qingcheng Zeng, Yu He Ke, Wanxin Li, Lechao Cheng, Qingyu Chen, James Caverlee, Yutaka Matsuo, Irene Li

arXiv:2403.05881v3cs.CL

TL;DR

LLMs’ medical factual consistency is a concern, and prior external-knowledge retrieval can introduce irrelevant or unreliable information. KG-Rank retrieves medical-KG triplets and ranks and re-ranks them for LLM answering; it reports ROUGE-L improvements across four medical datasets and in open domains.

  • Problem

    LLMs’ limited medical training raises concerns about consistency with established medical facts, while external-knowledge retrieval can introduce irrelevant or unreliable information.

    Source: §1 · p. 1

  • Method

    KG-Rank identifies question entities, retrieves related medical-KG triplets, and uses ranking and re-ranking techniques to order them for LLM inference.

    Source: §1 · p. 1, §2.3 · p. 3

  • Results

    Over 18% improvement in ROUGE-L across four medical QA datasets and 14% ROUGE-L enhancement in open domains.

    Source: §5 · p. 5

  • Takeaways & Limitations

    The reported medical and open-domain results indicate KG-Rank’s effectiveness and versatility across the evaluated settings.

    Source: §5 · p. 5

    Physician evaluations of answer factuality remain planned, and the ranking method may increase computational time.

    Source: Limitations · p. 5

Abstract

from arXiv · show

Large language models (LLMs) have demonstrated impressive generative capabilities with the potential to innovate in medicine. However, the application of LLMs in real clinical settings remains challenging due to the lack of factual consistency in the generated content. In this work, we develop an augmented LLM framework, KG-Rank, which leverages a medical knowledge graph (KG) along with ranking and re-ranking techniques, to improve the factuality of long-form question answering (QA) in the medical domain. Specifically, when receiving a question, KG-Rank automatically identifies medical entities within the question and retrieves the related triples from the medical KG to gather factual information. Subsequently, KG-Rank innovatively applies multiple ranking techniques to refine the ordering of these triples, providing more relevant and precise information for LLM inference. To the best of our knowledge, KG-Rank is the first application of KG combined with ranking models in medical QA specifically for generating long answers. Evaluation on four selected medical QA datasets demonstrates that KG-Rank achieves an improvement of over 18% in ROUGE-L score. Additionally, we extend KG-Rank to open domains, including law, business, music, and history, where it realizes a 14% improvement in ROUGE-L score, indicating the effectiveness and great potential of KG-Rank.

1 Introduction

The introduction identifies a gap in how external medical knowledge is integrated with LLMs and presents KG-Rank, which uses ranking and re-ranking to support long-form medical QA. It frames the framework as addressing relevance and reliability in retrieved knowledge.

1 Introduction

  • External medical knowledge may improve LLM-generated content, but retrieval can also introduce irrelevant or unreliable information.The introduction identifies better integration of external knowledge as an overlooked question.Source: §1 · p. 1
  • KG-Rank retrieves medical-KG relations for question entities, then uses ranking and re-ranking to order triplets for LLM inference.The KG is UMLS; the methods are intended to filter noise and redundancy.Source: §1 · p. 1
  • The paper presents KG-Rank as an approach to more accurate and reliable long-form medical question answering.The motivation is concern that LLMs’ limited medical training may leave generated content inconsistent with established medical facts.Source: §1 · p. 1
  • KG-Rank combines a medical knowledge graph with ranking techniques for long-answer medical QA, which the authors describe as a first application of this combination.The authors also report validation on medical and open-domain QA tasks.Source: §1 · pp. 1–2

2 Methodology

KG-Rank maps question entities to UMLS, retrieves their one-hop relations, ranks the resulting triplets, and supplies the re-ranked top-p triplets to an LLM for answer generation.

2 Methodology

  • The methodology presents KG-Rank as a four-step framework for long-form medical QA, ending with an LLM response.Source: §2 · p. 2

2.1 External Knowledge Graph

  • KG-Rank uses UMLS as its medical knowledge graph, where knowledge is represented as triples linking medical concepts through relations.The Metathesaurus contains over 3.8 million concepts and more than 78 million relations.Source: §2.1 · p. 2

2.2 Entity Extraction and Mapping

  • For each question, a medical NER prompt identifies related medical entities and maps them to corresponding entities in the knowledge graph.Source: §2.2 · p. 2

2.3 Relation Retrieval and Triplet Ranking

  • After mapping, KG-Rank retrieves one-hop relations; it ranks triplets using question similarity, an LLM-expanded question embedding, or MMR to balance relevance and redundancy.A medical cross-encoder, MedCPT, then re-ranks the triplets, and the top-p are supplied with the task prompt to the LLM.Source: §2.3 · pp. 2–3

3 Experiments

The experiments evaluate KG-Rank on four medical QA datasets using multiple answer-quality metrics and compare ranking settings across GPT-4 and LLaMa2-13b. RR increases ROUGE-L by over 18% on ExpertQA-Bio, while GPT-4 consistently surpasses LLaMa2-13b.

3 Experiments

  • Four selected medical QA datasets contain free-text answers; evaluation uses ROUGE, BERTScore, MoverScore, and BLEURT.Table 1 summarizes average sentence and word counts across the datasets.Source: §3 · p. 3, Table 1 · p. 3

3.1 Results

  • The experiments evaluate GPT-4 and LLaMa2-13b using zero-shot, Similarity, Answer Expansion, and MMR ranking, with re-ranking added to Similarity.Source: §3.1 · p. 3

3.2 Datasets

  • Over 18% increase in ROUGE-L score for ExpertQA-Bio; RR also excels on ExpertQA-Med and Medication QA.RR does not show steady improvement on LiveQA, whose automatically extracted answers may have semantic-coherence and formatting issues.Source: §3.2 · p. 3
  • KG and ranking techniques enhance performance in almost all benchmarks and evaluation metrics in the zero-shot setting.Source: §3.2 · p. 3
  • The three ranking methods show variable performance across datasets, reflecting different strengths and limitations across contexts.Source: §3.2 · p. 3
  • GPT-4 consistently surpasses LLaMa2-13b in zero-shot and ranking settings.Source: §3.2 · p. 3

4 Ablation Study and Analysis

The analysis compares medical-model and re-ranking alternatives, assesses answer factuality, and tests KG-Rank in open-domain QA. Its evaluations report gains across medical and nonmedical settings, alongside a case study of changed clinical advice.

Medical LLM

  • In zero-shot evaluation without KG-Rank, Baize-Healthcare, fine-tuned on LLaMa-7b using medical data, consistently outperforms LLaMa2-7b across all four datasets, as shown in Figure 2.Figure 2 compares BERTScore in the zero-shot setting with LLaMa2-7b and Baize-Healthcare.Source: §4 · p. 3, Figure 2 · p. 4

Re-ranking Models

  • MedCPT steadily outperforms Cohere’s re-rank model on all datasets in the re-ranking comparison.MedCPT is trained on PubMed articles, while Cohere is designed for broader-domain applications.Source: §4 · p. 3, Table 3 · p. 4

Case Study

  • For a patient with acute renal and hepatic failure, KG-Rank recommends 0.8–1 g/kg protein rather than 1.6–2.2 g/kg.The authors judge the lower-range recommendation more factually correct in this case study.Source: §4 · p. 4

LLM-based Evaluation

  • Across four medical datasets, KG-Rank outperforms GPT-4 zero-shot in the GPT-4 factuality evaluation.The evaluation criteria were designed by two resident physicians with over five years of experience.Source: Table 4 · p. 5, §4 · p. 5

KG-Rank in Open Domain

  • In the open-domain Mintaka evaluation, KG-Rank raises accuracy from 60.40% to 61.90%.The test randomly selected 1,000 pairs and replaced UMLS with Wikipedia through the DBpedia API.Source: §4 · p. 5
  • Across law, business, music, and history benchmarks, KG-Rank outperforms the GPT-4 baseline.These experiments use the ExpertQA dataset and evaluate ROUGE-L, BERTScore, and MoverScore.Source: §4 · p. 5

5 Conclusion

KG-Rank combines a medical knowledge graph with ranking techniques for factual medical QA, and the paper reports gains in medical and open-domain evaluations.

5 Conclusion

  • Over 18% improvement in ROUGE-L score is reported across four medical QA datasets.The conclusion also reports a 14% ROUGE-L score enhancement in open domains.Source: §5 · p. 5
  • KG-Rank integrates a medical knowledge graph and ranking techniques to improve the factuality of medical QA.The paper describes it as the first application of KG combined with ranking techniques for long-answer medical QA.Source: §5 · p. 5
  • Table 5 presents base and KG-Rank performance in the open domain.The conclusion reports a 14% ROUGE-L score enhancement for the open-domain application.Source: Table 5 · p. 5, §5 · p. 5

Limitations

The authors identify planned physician evaluation, testing additional medical-specific base models, and ranking-related computational time as limitations or future work.

Limitations

  • Physician evaluations are planned to validate the factual accuracy of KG-Rank’s answers.The authors state that physician evaluation has not yet been conducted.Source: Limitations · p. 5
  • The authors plan to assess more medical-specific base models on medical QA tasks.Source: Limitations · p. 5
  • The ranking method may increase computational time, motivating planned efficiency optimization.Source: Limitations · p. 5

Ethical Considerations

The study uses public medical datasets for academic purposes rather than practical application and reports that its datasets and outputs contain no individual privacy data.

Ethical Considerations

  • The research uses public medical datasets solely for academic purposes, not practical application.Source: Ethical Considerations · p. 5
  • The study reports that benchmark datasets and model outputs contain no individual privacy data.The text also states that GPT-4, LLaMa2-13b, LLaMa2-7b, and baize-healthcare were used for generation and that no harmful content was produced.Source: Ethical Considerations · p. 5

A Prompt Templates

The appendix presents prompt templates used as LLM inputs at different KG-Rank phases, along with physician-designed criteria used in GPT-4 evaluation.

A Prompt Templates

A.1 Medical NER Prompt

  • The Medical NER prompt is designed to extract medical terminologies from a question.Figure 4 depicts the prompt template.Source: §A.1 · p. 8, Figure 4 · p. 8

A.2 Answer Expansion Prompt

  • The answer expansion prompt supports the proposed ranking strategy in step 2 of Figure 1.Figure 5 depicts the prompt template.Source: §A.2 · p. 8, Figure 5 · p. 8

A.3 KG-Enhanced Prompt

  • The KG-enhanced prompt template is used to obtain final answers from LLMs in step 4 of Figure 1.Figure 6 depicts the prompt template.Source: §A.3 · p. 8, Figure 6 · p. 8

A.4 Physician-Designed Criteria for GPT-4 Evaluation

  • Two resident physicians with over five years of experience established criteria for evaluating medical long-form QA, included in GPT-4’s evaluation prompt.Table 6 presents the physician-designed criteria.Source: §A.4 · p. 9, Table 6 · p. 9

A.5 KG-Enhanced Prompt for Mintaka Task

  • A KG-enhanced answer prompt is specially designed for the Mintaka dataset.Figure 7 presents this prompt template.Source: §A.5 · p. 9, Figure 7 · p. 9

B Detailed Evaluation Results

The detailed evaluation reports GPT-4 as outperforming other LLMs in zero-shot medical QA and MedCPT as outperforming Cohere re-ranking in the GPT-4 similarity-ranking setting.

B.1 Zero-shot Performance of Different LLMs

  • GPT-4 performs better than the other LLMs on four medical datasets in the zero-shot setting.The comparison is reported in Table 7.Source: Table 7 · p. 10, §B.1 · p. 10
  • Tables 7 and 8 compare six automatic metrics: ROUGE-1, ROUGE-2, ROUGE-L, BERTScore, MoverScore, and BLEURT.Table 7 covers zero-shot LLMs; Table 8 covers re-rank models.Source: Table 7 · p. 10, Table 8 · p. 10

B.2 Performance of Different Re-rank Models

  • MedCPT outperforms the Cohere re-rank model on four medical datasets within GPT-4 with similarity ranking.The comparison is reported in Table 8.Source: Table 8 · p. 10, §B.2 · p. 10

C More Case Studies

The case studies contrast KG-Rank-enhanced and vanilla models on a breast-cancer prognosis question and an open-domain question about Robert De Niro's Godfather appearances.

C More Case Studies

  • KG-Rank's breast-cancer prognosis answer is described as more factually accurate in terms of medical evidence than the vanilla model's answer.Figure 8 compares vanilla and KG-Rank-enhanced LLaMa2-13b.Source: C More Case Studies · p. 11, Figure 8 · p. 11
  • In Mintaka’s open-domain QA case study asking how many Godfather movies Robert De Niro appeared in, vanilla GPT-4 answered 2, while KG-Rank-enhanced GPT-4 answered 1, matching ground truth.Source: C More Case Studies · p. 11
  • The examples cover both medical QA with LLaMa2-13b and open-domain QA with GPT-4.Figure 8 identifies LLaMa2-13b, while Figure 9 identifies GPT-4.Source: Figure 8 · p. 11, C More Case Studies · p. 11, Figure 9 · p. 12
Loading 2403.05881v3…