PaperlayerView exact revision on arXiv

Source-linked AI summary

Leveraging Medical Knowledge Graphs Into Large Language Models for Diagnosis Prediction: Design and Application Study

Yanjun Gao, Ruizhe Li, Emma Croxford, John Caskey, Brian W Patterson, Matthew Churpek, Timothy Miller, Dmitriy Dligach, Majid Afshar

arXiv:2308.14321v2cs.CLcs.AI

TL;DR

Complex, verbose EHR narratives can overwhelm providers and contribute to diagnostic-error risk, while LLM-generated diagnoses raise concerns about misleading or hallucinated information. The study augments LLM diagnosis generation with DR.KNOWS, a graph model that retrieves case-specific UMLS knowledge paths for prompting. The in-house dataset had more input CUIs and longer CUI outputs on average, while MIMIC-III had a wider range of abstractive concepts.

  • Problem

    Complex, verbose EHR narratives can overwhelm providers and contribute to diagnostic-error risk, while LLM-generated diagnoses raise concerns about misleading or hallucinated information.

    Source: Background · p. 2

  • Method

    The study augments LLM diagnosis generation with DR.KNOWS, a graph model that retrieves case-specific UMLS knowledge paths for prompting.

    Source: Architecture Overview · p. 5

  • Results

    The in-house dataset had more input CUIs and longer CUI outputs on average, while MIMIC-III had a wider range of abstractive concepts.

    Source: Data Overview · p. 5

  • Takeaways & Limitations

    DR.KNOWS enhances retrieval by leveraging graph structure and attention over path-concept relationships to select more accurate, contextually relevant paths.

    Source: Architecture Overview · p. 5

    UMLS concept extractors may miss indirect or nuanced medical concepts, and cosine-similarity path selection depends heavily on embedding quality.

    Source: Limitations · p. 13

Abstract

from arXiv · show

Electronic Health Records (EHRs) and routine documentation practices play a vital role in patients' daily care, providing a holistic record of health, diagnoses, and treatment. However, complex and verbose EHR narratives overload healthcare providers, risking diagnostic inaccuracies. While Large Language Models (LLMs) have showcased their potential in diverse language tasks, their application in the healthcare arena needs to ensure the minimization of diagnostic errors and the prevention of patient harm. In this paper, we outline an innovative approach for augmenting the proficiency of LLMs in the realm of automated diagnosis generation, achieved through the incorporation of a medical knowledge graph (KG) and a novel graph model: Dr.Knows, inspired by the clinical diagnostic reasoning process. We derive the KG from the National Library of Medicine's Unified Medical Language System (UMLS), a robust repository of biomedical knowledge. Our method negates the need for pre-training and instead leverages the KG as an auxiliary instrument aiding in the interpretation and summarization of complex medical concepts. Using real-world hospital datasets, our experimental results demonstrate that the proposed approach of combining LLMs with KG has the potential to improve the accuracy of automated diagnosis generation. More importantly, our approach offers an explainable diagnostic pathway, edging us closer to the realization of AI-augmented diagnostic decision support systems.

Introduction

Verbose EHR narratives can burden clinicians, while diagnosis generation with LLMs must meet high reliability demands. This study proposes DR.KNOWS, which retrieves case-specific UMLS knowledge paths to augment foundational LLMs and support an explainable diagnostic pathway.

Background

  • Complex, verbose EHR narratives may overwhelm providers, and clinicians may skip repetitive sections and rely on decisional heuristics that can contribute to diagnostic errors.EHR notes record patient health trajectories, diagnoses, and treatment plans; narrative complexity and verbosity are linked to diagnostic-error risk.Source: Background · p. 2
  • Automated diagnosis generation from progress notes is an emerging LLM task, but misleading or hallucinated outputs raise patient-safety concerns.Prior work includes diagnosis summarization from progress notes and the ProbSum shared task.Source: Background · p. 2
  • UMLS offers biomedical concepts and semantic relations, but diagnosis-oriented knowledge mining must address patient-specific narratives and a graph of 4.5 million concepts and 15 million relations.Prior UMLS KG applications include information extraction and question answering.Source: Background · p. 2
  • DR.KNOWS retrieves top-N case-specific knowledge paths related to disease pathology and feeds them into foundational LLMs for diagnosis prediction.The study examines fine-tunable T5 and zero-shot prompting with a sandboxed ChatGPT.Source: Background · p. 2

Objectives

  • The work combines a graph model that selects probable diagnoses with explainable paths and an additional module augmenting pretrained language models.It also proposes a human-evaluation framework adapting a diagnostic-safety survey instrument.Source: Objectives · p. 3
  • The study positions path-based KG prompts as its distinction from prior diagnosis summarization and KG integration work.The authors describe KG integration with LLMs as an emerging trend and identify their contribution as a novel graph model for diagnosis prediction.Source: Objectives · p. 3

