🤖 AI

Trajectory-Aware Clinical Risk Prediction via Severity-Grounded Knowledge Graphs and Retrieval-Augmented Generation

The paper proposes TRACER, a novel framework that integrates severity-grounded medical knowledge graphs with retrieval-augmented generation to effectively leverage heterogeneous external knowledge and unstructured clinical notes, achieving significant improvements in mortality and readmission risk prediction on MIMIC datasets.

Original authors: Kyunghoon Jeon, Youmin Ko, Woohwan Jung, Hyunjoon Kim

Published 2026-07-22

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Kyunghoon Jeon, Youmin Ko, Woohwan Jung, Hyunjoon Kim

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to predict the weather for a specific city, but you only have a few scattered, messy notes from the last time someone looked out the window. You don't have a full history of storms, just a few scribbles like "cloudy" or "windy." Now, imagine trying to predict a patient's health future using their Electronic Health Records (EHRs). These records are like those scattered notes: they contain lists of diagnoses, medications, and doctor's notes, but they are often incomplete, irregular, and full of gaps. Doctors need to know if a patient is likely to get sicker or return to the hospital, but the data is too sparse to see the whole picture.

To solve this, scientists have started using "Knowledge Graphs," which are like giant, interconnected maps of medical facts. Think of a knowledge graph as a massive library where every disease, symptom, and treatment is a book, and the shelves show how they connect. If you have a patient with "asthma," the map can show you that asthma often leads to "respiratory failure." However, traditional maps treat every book on the shelf as equally important. They don't realize that "a mild cough" is very different from "heart failure," even if both are just words on a list. This paper introduces a new way to read these maps and the messy notes, turning a chaotic pile of data into a clear, story-driven prediction.

The researchers behind this study, led by Kyunghoon Jeon and colleagues, propose a system called TRACER. They realized that existing methods were failing because they treated all medical conditions the same, ignoring how serious a specific illness really was. To fix this, they built a "severity-grounded" knowledge graph. Imagine this as a map where the most dangerous roads are highlighted in bright red and the safe paths are in soft green. They used a smart computer program (an AI) to read medical literature and assign a "severity score" to every disease, ranging from 1 to 20. This allows the system to know that a patient with a score of 18 is in much more trouble than one with a score of 3, even if they have the same number of recorded visits.

But TRACER doesn't just look at the map; it also reads the story. The team noticed that doctors write detailed notes about how a patient is feeling, which often contain clues that the structured lists miss. TRACER digs into these unstructured notes to find "red flags"—subtle warnings like a doctor noting that a patient's breathing is getting harder, even if the official diagnosis code hasn't changed yet. Furthermore, instead of looking at a patient in isolation, TRACER finds "look-alike" patients. It asks, "Who else had a similar history and a similar set of recent symptoms?" and uses their outcomes as a guide.

The results of this approach are impressive. When tested on two massive, real-world hospital datasets (MIMIC-III and MIMIC-IV), TRACER significantly outperformed the best existing methods. For predicting patient mortality (death), the system improved its sensitivity by up to 28.5% compared to the previous top method. For predicting hospital readmissions (patients coming back too soon), it saw a 19.7% increase in Macro F1 score. The authors suggest that this happens because TRACER can see the "trajectory" of a disease—how it moves from one visit to the next—rather than just looking at a static snapshot. By combining the severity-weighted map, the detailed story from the notes, and the wisdom of similar past cases, TRACER provides a much clearer, more reliable prediction of what might happen next, offering doctors a powerful new tool to catch critical risks early.

Technical Summary: Trajectory-Aware Clinical Risk Prediction via Severity-Grounded Knowledge Graphs and Retrieval-Augmented Generation

Problem Statement

Clinical risk prediction, specifically for mortality and hospital readmission, relies heavily on Electronic Health Records (EHRs). However, existing methods face three critical limitations in real-world scenarios:

Data Sparsity and Irregularity: Most patients have limited visit histories (often fewer than five visits) with heterogeneous time intervals. Standard sequential models struggle to capture meaningful temporal dynamics or identify critical turning points under these sparse conditions.
Uniform Treatment of Medical Concepts: Existing Knowledge Graph (KG) approaches typically treat all diagnoses equally, ignoring the varying clinical severity of different conditions. This leads to the failure to capture "red-flag" events or high-risk disease progressions, as critical diagnoses are not distinguished from less consequential ones.
Underutilization of Unstructured Data: Rich, nuanced information in clinical notes (e.g., physician assessments, suspected complications, treatment responses) is often discarded or oversimplified into structured codes, stripping away vital narrative evidence for early warning signs.
Lack of Cohort Context: Relying solely on a single patient's sparse history fails to leverage clinical experience embedded in historically similar patient cohorts.

Methodology: The TRACER Framework

The authors propose TRACER (Trajectory-Aware and Clinically Grounded prediction framework), a Retrieval-Augmented Generation (RAG) system designed to address these limitations through three core stages:

1. Severity-Weighted Medical Knowledge Graph (SMKG) Construction

To address the uniform treatment of diagnoses, TRACER constructs a novel Severity-Grounded Medical Knowledge Graph (SMKG).