Methods

The task predicts diagnoses in SOAP progress-note plans from the other note sections. The methods use UMLS concepts and semantic relations for multihop diagnosis reasoning, and evaluate across ProbSum and in-house EHR data with different annotation and concept characteristics.

Daily Progress Notes for Diagnosis Prediction

  • The task predicts the list of problems and diagnoses in the plan section from SOAP progress notes, using ProbSum annotations for training and evaluation.The in-house data has no human annotation; the graph model uses extracted disease-and-syndrome concepts as ground truth, and the final set contains 4,815 notes.Source: Daily Progress Notes for Diagnosis… · p. 3, Data Overview · pp. 4–5

Using UMLS KGs to Find Potential Diagnoses, Given Medical Narratives

  • The study builds its UMLS KG from SNOMED CT concepts and selected semantic types, networks, and relations.A physician selected 107 diagnostic relations from the UMLS repository’s 270 relations to limit complexity and support retrieval accuracy.Source: Using UMLS KGs to Find Potential… · pp. 3–4
  • Multihop reasoning starts from extracted patient concepts, explores relevant neighboring concepts, and selects the top N nodes at each hop to infer diagnoses.The UMLS KG represents concepts as nodes and semantic relations as edges; the example proceeds for up to two hops.Source: Using UMLS KGs to Find Potential… · p. 4, Figure 2 · p. 4

Data Overview

  • MIMIC-III covers ICU patients, while the in-house dataset spans adult patients across hospital settings; both use SOAP progress notes.MIMIC-III contains more than 38,000 patients from 2001–2012; the in-house set covers 2008–2021.Source: Data Overview · p. 4
  • The in-house dataset has more input CUIs and a longer CUI output on average, whereas MIMIC-III has a wider range of abstractive concepts.Table 1 reports average input and output CUIs and defines abstractive concepts as gold-standard diagnoses absent from the input.Source: Data Overview · p. 5, Table 1 · p. 5

Graph Model Development

DR.KNOWS retrieves and ranks UMLS knowledge paths relevant to a patient’s concepts and conditions. Its design combines graph-based representations with path selection informed by path–concept relationships.

Overview

  • DR.KNOWS integrates structured UMLS knowledge into patient-specific diagnostic predictions by retrieving and ranking clinically pertinent paths.A graph neural network incorporates UMLS topology into concept representations to estimate node relevance to patient conditions.Source: Overview · p. 5

Architecture Overview

  • The model encodes input concepts’ 1-hop UMLS subgraphs with SGIN, producing node embeddings that capture neighboring concepts and pretrained concept embeddings.These node embeddings form the basis for path embeddings processed by a path encoder; SGIN was chosen for its Weisfeiler-Lehman-level expressive power.Source: Architecture Overview · p. 5
  • The path encoder scores paths for semantic and logical alignment with the input, then uses top-N scores to select nodes for the next hop.If no suitable diagnosis node is found, exploration ends with a self-loop on the current node.Source: Architecture Overview · p. 5
  • DR.KNOWS adds graph structure and attention-based path–concept relationships beyond vector representations and cosine similarity for retrieval and ranking.The attention network is used to select paths based on their relationships to concepts, not solely their embeddings.Source: Architecture Overview · p. 5

Contextualized Node Representation

The method builds contextualized representations for input-linked UMLS concepts, then scores and ranks candidate paths against the patient’s text and concepts. It retains the highest-scoring concepts for the next reasoning step.

Contextualized Node Representation

  • SapBERT encodes concatenated concept names into contextualized concept representations, which SGIN updates with graph-topology information.The model constructs relation paths from source CUIs and their 1-hop relations; relation edges use one-hot embeddings.Source: Contextualized Node Representation · p. 6
  • A stacked GIN aggregates neighboring-node information across layers to produce each node’s final representation.The final representation concatenates the outputs of the stacked layers.Source: Contextualized Node Representation · p. 6
  • A TF-IDF-based scheme gives higher weights to CUIs and semantic types considered more relevant to predicting diseases.The paper notes that temporal concepts such as “recent” may lead to less useful disease-prediction paths.Source: Contextualized Node Representation · p. 6

Path Reasoning and Ranking

  • Multihead and trilinear attention score candidate paths using the patient’s text, concept list, and current path representations.Multihead attention explicitly models relationships among these inputs; trilinear attention learns them through inner products.Source: Path Reasoning and Ranking · p. 6
  • The path ranker aggregates attention scores and selects the top N concepts for the next iteration.Path embeddings are recursively aggregated as nodes and relations extend the path.Source: Path Reasoning and Ranking · p. 6

Loss Function

Training combines a CUI prediction objective with contrastive learning. The latter distinguishes paths that lead to gold-standard concepts from paths that end at incorrect concepts.

Loss Function

  • The total loss combines CUI prediction loss with contrastive learning loss.The CUI prediction component uses binary cross-entropy to assess whether predicted nodes are in the gold-standard labels.Source: Loss Function · p. 7
  • Contrastive learning compares anchor embeddings with positive and negative path features to encourage discriminative representations.Positive paths lead to the ground-truth concept; negative paths start from the source but end at an incorrect concept.Source: Loss Function · p. 7
  • The training algorithm iteratively selects and ranks relevant paths to extend rather than ranking all n-hop paths in one pass.The authors state that this reduces computational requirements.Source: Loss Function · p. 7

Selection of Foundational Models and Experiment Setup

The study tests graph-predicted paths as auxiliary knowledge for language models, using fine-tuned T5 variants and zero-shot GPT models. It selects a medical-persona prompt with structural path representations for reported results.

Selection of Foundational Models and Experiment Setup

  • The experiments use T5-Large and GPT-3.5-Turbo, evaluating T5 in fine-tuning scenarios and GPT models in zero-shot settings.The reported model sizes are 770 million parameters for T5-Large and 154 billion for GPT-3.5-Turbo.Source: Selection of Foundational Models and… · p. 7
  • Three T5-Large variants—vanilla T5, Flan-T5, and Clinical-T5—are fine-tuned on the ProbSum summarization dataset.Clinical-T5 was specifically trained on MIMIC; Flan-T5 was fine-tuned using an instructional approach.Source: Selection of Foundational Models and… · p. 7
  • To fit T5’s 512-token limit and compare it fairly with GPT’s 4096-token capacity, the input is restricted to subjective and assessment note sections.The paper contrasts these sections with objective sections, which mainly contain numerical values.Source: Selection of Foundational Models and… · p. 7

Prompting Foundational Models to Integrate Graph Knowledge

  • Prompt development expanded manually written task prompts through paraphrasing and backtranslation, then ranked them by perplexity.The five initial prompt sets included three from a nonmedical expert and two from medical experts.Source: Prompting Foundational Models to… · pp. 7–8
  • The selected prompt uses a medical persona and incorporates DR.KNOWS paths in structural form, which outperformed clause-style paths in preliminary experiments.The final prompt was paraphrased from a subject-matter-expert prompt and chosen from a template with minimal perplexity.Source: Prompting Foundational Models to… · p. 8

Evaluation Metrics

The study evaluates graph-model concept retrieval and language-model diagnosis predictions with automated metrics, and assesses generated reasoning using a structured human-evaluation framework.

Automated Evaluation Metrics for Quantitative Analysis

  • DR.KNOWS concept retrieval is evaluated against gold-standard disease CUIs using Recall@N, Precision@N, and F-score.The gold-standard set uses disease and syndrome CUIs from plan sections that connect to input CUIs within at most two hops.Source: Automated Evaluation Metrics for… · p. 8
  • Language-model diagnosis outputs are assessed with the concept-retrieval metrics and ROUGE-L and ROUGE-2.ROUGE-L evaluates longest common subsequence, while ROUGE-2 focuses on bigram matching.Source: Automated Evaluation Metrics for… · p. 8

Human Evaluation for Qualitative Analysis

  • Two medical professionals assessed outputs with and without knowledge paths across two defined events, using a framework adapted from the Safer DX Instrument.Two senior board-certified clinical informatics physicians advised, pilot-tested, and trained the evaluators.Source: Human Evaluation for Qualitative… · pp. 8–9
  • Human evaluation examines model-generated reasoning for comprehension, rationale, knowledge recall, omitted diagnostic reasoning, and abstraction.Comprehension, rationale, and knowledge recall use a 1-to-5 Likert scale; omission and abstraction measures are yes-or-no questions.Source: Human Evaluation for Qualitative… · pp. 8–9

Results

Across evaluations, DR.KNOWS outperformed QuickUMLS on diagnostic-concept precision and F-scores, and knowledge-path prompts improved diagnosis-prediction results. Human ratings showed mixed changes, while error cases included irrelevant paths and failures to use correct paths.

Intrinsic Evaluation of DR.KNOWS on Predicting Diagnostic Concepts

  • DR.KNOWS beat QuickUMLS in precision and F-score on both datasets: MIMIC precision 19.10 vs 13.59 and F-score 25.20 vs 21.13; in-house 22.88 vs 12.38 and 25.70 vs 20.09.Source: Intrinsic Evaluation of DR.KNOWS on… · p. 9, Table 2 · p. 9
  • TriAttn exceeded MultiAttn in F-score on MIMIC (25.20 vs 23.10) and in-house (25.70 vs 17.69), while QuickUMLS led recall, especially in-house.The concept-extractor baseline’s recall was highest because it identified input concepts overlapping with reference CUIs; the in-house set was largely extractive.Source: Intrinsic Evaluation of DR.KNOWS on… · p. 9