Integration: It unifies existing medical KGs (e.g., UMLS), PubMed abstracts, and Wikipedia summaries.
Severity Scoring: An LLM assigns a fine-grained severity score (1–20) to each diagnosis node based on retrieved medical literature and clinical guidelines. This score reflects clinical severity, distinguishing high-risk conditions from incidental events.
Weighting: Diagnosis nodes are assigned these severity scores, while other nodes (procedures, medications) are assigned a weight of 0, ensuring the graph prioritizes disease severity.

2. Patient Medical Profile (PMP) Retrieval

To mitigate data sparsity, TRACER augments the target patient's sparse EHR with a structured and unstructured medical profile consisting of:

Severity-Aware Trajectory Retrieval: Instead of treating visits as isolated points, the system models health evolution as paths in the SMKG. It retrieves Key Supporting Trajectories (KSTs) connecting concepts across consecutive visits.
- Dual-Query Strategy: It performs a dual retrieval for Risk Factors (e.g., progressive organ dysfunction) and Protective Factors (e.g., treatment response, stability).
- Scoring: Trajectories are scored based on semantic relevance to the query and their aggregate severity.
- Refinement: A Maximal Marginal Relevance (MMR) algorithm and Natural Language Inference (NLI) based stopping criteria are used to filter redundancy and ensure the selected trajectories are diverse, relevant, and clinically supportive.
Clinical Note Retrieval: Relevant passages from unstructured clinical notes (progress notes, discharge summaries) are retrieved using the same risk/protective queries to capture subjective assessments and nuanced cues.
Similar Patient Retrieval: The system retrieves the top- $m$ similar patients based on Recent-visit Jaccard similarity (overlap in medical concepts from the most recent visits) rather than broad sequence embeddings. This provides cohort-level context and outcome patterns.

3. Retrieval-Augmented Clinical Risk Prediction

The final stage employs an LLM (specifically GPT-4o-mini in the experiments) to generate predictions.

Prompt Construction: The prompt integrates the task description, the target patient's demographics, the retrieved KSTs (repacked to prioritize relevance), clinical note passages, and the context of similar patients.
Reasoning: The LLM is guided through a step-by-step reasoning process to weigh risk vs. protective signals, compare with similar cohorts, and output a binary prediction (Mortality: 0/1; Readmission: 0/1).
Repacking Strategy: To mitigate the "lost-in-the-middle" problem in long-context LLMs, the system dynamically repacks the retrieved evidence, placing the most relevant trajectories and notes closer to the task description.

Key Contributions

Severity-Grounded KG (SMKG): The proposal of a medical KG where diagnosis nodes are assigned fine-grained, literature-derived severity scores, enabling the model to distinguish critical clinical signals.
Trajectory-Aware Augmentation: A method to address data sparsity by augmenting patient data with severity-aware progression paths (KSTs) and clinically similar patient cases, allowing the model to reason about disease progression even with limited history.
Nuanced Textual Evidence: The integration of salient, unstructured clinical notes alongside structured data to capture red-flag events and personal experiences often missed by structured codes.
Interpretability and Performance: The framework provides path-level, clinically grounded explanations for predictions while achieving state-of-the-art performance on standard benchmarks.

Experimental Results

Experiments were conducted on MIMIC-III and MIMIC-IV datasets for mortality and 15-day readmission prediction tasks.

Performance Gains: TRACER significantly outperforms state-of-the-art baselines (including GraphCare, KARE, and DearLLM).
- Mortality Prediction (MIMIC-III): Achieved a 28.5% increase in Macro F1 and a 28.5% increase in Sensitivity compared to the best baseline (KARE).
- Readmission Prediction (MIMIC-III): Achieved a 19.7% increase in Macro F1.
- MIMIC-IV Results: Consistent improvements were observed, with TRACER achieving the highest Macro F1 and Sensitivity across both tasks.
Ablation Studies:
- Removing severity grounding ("w/o Sev.") caused a significant drop in sensitivity (from 0.7128 to 0.4950), confirming the importance of severity-aware retrieval.
- Removing trajectory retrieval ("w/o Traj.") resulted in the largest drop in Macro F1 for mortality prediction.
- Removing clinical notes and similar patient evidence significantly impacted readmission prediction, highlighting the value of unstructured text and cohort context.
Efficiency: TRACER demonstrated approximately 33% faster inference time compared to KARE, attributed to the efficiency of trajectory retrieval versus iterative multi-hop subgraph expansion.
Robustness: The framework showed robustness across different LLM backbones (including open-source models like LLaMA2 and Qwen) and different knowledge snapshots.

Significance and Claims

The paper claims that TRACER represents a substantive methodological advance beyond incremental model refinements. Its significance lies in:

Addressing Real-World Sparsity: By leveraging severity-weighted trajectories and similar patient cohorts, the model effectively handles the extreme sparsity of real-world EHR data where most patients have fewer than five visits.
Clinical Grounding: The integration of severity scores derived from medical literature ensures that the model prioritizes high-risk clinical signals, aligning more closely with clinical reasoning than methods that treat all diagnoses equally.
Interpretability: Unlike "black-box" deep learning models, TRACER provides transparent, path-level explanations derived from retrieved evidence, enhancing trust in clinical decision support systems.
Balanced Performance: The framework successfully balances sensitivity (identifying true positives) and specificity, which is critical in clinical settings to avoid both missed diagnoses and unnecessary interventions.

The authors conclude that capturing causal progression and amplifying high-risk signals through a hybrid retrieval-augmented design leads to more effective risk stratification than flattening medical history or relying solely on structured data.