Assessing the Impact of DR.KNOWS on Diagnosis Prediction

Human Evaluation Results

  • KG increased rationale agreement (55% vs 50%; P<.001) and effective abstraction (P=.002), reduced affirmative abstraction (78% vs 88%; P=.03), and did not significantly change omission (16% vs 10%; P=.16).Source: Human Evaluation Results · p. 11, Figure 4 · p. 11

Error Analysis

  • Retrieval errors included self-loop or extractive paths, irrelevant abstraction targets, and source CUIs producing unhelpful or contradictory paths.One drug-allergy path contradicted the input note’s “no known drug allergies” description.Source: Error Analysis · pp. 11–12, Figure 5 · p. 11
  • ChatGPT sometimes failed to use a correct retrieved path: the provided cirrhosis-of-liver path was not reflected in its diagnosis output.Source: Error Analysis · p. 12, Figure 6 · p. 12

Discussion

The study reports benefits from DR.KNOWS and KG-informed LLM diagnosis generation, while identifying limitations in knowledge-path relevance and model use of retrieved paths. It frames graph prompting and further refinement as necessary for more reliable diagnostic outputs.

Principal Findings

  • DR.KNOWS improved correct-concept extraction, and predicted paths enhanced T5 performance across metrics; ChatGPT’s reasoning also received strong agreement from human evaluators.Vanilla T5 with path prompts achieved top ROUGE-L and CUI F-scores.Source: Principal Findings · p. 12
  • UMLS paths may reflect biases in their source populations, and case-specific selection may not fully align retrieved paths with patient-specific needs.The authors suggest evaluating relevance with additional context, such as demographic details.Source: Principal Findings · p. 12
  • DR.KNOWS sometimes retrieves paths unrelated to the patient representation, while semantic-similarity ranking depends on embedding quality and accurate starting concepts.The authors identify improved representations, probabilistic modeling, and graph reasoning as future directions.Source: Principal Findings · p. 12
  • ChatGPT sometimes neglects beneficial paths, and its API model’s internal weights and training processes are inaccessible.The authors suggest exploring graph prompting and instruction tuning on open-source language models.Source: Principal Findings · p. 13
  • The authors conclude that graph prompting is needed to guide LLMs toward correct reasoning paths and comprehensive diagnoses, and present a diagnostic-safety evaluation framework for assessing them.They describe DR.KNOWS as a step toward knowledge-grounded AI in medicine that could potentially reduce factual errors.Source: Principal Findings · p. 13

Limitations

The authors identify concept extraction and path selection as limitations requiring continued refinement. The section also records data availability, contributor roles, funding, and a conflict-of-interest disclosure.

Limitations

  • UMLS-based concept extractors may miss indirect or nuanced medical concepts, leaving knowledge representations incomplete and potentially less accurate.The authors call for more sophisticated concept identification techniques.Source: Limitations · p. 13
  • Cosine-similarity path selection depends on embedding quality, and weak semantic representations may yield noisy paths that affect diagnostic suggestions.The authors call for more robust embeddings and reduced reliance on cosine similarity.Source: Limitations · p. 13

Data Availability

  • The study’s generated source-code knowledge graph is available on GitHub, and MIMIC-III is available through PhysioNet..Source: Data Availability · p. 13

Authors' Contributions

  • The listed authors held roles spanning conceptualization, methodology, analysis, writing, validation, and other study activities.Funding acknowledgments name NIH institutes and identify supported investigators.Source: Authors' Contributions · p. 13, Limitations · p. 13

Conflicts of Interest

Multimedia Appendix 1

The appendix listing covers data preprocessing, DR.KNOWS training details, ChatGPT prompt engineering, and T5 fine-tuning.

Multimedia Appendix 1

  • Multimedia Appendix 1 covers data preprocessing, DR.KNOWS training, ChatGPT prompt engineering, and T5 fine-tuning.Source: Multimedia Appendix 1 · p. 14

Abbreviations

The section defines abbreviations used in the paper and provides publication and licensing information.

Abbreviations

  • The abbreviation list expands terms including DR.KNOWS as Diagnostic Reasoning Knowledge Graph System and UMLS as Unified Medical Language System.Source: Abbreviations · p. 16
  • The publication record lists the journal, article identifier, DOI, and publication dates.Source: Abbreviations · p. 17
  • The article is identified as open access under a Creative Commons Attribution License, with reuse requiring attribution and bibliographic and license information.Source: Abbreviations · p. 17

Related papers

Loading 2308.14321v2…