Reasoning-Enhanced Healthcare Predictions with Knowledge Graph Community Retrieval
Authors: Pengcheng Jiang
Cao Xiao
Minhao Jiang
Parminder Bhatia
Taha Kass-Hout
Jimeng Sun
Jiawei Han
Venue: ICLR 2025 Poster (2025)
Keywords: EHR Prediction
Large Language Models
Knowledge Graphs
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: Large language models (LLMs) have demonstrated significant potential in clinical decision support. Yet LLMs still suffer from hallucinations and lack fine-grained contextual medical knowledge, limiting their high-stake healthcare applications such as clinical diagnosis. Traditional retrieval-augmented generation (RAG) methods attempt to address these limitations but frequently retrieve sparse or irrelevant information, undermining prediction accuracy. We introduce KARE, a novel framework that integrates knowledge graph (KG) community-level retrieval with LLM reasoning to enhance healthcare predictions. KARE constructs a comprehensive multi-source KG by integrating biomedical databases, clinical literature, and LLM-generated insights, and organizes it using hierarchical graph community detection and summarization for precise and contextually relevant information retrieval. Our key innovations include: (1) a dense medical knowledge structuring approach enabling accurate retrieval of relevant information; (2) a dynamic knowledge retrieval mechanism that enriches patient contexts with focused, multi-faceted medical insights; and (3) a reasoning-enhanced prediction framework that leverages these enriched contexts to produce both accurate and interpretable clinical predictions. Extensive experiments demonstrate that KARE outperforms leading models by up to 10.8-15.0\% on MIMIC-III and 12.6-12.7\% on MIMIC-IV for mortality and readmission predictions. In addition to its impressive prediction accuracy, our framework leverages the reasoning capabilities of LLMs, enhancing the trustworthiness of clinical predictions.

Published as a conference paper at ICLR 2025
REASONING-ENHANCEE HEALTHCARE PREDICTIONS WITH KNOWLEDGE GRAPH COMMUNITY RETRIEVAL
Pengcheng Jiang'
Cao Xiao
Minhao Jiang"
Parminder Bhatia
Taha Kass-Hout
Jimeng Sun
Jiawei Han
UIIC
1GE HealthCare
ABSTRACT
Large language models (LLMs) have demonstrated significant potential clin- ical decision support. Yet LLMs still suffer from hallucinations and lack fine- grained contextual medical knowledge, limiting their high-stake healthcare ap- plications such as clinical diagnosis. Traditional (RAG) methods attempt to address these limitations but frequently retrieve sparse irrelevant information, undermining prediction accuracy. We introduce KARE, a novel framework that integrates knowledge graph (KG) community-level re- trieval with LLM reasoning to enhance healthcare predictions. KARE constructs a comprehensive multi-source KG by integrating biomedical databases, clinical literature, and L-geeeeee insights, and organizes it using hierarchical graph community detection and summarization for precise and contextually relevant in- formation retrieval. Our key innovations include: (1) a dense medical knowl- edge structuring approach enabling accurate retrieval of relevant information; (2) a dynamic knowledge retrieval mechanism that enriches patient contexts with fo- cused, multi-faceted medical insights; and (3) a reasoning enhanced prediction framework that leverages these enriched contexts produce both accurate and interpretable clinical predictions. Extensive experiments demonstrate that KARE outperforms leading models by up to 10.8-15.09 on MIMIC-III and 12.7% MIMIC-IV for mortality and readmission predictions. In addition to its im- pressive prediction accuracy, our framework leverages the reasoning capabilities of LLMs,, enhancing the trustworthiness of clinical predictions.
generation
reasoning-enhanced
INTRODUCTION
Large language models (LLMs) (Touvron al., 2023a;b; OpenAI et al., 2024; Team al., 2024) has revolutionized natural language processing, offering unprecedented capabilities in understanding and generating human-like text. the healtheare domain, LLMs hold the potential to transform clinical decision-makinn by providing insights derived from vast amounts of medical data (Wornow et al., 2023; Yang et al., 2022). There has been many recent explorations on applying ML-based methods in healthcare domain (Choi al., 2016; 2017; Shickel al., 2018; Choi al., 2018; Ma et al., 2020a; Gao et al., 2020; Zhang et al., 2021: Wu et al., 2023; Jiang et al., 2024a; Zhu al., 2024a; Xu et al., 2024). However, deploying LLMs in clinical settings presents significant challenges, mainly because LLMs may produce hallucinations or incorrect information due to a lack of specialized medical knowledge. Traditional retrieval-augmeened generation (RAG) techniques (Lewis et al., 2021). which aim to mitigate hallucinations by retrieving external knowledge, often fall short in healthcare applications. They tend retrieve information that, while semantically similar in latent space, fails to provide meaningful clinical insights, leading to suboptimal outcomes for precise healthcare predictions (Shi al., 2024; Magesh et al., 2024; et al., 2024). For instance, when dealing with the diagnosis of heart failure, traditional RAG model might retrieve data on several conditions that are semantically similar, such as "acute coronary syndrome" or "ischemic heart disease" due to their close proximity in latent space. However, these conditions, while related, do not capture the specific nuances of heart failure, such as the impact of left ventricular ejection fraction specific biomarkers like NT-proBNP levels.
Knowledge graphs (KGs) offer promising solution by providing structured representations of med- ical knowledge, capturing complex relationships between clinical entities (Liu et al., 2019; Yasunaga

---

Published as a conference paper at ICLR 2025
et al., 2022; Zhang al., 2022). Integrating KGs with LLMs enhance the models' reasoning capabilities and provide Nomain-specifit knowledge essential for accurate healthcare predictions (Soman et al., 2024). However, previous studies have often lacked interpretability and failed to fully leverage the reasoning strengths of LLMS (Jiang et al., 2024a; Xu et al., 2024; Zhu et al., 2024a).
Graph community retrieval has been a proven technique in various domains, such social network analysis (Fortunato, 2010; Jin et al., 2021) and recommendation systems (Salha et al., 2019) for efficiently extracting relevant and contextual information from large-scale graphs. Recent work like GraphRAG (Edge et al., 2024) has demonstrated the superior performance of graph community retrieval compared to naive RAG in the query-focusee summarization task. However, the application of graph retrieval for LLM- baae- healthcare prediction remains largely unexplored.
In this paper. we introduce KARE (Knowledge Aware Geasoning-Enhaneee HealthCare Predic- tion), a new framework that combines KG community-level retrieval (e.g., retrieving relevant sub- graphs) with LLM reasoning to improve healthcare prediction.
Our technical contributions can be summarized as follows:
1. Multi-Source Medical Knowledge Structuring and Indexing: We develop novel method to construct and index multi-source medical concept KGs by integrating concept-specific knowl- edge derived from relationships among different concepts in patients' electronic health records (EHRs). We employ hierarchical graph community detection and summarization techniques to organize the KG into semantically meaningful communities, facilitating precise, fine-grained, and contextually relevant information retrieval.
2. Context Augmentation with Dynamic Knowledge Retrieval from KG: We propose a context augmentation technique that can dynamically enrich patient data with knowledge from relevant KG communities tailored to the patient context. By retrieving pre-summarized communities, enrich the input the LLMs with focused, multi-faceted medical insights, addressing the limitations of traditional RAG methods.
Reasoning-Enhannne Clinical Prediction Framework: We leverage the augmented patient context to enable LLMs to generate step-by-step reasoning chains, enhancing both interpretability and prediction accuracy in clinical tasks.
To evaluate the KARE framework, we conducted experiments on in-hospital mortality and hospi- tal readmission prediction tasks using the MIMIC-III and MIMIC-IV datasets(Johnson al., 2016; 2020). KARE sigmificantly outperforms the best baseline models. Specifically, KARE achieves improvements over best baselines up 10.8%, 15.0%, 12.6%, and 12.7% on the MIMIC-III mor- tality, MIMIC-III readmission, MIMIC-IV mortality, and MIMIC-IV readmission prediction tasks, respectively. By attaining higher prediction accuracy and leveraging reasoning capabilities, KARE enhances the trustworthiness of clinical decision support systems. The reasoning process incorpo- rates valuable evidence from relevant medical knowledge, facilitating more informed and explain- able predictions that are needed in clinical decision making.
RELATED WORKS
Clinical Predictive Models. Electronic health record (EHR) data have become invaluable in the medical field, supporting predictive tasks aimed improving patient care and clinical outcomes (Cai et al., 2016; Ashfaq et al., 2019; Bhoi et al., 2021). The development of deep learning mod- els (Hochreiter searchers to capture complex patterns within structured EHR data. Models such as RETAIN (Choi et al., 2016), GRAM (Choi et al., 2017), and others (Nguyen et al., 2016: Choi et al., 2018; Ma et al., 2020a;b; Gao et al., 2020; Zhang et al., 2021; Yang et al., 2023b) have shown promise in various predictive tasks. However, traditional predictive models are often inflexible, requiring spe- cific labeled training data and struggling to generalize beyond their original scope. This limitation is particularly problematic in the dynamic healthcare environment. To address this, there is growing interest in using LLMs for clinical predictive tasks. LLMs offer greater adaptability and potential to interpret diverse medical information, including unstructured text and knowledge graphs enabling more robust and versatile clinical decision support systems.
Schmidhuber, 1997; Chung et al., 2014; Vaswani et al., 2017) has enabled re-
2

---

Published as a conference paper at ICLR 2025
Step 1, Medical Knowledge Graph Construction and Indexing
Step 3. Brasoning-Enhaneee Precise Healthcare Prediction
Knowledge Source
Semantic Clustering
3.1 Training Sample Generation
ftr.md
A
Ressoning Chaina
Training Sample
Emb
Potient
Blomedical
Blomedical
Carpus
Community
Extraction
PRONFTI
Ground Truth
Run
Community Gurmmniitttion
Medical Concept Knowicdge Graph
Tuning
Step
Construction Augmentation
Reasening Instruction
Local LLM
Reasoning
Chain
B CF
Patlent
(5imlar Canes)
Content
Pattentts EHR
Community Summartes
Solaction Metrics
11: Noda Hits
W1
Coherence
Labet Prediction
13) Recency
Instruction
14) Tramm-Buuee Bales
Figure I: A conceptual illustration of our KARE framework. Step 1 constructs comprchensive medical concept knowledge graph by integrating information from multiple sources, organizing into hierarchical community structure. This structure allows the generation of community summaries that facilitate precise knowledge retrieval. Step 2 dynamically augments the patient's s EHR context with relevant summaries from the knowledge graph, offering the LLM focused and relevant medical insights. Step 3 generates training samples by employing expert LLM to create reasoning chains based on the augmented patient context and ground truth labels. then fine-tunes munetunes local LLM using multitask learning approach to produce interpretable reasoning chains and accurate predictions. By combining knowledge retrieval with LLM-driven reasoning. KARE significantly enhances the accuracy and reliability of clinical predictions.
LLMs for Healthcare Predictions. LLMs have revolutionized healthcare applications due to their advanced language understanding and generation capabilities (Xu al., 2023; Kim al., 2024; Bedi et al., 2024; Denecke et al., 2024). Recent works like MedRetriever (Ye et al., 2021), Graph- Care (Jiang et al., 2024a), RAM-EHR (Xu et al., 2024), EHR-KnowGar (Niu al., 2024b), and EMERGE (Zhu et al., 2024a) have attempted inject knowledge from retrieved literature LLMs into patient representations, but they still lack interpretability and do not fully exploit the reasoning capabilities of LLMs. On the other hand, when directly applied to domain-specifit tasks like EHR prediction, L.MSs can produce significant errors and hallucinations due to the lack of integration of specialized domain knowledge (Zhu et al., 2024b; Xu al., 2024; Cui et al., 2024; Shi al., 2024; Chen et al., 2024). Recent works like KG-RAG (Soman et al., 2023) demonstrate the broader value of KG integration with LLMs in biomedical applications. Therefore,
work integrates KG com- munity indexing and dynamic graph retrieval, compared to traditional RAG (Lewis et al., 2021: I: Niu al., 2024a) and KGs, to construct and query fine- tne-grained, precise knowledge, enhancing patient context. Furthermore, existing M-LM-based methods often fail to fully harness the inherent reasoning capabilities of L.Mss. Recent efforts (Cui al., 2024; Shi al., 2024) explored agentic frameworks for EHR prediction but rely on prompting that does not deeply engage with underlying EHR data patterns, resulting in suboptimal performance. Our approach distinguishes itself by fine-tuning a specialized, smaller LLM that incorporates reasoning abilities distilled from larger models.
KARE: KNOWLEDGE A WARE REASONING- ENHANCED REASANENG-ENIIANEED FRAMEWORK
Our KARE framework (Figure 1) aims to improve healthcare predictions by combining relevant medical knowledge along with reasoning capabilities with LLMs. The following steps achieve this: (1) medical concept knowledge graph construction and indexing, (2) patient context construction and augmentation, and (3) reasoning-enhancee enhanced precise healthcare prediction.
3.1 STEP 1: MEDICAL CONCEPT KNOWLEDGE GRAPH CONSTRUCTION AND INDEXING
Objective of Step 1 is to create medical knowledge base that is specifically tailored to electronic health record (EHR) data. Unlike most existing medical knowledge graphs, which static and not
3

---

Published as a conference paper at ICLR 2025
connected to the EHR data, KARE dynamically generates a high-quality knowledge base that can be used for retrieving and predicting information in later stages.
3.1.1 MEDICAL CONCEPT-SPECITE KNOWLEDGE GRAPH EXTRACTION
For each medical concept & in the EHR coding system, we extract a concept-specific knowledge graph Go. (Vo, Ee,)) globally tailored to the EHR datasets from three sources:
(a) Biomedical KG (e.g.. UMLS (Bodenreider, 2004)): For each medical concept o in EHR data, extract subgraph G by first iterating through the patient EHR dataset collect the top X co-existing concepts appeared in each patient's data, forming the set of related concepts Re We then find the shortest path P.j in the KG for each pair (c, Cy) in Re with specified maximum path length. G (VKG, _x()) is constructed by combining all these shortest paths, where VKG and EKC are the union of nodes and edges in all Pij. respectively. See more details in Appendix B.I.
(b) Biomedical Corpus (e.g., PubMed (Canese & Weis, 2013)): We iterate through the EHR dataset and, for each visit of the patients, and collect all the involved medical concepts. We then retrieve the top documents from the corpus based on these medical concepts. For each retrieved document, we perform entity extraction and relation extraction extract KG triples. The extracted triples are then added to the KG of the medical concepts mentioned in the document. By doing so, G is built for each concept Cj We showcase more details in Appendix B.2.
(c) LLMs; We iterate through the EHR dataset and prompt the LLM to identify the relationships among the concepts that are helpful to the clinical predictions, where we allow the LLM to add intermediate relationships within two concepts. The process is detailed in Appendix B.3.
The final concept-specifie KG Ge, the union of the subgraphs from each source:
G. = =G U GB, U UGI
Finally, we integrate all concept-specific KGs for the medical concepts in our EHR coding system. The resulting knowledge graph = (V',', of all medical concepts in the specified EHR coding system.
E') ,', is defined as G' = Uc.ec Ge, where C is the set
Note:
Different from the KG construction method introduced by GraphCare (Jiang al., 2024a), which retrieves sparse and random relationships from LLMs and biomedical KGs, our approach utilizes the EHR dataset to anchor the relevant relationships and interactions among the medical concepts present in patient data. This targeted focus allows us to construct a more relevant and context-tailored KG for clinical predictions.
3.1.2 SEMANTIC CLUSTERING
Semantic clustering in our KG addresses the challenge of differently named entities and relations from various sources that may refer to the same concept. We employ agglomerative clustering (Müllner, 201 with an automatically determined optimal threshold. First, we generate text em- beddings e; = TextEmbed((;) for each entity v E V' and e, TextEmbed(r,)) for each re- lation r, R' using an LLM. To determine the optimal clustering thresholds 0g and 0, for en- tities and relations, we refer to the silhouette score (Shahapure & Nicholas, 2020; Jiang et al., 2024b), which considers both intra-cluster similarity and inter-cluster dissimilarity. We sample subset of entities and relations, perform agglomerative clustering with varying distance thresholds, and select those yield the highest scores. We then cluster entities and relations using their respective optimal thresholds. Each cluster is represented by the element closest to the cluster ter, determined by the average embedding of all elements within the cluster. We create mappings de : V + V and or R' - between the original entities/relations and their cluster representa- tives. Each triple (h', ((,')) in the original KG is mapped to its corresponding cluster representatives rh, (h)()') (, (r')), (.))).), resulting in refined knowledge graph G (V, R, E).
3.1.3 HIERARCHICAL KG COMMUNITY DETECTION AND INDEXING
We organize the refined knowledge graph (KG) into a hierarchical structure of communities using the Leiden algorithm (Traag et al., 2019). This is done at multiple levels of granularity, from coarse to fine. We run the algorithm multiple times with different randomness parameters to explore diverse
4

---

Published as a conference paper at ICLR 2025
community structures, and generate multi-theme summaries for each community, providing more comprehensive understanding of the knowledge contained withtin the KG.
To keep computation manageable, limit the maximum number of triples per community (Z,) and the maximum number of triples per initial summary (Z,).
For each community, we generate types of summaries using LLM (prompts used: Figure 12):
General summary: A concise summarization of the medical concepts and relationships in the community without focusing on any specific theme.
Theme-specipy summary: A summary that highlights how the knowledge in the community relates to specific theme (e.g., mortality prediction), if relevant. Figure 14 shows an example.
The summarization process depends on the community size:
For small communities (size I Zz), we directly summarize all triples.
For large communities (Z, < size < Ze), we shuffle and split the triples into subsets, summarize each subset, and then iteratively aggregate the summaries until get single comprehensive summary (prompt shown Figure 13).
For extremely large communities (size Zc), we not generate summaries due to the limit of the LLM context window.
As we move up the hierarchy from fine to coarse levels, triples from small communities get merged into larger ones, which are then summarized using the same process.
The result is a hierarchical structure of communities at different granularities, each with theme- specific summaries. Running the Leiden algorithm multiple times with different randomness pa- rameter gives us diverse communities, allowing entities to contribute to multiple summaries. This rich, multi-level representation of the KG the foundation for later steps.
3.2 STEP 2: PATTENT CONTEXT CONSTRUCTION AND AUGMENTATION
Objective of Step This step constructs patient's EHR context with the highly relevant and fine- grained medical knowledge attached.
Base Context Construction. For a patient p. we construct a base context B, with their EHR data with a standardized template. This context focuses on (1) task description, (2) the patient's is con- ditions, procedures, and medications across different visits, and (3) similar patients to the target patient. For (3), two most similar patients are retrieved from the reference set (i.e., training data) based on the EHR similarity where one has the same label as patient and the other has a different label (Cui et al., 2024). Figure 11 shows example of the base context and the template used.
Context Augmentation. To enrich the patient's base context with relevant information from the knowledge graph, first construct patient-specific knowledge graph Gp by aggregating the concept-specific graphs Ge, (defined in Eq. 1) for all medical concepts e; in the patient's EHR, using the mappings de and o, from $3.1.2:
Gp Ue,EEHR, (de (h), (c(() (h,r,
G..)
(2)
From Gp. we derive two sets nodes: Vidireet, rrepreeentin medical concepts that directly appear in the patient's 's EHR, and lyindirect, containing the remaining nodes in Gp.
We then introduce a combined relevance score for each community C₁ select the most relevant summaries for context augmentation:
Relevance(C*) (H(Ck, Vydirest + M(Ckk V vrindirect)) x Decay(Cre Vidinery
x Coherence(Sc,, B,) Recency (Ct, Vidrrct x ThemeRel,
In Eq. 3, H(Ck, Vddirett and H(Ck,, V Vindirect) calculate the normalized counts of direct and indirect node hits by comparing the nodes in community C₁ with the corresponding sets direct and indirect nodes. The parameter E [0, 1) weights the importance of indirect hits relative direct hits. The
5

---

Published as a conference paper at ICLR 2025
decay function Decay (eka V Vplinett) reduces the contribution of previously hit nodes in community C& by a factor 3 where 3 E (0, 1] is a decay constant and (e) is the hit count of node v in previous selections, considering only the direct nodes Vyliress Additional factors are defined as:
Coherence(Sc.: Bp) = 1 + A1 cos (e(Sec)).
(4)
EveVor Vydess visit(e) IVc, n Vilinee
Recency(Ck, y Vddrect) = -1- A
(5)
= (Ct)
Aa
Σ
max cos (e(v), SE,,
(6)
IVc.
veVo,
Here, e(-) denotes text embedding function, cos(.,-) is the cosine similarity between embeddings, visit(v) returns the visit index of node v, and X1, A2 A3 E [0, 1] control the weights of the metrics. The set T- contains representative terms for the theme T (e.g., (end-stage, life-threateningg... for mortality prediction), same those used for attention initialization in (Jiang et al., 2024a).
The proposed metrics in Eq. 3 serve different purposes: node hits H ensure specificity to the pa- tient's conditions, decay factor promotes diversity, coherence aligns the selected summaries with the patient's overall context, recency prioritizes more recent information, and theme relevance main- tains trik-oriented selection. In addition, propose Dynamic Graph Retrieval and Augmentation (DGRA) method to iteratively select the most relevant summaries to augment the patient's context. At each iteration, performs
Algorithm Dynamic Graph Retrieval and Augmentation
(1) Compute the relevance scores for all candidate communities C2 C using Eq. 3.
Input: Set of communities C, patient graph Gp. base text B,,, desired number of summaries N Output: Augmented patient context Ap
(2) Identify the community Chest with the highest relevance score and add its summary Sc... the of se- lected summaries Sp.
Initialize Sp +- 0
Initialize hit counts H(v) +- 0 for each node v E / while |Spl < N do
Compute Relevance(Ca, for all Cs E using Eq.
Select Cbea f arg max EC Relevance Add Sci. to Sp: Sp f -Sp Scau
(3) Inerement the hit count (v) for each node in Vow.. which will im- pact the decay future relevance cal- culations.
For each v E Vor. n Vdlined H(()) + H(v) +1
Remove Ohest from C: C +-C C\ Chest
end Augment patient context: Ap Bp + Sp return Ap
(4) Remove Crest from the candidate communities C, ensuring it is not re- considered subsequent iterations.
The process continues until N summaries have been selected. The final augmented patient context Ap is obtained by concatenating the base context B, with the selected summaries.
By dynamically updating the node hits and recalculating relevance scores at each iteration, we pri- oritize communities that contribute new and valuable information. This ensures that the augmented context includes the most relevant and diverse information from the KG, tailored to the patient's specific conditions and the prediction task.
3.3 STEP 3: REASONING- -NHANNE PRECISE HEALTHCARE PREDICTION
Objective of Step 3: This step trains an LLM capable of predicting healthcare outcome while gen- erating a reasoning process, using the augmented patient context constructed in the previous step.
3.3. TRAINING SAMPLE GENERATION
Inspired by recent rationale distillation approaches (Kang et al., 2024; Kwon et al., 2024; Jiang et al., 2024b), we employ an LLM to generate reasoning chains in unified format for each patient p and
task T. This process involves entering (1) the task description D, (e.g., Figure 15), (2) the augmented patient context Ap, and (3) the corresponding ground truth label yf.r into the LLM. The specific prompt utilized for the reasoning chain (training sample) generation is showcased in Figure 16 in Appendix. The L.M generates K reasoning chains Pp.r.k along with confidence levels. We select

---

Published as a conference paper at ICLR 2025
Table 1: Statistics of pre- preeerscesed EHR datasets. "#"; "the number of", "/ patient": "per patient".
MIMIC-IIl-Mot.
MIMIC-III-Rad.
MIMIC-IV-Mort.
MIMIC-IV-Rea.
Train
Valid
Test
Train
Valid
Test
Train
Valid
Test
Train
Valid
Test
Patients (Samples)
7730
991
996
7730
991
996
8018
996
986
8029
958
1013
Visits / Patient
1.56
1.60
1.61
1.56
1.60
1.61
1.26
1.30
1.21
1.26
1.28
1.25
Conditions/ / Putient
23.27
23.92
25.89
23.27
23.92
25.89
14.34
15.30
13.59
13.62
14.21
13.21
Procedures/ Patient
6.22
6,56
7.17
6.22
6,56
7.17
2.96
3.08
2.84
2.89
2.96
2,81
Medications/ / Patient
54.79
55.77
63.73
54.79
55.77
63.73
30.66
32.86
28.40
28.74
30.61
27.59
the reasoning chain with the highest confidence, ensuring that only the most reliable explanations are used. The final set of training data for each patient-task pair is then {(D,, Ap, p v2,r),, where Pp. is the reasoning chain with the highest confidence level.
3.3.2 MULTITASK-BAN BASED TUN-TUNING AND PREDICTION
We fine-tune a relatively small local LLM (e.g., a 7B-parameter model) to perform both reasoning chain generation and label prediction for each patient p and healthcare prediction task (such as mortality or readmission prediction). The model is trained using inputs that consist the task description D, and the augmented patient context Ap. with an prepended instruction indicating whether to generate a reasoning chain or predict the label. These inputs and outputs are formatted according to the templates shown in Figure 17 in the Appendix.
During fine-tuning, when instructed to generate a reasoning chain (with the prefix [Reasoning], ) the model aligns its output with the reasoning chain p obtained from the previous step. When instructed to predict the label (with the prefix [Label Pred ct i onn]. it aligns its output with the ground truth label vp.r- We minimize the cross entropy loss across both tasks, encouraging the development of shared representations that enhance performance in both reasoning and prediction.
In the prediction phase, given new patient Phew and task r. we provide Ares T, and the appropriate instruction to the fine-tuned model. Based on the instruction, the model can either generate the reasoning chain Pp or predict the label yp This flexible approach allows us obtain detailed reasoning when necessary or perform efficient label prediction, leveraging the multitask training to effectively handle both tasks during inference.
4 EXPERIMENTS
4.1 EXPERIMENTAL SETTING
Tasks. In this work, we focus on the following EHR-bbsee prediction tasks:
Mortality Prediction. This task estimates mortality outcome for next visit, defined as ʃ : (21, 24-11) y [e), where y(жx] {0. 13 is patient's survival status during visit Zt.
Readmission Prediction. This task predicts if patient will be readmitted within a days, defined as f ;: (I1, T2,,........ 21-1) ylp(0;) (((x-1)), where E {0,1], 1}. (0(x)) is timestamp of visit øt, and y vx)())) - (p(++1)) if ((xx)) (((x-1) S a, else 0. is set to in this study.
Datasets. We utilize the publicly available MIMIC-III (Johnson et al., 2016) (v1.4) and MIMIC-IV (Johnson et al., 2020) (v2.0) EHR datasets, and use Py (Yang et al., 2023a) for preprocessing. MIMIC-III is processed (full set) using the same approach as GraphCare (Jiang et al., 2024a). For MIMIC-IV mortality prediction, we retain 2,152 patients with a label of (mortality), excluding 54 patients with more than 10 visits. We then randomly (seed=42) sample unique patients with label of 0, each having no more than 10 visits, until reaching sample size of 10,000. For MIMIC-IV readmission prediction, we randomly (seed=42) select 5,000 unique patients with label of (will be readmitted) and 5,000 with a label of O. Both datasets are split into training, validation, and test sets a 100/ ratio by patient, ensuring that all samples from the same patient are confined to single subset, preventing data leakage. We use Clinical Classifications Software (CCS) for mappings and the Anatomical Therapeutic Chemical classification system at the third level (ATC3) for medication mapping, with the resulting statisties presented in Table 1.
Evaluation Metrics. We employ four standard binary classification metrics: (1) Accuracy mea- suring overall correct predictions; (2) Macro-F1, providing a balanced measure for imbalanced
7

---

Published as a conference paper at ICLR 2025
Table 2:
Comparative analysis of mortality and readmission predictions using MIMIC-III III and MIMIC-I datasets. "pos" indicates the proportion of positive samples (label and Sensitivity are emphasized by an asterisk (*) due to their importance handling imbalanced datasets. Results are averaged by multiple runs: 30 for MI.- 10 for LMMMI--haeee and 3 for LLM-hased methods with different random seeds, We highlight the highest value for each metric.
the test set. Metrics like Macro
MIMICHII
Mortality Prediction 5.42%)
Readmission Prediction (pos 54.89%)
Type
Models
Accuracy
Sendlivity
Specilicity
Accuracy
Macro
Scaditivity
Speciliciity
GRU (Clung 2014)
50.7
3.7
97.8
62.2
61.5
Transfomorr (Vanwani 2017)
07.6
RETAIN
92.4
97.6
56.0
(Chel
Deepe (Nggyen
3.7
ICN cial. 2015)
91.6
96.4
al., 20206)
94.5
1000
59.2
59.0
AdhCrre (Maee 70200)
90.6
54.1
97.6
61.6
60.5
70.8
GRASP (2hng et 2021)
93.7
49,9
98.9
61.3
59.5
74.9
64.8
SeageNet (Gro et al. 2020;
90.5
50.5
5.6
95.4
60.5
80.0
B5.1
$4.9
KesPrat (Yang at 2023b)
92.4
52.2
933
94.7
03.5
62.1
68.0
56.1
GraphCare Chang et 2024a)
94.9
58.3
17.2
97.1
65.4
64.1
20.3
57.8
LM
RAM-EHR (Xu al. 2024)
94.4
29.5
14.8
98.9
64.8
63.5
74.7
$2.4
EMERGE (2hu et dl., 202
94.1
57.7
13.2
98.4
03.7
62.0
65.0
55.9
Zere-shat (LLM: Claude 3.5 Somee)
wi FHR contest only
89.5
50.4
644
94.4
$4.3
35.4
98.9
0.2
w/ Classic RAG
89.9
51.2
10.2
92.8
53.2
34.6
91.2
1.4
KARE augmented contec!!
92.3
54,5
14.2
94.6
56,3
43.8
93.9
10.6
Fow-Shon (LLM: Clands 3.5 Sannet)
exemplar only (N=236
LLM
esemplar only
44.7
15.7
w/ FHR-CGAgen! (Cui al. al. 2014)
87.4
51.7
91.8
55.2
46.1
78.2
w/ KARI augmented context
91.5
53.5
13.7
94.0
57.1
49.3
75.5
27.2
Finc-tined Mistrol (B-Iistrrettt--3)
90.4
57.6
w/ Classic RAG
90.
51.4
12.5
91.6
60.2
59.9
56.
64.5
KARE (ours)
95.3
64.6
24.7
98.3
73,9
73.7
76.7
70.7
MIMIC-IV
Mortulity Prediction (pos 19.165))
Readmision Prediction (pos 46.50%)
Type
Models
Macro
Sensitivity
Specificity
Accuracy
Macro F1
Sensitvvity
Specificity
GRU (Chung et al. 2014)
88.7
76.4
42.9
99.6
67.4
62.2
68.3
56.2
Transfomer (Vaswani al... 2017)
83.7
71.3
67.1
92.3
61.3
61.3
63300
59.5
RCTAIN (Chok 2016)
2016)
84.8
73.8
52.4
92.4
62.8
62.6
68.7
56.6
GRAM (Cheei 2017)
H6.4
74.4
50.6
93.9
62.5
Deepr (Npryan
$9.2
ConCare 20200)
GRASP 22hang 2011)
79.1
99.5
59.6
53.1
StageNet et 2020)
77.8
96.7
62.7
KeiPriat (Nang al. 20236)
88.7
79.8
$3.1
98.0
63.5
63.3
67.0
GraphCam Cliang
7024s)
91.5
96.6
66.2
1M4MI
RAM-EHR (Xu 3024)
90.5
78.4
97.0
65.5
64.0
EMERGE (2hu et al.. 2025a)
90.7
78.3
53.4
96.5
63.3
63.2
61.5
51.9
Zemoshat MI.MM Clande Sanne))
w/ EIRR context only
80.5
47.0
2.7
98.7
49.4
45.7
81.8
21.5
w/ Classic RAG
81.0
49.9
8.1
94.6
49.0
44.2
83.2
18.8
KARE sugmented contegt
83.2
54,3
12.7
96.3
52.3
49.7
80.6
23.7
Fow-Shun (l.LM: Clande 3.5 Sumner)
exemplar unly (N-2j4
LLM
exemplar only
99.5
68.1
wi EHR CoAgent (Cu al.. 2024)
81.0
55.5
13.8
97.0
51,2
46.3
75.4
24.0
w/ KARE augmetted context
84.5
57.4
15.8
97.6
54.1
31.9
75.2
34.1
Fine-tuned (1LM Mistral- 76- -.--m
Backbone
92.2
$3.1
65.0
96.2
56.1
46.7
23.1
76.2
wi Classic RAG
92.5
$3.8
63.2
97.6
58.8
$2.1
46.7
$7.5
KARK (ours)
94.1
90.4
73.2
99.5
73.9
73.8
85.6
63.7
b)) We nomrowe umn docimmnts from M0 PhsMMeed shoiracss thas are mosi similer as the hoc contust Denne setrional applicd with Nomnic (Nteohoum al. 2004) (dimm:766) Centea is augmental dearthed in Saxtion using our comeracted KG communiccs. Up tu community cammerio. arr eppanded suppkementary information. froz legathee samples), similar the strutegy imsodieed (Cui ot 2004). implemenie
N-2 mexss that uesiive
poticets froms positive paper
CoAkgaa upcn-anneeed
Appendis C
datasets: (3) Sensitivity, quantifying the model's s ability to identify high-risk patients; and (4) Speci- ficity, assessing accuracy in identifying low-risk patients. Detailed discussion of metric selection and computation is provided in Appendix E.
Baselines. We compare to three categories of baselines: (1) ML-based methods: GRU (Chung et al.,
2014), Transformer (Vaswani et al., 2017), RETAIN (Choi et al., 201 6). GRAM (Choi et al., 2017), Deepr (Nguyen et al., 2016), TCN (Bai et al., 2018), StageNet (Gao al., 2020), ConCare (Ma et al., 2020b), AdaCare (Ma et al., 2020a), GRASP (Zhang et al., 2021), and KerPrint (Yang et al., 2023b); (2) LM + ML-based methods: GraphCare (Jiang et al., 2024a), RAM-EHR (Xu et al., 2024), and EMERGE (Zhu et al., 2024a); and (3)
methods: zero-shot and few-shot prompting with the advanced LLM Claude 3.5 Sonnet (Anthropic, 2024), and a few-shot-based method EHR- CoAgent (Cui al., 2024). We showcase the details of baseline implementation in Appendix C.
Implementation Details. We utilize Scikit-learn (Pedregosa et al., 2018) for agglomerative cluster- ing and Graspy (Chung al., 2019) for the hierarchical Leiden algorithm. Semantic similarity cal-
8

---

Published as a conference paper at ICLR 2025
Table 3: Ablation study of fine-tuning components. Results averaged by with different seeds.
Sämilar
Retrieved
Rensoning
MIMIC-III-Mortalle
MIMIC-III-Rese Readmission
Patients
Knowledge
Accurncy
Macro
Sensitivity
Specificity
Accurncy
Macro
Sensitivity
Specificity
90.4
53.0
11.4
94.3
57.6
57.6
50.5
66.3
93.1
58.4
15.8
97.5
65.5
64.7
62.3
67.7
95.3
64.6
24.7
98.3
72.8
74.7
70.6
93.6
61.3
18.4
98.6
73,9
73.7
76.7
70.7
Similar
Retrieved Knowledge
Reusoning
MIMIC-NV-Mortality
MIMIC-IY-Rcddmisoon
Patients
Accuracy
Macro F1
Sensitivity
Specificity
Aecuracy
Macro F1
Sensitivity
Specificity
X
X
92.2
83.1
65.0
96.2
56,1
46.7
23.1
76,2
93.3
93.8
89.6
74,5
98.8
72.2
71.9
81.1
64.0
94.1
90.4
73,2
99.9
73,9
73.8
85,6
63.7
culations performed using the Nomic embedding (Nussbaum et al., 2024) for PubMed abstracts and the sutt--mbeddinng3-1e 3-large model from Azure OpenAl for all other purposes. The optimal thresholds for semantic clustering
determined to be 0g 8, 0.14. We generate community summaries using Z, = 20 and Z = 150, and employ hyperparameters a = 0.1, 0.7, A = 0.2, A2 = 0.2, and \3 0.3 for patient context augmentation. Claude 3.5 Sonnet, accessed via the Ama- zon Bedrock platform', , is used as our expert LLM for generating reasoning chain training samples. Our fine-tuning framework is implemented using the TRL (von Werra al., 2020), Transformers (Wolf et al., 2020), and FlashAttention-2 (Dao, 2024), with 2023) as our local L.M.. We provide step-by-step implementation details Appendix
I-nnall7-tt-ttut---O..... (Jiang et al.,
4.2 EXPERIMENTAL. RESULTS
Main Results. Table 2 presents the main results and highlights several key observations: (1) KARE consistently outperforms all other methods across every dataset and task; (2) the naïve RAG model sometimes fails to enhance zero-shot performance, while our method effectively augments context, leading to improved zero-shot predictions; (3) our context augmentation method is comparable to the state-of-the-aa EHR-Co. CoAgent in few-shot scenarios; and (4) our approach identifies more unique patterns, particularly excelling in correctly predicting true positives for mortality prediction, which other supervised models struggle to capture. We place some case studies in Appendix F.
Note: In mortality prediction using MIMIC-II//V, sensitivity crucial because positive cases significantly fewer than negative ones, increasing the risk of overfitting. Accurately predicting posi- tive cases is essential. Our model's specificity is not always the highest, as efforts to enhance model's overall capability can sometimes lead to misclassification of negative cases as positive. This is well-known trade-off between sensitivity and specificity (Zweig & Campbell, 1993; Powers, 2020). Conversely, for readmission prediction, where datasets are balanced, the model is expected to per- form equally well on both positive and negative samples.
Ablation Study of Fine-tuning Components. We perform an ablation study assess the in- dividual contributions of each component in boosting the performance of our fine-tuned model, as illustrated in Table 3 and Figure 2. The results Table 3 show that all components (similar patients, retrieved knowledge, and rea- soning) contribute positively to performance in most cases. However, in highly imbalanced datasets like MIMICIII Mortality, including similar patients can degrade performance. This is likely because the retrieved patients for pos- itive cases (label = 1) tend to be less similar when positive samples are scarce. Addition- ally, the absence of these components makes the fine-tuned model more prone to label bias, as seen MIMIC-III Mortality and MIMIC- IV Readmission. Figure further illustrates
0.28
0.24
0.22
0.18
0.16
0.50
1.50
Figure 2: Validation loss of the label prediction dur- ing the fine-tuning with different settings. Loss is puted every 1/4 epoch. Task; mortality prediction on MIMIC-IV. "Basc" and "Aug." denote base context and augmented context, respectively.
'The use of Amazon Bedrock is authorized by MIMIC: http:///wrriannooooooo////ggoggr Our code is available https://gituuh.co/paajjjj//

---

Published as a conference paper at ICLR 2025
w/o 6
We G
W/o G
Figure 3:
Ablation study of (left) the metrics we proposed for patient context augmentation, and (right) the KG used as the knowledge source. N.H., Coh., Rec., and T.R. denote node hits, coherence, recency, and theme relevance, respectively. Tested task: MIMICCIV-Readmisssonn
that fine-tuning with base context leads to early overfitting compared to fine-tuning with augmented context. Moreover, adding reasoning as multitask objective accelerates convergence for models using base context, whereas it slows convergence when applied to models with augmented context. This suggests that learning reasoning over more information-rion contexts is more challenging, but ultimately results in lower final loss once mastered.
Effeet of Context Augmentation Components. The LHS of Figure 3 compares the contribution of each metric proposed for community summary selection in patient context augmentation. The study shows that node hits is the most critical factor, followed by the DGRA algorithm, theme relevance, coherence, and recency, with each playing distinct role in enhancing the final performance.
Effeet of the Knowledge Source. The RHS of Figure 3 shows how removing individual knowledge the MIMIC-IV readmission task. While all KGs im-
sources affects the model's performance prove predictions, removing G causes the smallest performance drop. whereas removing L.LMM leads the largest decline. This highlights the importance of the LLM-extracted KG. likely due its contextually relevant, clinically specific relationships. The UMLS-derived KG contributes less, likely because code mapping introduces sparsity by generalizing fine-grained concepts into more abstract categories (c.g., mapping "acute myocardial infarction" " to "cardiovaseular diseases"). " This generalization limits the exploration of detailed relationships within the large KG. Future work will explore methods for retrieving knowledge with more fine fine-grained concepts from biomedical KGs.
Benefit of Multitask Learning. We compare our multitask learning approach, which treats reasoning chain generation and outcome prediction as separate tasks, with "Two--n-Onee method that only outputs the concatenated reasoning chain and ground ground-truth label. As shown in Table 4, multitask learning sig- nificantly outperforms the "Two-In-One approach
Table 4: Comparison of strategies for fine- tuning LLM with reasoning chain and label.
MIMIC-IV-Mortaii
MAMIC-IV-Renmmisiin
Surtegy
Macro F1
Seasatvity
Accuncy
Macro F1
Multiask "Two-In-Ope"
90.4
73.2
73.9
73.8
86.3
for both mortality and readmission prediction MIMIC-IV. This demonstrates that decoupling tasks allows better capture of each component's 's nuances, yielding more robust patient representa- tions. This framework enables the LLM to specialize in generating quality reasoning chains while making accurate predictions, resulting in more effective and interpretable model.
5 CONCLUSION
We propose KARE, a novel framework that combines community-based knowledge graph retrieval with large language model reasoning enhance healthcare predictions. KARE constructs compre- hensive knowledge graph, employs hierarchical community detection, and dynamically augments patient context with fine-grained, contextually relevant information. By fine-tuning specialized smaller LLM, KARE generates interpretable reasoning chains for accurate predictions. Experi- ments on MIMIC-II and MIMIC-IV datasets demonstrate KARE's superiority over state-of-the-arn methods for mortality and readmission prediction tasks. Future work will focus scaling KARE to more challenging healthcare tasks and exploring its applicability to other scientific domains, where integrating knowledge graphs and powerful language models potentially drive groundbreaking scientific progress. We discuss ethics, broader impacts, and limitations in Appendix A. Human evaluation of KARE-generated reasoning chains presented in Appendix H.
10

---

Published as a conference paper at ICLR 2025
6 ACKNOWLEDGEMENTS
Research was supported in part by US DARPA INCAS Program No. HR0011-21-CO1GE and BRIES Program No.
National Science Foundation I1S-19-56151, the Molecule Maker Lab Institute: An AI Research Institutes program supported by NSF under Award No. 2019897, and the Institute for Geospatial Understanding through an Integrative Discovery Environment (I- GUIDE) by NSF under Award No. 2118329.
REFERENCES
Anthropic. Introducing claude 3.5 sonnet.
claude - sonsonett,200 2024. Accessed: 2024-10-01.
Awais Ashfaq, Anita Sant' Anna, Markus Lingman, and Stawomir Nowaczyk. Readmission pre- diction using deep learning on electronic health records. Journal of Biomedical Informatics, 97: 103256, 2019. ISSN 1532-0464. doi: wttp://iiiieeearrttat///////2222 //www. www.:sssieececceceeemeemoeec com/ wttpeec...ccc.eeiieeeeeeeeeee
URL https:
Shaojie Bai, J Zico Kolter, and Vladlen Koltun. An empirical evaluation of generic convolutional and recurrent networks for sequence modeling. arXiv preprint 1.03:00000 2018.
Suhana Bedi, Yutong Liu, Lucy Orr-Ewing, Dev Dash, Sanmi Koyejo, Alison Callahan, Ja- son A. Fries, Michael Wornow, Akshay Swaminathan, Lisa Soleymani Lehmann, Hyo Jung Hong, Mehr Kashyap, Akash R. Chaurasia, Nirav R. Shah, Karandeep Singh, Troy Tazbaz, Arnold Milstein, Michael A. Pfeffer, and Nigam H. Shah. A systematic review of testing and evaluation of healthcare applications of large language models (Ilms). medRviv, 2024. doi: 10. 10.1101/2024.04 15.2430566. URL https: http///wwwwww..iio/eeeeet//ttttt www..mv.ddivooooommmm. www./onteetttttoo///eoo///o// coot/ot///tttteeeeeeee /early/ 2024 2024/04/ 00224/04/16/222 04 15. 24305869.
Suman Bhoi, Mong Li Lee, Wynne Hsu, Hao Sen Andrew Fang, and Ngiap Chuan Tan. Personaliz- ing medication recommendation with a graph-based approach. 40(3), nov 2021. ISSN 1046-8188. doi: 10.1145/3488668. URL https: //o...i..////// org/ 10. 145 3488668.
Olivier Bodenreider. The unified medical language system (umls): integrating biomedical terminol- ogy. Nucleic acids research, 32(supp..1::1267-D2770 2004.
X Cai, O Perez-Concha, E Coiera, F Martin-Sanchez. R Day, D Roffe, and B Gallego. Real- time prediction of mortality, readmission, and length stay using electronic health record data. Journal of the American Medical Informatics Association, 23(3):553-566 5611, May 2016. doi: 10. 1093/jamia/ocy 100. Epub 2015 Sep 15.
Kathi Canese and Sarah Weis. 2013.
Pubmed: the bibliographic database. The NCBI handbook, 2(1),
Canyu Chen, Jian Yu, Shan Chen, Che Liu, Zhongwei Wan, Danielle Bitterman, Fei Wang, and Kai Shu. Clinicalbench: Can lims beat traditional ml models in clinical prediction?, 2024. URL.
httpt:: htttps://axxxv..oo org/ wwwgg//bbs.. abss 2211.1114 6466.
Edward Choi, Mohammad Taha Bahadori, Jimeng Sun, Joshua Kulas, Andy Schuetz, and Walter Stewart. Retain: An interpretable predictive model for healthcare using reverse time attention mechanism. Advances in neural information processing systems, 29, 2016.
Edward Choi, Mohammad Taha Bahadori, Le Song, Walter F Stewart, and Jimeng Sun. Gram: grsph attention model for healthcare representation learning- Proceedings of the 23rd ACM SIGKDD international conference on knowledge discovery and data mining, pp. 787-795, 2017.
Edward Choi, Cao Xiao, Walter F. Stewart, and Jimeng Sun. Mime: Multilevel medical embedding of electronic health records for predictive healthcare, 2018. URL https: / /aw.iaiivv.o..//... /rrg/ abs /1ss/1110 09593.
11

---

Published as a conference paper at ICLR 2025
Jaewon Chung, Benjamin D Pedigo, Erie w Bridgeford, Bijan K Varjavand, Hayden S Helm, and Joshua T Vogelstein. Graspy: Graph statistics in python. Journal of Machine Learning Research, 20( 20(158):1-7,20 2019.
Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint arXiv:1412.3555, 2014.
Hejie Cui, Zhuocheng Shen, Jieyu Zhang, Hui Shao, Lianhui Qin, Joyce C Ho, and Carl Yang. Lims-based few-shot disease predictions using ehr: A novel approach combining predictive agent reasoning and critical agent instruction. arXiv preprint arXiv:2403.I 15464, 2024.
Tri Dao. FlashAttention-2:: Faster attention with better parallelism and work partitioning. In Inter- national Conference on Learning Representations (ICLR), 2024.
K. Denecke, R. May, LLMHcalthGroup. and O. Rivera Romero. Potential of large language models in health care: Delphi study. Journal of Medical Internet Research, 26:e52399, May 2024. doi: 10.2196/52399.
Darren Edge, Ha Trinh, Newman Cheng, Joshua Bradley, Alex Chao, Apurva Mody, Steven Truitt, and Jonathan Larson. From local to global: A graph rag approach to query-focused summariza- tion. arXiv preprint arXiv:2404.1 100022444..633, 2024.
Santo Fortunato. Community detection in graphs. Physics reports, 486(3-5):75-174,200 2010.
Junyi Gao, Cao Xiao, Yasha Wang, Wen Tang, Lucas M Glass, and Jimeng Sun. Stagenet: Stage- aware neural networks for health risk prediction. Proceedings of The Web Conference 2020, pp. 530-540, 2020.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Comput., 9(8): 1735-1780, nov 1997. ISSN 0899-7667. doi: (1116///ccoo.199......... URL https:
/ //dpi. 11 wwwp:///ecc...eoe/ 19...9.8............... 8.1.1.55.
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint urXiv:2310.0682) 2023.
Pengcheng Jiang, Cao Xiao, Adam Richard Cross, and Jimeng Sun. Graphcare: Enhancing health- care predictions with personalized knowledge graphs. In The Twelfth International Conference on Learning Representations, 2024a.
Pengcheng Jiang, Cao Xiao, Zifeng Wang, Parminder Bhatia, Jimeng Sun, and Jiawei Han. TriSum: Learning summarization ability from large language models with structured rationale. In Kevin Duh, Helena Gomez, and Steven Bethard (eds.), Proceedings of the 2024 Conference of the North American Chupter of the Association for Computational Linguistics: Human Language Tech- nologies (Volume 1: Long Papers), pp.
Mexico City, Mexico, June 2024b. Asso- ciation for Computational Linguisties. doi: 10. (0.18653/v1/2224.mo www.aa--lllo...... URL https: //aclanthoi acllanhology.... org/ 2024..... pataclllonge 154.
Di Jin, Zhizhi Yu, Pengfei Jiao, Shirui Pan, Dongxiao He, Jia Wu, S Yu Philip, and Weixiong Zhang. A survey of community detection approaches: From statistical modeling to deep learning, IEEE Transactions on Knowledge and Data Engineering, 35(2):1 149-1170. 2021.
Alistair Johnson, Lucas Bulgarelli, Tom Pollard, Steven Horng, Leo Anthony Celi, and Roger Mark. Mimic-iv. PhysioNet. Available online at: https:/ppyssonet. org/content/mimiccil..cc (ttpi/mmiiiii/...///eeeaeee accesssed August 23, 202/), 2020.
Alistair EW Johnson, Tom J Pollard, Lu Shen, Li-wei H Lehman, Mengling Feng, Mohammad Ghassemi, Benjamin Moody, Peter Szolovits, Leo Anthony Celi, and Roger G Mark. Mimic-iii, freely accessible critical care database. Scientific data, 3(1):1-9, 2016.
Minki Kang, Seanie Lee, Jinheon Back, Kenji Kawaguchi, and Sung Ju Hwang. Knowledge- augmented reasoning distillation for small language models in knowledge-intenee intensive tasks. Ad- vances in Neural Information Processing Systems, 36, 2024.
12

---

Published as a conference paper at ICLR 2025
Yubin Kim, Xuhai Xu, Daniel McDuff, Cynthia Breazeal, and Hae Won Park. Health-Ilm: Large language models for health prediction via wearable sensor data, 2024. URL https: /aapps:///rxvv.. arxiv. org/abs/ 2201. 06866.
Taeyoon Kwon, Kai Tzu-iunn Ong. Dongjin Kang, Seungjun Moon, Jeong Ryong Lee, Dosik Hwang, Beomseok Sohn, Yongsik Sim, Dongha Lee, and Jinyoung Yeo. Large language models are clinical reasoners: Reasoning-awann diagnosis framework with prompt-generated rationales. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 18417-18425, 2024.
Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen tau Yih, Tim Rocktäschel, Sebastian Riedel, and Douwe Kiela. Retrieval-auggmented generation for knowledge-intensiee nlp tasks, 2021. URL https: hacpgg/bss/2000 11401.
Jiarui Li, Ye Yuan, and Zehua Zhang. Enhancing Ilm factual accuracy with rag to counter hal- lucinations: case study on comain-speccifcc queries in private knowledge-bet bases, 2024. URL
https://arxis org/ alb/ab////// 2403. 12.3..0440 0446.
Weijie Liu, Peng Zhou, Zhe Zhao, Zhiruo Wang. Qi Ju, Haotang Deng, and Ping Wang. K-bert: Enabling language representation with knowledge graph, 2019. URL https://///.. /ttp///axvvv.//////.o abs/1909 07606.
Liantao Ma, Junyi Gao, Yasha Wang, Chaohe Zhang, Jiangtao Wang, Wenjie Ruan, Wen Tang, Xin Gao, and Xinyu Ma. Adacare: Explainable clinical health status representation learning via scale- adaptive feature extraction and recalibration. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 825-832, 2020a.
Liantao Ma, Chaohe Zhang, Yasha Wang, Wenjie Ruan, Jiangtao Wang, Wen Tang, Xinyu Ma, Xin Gao, and Junyi Gao. Concare: Personalized clinical feature embedding via capturing the healthcare context. In Proceedings of the AAAI Conference Artificial Intelligence, volume 34, pp.
2020b.
Varun Magesh, Faiz Surani, Matthew Dahl, Mirac Suzgun, Christopher D. Manning, and Daniel E. Ho. Hallucination-free? assessing the reliability of leading ai legal research tools, 2024. URL https:: /tt.xvaxivvvv//////////////.avla htsp/bb//2000 20362.
Daniel Müllner. Modern hierarchical, agglomerative clustering algorithms. arXiv preprint
arXiv: 109.2378, 2011.
Phuoc Nguyen, Truyen Tran, Nilmini Wickramasinghe, and Svetha Venkatesh. Deepr: convolu- tional net for medical records. IEEE journal of biomedical and health informatics, 21(1):22-30, 2016.
Cheng NIu, Yuanhao Wu, Juno Zhu, Siliang Xu, Kashun Shum, Randy Zhong, Juntong Song, and Tong Zhang. Ragtruth: A hallucination corpus for developing trustworthy retrieval-augmenner language models, 2024a. URL https: :ttp:p:://arixvv org/ aba/aaas/2/00 /2....40.10000 00396.
Shuai Niu, Jing Ma, Liang Bai, Zhihua Wang, Li Guo, and Xian Yang. Ehr-knowgen: Knowledge- enhanced multimodal learning for disease diagnosis generation. Information Fusion, 102:10206). 2024b.
Zach Nussbaum, John X Morris, Brandon Duderstadt, and Andriy Mulyar. Nomic embed: Training reproducible long context text embedder. arXiv preprint arXiv:2402.01613. 2024.
OpenAI, Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Floren- cia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, Red Avila, Igor Babuschkin, Suchir Balaji, Valerie Balcom, Paul Baltescu, Haiming Bao, Moham- mad Bavarian, Jeff Belgum, Irwan Bello, Jake Berdine, Gabriel Bernadett-Shapiro, Christopher Berner, Lenny Bogdonoff, Oleg Boiko, Madelaine Boyd, Anna-Luisa Brakman, Greg Brock- man, Tim Brooks, Miles Brundage, Kevin Button, Trevor Cai, Rosie Campbell, Andrew Cann, Brittany Carey, Chelsea Carlson, Rory Carmichael, Brooke Chan, Che Chang, Fotis Chantzis, Derek Chen, Sully Chen, Ruby Chen, Jason Chen, Mark Chen, Ben Chess, Chester Cho, Casey
13

---

Published as a conference paper at ICLR 2025
Chu, Hyung Won Chung, Dave Cummings, Jeremiah Currier, Yunxing Dai, Cory Decareaux, Thomas Degry, Noah Deutsch, Damien Deville, Arka Dhar, David Dohan, Steve Dowling, Sheila Dunning, Adrien Ecoffet, Atty Eleti, Tyna Eloundou, David Farhi, Liam Fedus. Niko Felix, Simón Posada Fishman, Juston Forte, Isabella Fulford, Leo Gao, Elie Georges, Christian Gib- son, Vik Goel, Tarun Gogineni, Gabriel Goh, Rapha Gontijo-Lopes, Jonathan Gordon, Morgan Grafstein, Scott Gray, Ryan Greene, Joshua Gross, Shixiang Shane Gu, Yufei Guo, Chris Hal- lacy, Jesse Han, Jeff Harris, Yuchen He, Mike Heaton, Johannes Heidecke, Chris Hesse, Alan Hickey, Wade Hickey. Peter Hoeschele, Brandon Houghton, Kenny Hsu, Shengli Hu, Xin Hu, Joost Huizinga, Shantanu Jain, Shawn Jain, Joanne Jang, Angela Jiang, Roger Jiang, Haozhun Jin, Denny Jin, Shino Jomoto, Billie Jonn, Heewoo Jun, Tomer Kaftan, Lukasz Kaiser, Ali Ka- mali, Ingmar Kanitscheider, Nitish Shirish Keskar, Tabarak Khan, Logan Kilpatrick, Jong Wook Kim, Christina Kim, Yongjik Kim, Jan Hendrik Kirchner, Jamie Kiros, Matt Knight, Daniel Kokotajlo, Lukasz Kondraciuk, Andrew Kondrich, Aris Konstantinidis, Kyle Kosic, Gretchen Krueger, Vishal Kuo, Michael Lampe, Ikai Lan, Teddy Lee, Jan Leike, Jade Leung. Daniel Levy, Chak Ming Li, Rachel Lim, Molly Lin, Stephanie Lin, Mateusz Litwin, Theresa Lopez, Ryan Lowe, Patricia Lue, Anna Makanju, Kim Malfacini, Sam Manning. Todor Markov, Yaniv Markovski, Bianca Martin, Katie Mayer, Andrew Mayne, Bob McGrew, Scott Mayer McKinney, Christine McLeavey, Paul McMillan, Jake McNeil, David Medina, Aalok Mehta, Jacob Menick, Luke Metz, Andrey Mishchenko, Pamela Mishkin, Vinnie Monaco, Evan Morikawa, Daniel Mossing, Tong Mu, Mira Murati, Oleg Murk, David Mély, Ashvin Nair, Reichhro Nakano, Ra- jeev Nayak, Arvind Neelakantan, Richard Ngo. Hyeonwoo Noh, Long Ouyang, Cullen O' Keefe, Jakub Pachocki, Alex Paino, Joe Palermo, Ashley Pantuliano, Giambattista Parascandolo, Joel Parish, Emy Parparita, Alex Passos, Mikhail Pavlov, Andrew Peng, Adam Perelman, Filipe de Avila Belbute Peres, Michael Petrov, Henrique Ponde de Oliveira Pinto, Michael, Pokorny, Michelle Pokrass, Vitchyr H. Pong, Tolly Powell, Alethea Power, Boris Power, Elizabeth Prochl, Raul Puri, Alec Radford, Jack Rae, Aditya Ramesh, Cameron Raymond, Francis Real, Kendra Rimbach, Carl Ross. Bob Rotsted, Henri Roussez, Nick Ryder, Mario Saltarelli, Ted Sanders, Shibani Santurkar, Girish Sastry, Heather Schmidt, David Schnurr, John Schulman, Daniel Sel- sam, Sheppard, Toki Sherbakov, Jessica Shieh, Sarah Shoker, Pranav Shyam, Szymon Sidor, Eric Sigler, Maddie Simens, Jordan Sitkin, Katarina Slama, Ian Sohl, Benjamin Sokolowsky, Yang Song, Natalie Staudacher. Felipe Petroski Such, Natalie Summers, Ilya Sutskever, Jie Tang, Nikolas Tezak, Madeleine B. Thompson, Phil Tillet, Amin Tootoonchian, Elizabeth Tseng, Pre- ston Tuggle, Nick Turley, Jerry Tworek, Juan Felipe Cerón Uribe, Andrea Vallone, Arun Vi- jayvergiya, Chelsea Voss. Carroll Wainwright, Justin Jay Wang, Alvin Wang, Ben Wang. Jonathan Ward, Jason Wei, CJ Weinmann, Akila Welihinda, Peter Welinder, Jiayi Weng, Lilian Weng. Matt Wiethoff, Dave Willner, Clemens Winter, Samuel Wolrich, Hannah Wong, Lauren Work- man, Sherwin Wu, Jeff Wu, Michael Wu, Kai Xiao, Tao Xu, Sarah Yoo, Kevin Yu, Qiming Yuan, Wojciech Zaremba, Rowan Zellers, Chong Zhang, Marvin Zhang, Shengjia Zhao, Tianhao Zheng, Juntang Zhuang, William Zhuk, and Barret Zoph. Gpt-4 technical report, 2024. URL https: //tasi///avvvv./ org/ abs// / 2303 08774.
Fabian Pedregosa, Gaël Varoquaux, Alexandre Gramfort, Vincent Michel, Bertrand Thirion, Olivier Grisel, Mathieu Blondel, Andreas Müller, Joel Nothman, Gilles Louppe, Peter Prettenhofer, Ron Weiss, Vincent Dubourg, Jake Vanderplas, Alexandre Passos, David Cournapeau, Matthieu Brucher, Matthieu Perrot, and Édouard Duchesnay. Scikit-learn: Machine learning in python, 2018. URL wwtphhtt//// //tp:i//aiivv.//. org/ abs 201.001.55
David MW Powers. Evaluation: from precision, recall and f-measure to roc, informedness, marked- ness and correlation. arXiv preprint arXiv:2010. 2020.
Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He. Deepspeed: System opti- mizations enable training deep learning models with over 100 billion parameters. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery Data Mining, KDD '20, pp. 3505-3506, New York, NY, USA, 2020. Association for Computing Machin- ery. ISBN 47814503799 doi: 00.1145/3994486.4406777 URL https: aedoi..or//1 1145/33 338333399443494 3406703.
Guillaume Salha, Stratis Limnios. Romain Hennequin, Viet-Anh Tran, and Michalis Vazirgiannis. Gravity-inspired graph autoencoders for directed link prediction. In Proceedings of the 28th ACM international conference on information and knowledge management, pp. 598, 2019.
14

---

Published as a conference paper at ICLR 2025
Ketan Rajshekhar Shahapure and Charles Nicholas. Cluster quality analysis using silhouette score. In 2020 IEEE 7th international conference on data science and advanced analytics (DSAA), pp- 747-748. IEEE, 2020.
Wenqi Shi, Ran Xu, Yuchen Zhuang. Yue Yu, Jieyu Zhang, Hang Wu, Yuanda Zhu, Joyce Ho, Carl Yang, and May D. Wang- Ehragent: Code empowers large language models for few-shot complex tabular reasoning on electronic health records, 2024. URL https: http///axxvvv///////ee abs/2401. 07128.
Benjamin Shickel, Patrick James Tighe, Azra Bihorac, and Parisa Rashidi. Deep ehr: survey of recent advances in deep learning techniques for electronic health record (ehr) analysis. IEEE Journal of Biomedical and Health Informatics, 22(5)) 189-11044, September 2018. ISSN 2168- 2208. doi: 10.1 1010099hhi 2017.27677 2767063.
URL http: : /xx. doi wo..ooo.. oooiioogg/000 11.0..100 10/9/09/000 JBHI.
Karthik Soman, Peter W Rose, John H Morris, Rabia E Akbas, Brett Smith, Braian Pectoom, Catalina Villouta-Reyes Gabriel Cerono, Yongmei Shi, Angela Rizk-Jackson, et al. Biomedi- cal knowledge graph-enhance prompt generation for large language models. arXiv e-prints, pp- arXiv-2311, 2023.
Karthik Soman, Peter W Rose, John Morris, Rabia E Akbas, Brett Smith, Braian Peetoom, Catalina Villouta-Reyes, Gabriel Cerono, Yongmei Shi, Angela Rizk-Jackson, Sharat Israni, Charlotte A Nelson, Sui Huang, and Sergio E Baranzini. Biomedical knowledge graph-optimized prompt generation for large language models, 2024. URL https: /ttp/v/aavvvv//v////////e//aaaaa org/ abs/ 2311.17330.
Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre. Morgane Rivière, Mihir Sanjay Kale, Juliette Love, Pouya Tafti, Léonard Hussenot, Pier Giuseppe Sessa, Aakanksha Chowdhery, Adam Roberts, Aditya Barua, Alex Botev, Alex Castro-Ros, Ambrose Slone, Amélie Héliou, Andrea Tacchetti, Anna Bulanova, An- tonia Paterson, Beth Tsai, Bobak Shahriari, Charline Le Lan, Christopher A. Choquette-Choo Clément Crepy. Daniel Cer, Daphne Ippolito. David Reid, Elena Buchatskaya, Eric Ni, Eric Noland, Geng Yan, George Tucker, ChhiiteChrristian Muraru, Grigory Rozhdestvenskly, Hen- ryk Michalewski, Ian Tenney, Ivan Grishchenko, Jacob Austin, James Keeling, Jane Labanowski, Jean-Baptiste Lespiau, Jeff Stanway, Jenny Brennan, Jeremy Chen, Johan Ferret, Justin Chiu, Justin Mao-Jones, Katherine Lee, Kathy Yu, Katie Millican, Lars Lowe Sjoesund, Lisa Lee, Lucas Dixon, Machel Reid, Maciej Mikula, Mateo Wirth, Michael Sharman, Nikolai Chinaev, Nithum Thain, Olivier Bachem, Oscar Chang, Oscar Wahltinez, Paige Bailey, Paul Michel, Petko Yotov, Rahma Chaabouni, Ramona Comanescu, Reena Jana, Rohan Anil, Ross Mcllroy, Ruibo Liu, Ryan Mullins, Samuel I. Smith, Sebastian Borgeaud, Sertan Girgin, Sholto Douglas, Shree Pandya, Siamak Shakeri, Soham De, Ted Klimenko, Tom Hennigan, Vlad Feinberg, Wojciech Stokowiee, Yu hui Chen, Zafarali Ahmed, Zhitao Gong. Tris Warkentin, Ludovic Peran, Minh Giang, Clément Farabet, Oriol Vinyals, Jeff Dean, Koray Kavukcuoglu, Demis Hassabis. Zoubin Ghahramani, Douglas Eck, Joelle Barral, Fernando Pereira, Eli Collins, Armand Joulin, Noah Fiedel, Evan Senter, Alek Andreev, and Kathleen Kenealy. Gemma: Open models based on /t///axxiiiv///e gemini research and technology, 2024. URL https: http://aawvv.r.././/////// www.oorgg/bbb/eeeeee 2403 08295.
Brandon Theodorou, Cao Xiao, and Jimeng Sun. Synthesize high-dimensional longitudinal elec- tronie health records via hierarchical autoregressive language model. Nature communications, 14 (1):5305, 2023.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Ar- mand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models, 2023a. URL https: : /ttxiaivvvva////a//eee
299bb/2200..99... 13971.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi. Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull. David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshom, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel
15

---

Published as a conference paper at ICLR 2025
Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuming Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023b. URL https: http//axxvvv..////.//.e/e..a org/ abs/ 2307 09288.
Vincent A Traag, Ludo Waltman, and Nees Jan Van Eek. From louvain to leiden: guaranteeing well-connected communities. Scientific reports, 9(1):1-12, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- fion processing systems, 30, 2017.
Leandro von Werra, Younes Belkada, Lewis Tunstall, Edward Beeching, Tristan Thrush, Nathan Lambert, Shengyi Huang, Kashif Rasul, and Quentin Gallouédec. Trl: Transformer reinforcement learning. https: //gi/gittuu//t com http///ugggggggggeeeeeeec /t rrr, 2020.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma. Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gug- ger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-the--a natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pp. 38-45, Online, October 2020. As- sociation for Computational Linguistics. URL https:////// htwp::://wwwww. /wwww a acl web. org/ / henttolloy.. 2020. emnlp-aae demos. 6.
M. Wornow, Y. Xu, R. Thapa, et al. The shaky foundations of large language models and foun- dation models for electronic health records. NPJ Digital Medicine, 6:135, 2023. doi: 10.1038) 341746-023-00879-. URL https:ttp/// otg/0............0 1038/ 004001000001000000000000009
8.
Zhenbang Wu, Cao Xiao, and Jimeng Sun. Medlink: De-identified patient health record linkage. In Proceedings of the 29th ACM SIGKDD Conference Knowledge Discovery and Data Mining, pp. 2672 2682, 2023.
Ran Xu, Wenqi Shi, Yue Yu, Yuchen Zhuang, Bowen Jin, May D Wang. Joyce C Ho, and Carl Yang. Ram-chr: Retrieval augmentation meets clinical predictions on electronic health records. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics, 2024.
Yongxin Xu, Xu Chu, Kai Yang, Zhiyuan Wang, Peinie Zou, Hongxin Ding, Junfeng Zhao, Yasha Wang, and Bing Xic. Seqcare: Sequential training with external medical knowledge graph for diagnosis prediction in healthcare data. In Proceedings of the ACM Web Confer- ence 2023, WWW '23, pp. 2819-2830, New York, NY, USA, 2023. Association for Com- puting Machinery. ISBN 9781450394161. doi: 0..145/3543507.35883383 //doi.org/10 10. 0445/5543507 3543507.3383543
URL https:
Chaoqi Yang, Zhenbang Wu, Patrick Jiang. Zhen Lin, Junyi Gao, Benjamin Danck, and Jimeng Sun. PyHealth: A deep learning toolkit for healthcare predictive modeling. In Proceedings of the 27th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining (KDD) 2023, 2023a.
URL https:://// / / ihubb. com/
Kai Yang, Yongxin Xu, Peinie Zou, Hongxin Ding, Junfeng Zhao, Yasha Wang, and Bing Xie. Ker- print: local-gh local-global knowledge graph enhanced diagnosis prediction for retrospective and prospec- tive interpretations. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp. 5357-5365, 2023b.
X. Yang, A. Chen, N. PourNejatian, H. C. Shin, K. E. Smith, C. Parisien, C. Compas, C. Martin, A. B. Costa, M. G. Flores, Y. Zhang, T. Magoc, C. A. Harle, G. Lipori, D. A. Mitchell, W. R. Hogan, E. A. Shenkman, J. Bian, and Y. Wu. A large language model for electronic health records. NPJ Digital Medicine, 5(1):194, December 2022. doi: 003866447746-02200077444
16

---

Published as a conference paper at ICLR 2025
Michihiro Yasunaga, Hongyu Ren, Antoine Bosselut, Percy Liang, and Jure Leskovec. Qa-gnn: Reasoning with language models and knowledge graphs for question answering, 2022. URL
https:/ /arxiv..vv/. www./rrggss...oo.oo....
06378.
Muchao Ye, Suhan Cui, Yaqing Wang, Junyu Luo, Cao Xiao, and Fenglong Ma. Medretriever: Target-driven interpretable health risk prediction via retrieving unstructured medical text. In Pro- ceedings of the 30th ACM International Conference on Information Knowledge Management, pp. 2413-2423, 2021.
Chaohe Zhang, Xin Gao, Liantao Ma, Yasha Wang, Jiangtao Wang, and Wen Tang. Grasp: generic framework for health status representation learning based on incorporating knowledge from sim- ilar patients. In Proceedings of the AAAI conference artificial intelligence, volume 35, pp-
715-723, 2021.
Xikun Zhang, Antoine Bosselut, Michihiro Yasunaga, Hongyu Ren, Percy Liang, Christopher D. Manning, and Jure Leskovec. Greaselm: Graph reasoning enhanced language models for question answering, 2022. URL
orgp./////bbbb/11111////// 08860.
Yinghao Zhu, Changyu Ren, Zixiang Wang, Xiaochen Zheng, Shiyun Xie, Junlan Feng, Xi Zhu, Zhoujun Li, Liantao Ma, and Chengwei Pan. Emerge: Integrating rag for improved multimodal ehr predictive modeling. arXiv preprint arXiv:2406.00036, 2024a.
Yinghao Zhu, Zixiang Wang, Junyi Gao, Yuning Tong, Jingkun An, Weibin Liao, Ewen M Harrison, Liantao Ma, and Chengwei Pan. Prompting large language models for zero-shot clinical predic- tion with structured longitudinal electronic health record data. arXiv preprint arXiv:2402.01713. 2024b.
Mark H Zweig and Gregory Campbell. Receiver-operating characteristic (roc) plots: fundamental evaluation tool in clinical medicine. Clinical chemistry, 39(4):561-577 1993.
17

---

Published as a conference paper at ICLR 2025
Contents of Appendix
A Ethics, Broader Impacts, and Limitations
19
B Details of Knowledge Graph Construction
19
B.1 KG Extraction from Large Biomedical KG
19
B.2 KG Extraction from Large Biomedical Corpus
20
B.3 KG Extraction from Large Language Models
22
B.4 Case Studies of KG Construction
24
Baseline Implementations
26
D Implementation Details of KARE
29
D.1 Step Medical Concept Knowledge Graph Construction
29
D.2 Step 2: Patient Context Construction and Augmentation
29
D.3 Step Reasoning-Enhancee Precise Healthcare Prediction
................................
29
Evaluation Metrics
31
E.1 Primary Metrics
31
E,2 Choice of Metrics
................................
31
F
Case Study
32
F.1
Reasoning-Ennaane Prediction by Our Tune-Tuned Model
................................
32
F22 Zero-Shot & Few-Shot Comparison
33
Templates, Prompts, and Examples
38
H Human Evaluation of Reasoning Chains
43
I Parameters & Training Time of Models
45
Notations
46
18

---

Published as a conference paper at ICLR 2025
A ETHICS, BROADER IMPACTS, AND LIMITATIONS
Our work involves analysis of electronic health record (EHR) data, which contains sensitive personal medical information. To ensure the ethical handling of this data, we conducted all interac- tions between the language models and the EHR data through Amazon B which provides rigorous compliance standards and privacy protection measures. This allowed us fully leverage the capabilities of the LLMS while maintaining strict confidentiality of the patient data.
Ethics:
Our framework, KARF, demonstrates the potential for integrating knowledge graphs and language model reasoning to enhance clinical decision support systems. By providing more accurate and interpretable predictions for critical outcomes like mortality and readmission, KARE could assist healthcare providers in identifying high-risk patients who may require additional interventions or closer monitoring. This could ultimately lead to improved patient outcomes and more efficient allocation of healthcare resources. However, it is important to recognize that our models are intended to augment, rather than replace, the judgment of healthcare professionals. The predictions should be considered as additional data points to inform clinical decision making, not as definitive diagnoses or treatment recommendations.
Broader Impacts:
Limitations: While KARE achieves promising results, there are several limitations to consider. First, our evaluation is based on the MIMIC-III and MIMIC-IV datasets, which represent a specific patient population from single hospital system in the United States. The generalizability of our findings to other patient populations or healthcare settings may be limited. Second, our knowledge graphs are constructed from a subset of biomedical databases, literature, and language model out- puts, and may not capture the full breadth of medical knowledge. Expanding the knowledge sources and improving the knowledge extraction and integration processes could further enhance the perfor- mance of our models. Third, our framework relies on the outputs of large language models, which are known have biases and can generate hallucinations. While we have taken steps mitigate these issues, such as collaborating with medical experts to validate the extracted knowledge, there remains risk of the models producing incorrect or biased predictions
some cases. Fourth, our work is based on English biomedical literature and EHR data, and we did not evaluate the perfor- mance of KARE on data other languages. The applicability and effectiveness of our approach for non-English clinical settings requires further investigation. Ongoing research is needed to de- velop more robust and reliable language models for clinical applications across diverse languages and populations.
In conclusion, KARE represents an important step towards leveraging knowledge graphs and lan- guage model reasoning for improved clinical predictions. However, further research is needed to address the limitations and ensure the safe and responsible deployment of such models in roal-world healthcare settings. We encourage future work to focus on enhancing the generalizability, inter- pretability, and robustness of these approaches, as well as engaging with healthcare stakeholders to develop guidelines for the ethical and effective use of Al in clinical decision support.
B DETAILS OF KNOWLEDGE GRAPH CONSTRUCTION
B. KG EXTRACTION FROM LARGE BIOMEDICAL KG
To construct concept-specific knowledge graphs from the Unified Medical Language System (UMLS), we follow multi-step process:
1. Extracting Co-existing Concepts from EHR Data: We iterate through the patient EHR dataset and collect the top X (X = 20 in our implementation) most frequently co-existing concepts for each unique medical concept based on their co-occurrence in patient records.
2. Mapping Concepts to UMLS CUTs: corresponding Concept Unique Identifiers (CUIs) in UMLSS This involves mapping condition and procedure concepts to CCS codes, then to ICD-9 codes, and finally to UMLS CUIs. Drig concepts are directly mapped to ATC codes and then to UMLS CUIs.
We map the medical concepts from the EHR data to their
'https://dccs" aws. amazon. http:///woooo.oo..ooooooooocoocr html
de wwtp//hhtitt i bedrook.
onet http:///ss.tttttt..ee////////ss/
19

---

Published as a conference paper at ICLR 2025
EHR
Concepts in Each Visit
Top Coexisting Concepts
Pairwise Shortest Path Finding Length -
Concept-Specite
concept 1:f concept concept_s,
[concestL) []
G&G concept,
UMLS
[concepe 3]
[conceoL4) [conceot_5]
Figure 4: Our pipeline to construct concept-specific KG OKG with bio KG (UMLS) and EHR.
For each medical concept 6 and its top co-existing concepts Rey extract concept-specific subgraph G from UMLS using bidirectional shortest path finding algorithm. The algorithm parameters are set as follows:
Extracting Subgraphs from UMLS:
max max_length=7: The maximum length of the shortest paths considered between concepts.
max paths=40: The maximum number of shortest paths to retrieve for each pair of con- cepts.
max mexx_ndds==12000 The maximum number of nodes to explore during the bidirectional
search.
The bidirectional shortest path finding algorithm is implemented as follows:
Algorithm Bidirectional Shortest Path Finding
Input: Graph G, start node 8, end node f, max length 1, max paths p. max nodes n
Output: Set of shortest paths p
Initialize forward queue 2 + [[s}} and backward queue Qo + {[}}} Initialize forward visited dict Vy < {s [s|} and backward visited diet Vú + {t [t] Initialize paths + and nodes explored N + while 2 and Q₁ and |P| < p and N do
Expand forward path Ty + Q2 popleff() and last node a + x-[-1]
u E V then
P - PU (T + Vu[uI[:: [1] :}}
end
Expand backward path Ti Q Qo-popleft( and last node +- Tol-1]
uf E V then
m|:: - [][1 }}} P P (Vylu) m
end
if (en(TTT) ) 21 (en(T)) < 2l then
Expand neighbors and update Q,, Qs Vy. Vi
end end return
The extracted shortest paths for each pair (c,, 0j) are combined to form the concept-specific subgraph (KG The nodes yKG and edges EKG of the subgraph are the union of all nodes and edges in the extracted paths. c
By following this process, we construct set concept-specific knowledge graphs (GKG) that capture the relevant relationships and contextual information for each medical concept G based on the UMLS knowledge graph and real-world co-occurrence patterns in patient EHR data.
The number of the resulting KG triples from UMLS is 29,434.
B.2 KG EXTRACTION FROM LARGE BIOMEDICAL CORPUS
To construct concept-specific knowledge graphs, we process the EHR dataset and the PubMed Ab- stracts corpus using the following pipeline (Figure 5):
20

---

Published as a conference paper at ICLR 2025
EHR
Concepts
Concept KGs
G
Dack:" "... (c..cccnnttt...... 11 [concept 2]... [concept 4]...
(concept
[concept 1]
PublMed
GeG
iconcept_2 2]
RBC2: ".... 2ooccpp.2]...
Entity/Retation Extraction
Gie (ounoppt. 3]
1oal
[ooncept
[concopt "]
[concept_]]
(query set of concepts)
G (ooneppt
Doc ".. (...Iconceett_B [concopt-]..."
GBC (onerpt
Figure 5: Our pipeline to construct concept-specífic KG GBC from biomedical corpus with EHR.
Concept Set Extraction from EHR Data: We iterate through the EHR dataset and, for each patient visit, we collect all the involved medical concepts (conditions, procedures, and drugs) into a concept set. This results in a list of concept sets, where each set represents the concepts associated with specific patient visit.
2. Filtering Similar Concept Sets:
To reduce redundancy and computational overhead, we filter highly similar concept sets based on their concept multi-hot vector representation. We use the Count Vectorizee from scikit-learn (Pedregosa et al., 2018) to create a vocabulary of unique concepts and transform each concept set into multi-hot vector. We then compare the symmetric difference between pairs of concept sets and filter out sets that differ by fewer than specified similarity threshold (5 in our case). This step helps to eliminate nearly duplicate concept sets while retaining diverse range of concepts for knowledge graph construction.
This process results in 26, 134 concept sets in experiment.
3. Dense Retrieval of Relevant PubMed Abstracts: For each filtered concept set, we retrieve the top n n = our case) most relevant documents from PubMed Abstracts using dense retrieval. Considering the expensive time consumption to retrieve abstracts from the full corpus (with 30 million abstracts), we randomly (sced=42) select 1/10 abstracts as the subset corpus to proceed. We employ the nomic waii//omico ee-e---beeeeeteeee -vl model to embed both the concepts and the PubMed abstracts into a dense vector space (dim=768). We then compute the cosine similarity between the concept set embedding and the abstract embeddings to identify the most relevant documents. The dense retrieval process optimized for efficiency by processing the embeddings in chunks and utilizing GPU acceleration when available.
4. Triple Extraction from PubMed Abstracts:
For each retrieved PubMed abstract, we perform entity extraction and relation extraction to extract knowledge graph triples. use large lan- guage model (Claude 3.5 Sonnet in our case), to identify relationships between the concepts mentioned in the abstract. The LLM is provided with a prompt that includes the abstract text, the list of relevant concepts, and example of the desired triple format. The prompt instructs the LLM to extract at most 10 informative and logically sound triples for each abstract, focusing relationships closely related the provided concepts. The extracted triples follow the format (ENTITY1, RELATIONSHIP ENTITY2], where the entities are replaced with the exact concept terms when applicable. The prompt we used for the triple extraction is:
Given a medical text and list of important concepts, extract relevant relationships between the concepts from the text (if present). For each triple, if an entity matches one of the given concepts, replace the entity with the exact concept term.
Focus on generating high-quality triples closely related to the provided concepts. Aim to extract at most 10 triples for each text. Each triple should follow this format: [ENTITYI, RELATIONSHIP ENTITY2). Ensure the triples are informative and logically sound.
Example:
Text:
Asthma is chronic respiratory condition characterized by inflammation and narrowing of the airways, leading to breathing difficulties. Common symptoms include wheezing. coughing, shortness of breath, and chest tightness. Triggers can vary but often include allergens, air pollution, exercise, and respiratory infections. Management typically involves combination of long-term control medications, such as inhaled corticosteroids, and
21

---

Published as a conference paper at ICLR 2025
quick-relief medications like short-acting beta-agonists. Recent research has focused personalized treatment approaches, including biologics for severe asthma and the role of the microbiome in asthma development and progression. Proper inhaler technique and adherence to medication regimens are crucial for effective management. Asthma action plans, developed in partnership with healthcare providers, help patients manage symptoms and exacerbations.
Concepts:
[asthma, inflammation, airways, wheezing, coughing, inhaled corticosteroids, short-acting beta-agonists, allergens, respiratory infections]
Extracted triples:
[[asthma, is chronic respiratory condition], [asthma, characterized by, inflammation of airways], [inflammation, causes, narrowing of airways], [narrowing of airways, leads to, breathing difficulties], [wheezing, is symptom of, asthmal, [coughing, is symptom of, asthma]. [allergens, can trigger, asthma), [respiratory infections, can trigger, asthma), [inhaled corticosteroids, used for, long-term control of asthma], [short-acting beta-agonists, provide, quick relief in asthma]]
Text:
ftext)
Concepts:
[concepts)
Extracted triples:
Knowledge Graph Construction: The extracted triples from each abstract are added to the knowledge graph of the medical concepts mentioned in the document. This process builds concept-specific knowledge graph G for each concept G. incorporating relevant information from the PubMed corpus. The resulting knowledge graphs capture the relationships and contex- tual information associated with each medical concept.
By following this pipeline, construct comprehensive of concept-specific knowledge graphs that integrate information from both EHR data and the PubMed corpus. These knowledge graphs serve valuable resource for EHR-based downstream tasks, such as patient representation learn- ing and predictive modeling.
The number of resulting KG triples from the PubMed Abstracts is 259.938.
B.3 KG EXTRACTION FROM LARGE LANGUAGE MODELS
Instruction
"Plesise identiff relationahio among these concepts that can be potentially helpru clinccal predictions te gu, mortality predictions madm esion prediction).
Concapt-Speciffic
Concepts
You con intrdduce intormediate relationships with other ansiliw besed your knnosladge.
Identified Relationships
KGs
There's
the relationship connected.
[oonGRpE Dincappi
G
GLL
conoupt.
[concopt_L 1 jooncapt 2]
Output
[concopt E.
GLUM
Comoapt.2
[ENTITT, RELATIONBHIP_ ENTTY2), ENIIYY, RILATIONSSP ENITYY].... ..
GLUM
[concept.]
I
couneps.
Iconcept_4 [concapt 9
concopt.4]
GLUM
conncapt concept 3]
(conorpt22)
Figure 6: Our pipeline construct concept-specific KG with LLM and EHR.
To extract concept-specific knowledge graphs from large language models (LLMs), we follow process similar to the initial steps of the "KG extraction from corpus" pipeline (Appendix B.2):
22

---

Published as a conference paper at ICLR 2025
1. Concept Set Extraction from EHR Data: We iterate through the EHR dataset and, for each patient visit, we collect all the involved medical concepts (conditions, procedures, and drugs) into concept set. This results in a list of concept sets, where each set represents the concepts associated with a specific patient visit.
Filtering Similar Concept Sets: reduce redundancy and computational overhead, we filter out highly similar concept sets based their concept multi-hot vector representation, de- scribed in the previous subsection. This step helps to eliminate nearly duplicate concept sets while retaining diverse range of concepts for knowledge graph construction.
3. Prompting LLMs for Relationship Extraction:
For each filtered concept set, we prompt large language model (Claude 3.5 Sonnet in our case) to identify relationships among the medi- cal concepts that can be potentially helpful for clinical predictions, such mortality prediction and readmission prediction. The LLM is encouraged to introduce intermediate relationships with other entities based its knowledge, and there is requirement keep all the relationships connected. The prompt instructs the LLM to use the original names the provided concepts in the outpur, which should follow the format [ [ENTITY], RELATIONSHIP 1 ENTITY2] [ENTITY2, BELATTONSHTP 2,, ENTITY3], ...]. J. The prompt we used for the rela- tionship extraction is:
Please identify the relationships among these medical concepts that can be potentially helpful to clinical predictions (e.g., mortality prediction, readmission prediction) as many as possible.
You can introduce intermediate relationships with other entities based on your knowledge.
Consider how these concepts would interact with others to be useful for clinical predictions. There's no need to keep all the relationships connected.
For the concepts provided in the list, you MUST use the their original name without any changes. Please output only the list of triples without any other information.
Output format:
[[ENTITYL, 1, RELATIONSHIP 1, ENTITY2), (ENTITY2, RELATIONSHIP ENTITY3). ...]
Medical Concepts:
concepts
Output:
4. Knowledge Graph Construction: The extracted triples for each concept set used to struct concept-specific knowledge graph GLLM for each concept q. For each concept e; in the concept set, store the connected 3-hop subgraph sourced from e; its corresponding concept-specific knowledge graph GLLM This step ensures that only the triples directly or in- directly connected the concept c; are included in its concept-specific knowledge graph. This process is iteratively performed for all the concepts in the concept set. The resulting knowledge graphs capture the relationships and contextual information associated with each medical concept based on the knowledge embedded in the large language model.
By leveraging the knowledge embedded in large language models, this process allows us to construct concept-specific knowledge graphs that incorporate a broad range of information beyond what is explicitly stated in the EHR data biomedical literature. These knowledge graphs can provide valuable insights and support various downstream tasks in the clinical domain.
The number of resulting KG triples from the LLM is 315.492.
23

---

Published as a conference paper at ICLR 2025
B.4 CASE STUDIES OF KG CONSTRUCTION
Target Concept: oplold analgesios
(Top-20 Co-existing Concepts for Each Concept)
Retonionship Triple::
jespects afaveer effects, Alloved qualifier, opio op ansigesics) [spinal fuston, Allowed quslitier, aspazts of soverse effechs]
Corraction Paths:
tustor (Allowed quatifee)
effects (Allowed
analgesics
"spinal Fusion":
Target Concept:
antibacteriats
"oploid analgesics",
"other seta-lactas antibacterials
Retorionship Triples:
"antithroebotic agents",
(espects adveess effects, Alloved qualifier, Carbaperam-cmmanii product] [Cartspenem containing product. iss. other bate-lnctam antibactetsis in atte [spinol fusion, Alloved qualifior, aspeces of soverse effects]
supplesents",
"drugs constipation",
"other analgesics antipynetics
Connection Pathi:
"1.v. solution additives", "drugs for peptic ulcer and gatranorsqphaaa reflux disease (gord)",
spinal fusion
(Allowed
quzlifler) Carbape
antibacterist atc
"Iv solutions used in parenteral adninistration flulds,
Target Concapt: antithrombotic agants
lecctrolytes and nutrients", "spendyloris; Intervertebrai disc prablees", disorders; antinauseents"" "anxolytics",
(Alg. 2)
Relerionship Tripls:
Jospects of adverse offects, Allowed quelifirr Fibrinclytic Agonts) (Fininulyic Apents, meliiiiccniii oniibro agentss) (epinal lusion, Allowad qualifier, saparts adverse allects]
"other nineral supplemets atc"
"antiinfectives and antisaptics, exel. combinations with corticasteraids",
Connection Paths:
[Allowed quslitter Ribtnolytic Agents thas relotionshipi antthrombettc agents
"antlepilaptics""
"other nutrients In atc", "antivaricose therapy drugs",
Target Concept: potassium supplements
blocking
"lipid modifying plain", "ecooppesstoon peripheral
Relaskonchip Triples:
|FOTASSILM CITRATE, member of, potiesiim supplerentss
[FCTASSIUM CHLORIDE, member, ol, polassium susplements)
eftects, qualtted PMTASS CITRATE]
POTASSIUM CHLORIDE
aspects
Connection
spinal fusion (Allnwed q qualiiery aspects aclverss effects a (Allowed
qualiflen) - POTAESIUM CHLORIDE (membor, off potassium supplaments spánel faskar (Allowed qualfie)) aspectt adverse oflecta fean be qusliteed byp POTASSIUMCTRAATT (membee of) potassium supptements
Figure 7: An example of concept KG extraction from biomedical KG (UMLS).
24

---

Published as a conference paper at ICLR 2025
(Concept set in a patient visit)
(Retrieved Top-10 PubMed Abstracts)
(Triples)
Rank Similarty. 0.786074030788008
PHID: 310377
Thie: Mf soute trouma. Authors: CNN Shozly
"pleurisy; peuuntoroas; pulmonsry
collapse",
Abstract In azute poin, TENS, los pocks, ond calmm, resssusing utitlute and voice csuiful reddusing pair. Narootic recui uments be roduced, and chronic pain may be provented. Complication such os paralyti iluu and steloctasis can olso be roduced. These techniques can be used the umergenicy ward, the recovery doctor's office.
[opiaid analgesios, induned by, TENS]
"coronary atneroscierosts and other heart disease",
"potasstin spplesents", "other fractures",
[opioid analgrsics, reduced by. packs]
"anddolytics",
Triple Extraction
"estodd analgestes", "other pastrointestim dtsorders",
and the
leptlepsy: convalstons", "i.v. solution additives", "entibiatics for topical sse', "e coces: Iall",
Publication Year: 1979
Bank Simelar 0.724/ 7850642296622
"adrenengics, 1sholants", "entteptleptics""
PHI:: 235158
"other diagrostic agents in ate", "diabetes mcllitus without
Doc. Retrieval
Tiie: Intoraction between local anesthetiis and analeptic drugo.
[epllepsy convulsions, by. large ansleptic drugd
Authune B&G Thumpaon, J Aidweie
complscation",
Atsemeet: Allhough Aibhuughhlliii fortuiioua pharmazologi interactions betweon local enesthetic agents anslaptic drugs may be protectiv when coses both agants are used. Mice pretrested with procains, lidocaine, ietracaine had lowar incidence seizures when convutaive doses oither rikethamde de dosapram hydrecntoride were given atosporitoncolly. des wordd groupe animale
clivicad undeainatea,
"drugs for peptic ulce- and gastro- assophageal retlux disease (gond)",
[local anestnetic reduce, apllapsy: convulalons]
"other streral sppllments in ate", "soinal fuxion",
"beto blocking agents",
"other ansigestts and intipyretics in atc",
[idocaine, used as, antiepleptico]
"Iantrectoyy; excision
Triple Extraction
Intervertebral disc",
Murbelly
[lioccanne, most effective in reducing, mortality!
"Antrocrantal injury", "ether beta-lactas ant:bactertols
gieen metthhrmdde and treated with conspram, All animata treated partylerstetruzsl consulsed only idocans no (also used to an convusannt was oble redupe mortalityin this group. Of local ameathatic agantia, tetraca affented lassi protection from diatth, whereus idoca spemed be most effectivo. Publication Year: 1975
in atc"
[pertylenctatnzzal, causes,
upilspey: convulsions]
[lidocaine, reduces mortality from, epilepay; comvulsions
Figure 8: An example of concept KG extraction from Corpus (PubMed Abstract).
(Concept set in a patient visit)
(Triples)
"pleurisy; pneunothorax; pulmorary collaase",
"coronory atheroscierests and other
heart disease",
"potassiur supslemens"", "other Fractures",
[pisunsy; pneumothorar, pulmonary collapse, treated
adrenergics, inhalants],
"ansiclyties""
[diabetes mellitus without complication, may require monitoring of,
"opiotd amsigestes", "other castrotrtestinal disorders",
potassium supplements].
(epilapay:
antispliaptica.
"epilepsy: convulstons", "i.v. solution additives",
O
lepilepsy convulsions, may be treated with, amuiolyiiss. [other gastrointestinal discrders, treated with, drugs for postic ulcer and
"antibiotics for topscal use" "e codes: fall",
reflux disease (sord)).
"adrenengies, inholants", "antiopiieptics",
[e codes: commonty causes, other tractures), [e codee: full, can result in, intracranisl injury]. [other fractures, often treated with, opiaid anaigasics), [intracrenial injury, may monaged opicid analgssics), [spinal tusion, sungical alternative to, laminectomy; excision intervertebral disc).
"othar diagrostic agents in atc',
coeplicotion",
"elabetes mcilitus without
Triple Extraction
"drugs for peptis ulter and gastro-
oesophageal reflus disease (gord)", "other nineral supplaments in ster. "spinal fusion',
[coronary aheroscierosis with, blocking agants],
treated
"beta blocking agants",
[other beta- iete-icctam antibactorials atc, antibiotics for topical use]
balongs to same class 6s,
"other onalgestcs and antipyretics
in atc",
"laminectowy; excision
intervertebral disc",
"intracronial intury'.
"other bele-lacten anlibacteriels
ate"
Figure 9: An example oncept KG extraction from LLM (Claude-3.5),
25

---

Published as a conference paper at ICLR 2025
C BASELINE IMPLEMENTATIONS
ML-based Models. To ensure fair performance comparison, we implemented all machine learn- ing (ML)-based Electronic Health Record (EHR) models using PyHealth (Yang et al.. 2023a). Since GRAM (Choi et al., 2017) and KerPrint (Yang et al., 2023b) were not yet integrated into PyHealth, we separately implemented these models within the pyhealts th models module. Our implemen- tations were based on the original codebases for GRAM and KerPrint For consistency across all ML-based models, we set the embedding size to 256. We trained the models using a learning rate of le-4 and employed an early stopping mechanism based on validation loss to prevent overfitting.
LM+ML Models. We implement GraphCare (Jiang et al., 2024a) using their official codebase" with their default setting for each component. We use =ext sembeddingg333 version of eext- seatt-mbeddinng-dae 022 used in the original implementation) as the embedding model for the embedding initialization, and use their proposed BAT as the base GNN model. We implement RAM-EHR (Xu et al., 2024) using their codebase" with the settings mentioned in the implementation details in their paper. We implement EMERGE (Zhu al., 2024a) fully following the implementation details provided in their paper (with the LLMs Clinical-LongFormen BGE-M3, Qwen 1.5-7B Chat, and DeepSeek-V2 Chat used for different purposes in the pipeline).
(an improved
LLM-based Methods. For zero-shot and few-shot wompping-basee EHR prediction with the LLM, we utilize the template presented in Table which includes task description, task-specific instruc- tion, patient base context, supplementary information (based on retrieval), and Patient References (similar patients). Unlike the structured format used for reasoning-chie here is presented in free-style manner, which, as our study indicates, results in better performance.
generation, the reasoning
We implement the EHR-Cot CHR-CoAgent (Cui et al., 2024) approach as described in their paper', which combines the strengths of predictive agent reasoning and critical agent instruction to create accu- rate few-shot prediction system for our tasks. The implementation consists two main components: predictor agent and a critic agent.
The predictor agent is responsible for generating predictions and providing explanatory reasoning based the input EHR data. Given patient's ' medical history, the predictor agent analyzes the relevant information and generates the most likely prediction along with a step-by-step explanation of its reasoning process. The prompt used for the predictor agent is as follows:
Given the following task description, patient EHR context, task instructions, and similar patients. please make a prediction with reasoning.
Task
[Task Definition] [Regulator]
# Patient EHR Context # [Patient's ' Context (Base)]
# Task Instructions (Guidelines) [Refined Guidelines if iteration else Initial Guidelines]
# Similar Patients #
['Top-K similar patients" contexts]
Give the prediction and reasoning in the following format:
Reasoning #
hhttps://ithhb: wtt.h//..o/rttttaa/////oo/ao httpr
Note: While EHR-CaAgent was not originally designed for mortality readmission prediction tasks,
have made minor custom modifications to adapt to our specific use cases.
26

---

Published as a conference paper at ICLR 2025
(Your reasoning here]
# Prediction #
[Your prediction here (1/0))
Output:
where Task Definition is a brief definition of the task (e.g.. Mortality Prediction Task: Objective: Predict the mortality outcome for a patient's subsequent hospital visit based solely conditions, procedures, and medications. Labels: = mortality, 0 = survival). We introduce a Regulator to import prior knowledge of the dataset avoid the LLM over-focus on improving true positive or true negative. For example, we set ***t to Notice: Only the patients with extremely very high risk of mortality should be predicted I. as the regulator for the mortality prediction task. The existence of the regulator significantly affect the final result for imbalanced datasets, as shown in Table This is because that the instruction-upe updating approach used by EHR- CeR-CCoggent tends to excessively penalize false positives to produce instructions that boost true positíves, especially in cases of imbalanced data like mortality prediction the two datasets. However, this can compromise true negatives and accuracy in such scenarios.
Table 5: Significant performance difference between EHR-CoAgen w/ regulator and w/o regulator.
MIMIC-III Mortality
MIMIC-IV Mortality
Accuracy
Macro F1 Sensitivity
Specificity
Accuracy
Mucra Fl
Sensitivity
Specificity
EHR-CoAgeet w/ Regulator FHR CoAgent wla Regulater
87.4
51.7 39.7
13.0
91.7
81.0
55.5
13.8
97.0
53.6
51.9
53.6
72.4
61.8
51.3
77.4
Refined Guidelines are the instructions refined by the critic agent, while Initial Guidelines the seed instructions, which is as same as the task descriptions we used for other methods.
The critic agent, the other hand, plays different role in the CHA-Ctt framework. It ob- serves the predictor agent's outputs alongside the ground truth labels and identifies error patterns and discrepancies in the predictor's reasoning process. Based this analysis, the critic agent refine the instructions to improve the reasoning process of the predictor. The prompt used for the critic agent is follows:
You assistant who is good at self-reflection, gaining experience, and summarizing criteria. By reflecting on failure predictions that are given below, your task is to reflect on these incorrect predictions, compare them against the ground truth, and formulate criteria and guidelines to enhance the accuracy of future predictions.
The original instructions are provided under "# Task Instructions (Guidelines) #". Your task is to refine the instructions based on the discrepancies between the predictions and the ground truth.
# Input Data #
(input data batch with prompts, predictions, and ground-truth labels]
Instructions
1. Please always remember that the predictions above all incorrect. You should always the ground truth as the final basis to discover many unreasonable aspects the predictions and then summarize them into experience and criteria.
2. Identify why the wrong predictions deviated from the ground truth by examining discrepancies
in the medical history analysis.
3. Determine key and potential influencing factors, reasoning methods, and relevant feature combinations that could better align predictions with the ground truth.
4. The instructions should be listed in distinct rows, each representing a criteria or guideline. The instructions should be generalizable multiple samples, rather than specific individual samples.
Conduct detailed analysis and write criteria based on the input samples, rather than writing some criteria without foundation.
27

---

Published as a conference paper at ICLR 2025
7. Please note that the criteria you wrote should not include the word "ground truth".
Your output should be the new set of guidelines under "# Task Instructions (Guidelines) #" that can be used to improve the predictor's reasoning process.
Output:
which is mostly the same as the one in their paper (Cui et al., 2024).
Our implementation of EHR- CeAeett follows an iterative refinement process, where the predictor agent generates predictions, the critic agent analyzes incorrect predictions and refines the instruction, which is consolidated and integrated into the predictor's prompts for the next round. The incorrect predictions are divided into batches, and the critic agent refines the instructions for each batch. The instructions from all batches then consolidated using the LLM to identify the most important and recurring insights across the entire refined instruction list. This consolidated new guidelines integrated into the predictor's ' prompts for the next round, allowing the system to effectively improve its prediction performance. We iterate the process 5 times.
The consolidated instructions are generated using the following prompt:
Given the following set of guidelines, please consolidate the insights into a concise and coherent of guidelines for refining the predictor's reasoning process.
Set of Guidelines #
[A Batch of Guidelines]
# Instructions #
1. Analyze the provided guidelines and identify common themes, patterns, and key insights.
2. Synthesize the insights into a consolidated set of guidelines that capture the most important and recurring aspects.
3. Ensure that the consolidated guidelines clear, concise, and actionable to refine the predic-
tor's reasoning process.
4. Create numbered list of the consolidated guidelines in the same format the original guidelines.
Output:
The consolidated instructions are then recursively consolidated until the final list size is smaller than 10. This is done ensure that the consolidated instructions can be effectively integrated into the predictor's prompts for the next round, considering the limited context window size of the L.M.M
By incorporating the consolidated instructions into the predictor's prompts, the EHR- EHR-CoAgem ap- proach enables an iterative refinement process to improve the accuracy predictions.
28

---

Published as a conference paper at ICLR 2025
D IMPLEMENTATION DETAILS OF KARE
D.1 STEP 1: MEDICAL CONCEPT KNOWLEDGE GRAPH CONSTRUCTION
Step 1.1: Medical Concept-Specific Knowledge Graph Extraction
We use UMLS, PubMed Abstracts, and Claude 3.5 Sonnet as the sources for knowledge graph extraction from Biomedical KG, Biomedical Corpus, and Large Language Model, respectively. The extraction details are showcased in Appendix B.1, B.2, and B.3, respectively.
For UMLS, we utilize the "Full Release" version under "2024AA Full UMLS Release Files" For dense retrieval from PubMed abstracts, we utilize the local embedding model Nomic (dimension 768) (Nussbaum et al., 2024). We use Amazon Bedrock" to access the Claude model.
The resulting KG triples from UMLS, PubMed Abstracts, and Sonnet are 29,434, 259,938, and 315,492, respectively.
Step 1.2: Semantic Clustering
For semantic clustering of entities and relations in the KG we build above: (1) first use the text- embedding-33-art model (dimension = 1024) from Azure OpenAl to retrieve the text embeddings of entities and relations; and (2) use Seikit-learn (Pedregosa al., 2018) to perform agglomerative clustering based on those emheddings. The optimal cosine distance thresholds 0g and 0,. are both found to be 0.14, resulting 513,867 triples in total after clustering.
Step 1.3: Hierarchical KG Community Detection and Indexing
We employ Graspy (Chung al., 2019) implement the hierarchical Leiden algorithm, setting the maximum size for each top-level community kmax_cluster: to 5.
To enhance community diversity, the algorithm is 25 times with different randomness each iteration, resulting in unique 59,832 communities (with different combinations of triples) where there are 40,934 communities with the size smaller than 20 (Z,), and 57,247 communities with the size smller than 150 (Z.).
Using Claude 3.5 Sonnet as the LLM, we generate 147,264 community summaries (including both general and theme-specifie summaries) with the prompts shown in Figure 12 and 13.
D.2 STEP 2: PATIENT CONTEXT CONSTRUCTION AND AUGMENTATION
We use the template as shown in Figure 11 to construct patient's base context based on their EHR.
To retrieve the relevant medical knowledge for context augmentation, set 0.1, B 0.7, A1 = 0.2, A2 = 0.2, and Ag = 0.3 as the hyperparameters, each tuned in the range of [0, 1].
D.3 STEP 3: REASONING- REASONING-ENHANCEE ENHANCED PRECISE HEALTHCARE PREDICTION
Step 3.1: Training Sample Generation
To generate reasoning chain training samples, we leverage Claude 3.5 Sonnet as our expert LLM. Ensuring EHR data protection and ethical use is paramount; therefore, LLM interactions are conducted via the Amazon Bedrock platform a cloud infrastructure that allows to fully harness LLM capabilities while maintaining strict privacy measures. The maximum output length for Sonnet set as 4,096 tokens. We use the prompt in Figure 16 for the reasoning chain generation here.
Step 3.2 Multitask-Based Fine-Tuning and Prediction
Our fine-tuning framework is implemented using the TRL (von Werra et al., 2020), Transformers (Wolf et al., 2020), and FlashAttention-2 (Dao, 2024) Python libraries. use Mistral-7B-Insttnnn
nih htt../////eeeerc///mom.em s hot.a//aaacaaaaoooo/ooccccocoene
ummlskrooleedeeesuuccss.s html
"https: /ttp///occ/ooo///o///////coc aws, www.:/o.ccoo//eeeeeeeeeeeeeeee
1he use of Amazon Bedrack is authorized by MIMIC: http://:www.iioootoo.//////////a
29

---

Published as a conference paper at ICLR 2025
v0.3 (Jiang al., 2023) our local LLM, full full-parameter fine-tuned using DeepSpeed (Rasley al., 2020) with the following configurations:
Parameter
Value
model caapee__ame_orppppe
misttalli/Miisrra-----nnsceeee
torch_dtype
bfloat16
use_flash.attention
true
preprocessing num workers bf16
12
true
gradient accumulationn _ttps steps
4
gradient_check checkpointing
true
learning_rate
5.0e-06
max max_seq_eentth
6000
num_train_epp
3
per device train batch size
Ir scheduler type warmup_rand praioo
cosine
0.1
Table 6: LLM Fine-tuning Configuration Parameters
Hardware Information: The experiments were conducted on system with AMD EPYC 7513 32-Core Processor and 1.0 TB RAM. The setup includes eight NVIDIA A100 80GB PCIc GPUs, each with 81920 MiB of memory, providing total of 640 GB GPU memory. The system's root partition has 32 GB of storage.
Training of each model runs eight NVIDIA A100 GPUs and typically completes within five hours. After the training process, we select the best performing model checkpoint based on valida- tion loss perform the prediction.
30

---

Published as a conference paper at ICLR 2025
E EVALUATION METRICS
In this work, we employ four key evaluation metrics to assess model performance:
E.1 PRIMARY METRICS
Accuracy: Measures the overall proportion of correct predictions:
TP + TN
Accuracy
TP TN FP+ FN
2. Macro-F1: Provides balanced measure that is particularly important for imbalanced datasets:
Macro-F1 2 x Precision x Recall
(8)
Precision + Recall
where Precision TP/(PP FP) and Recall TP/(TP FN)
Sensitivity (True Positive Rate): Quantifies the model's ability to correctly identify high-risk
patients:
TP
Sensitivity
TP + FN
Specificity (True Negative Rate): Assesses the model's s accuracy in identifying low-risk patients:
TN
Specificity
(10)
TN +
where TP True Positives, TN = True Negatives, FP False Positives, and FN False Negatives.
E.2 CHOICE OF METRICS
While metrics like AUROC (Area Under the Receiver Operating Characteristic) and AUPRC (Area Under the Precision-Recce Curve) are commonly used for imbalanced classification tasks, they are suitable for our LLM-based approach for several reasons:
1. LLM Probability Limitations: LLMs compute next-token probabilities over their entire vocab- ulary rather than binary class probabilities. These probabilities:
Are distributed across the full vocabulary rather than just binary classes
Depend on how different LLMs encode the same label ("0"/"1") using different tokens
Are not directly comparable to class probabilities output by traditional ML models
Clinical Relevance: Our chosen metrics provide direct clinical interpretability:
Sensitivity is crucial for identifying high-risk patients who require immediate attention
Specificity helps avoid unnecessary interventions for low-risk patients
Together, they effectively evaluate performance on imbalanced datasets using only final predic- tions
Our metric selection aligns with recent L--M-based EHR prediction works, such as EHR-CoAgent (Cui et al., 2024).
The effectiveness of our metric choice is demonstrated in the MIMIC-II/V mortality prediction tasks, where despite high class imbalance (positive rates of 5.42% (12%%/19.166). 16%), KARE achieves signif- icantly higher sensitivity (24.7% compared to baselines while maintaining high specificity 108.3%
31

---

Published as a conference paper at ICLR 2025
F CASE STUDY
F.1 REASONING- PREDICTION BY OUR TUR-TUNED MODEL
INPUT (Reasoning Task)
follwwing polior Bei comoott in pollerss and icfiirodd mezicall roweedee, "eecee pervice soyitep that
predictin
cofco bouod or The portont contost and
Arter the secrnning propose, providde the preciction
Lrooor curtnnt lbblkf loceral on
Roblored Medical knowlodge
lubal (orl).
2 prododurcs dailey
fett membirnne anid dellory
Thly by droomiiin cooditoo eocee comalnti prenee ne prim of membrores we This may and more corininng poremtot fill readmasioon
Medicetiona
Bhymin proporations
On ective medirt The portent will revreo Red
ta tookccal use
tosphro 13 of decherge.
campputiiio arthatomins onsothttics anc.
"thysid dto dan dling prognancy conimpoct managed ad thycld conde ors may seut caivene postooium comple anoons necessko lingng recomiiion.
Lobwl: recommoddgn within recdinision within doys
5. colckum diemarrs
amtoynet
Note: drut the moto on comorehensheey to
2. <lugg for cord parton
denemiiee
of icodiiision gool
trroll propory
ceterrreelyyyyyy deringusie betwyen pallent Fon to mecimmieedd and there noxt
contentommmto ontin sumorte products, con sloriids
combinatio of abewns perinel troumn reoci o uu rapical nomon aod toolmennts iseliood of wecemedes
combinations
espectorsnte
coppooooaans oud. combinatiooe cough
colloto puilon tam lypcco for monaging poi and decommton mownver cowe mon and cimple
Potiont ouooso
oweld
bobet
adiiasee offects that could proting recovvny
I polmyytiinniin and other pronlem of covaicte
"Crai um upp memer or otwn pownotoed dur tinn progrn ancy health Adeuuute touuppor intake
Paterrt cuuzz
2 other complicatooa of pregeencc
3 thurso condde
Vist Dr
and moethrnng ng con helo prnvent mosicoiion cniooss co cou del delicancyy imbolonnes co combined with fhydld lorlio indicoless procetive reprrcten ongmeett the parkents thecols owee iox rocaant on during pregraancy and other vegiinal deliven The monogammnt
reloted
oe infinum winn
5 other pregnoncy and de vary including nammal
wy.damydraannicc. and other preclem of camiottc
Procedures
-b relcted hourooll to ood lourn and vuvvo
regar cureen obrietiic ioceration
3 camp cottons pregrancy
5. officr comp collors of bifh puoperiom affocling other pregrimncy and dellerr incluttng normal
mhool preporctios
arimatory lifoorra oik producti.
ogants tour uniss help
nin slooks
Pooedine eec
othee ordigwite cntipyytec ntc instmmnnt
reoe curent obotetic lnowrdl
prevent potemtoi comp comp and seduc
fasures for los calure
5. "Indhiaao locing fhyoldd proporatione and lio'l peeproncy theddd condilg lon hoose thyold cemons regulal coucloli the postportu pattod. meopes
anwothenin,,
à cnuge for comtipttion
Byydddd proparations
volciiie noi combiieed
cough suporsscnts, soct comtinattorann
8 entoc do
homonees
colcium succlements
3. Origi for cond pation
co complco for rec hopplo reodiniss
5. celikk onalgeescs
"CCon'iprtoo" is bequuit laue porpoot and drugs
ono
in
Sievker Polients a
trac hamosholaa and anai
monagaamen uote bowel kunction and prevent on of commitation can avoid repcmision"
to repical
rrnowy
ontiiaa onisiholiss
7. contocies
isowotone
Concttoon
Trool prppprttoos
dsinted dulivar roculcc oodll qullii caro and monihorin potectilly y leading warmmidees
sobrtrccainnii and other procioms amniccic
wwweertommnooooooooooo and ontineumeto paoducts
aibemicted trturn tn oetinuum wile complica pregeanccy
13. un og op/
8 theiii conden
lobel
pregponcty kee loctivinng
OUTPUT
Reasaninge
1. "Palleel Ovandasmt:
repeil ou rent 50 dodeelic looerd for and other pioce durss to osli delivey, Forpart medio medicrtinn includes thoof pipportin fortoss ondinng formmmtay and conin to produets anolgesics. pilyveiio agents nemommed
four orripuutics cougs polion, boclalll voochess crro co supolements ono fhrnodd
Refreved Medical tooweedgert
Tokkvorrmmed oon noreose the nik preter inbor and pie premonume rupture et membrones ponentitly woding ng compltont om other ded very that dloee monthring
Rlsudd diockes during fancy. il not croparle mer agod may rouut advase oulcomes and tn moce the lidl budroodd puttocrum comptica fore rocosalaiing
recommion
inpeention compicaatin from ded dellvery
aeditional
montoring potent weding rendmmsion
hoppo fhyold mone logu d lo oudial polloe fum and int alancco may lcod complicatione rquaingroot on mot clouchy moni urod and utlue
a "CCmppatton with Sandirr hotlereee
cen and severdliie artice the teail patient, ni paly mydramiiss chedl toura to and other compicca throks clerdaa resei curnnt sbateffic
medicanions towevee of otherr compricot om pangrim bnny comron sordd to the cond piiito oct acttert firt siill wet whovon van sea within is onyy ollo hoo this oor dilon, along with management
other eroceddess
aid
has oc
puerperum
Newards The polion's hoylyyy odlyt rofammiss hindid diordes, and other complkorton of oregnanco, compeied with the need to repoil of obste kic coriioiion and osdi and delivery proceduues t-ppeer a higher mie perpartum compleett coe fomm that mar recamre saten, The ponents smitarty to the patern who
wither dere and medicaltonn petveill live mossurrs the ciionco of mud spo isk footon and the nord for dies fooi of the fod lund and coreniic comp looo is tiom cilatoo dedivdyyy indicotoo
odkilionll
"CCunchoios"
Bosed potiert of podytty thyrii dioodder other lection oregronou and the penceduows clong with dimiiaily in.n potter who witk rectmitte whhh etyy ther 3 sigs toveoood the poneett well be recdentee to herppter within
conteiic lnoertton ordtreed dellveey
Figure 10: Case Study of the Fine-tuned KARE Model.
32

---

Published as a conference paper at ICLR 2025
F.2
ZERO-SHOR SHOT & FEW-SHOW COMPARISON
Table 7: Case Study of zero shot and few shot EIR prediction with LLM (Somnet-3.5 Ground Truth: (Prediction 1).
parient will die from syothetic MIMIC-III by HALO (Theodorou al,, 2023).
concern, parients
Case
Description
Input Prompe with EHR
Given the following task description and patient please make prediction with reasoning based the paticat's context.
Mortality Prediction
Objective: Predict the mortality outcome for paticnt's subsoquent hospital visit based solely on conditions,
procedures, and medications.
Labels:
mortality, 0 survival
Key Considerations:
Conditions:
Severity diagnosed conditions (e-g.. advanced
heart failure, sepsis)
Presence of multiple comorbiditios
Acute vs. chronic nature of conditions
Procedures:
Invasiveness and complexity of recent procedures
Emergency vs. elective procedures
Frequeccy sustaining procedures (e.g.. diulysis, mechanical ventilation)
Medications:
Use of high-risk medications (e.g.. chemetherapy drugs, immunsssuppressanns)
Multiple medication use indicating complex health issues
Presence of medications typically used in end-of-life care
Note: Focus on combinations of conditions, procedures, and medications that indicate critical illness or high mortality, Consider how these factors interact and potentially exacerbate each other. Only the patients with extremely very high risk of mortality should be predicted as 1.
Patient Context:
Patient ID: 29488
Conditions:
Deficiency
anemia
Essential hypertension
Complication of device; implant or graft
Congestive heart failure; nonhypertensive- Cancer prostate
Anxiety disorders
Thyroid disorders
Disorders lipid metaholism
Conduction disorders
Mycoses
Other diseases of kidney and ureters
Cancer of csophagus
Procedures:
Diagnostic catheterization; commary arteriography
Other or procedures on vessels other than head and nock
Colorcctal resection
Drugs:
Other drugs for obstructive airway discases, inhalants in ATC
Lipid modifying agents, plain
Antithommbotic agents
Angiotensin receptor
(ARRs), plain
Visit I:
Conditions:
Congestive heart failure; nomhypertensive
Cardiac dysthythmias
Shock
Alcohol-related disorders
Diabetes mellitus complications
E. codes: adverse effects of medical care
Disorders tooth and jaw
Procedures::
Other vascular catheterization: not heart
Blood transfusion Drugs:
Other drugs for obstructive airway diseuses, inhulants in ATC
analgesics antipyretics ATC.
Drugs for peptic ulcer and aastro-cesopageal refiux disease (GORD)
Beta blocking agents
Potassium supplements
Continued page
33

---

Published as a conference paper at ICLR 2025
Table 7 continued from previous page
Case
Description
Supplementtury Information: |PLACEHOLDER]
References: [PLAACTHOLDDER
Format:
Reasoning
YYour reasoning here]
Prediction
(Your prediction here (1/0))
Case (1): Zero-Shot Prompting with EHR Base Context Only
As provided, Sonnet 3.5 gave the following incorrect prediction:
Model Output:
Reasoning
While this patient's is condition is scrious and deteriorating, there are several factors to consider: Severity conditions:
The patient has multiple conditions, including congestive heart failure, cancer (prostate and esophugus), and cardiac issues.
New complications have arisen in visit I, such as shock and cardiac dysrhythmias, which are concerning
However, some of the most immediately life- sie-chreatening conditions (e.g.. sepsis, advanced stage cancer, multi mulli-urgam failure) are not explicilly mentioned.
Procedures:
transfusion and vascular catheterization) suggest decline health, but they are also interventions aimed at stabilizing the patient.
procedures
There's no mention of extreme life-sustaining procedures like mechanical veatilation or dialysis.
Medications:
The medications listed are primarily for managing chronic conditions rather end-of-life care.
mention high-risk medications chemotherapy drugs strong apioids for palliative
Disease progression:
While there's a clear decline between visits and the conditions in visit I, while serious, are potentially manageable with appropriate medical intervention.
Absence of certain High-morfalite indicators:
There's no mention of conditions like sepsis, acute respiratory failure, advanced stage metastatic cancer, which would significantly increase mortality risk.
only predict mortality (I) for patients "extremely very high mortality", this patient's condition, while serious, does not meet that extremely high threshold. The presence of active interventions (blood transfusion. modications for managing conditions) suggests ongoing treatment rather than purely pulliative
task description's
Prediction # (incorrect)
(2): Zero-Shot Prompting with Classic
Supplementary Information:
RAG-Augmmened Context (Knowledge Souree:
Rank 1, Similarity: 0.200470362854000 PMID: 314160 Tide: Resection of posterior ventricular uncurysm. replacement of mitral valve, and coronary bypuss. Authors: C Chavez, G Rodriquez, H Conn Abstract: A 58-r old man had sudden and progressive heart failure after a severe myocardial infarction. Aggressive medical teeament consisting of diuretics, vasopressors, and digitalis failed to improve his condition significantly Cardiac catheterization disclosed critical stenosis in ventricul ancurysm, and severe mitral insufficiency. Intermittent third degree heart block developed after admission. Surgical correction resulted in dramatic recovery, and three years after operation is fully recovered and usmplommatic, Publication Year: 1979
PubMed Abstract)
anterior descending branch the coronary artery, a large posterior left
1154419942 994336 PMID: 80488I Tirle [Acute coronary thrombosis in a 28 year-old woman]. Authors: Mottč Vogel. G Coatantise, Marictte Abstract: A 28-ycar-old with
Rank Similarity: 0.7 27123003599
hospitlized myocardial infarction complicated by bifascicular block followed by complete atrio-ventriculat block with regressive course. coronary arteriography performed the 10th day demonstrated thrombosis the unierior interventricular artery, the rest of the coronay network being normal. The influence of dyslipiduemia taking contraceptives was discussed as actiology. Publication Year: 1975
cardiovaseular lar history.
Continued on next page
34

---

Published as a conference paper at ICLR 2025
Table 7 continued from previous page
Case
Description
Rank Similarity: 0.7 0.71020050148777 PMID: 42254 Title: [Non artyythmoeenic sudden death as complication of coroniay heart discase). Authers: H Zilcher. Glogar Abstract: In cohort of 417 palients admitted consecutively the Corenary Care Unit for acule myocardial ischemia (unstable angina pectoris in 121, acute myocardial infarction in 296 patients) 21 cases non arrhythnogenic sudden death occurred within 24 hours after admission. 16 of these patients suffered from acute myocardial infarction and 5 from unstable angina pectoris. Cause of death was cardiac rupture in 12 and pump failure paients with acute myucardial infarction. whereas patients with unstable angina pectoris died from pump failure. Patients with cardiac rupture within 24 hours after admission, had signiticantly higher systolic and diastolic blood pressure in comparison with the other groups and with patients dying from cartiac rupture the third day, later. All patients dying from pump failure with unstable angina pectoris and one of the patients dying from pump failure with acute myocardial infarction had blocker therapy ... Publication Year: 1979
(..)
Rank Similarity: 0.6921 10066645777 PMID: 443259 Tide: Chest pain, shock, arrhythmmias and death in a young woman. Authors: Abstract 9 figures form the core of this article deseribing and discussing a case sudden death. hours after year old
presented
hespital
emergency with chest pains. She had taken no medications other than oral conrrecptives (OCs) for years. The patient was admitted the cornnayy care unit where findings included a palpable blood pressure of 94 mm of Hg. beart rate of 128/min and a resgiralory rate of 28/minute. were
Cyanosis was noted, jugular veins were distended, and bilaterally; cardiac sounds were soft and a third heart sound was audible. Arterial oxypen tension was 15 mm of Ugg and curbon dioxide tension was 42 mm of lgg pll was 7.2. Ventricular tachycardia developed and ventricular fibrillation ensued. The patient oxygensted, external cardiac compression was performed. sodium bicarbonate, epinophrine, and caleium were administered.
over the lung bases
intubated and well
electrical detibrillation was performed. After several attempts, latter resulted in a slow idioventricular thythm on the elecimcardiogran, but neither the blood pressure nor pulse was detectable. Publication Year: 1979
Rank 10, Similarity: 0.692069 132545471 PMID: 938220 Title: Left ventricular function during nortic surgery. Authors: R M Carrol, R B Laravuso, JF Schauble Abstrnct: Fourteen patients undorgoing surgery for ancurssm or ocelusive discase the abdominal aorta were studied. Thirtoon patients had history of hypertession myocardial infarction: pulmonary disease. Tachyeardia, hypertension, and elevated pulmonary artery occluded. (PAo) pressure occumed in response laryngoscopy and intuhation in two patients; elevation of PAo pressure in response to nortic cross-clamping occurred in two patients. In three of these four patients, electrocundingrphiic eviderce of myocardial ischemia appeared. These events are important consideration of the occurence of myocardial inforction in patients undergoing abdominal aortic surgery. Satisfactory treatment myocardial ischemia has been accomplished with the use of propranokol hydrochloride and sodium mitropusside. Publication Year: 1976
patients had chronic obstructive
Model Output:
Reasoning #
Based on the patient's context and given task description, analyze key factoes to predict mortality risk for the subsequent hospital visit: Conditions:
The patient multiple serious conditions, including congestive heart failure (nonhypertensive) cardiac dysthythmias, and shock. These severe cardiovascular issues that significantly increase mortality risk.
The presence of diabetes mellinus with complications adds overall complexity of patient's 's health starus.
Alcohol-relate disorders may complicate treatment and recovery.
Procedures:
The putient underwent bloud transfusion, which might indicate severe blood loss or anemiaa, patentially linked the shock condition.
Other vascular catheterization (not heart) was performed, possibly related to managing circulatory issues.
The supplementary materials discuss these procedures but focus on their use in orthopedic surgerics.
Medications:
The patient is multiple medications, including heta-blocking g agents and drugs for obstructive airway discases, indicating complex health issucs.
Polassium supplements might be failure or other medications.
linked to managing electrolyte imbulances, possibly due to heart
Progression from previous visit:
The patient's condirion appears to have worsened since the previous visit, with new serious conditions like shock and cardiac dyarhythmias appearing.
Comorbidities:
The combination cardiovascular issues, diabetes complications, and alcobool-clated discrders creates comples clinical picture with multiple interacting risk factors.
Continued on next page
35

---

Published as a conference paper at ICLR 2025
Table 7 continued from previous page
Case
Description
However, there are some factors that suggest the patient not be at the highest risk of morrialit::
The patient receiving active treatment, including medications
conditions.
There's 's no mention of advanced cancer or other terminal illnesses.
The procedures mentioned are not highly invasive indicative end-stage discase.
There's indication of mechanical ventilation other life suslaining procedures.
Prodiction #
(incorrect)
(3): Zero-Shol Prompting with
Supplementary Information:
Patients with congestive heart failure combined with cardiac dysthythmias and shock have significanly increased risk of mortality duc to impaired cardinc output and hemodynamic instability.
KARE-Augmented Context (Knowledge Sourec: KG
Communities)
The presence of esophugeal cancer, especially when co-occurring with prostale cancer, indicates a higher hunden malignancy, which exucerbate overall health decline.
Complicatoos of devices, implants, gralls can lead which is life- threatening, expecially in patients with weakened cantine function.
infections oegan dysfunction.
The combination anemia thyroid disorders can exacerbatc cardiovascular conditions, leading poorer outcomes beatt failure patients.
Alcohol related disorders can negatively impact medication adherence and exacerbate conditions like hypertension and diabetes, increasing the of death.
use of hlord transfusions vascular catheterization procedures indicates severe clinical interventions often linked to critical conditions and higher risk of death.
Polypharnacy, including beta-blockers, potassium supplements, and medications for obstructive airway diseuses, reflects complex medical managemenn, which can associated increased of adverse events.
Patients with diabctes mellitus with complications face a greater risk of scrious outeomes due to potential organ dumage delayed healing processes.
Model Output:
Reasoning #
on the patient's context and supplementary information: Severity of Conditions:
The patient has congestive heart failure, cardiac dysrhythmias, and shock, which together increase mortality risk.
signilicanly impuir cardiac
The presence of cancers (esophageal prostate) indicates high burden malignancy. wursening the palient's 's overall prognosis.
Complications of devices, implants, or grafts may lead to severe infections or organ dysfumction.
Interactions Between Conditions:
Anemia thyroid disorders can worsen
outcomes.
Alcohol-related disorders may impair medication adherence liver function, complicating
managementt of other conditions.
Diabetes mellitus with complications further increases mortality risk due to potential organ damage.
Procedures:
hlord transfusions vascular cutheterization suggests interventions severe anemia bemodynamic instability.
These procedures are associated with higher mortality rates in critically ill patients.
Medications:
Polypharnacy with beta- stta-blockers, polassium supplements, and other medications indicates complex managerment and high risk of adverse events.
Medications for obstructive airway disemes suggest respiratory complications.
The provided information highlights that this combination of severe conditions, multiple comorbidities, and invasive procodures is associated with a very high risk of mortality.
Given the critical combination of severe cardiac conditions, multiple cancers, complex comorbidities. and the necessity for invasive procedures, the patient al an extremely high risk of moriality and likely to die in the visit.
Prediction # (correct)
Continued on next page
36

---

Published as a conference paper at ICLR 2025
Table 7 continued from previous page
Case
Description
Case (4): Few Tew-Shot Prompting Supplementury Information: (Same in Case 3)
References:
Negative Patient References and KARE-Angmented Context
Example (Mortalit=
Patient ID: 10269
Visit Information:
Conditions:
Congestive failure; nonlippertensive
Sepiis
Acute respiratory failure
Shock
Procedures:
Mechanical venrilation
Blood transfusion
Vascular cathoterization
Medications:
High-dose vasopressors
Broad spectrum antibiotics
Example (Mortality
Patient ID: 11445
Visit Information:
Conditions:
hypertension
Type mellitus without complications
Hyperlipidemia
Procedures:
Routine blood tests
Electrocardiogram
Medications:
Metformin
ACE inhibitors
Model Output:
Reasoning
Comparing the new patient to the provided references: Similarity Example
Both patients have congestive heart failure and shock.
Both require invasive procedures like blond transfusions and vascular catheterization.
The patient 10269 has acule respiratory Cailure and sepsis, while the new putient has multiple cancers
dysrhythmias.
Both complex comorbidities and undergoing critical interventions, indicating high risk mortality.
Difference from Example
The negative reference has manageable chronic conditions without severe complications.
procedures or conditions present.
The patient bas significantly more conditions and interventions.
Patient's Conditions:
The combination congestive heart failure, cardiac dysrhythmias, shock, and multiple places the patient very high risk of mostality.
Disorders linked to alcohol use may reduce compliance with preseribed medications and wersen chronic illoesses such high blood pressure and diabetes, elevating mortality
Diabetes mellitus with complications and alcohol related disonlers further exacerbate health risks.
Procedures and Medications:
Invasive procedures and polypharnncy indicate critical illness and complex management.
Similar to the positive reference, these factors contribute to higher mortality risk.
The provided knowledge and patient references highlight that patients with combination severe cunditions and interventions likely die. Given the strong similarities the posirive reference and the critical nature the paticnt's 's health status, it highly likely that the patient will die in the next visit.
Prediction (correct)
Analysis
Case (2) demonsirates that when using traditional RAG methods with sparse retrieval from PubMed abstracts, the supplementary information contains several keywords matching the patient's conditions but fails provide relevant information regarding the patient's high mortality risk.
Continued page
37

---

Published as a conference paper at ICLR 2025
Table 7 continued from previous page
Case
Description
Case (3) shows that KARE's dense and relevant knowledge retrieval provides comprehensive insighis into the comorhidities, and implications procedures medications,
patient's condition, interactions
leading to correct prediction.
Case (4) illustrates including paticat references with KARE-ougmentd context forther aids the model by providing concrete examples compure aguinst, eahancing remsoning and ensuring
prediction.
Conclusion
This study emphisizes importance KARE's approsch providing dense, relevant knowledge
significanlly improves the model's ability to make accurale predictions. By integrating compechensive Ledwledge medical concepts, KARE beyond what achievable with base context traditional RAG methods.
gruphs and considering complex interactions
reasoning capubilities
G TEMPLATES, PROMPTS, AND EXAMPLES
Readrission Prediction Task:
Patient I0: иыУпУу
Objective: Predict if the patieett will be readritted the hospital within days of discharge.
Visit 8:
Conditions:
days, readrission within 15 days
pleurisy; pneunothorax; pulnonary
2. cardiac dysrhythhiass
3. congestive heart failure; monhypertensive
Note: Analyze the information imprehensivell to deternine the tikelshode of readnission. The geal accurately distinguish between patients who readritted
4. other liver diseases
Procedures:
incision of pleura; thoracentesis; chest
drainage
2. other
procedures
3. other diagnostic cardiovascular procedures
# Patient EHR Context
Medications:
Patient
antithronbotic
other analgesics antipyretics in atc high-ceiling diuretics antacids
Visit
Canditians:
1. pleurisy; pneunothorax; pulronary collapse
5. drugs for constipation
2. pneuronta
congestive heart failure; nonhypertensive
nutritional deficiencies
Label:
circulatory disease
cardiac dysrhythnias
screening and histary of nental health and substance abuse codes
Patient 10: 22222222
heart valve disorders
Visit 8:
Conditions:
Procedures:
1. pleurisy; pneunothorak: pulronary col lapse failure; nonhynertensive complications surgical procedures nedical care
incision pleura; thoracentesis; drainage
congestive
Medications:
4. screening and history of nental health and substance abuse codes
high-ceiling diuretics
adrenergics, inholants
5. cardiac dysrhythnias
sunp
drugs abstructive airway diseases,
inhalants atc
Procedures:
5. beta blacking agents
diagnostic procedures respiratory and nediast
6. viral vaccines
other analgesics and antipyretics in drugs for constipation
2. other non-or therapoutic procedures on respiratory system
antithronbotic agents
3. incision of pleura; thoracentesis; chest drainage
opinid ana lgesics
11. antiemetics and antinauseants
12. other rineral supplerents in atc
13. L.v. solution additives
1. other agents acting on the renin- angintensin system in atc
14. other diagnostic agents in atc antiarrhythmics, class iii
2. antiadrenergic agents, peripherally acting
(continue)
Label:
Figure 11: An example of the patient base context.
38

---

Published as a conference paper at ICLR 2025
General Summarization
Theme-Specific Summarization
You are knowlndigebble medical assistant fosked with generoting comprehersive summary of the medical concepls relafionships provided
You are knowledgecble medica assitonnt tasked with generating theme speciiic surmmary of the medical concepts relationships provided
Given list medical triples in the format (enfityl, relationahip, entily2), please
Given a isl of medical hiples in the formal [eriliyl, relationship, entit?) summory that focuses on how the knowledge represented by these Hoples relevani le the given lheeme. The summary should highight the concepti, relctiorshiiss and implications
single, summmar that caplures the key medica knowledie represented by these concepts and their relationsnips. the summary should written the third al the entity names for full context.
specific
plense
and include
mest perlinent
Tocus how these medical concepts intercennected and their relevance to medical understanding pattent care. If there are any contradictions in the trples, please resolve them the summmray.
The summary should be written the third person ond eniily names context Pleuse provide only the summary without any starting words or phroses, and without mentioning the individuull triples.
Please providee only the summary without any starting words phroses, and without mentioring the individual triples.
Exomples:
Martality Prediction:
The sumnmory should an infegrated representation of the medical inn enowedge containec the trples.
Triples:
(Diubeles, is a risk factar for, Cardiovasculer Disease] [Hypertension, assaciated with, Dinbetes) (Obesity. contributes to, Diobetes)
Example:
Triples:
!Diabetes, risk
Cardiovascular Diseose)
Theme: Martaliy prediction
|Hypertensin, ossociated with, Diaberes) [Obesily. contributes to, Diaberes)
(Theme Description)
Summory:
Summary:
Diabeles, hyperlension,
are sigr sigiliiaal factors that can incrense the Ikelhone of mortaiity. of
Diabetes, hypertension, and obesity are closely interconneciee medical conditions. Diabeles signilicanl foctor developing cordiavescular disecse. Hypertension, or high blood pressure. often assoclated with clobeles. Obesily contribules the development diabetes. excess body weight
Diaiette directly associated ol with on increased which leading couse death. Hypertention and obesity, which offen ca occur with dlabetes, further compound these risks, Patients with these conditions close monitoring aggressive
cardiovascular
resislance. Managing lhese condilions logelher crucial for reducing the rikk of serious complicctions and improving overail pattent
management to mitigate the mortality.
Triples: (Formalted Triples)
Triples:
(Fomatted Triples)
Iheme: (Theme)
Summary:
Summmory:
Figure 12: The prompts for (left) general and (right) theme-specific KG community summarization.
Summary Combination for Large Community
You ore knowledgeable medical ossistant tosked with generating comprehensive summary of the medical concepts and relationships based on the individual summarres provided below.
Given follwing summaries of community, pisase create single, conerent summary that captures the key medicol knowledge represented by Ihis communily. The combined summary should include al the imporiani enfilies and relationssips mentioned the individual summaries.
Focus on the main concepts, their relarionshios, and their relevance to medical understanding or palient care. The summarr should be concise yet informalive, provicing comprehensive averview of the community's medical knowledge.
Please ensure that the combined summory wiitten in the third person and includes all the relevant entily nomes for ful conlext.
Pleass provide only the summary without any starting words or phrases.
(summories)
Combined Summory:
Figure 13: The prompt for summary combination for large communities.
39

---

Published as a conference paper at ICLR 2025
KG Community
Theme- Specifiic Summarization
Hypariataizo
Dishotns
Mortality
"Diabetes and Hyperionsion are majer risk factors Chronic Disease propresses, such es Anemma, which can furtterr incresse mortidity and mortality. Early detection and management of Diabetes and Hyperiension are crucial preventing develoomeet and progeesionn of premature death this patient peputations."
"Chronic Kidney Disease (OKD), commonly coused by Hypertension, lead to Anemia, characterized by
Chronic Hidiney Disensn
complicationn
number Anemia may cause symproms such as fetigue shortness of breath, increasing the risk hospital readmissions. to prevent reacmissions CKD patients, essental
tnsite
underlying
reducing
hemoglobin
Supplementarion
Anumiaa
Supplementaatoo,, ensure treatment adherence."
Figure 14: An example of theme-specific summarization of knowledge graph community.
Task Description Mortality Prediction
Mostality Prediction Task:
Dojective: Predict the mortailly outcome for patient's subseqqunt hospital viil bosed solely on condfioon, procedures, and medicalions.
labels: mortality. C survival
Consideralion:::
Conditions:
Severily tiagnased conditions (e.g.. advanced cancer, severe heart failure, sepiis)
Presence of multiple comorbiditles
"Key Considerations" section only for zero/few-shof prompling & fraining sample generation
Procedlurs:
Invosiveness and complexity of recent procedures
Emergency vs. electiv procedures
Frequency of lfe-sustaining procedures (e.g. dialysé, mechanical ventilation)
Medicalions:
Use of high risk medicalions (e.g. chemotherapy crugt, immuncuppresannt
Multiple medicorion indicating complex healt issues
Presence
lypically
Nole: Focuss un combinations of condilions, procedures, and medicalions that indicale crilical illness high mortally, Consider how these factors internet potenticily exccerbate each other. Only the pottents with extremely very high risk of monoliy should be predicted as I,
Task Description Readmission Prediction
Readmission Prediction Task:
Objeclive: Predicl the palient will readmilled the hospilal within 15 daass discharge based solety condition, procedures,
Labels: = readmission within 15 days. 0= no readmisiion within 15 days
Key Conscerations:
Condilions:
"Key Considerations" section only for zero/few-shot prompting & training sample generation
Chroric ciseases with high risk of exocerbation
frequent adjustments Recent acute conditions potenticl for complications
Conditions requiing
Procedures:
Recent major surgeries interventions high complicatio rates
Procedures that require extensive followup care
incomplete partiilly successful procedures
Medicalioon:
New medication regiment that may require adjustment
high side Complex medication :chedues that leyd lead to adherence issues
with therapeutic
Note:
Analyze the informatin comprehensively determine the Ikelhood of readmision. The goal accurately distinguish between porients who are Ikkly to readmitted those who not.
Figure 15: Task descriptions of the mortality prediction and the readmission prediction tasks used the paper. Note: For the fine-tuning process, we do not include the "Key Considerations" section in the input template.
40

---

Published as a conference paper at ICLR 2025
Civen the following fask description, patient EHR context, similar patients, refrieved medical knowledge, and ground Iruth label. provide a step-by-s Sep reasoning process that leads to the correct prediction:
Task # description)
Patient EHR Context #
(context)
Similar Patients #
(similar patients)
Retrieved Medical Knowledge #
(medical knowledge)
Ground Truth #
(ground fruth)
Please provide a step-by-step reasoning process that leads to correct prediction bosed on the palient's context, similar patients, and the retrieved relevant medical knowledge.
The reasoning choin should follow structured format:
1. Patient Overview: Check the key information in the patlent's context, with the Key Considerations from task description mind.
Relevant Retrieved Medical Knowledge: Highlihtt the retrieved medical knowledge pertinent the patient's condition.
Comparison with Similar Patlents: Analyze the similarifies and differences between the patient and similar patients, exploining how these factors influence the prediction. 4. Reasoning Towards Prediction: integrate the above information to logicolly reason towards the predicted outcome.
Conclusion: Summarize the reasoning and state the prediction without mentioning ground fruth.
The reasaning should be comprehenalve, medically sound, and clearly exploin how the patient's information leads the predicted cutcome.
Important Notes:
Do not mention ground truth label the reasoning process.
Use the relevant knowledge as needed.
Anolyze the similarilies and differences between the patient and similar patients to justily the prediction.
After generating the reasoning chain. please review it and indicate confidence in the reasoning chain at the end.
Options of confidence: [Very Conficent, Confident, Neutral, Not Confident, Very Not Confident.]
Output Formal:
Reasoning Chain
Patient Overview: (YOUR OUTPUT]
2. Relevant Retrieved Medical knowledge: [YOUR OUTPUT]
3. Comparison with Simiar Patients: [YOUR OUTPUT]
4. Reasoning Towards Prediction:
[YOUR OUTPUT]
5. Conclusion: [YOUR OUTPUT]
Confidence
[CONFIDENCE ("Very Confident", "Confident", 'Neutral", "Not Confident", "Very Not Confident"]]
Figure 16: Prompt used for reasoning chain generation for training sample.
41

---

Published as a conference paper at ICLR 2025
Fine-Tuning Input (Reasoning)
Fine-Tuning Input (Label Prediction)
[Reasoning] Given the following losk description, pallent EHR context, smiar potients, and retreved medical Inowledge. please provide a slep-by-sep reasoring process thatl leads to the preciction outcome based on the patient's context and relevon medical knoweddge. Alter the reasoning process. providee the predict label (0/))
[lobel Prediction] Given following lask description, polient EHR context siilrr patients, and retneved medical inowledge, please direcily predil the lobel [C/]]
# Task
(lask description)
* Patient EHR Context (context)
(lask, description)
(context)
(similor, pallents)
Similar Palienis
Rehieved Medical Knowledge
(similar patilets_
(medical ol [nkwwledee
Retrieved Medical Knowledge (medical knowledge)
Fine-Tuning Output (Reasoning)
Fine-Tuning Output (Label Prediction)
Reasoning (reasoning)
(Label (0/1))
Prediction (0/1))
Template used for the input and output of fine-tuning, (To improve the reprodueibility of KARE, we will publicize the processed data for fine-tuning the local LLM through PhysioNet)
Figure 17:
42

---

Published as a conference paper at ICLR 2025
H HUMAN EVALUATION OF REASONING CHAINS
We hired three MD students and an MD to conduct expert evaluation of the reasoning chain gen- erated by KARE. We instructed them to use same metrics introduced by Kwon et al. (2024) to evaluate 100 reasoning chains. All the evaluated samples are from test sets. Half of the chains were for mortality prediction, and the other half were for readmission prediction. For each task, 35 chains positive (correctly predicted), and 15 chains were negative (incorrectly predicted). The definitions of the metrics and rating scales are described as follows:
CONSISTENCY
Definition: Consistency measures how well the generated rationale aligns with the presented data and the model's prediction. evaluates whether the reasoning contains contradictions.
(Excellent): The rationale is entirely consistent with the provided data and prediction. No contradictions are present.
4 (Good); Minor inconsistencies are present but do not significantly detract from the overall coherence.
(Fair): Some inconsistencies exist that partially undermine the rationale's reliability. (Poor): Significant inconsistencies exist, making the rationale difficult to trust.
(Very Poor): The rationale is fundamentally contradictory to the data or prediction.
CORRECTNESS
Definition: Correctness assesses the medical validity of the knowledge and reasoning presented in the rationale.
(Excellent): The rationale is entirely medically accurate and reflects evidence- based clini- cal knowledge.
(Good): Minor inaccuracies are present but do not affect the overall clinical validity.
(Fair): The rationale contains some medically incorrect statements that could mislead clinicians.
(Poor): Multiple inaccuracies significantly reduce the credibility the rationale.
(Very Poor): The rationale is largely or entirely medically incorrect, making unusable.
SPECIFICITY
Definition: Specificity evaluates how detailed and precise the reasoning is in addressing the clin- ical scenario.
(Excellent): The rationale provides highly detailed and tailored insights specific to the patient case.
(Good): The rationale detailed but occasionally includes generalities.
- (Fair): The rationale is moderately specific, with noticeable generalizations.
(Poor): The rationale is vague and lacks sufficient detail to guide clinical decision-making.
(Very Poor): The rationale is overly generic and lacks any meaningful detail.
HELPFULNESS
Definition: Helpfulness measures the extent to which the rationale aids the prediction toward the correct diagnosis.
(Excellent): The rationale strongly supports the correct prediction and adds valuable clini- cal insights.
(Good): The rationale is helpful overall but lacks some critical insights.
3 (Fair): The rationale provides some useful guidance but is incomplete not compelling.
2 (Poor): The rationale adds minimal value to the prediction and lacks actionable insights. (Very Poor): The rationale is unhelpful and does not contribute meaningfully to the diag- nosis.
HUMAN-LIKENESS
Definition: Human--lieeness likeness measures how well the clinical rationale demonstrates insight and understanding of the patient description or diagnosis in a way that resembles human clinical rea- soning.
The rationale reflects deep clinical insight, contextual understanding, and reasoning that fully mimics human behavior.
(Excellent):
43

---

Published as a conference paper at ICLR 2025
(Good): The rationale generally matches human reasoning but lacks minor elements of nuanced understanding.
3 (Fair):
The rationale demonstrates basic human-like reasoning but misses several critical elements of insight.
(Poor): The rationale shows limited resemblance to human reasoning and lacks essential clinical insight.
(Very Poor): The rationale fails resemble human clinical reasoning and demonstrates no meaningful understanding of the patient description or diagnosis.
The evaluation results are shown in Fig. 18.
Mortality Predictions)
Mortality (Incorrect
4.3
4.25
3.87
3.60
3.5
3.5
1:33
3.0
3.0
2:80
2.5
2.5
Condistency
Speciiitty
Haptuness
Consistoncy
Specffitty
Readmission (Correct Predictions)
Readmission (Incorrect Predictions)
5.0
5.0
4.49
4.37
3.53
3:33
3.0
2.73
2:73
2.5
2.5
Condetencc
Conrectness
Specticity
Helpuiesss Human-Lkeresy Leeness
Consistency
Correctress
Specitctty
Helpfuiness Human-UUeeeet
Figure 18: Evaluation of reasoning chains generated by KARE.
Discussions:
1. Reasoning chains leading incorrect clinical predictions (negative cases) consistently score lower across all metrics. This highlights the critical role of high-quality reasoning chains in ensuring prediction accuracy. Enhancing the consistency and correctness of these chains could potentially improve the overall prediction performance.
2. Human-Likeness scores are notably lower for the readmission prediction task. According the experts we consulted, this is primarily because determining whether patient will be readmitted within 15 days is inherently challenging, even for experienced clinicians, given the limited information provided in the patient context. Factors like the absence of basic demographic details (c.g., gender and age) further complicate this task. Despite these limitations, KARE demonstrates remarkable ability to outperform clinicians predicting readmissions, showcasing its potential in information-scane scenarios.
3. For the mortality prediction task, inconsistencies between the reasoning chains and the final predictions were observed in certain cases. For instance, some reasoning chains concluded that the patient would survive the next visit, yet the final predicted label was "1" (indicating mortality). These discrepancies negatively impacted the consistency scores. Incorporating an additional verification step to align the reasoning chain with the final prediction may help address this issue and enhance overall reliability.
44

---

Published as a conference paper at ICLR 2025
1 PARAMETERS & TRAINING TIME OF MODELS
Performance vs Training
ConCare
GRASP
400
Monal-7B
Training Time (ninutes)
800
Parametors
1M
KAR!
10M
100M
1000M
7000M
GraphCare
EMERGE
Tranffomer RETAI FRAM
StageNat
AdaCamn Deepr
700
725
15.0
Th's
20.0
82.5
15.0
87.
93B
Performance (MIMIC-IV-Mortaitt Macro-F1)
Figure 19: Performance vs. Training Time. Training time refers to the duration required for the model achieve its optimal performance the validation set.
Performance Model
Mistru-78
Mitrolins RAG
KARE
107
Model Parameters (M)
RAM EHR
GraphCare
60 min
200 min
GRASP
101
GRAM
SeageNel
Transtormer
Despr
77.
806
Performance (MIMiC-IV-Mortalllyy Macro-F1)
Figure 20: Performance vs. Model Parameters.
45

---

Published as a conference paper at ICLR 2025
J NOTATIONS
Table 8: Notations used in our paper.
Notation
Description
G¹, V.V' RRR E'
Knowledge gruph belore semantic clustering and its components
G,V,R,E
Refined knowledge graph after semantic clustering and its components Concept-speed specific knowledge graph concept e; and its components
G Ea
GKG
Subgraph of Ge. from a biomedical knowledge graph
Gal
Subgraph of Ge, extracted from biomedical corpus Subgraph of G., extracted using large language model Patient-specith knowledge graph for patient and its entities Set of communities and the k-th community at level in run m Set of all communities, a community, and the best community Entities and summary of community C₁
GiLM
Gp Va,
Ci Ck
C, Cr. Obet
VoL Sc,
Bp
Base and augmented context for patient Selected community summaries for patient p
C, ej,e())
Text embedding of entityi / relation and embedding function Hit count of node previous selections
H(w)
T,
Healtheare prediction task (theme) and its representative Ground truth and predicted labels patient p and task
Uprrr UniT
Pp.zkk Pot
best
The k-th and best reasoning chains for patient and task r A medical concepl and the set of all concepts Top X co-existing concepts for concept Mappings from original entities and relations to cluster representatives Shortest path between concepts and çj Number of levels and in community detection
G.C C
Re:
der 9.
Pij
L.M
Ze, Ze
Maximum triples per community. triples per summary Clustering thresholds for entities and relations Hyperparameters for context augmentation
ee.
a, R,O, A1, A2,
46


Reviews:
[{"review_id": "H7DDegIk69", "reviewer": "Reviewer_83dM", "summary": "The paper presents KARE, a framework designed to improve healthcare predictions by combining KG community-level retrieval with LLM reasoning. It builds a medical KG from diverse sources, organizes it into meaningful communities, and dynamically augments patient data with relevant information. Extensive tests on MIMIC-III and MIMIC-IV datasets demonstrate significant performance improvements for mortality and readmission prediction.", "questions": "1. In Section 3.1.1, three different KGs are generated for each medical concept, and the final KG is to integrate the three KGs together. How do you handle the conflicts among the three KGs?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The proposed method improves the interpretability and trustworthiness of clinical predictions by generating reasoning chains.\n2. The experiments are extensive. The framework is compared to many baselines and shows a sufficient performance gain in both tasks.\n3. The visualizations are well-structured and the writing is easy to follow.\n4. The problem setting is clearly defined.", "weaknesses": "1. The framework is overly complex, making it difficult to optimize and reproduce. Moreover, any noise introduced in the earlier steps (e.g., entity and relation extraction, LLMs for relationship suggestions and reasoning chains generation, etc.) could affect subsequent stages and degrade the final performance.\n2. Given its complexity, the framework's efficiency should be evaluated. It involves building three different KGs for each medical concept, conducting community detection 25 times, and generating multiple summaries for each community (~60k communities in total). This process likely requires significant time and resources, making it inefficient.\n3. The contribution is a bit limited compared to GraphRAG; the primary difference is mentioned in line 213 that they run GraphRAG multiple times for better diversity.\n4. Acronyms like LLM, KG, EHR, and RAG are introduced multiple times throughout the paper.\n5. It would be better to include case studies for the three KGs generated by the biomedical KG, biomedical corpus, and LLMs.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5NJ86C1In8", "reviewer": "Reviewer_tZi7", "summary": "The paper presents KARE, a framework that enhances healthcare predictions by combining knowledge graph (KG) community-level retrieval with large language model (LLM) reasoning. It addresses LLM limitations like hallucinations and inadequate medical knowledge, which can affect clinical diagnosis. KARE builds a comprehensive knowledge graph from biomedical databases, clinical literature, and LLM insights, organized through hierarchical community detection and summarization to improve retrieval precision and relevance. Key innovations include: (1) a dense medical knowledge structuring approach for accurate information retrieval; (2) a dynamic retrieval mechanism that enriches patient contexts with multi-faceted insights; and (3) a reasoning-enhanced prediction framework that produces accurate and interpretable clinical predictions.", "questions": "please refer to the weaknesses", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "1. I really enjoyed the LLM-Driven Reasoning, producing both accurate and interpretable clinical predictions that enhance trust.\n2. The model's multi-source knowledge graph ensures relevant information retrieval, addressing LLM limitations like hallucinations and sparse data.\n3. KARE's retrieval mechanism enriches patient data with multi-faceted insights, enhancing EHR representation learning", "weaknesses": "1. Methodology Clarity: Some aspects of the methodology lack clarity, such as how the different knowledge graphs (KGs) are connected in Equation (1). Specifically, the approach for creating edges between nodes in different KGs, like G^KG and G^BC, is not well explained.\n\n2.  The experimentation could be broadened to include more general tasks, such as diagnosis prediction or drug recommendation. Was there a reason these broader tasks were not considered?\n\n3.  The results lack standard deviations or confidence intervals, which would help indicate the reliability of the reported performance.\n\n4. Ablation Study Design: The ablation study could be more informative if it involved removing each feature individually, rather than adding features one at a time.\n5. The provided anonymous GitHub code link cannot be opened.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1w63WxqMCc", "reviewer": "Reviewer_pYH3", "summary": "This paper introduces KARE, a framework integrating LLM reasoning with KG retrieval to improve healthcare predictions. KARE combines structured multi-source medical knowledge with dynamic, patient-specific context augmentation to provide accurate and interpretable clinical predictions. Evaluated on MIMIC-III and MIMIC-IV datasets for mortality and readmission predictions, KARE demonstrates improved accuracy and interpretability over conventional models.", "questions": "See Weaknesses Above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 2, "strengths": "- Pros\n  - Interesting topic for enhancing healthcare predictions by combining the reasoning capabilities of LLM.\n  - Clear and well-motivated reasoning in the paper. \n  - Comprehensive ablation studies validate the contributions of each model component.\n  - Well-written and structured.", "weaknesses": "- Cons\n  - The novelty of this paper is relatively limited, appearing to be incremental compared to GraphRAG [1].\n  - In Section 3.3.1, the authors select the reasoning chain with the highest confidence as training data. However, according to conclusions from some existing studies [2,3], the reasoning chain with the highest confidence is not necessarily the most reliable.\n  - MedRetriever [4] also adopts a retrieval-augmented approach for healthcare prediction, but this paper lacks a comparative analysis with MedRetriever.\n  - The definitions and calculation methods for Sensitivity and Specificity need to be clarified more thoroughly, and metrics such as AUROC and AUPRC should be added.\n  - Although Amazon Bedrock provides strict compliance standards and privacy protection measures, relying on it to generate reasoning chains for distillation may limit the generalizability of this approach in real healthcare scenarios with high privacy protection requirements.\n\n[1] Edge D, Trinh H, Cheng N, Bradley J, Chao A, Mody A, Truitt S, Larson J. From local to global: A graph rag approach to query-focused summarization. arXiv preprint arXiv:2404.16130. 2024 Apr 24.\n\n[2] Yang, Haoyan, et al. \"Can We Trust LLMs? Mitigate Overconfidence Bias in LLMs through Knowledge Transfer.\" arXiv preprint arXiv:2405.16856 (2024).\n\n[3] Xiong, M., Hu, Z., Lu, X., LI, Y., Fu, J., He, J. and Hooi, B., Can LLMs Express Their Uncertainty? An Empirical Evaluation of Confidence Elicitation in LLMs. In The Twelfth International Conference on Learning Representations.\n\n[4] Ye M, Cui S, Wang Y, Luo J, Xiao C, Ma F. Medretriever: Target-driven interpretable health risk prediction via retrieving unstructured medical text. InProceedings of the 30th ACM International Conference on Information & Knowledge Management 2021 Oct 26 (pp. 2414-2423).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "f23XN0MyCw", "reviewer": "Reviewer_thia", "summary": "The paper introduces KARE, a novel framework designed to enhance clinical decision support by addressing the limitations of Large Language Models (LLMs) in healthcare. While LLMs show potential, they suffer from hallucinations and lack the fine-grained medical knowledge necessary for high-stakes applications like diagnosis. Traditional retrieval-augmented generation (RAG) methods often retrieve sparse or irrelevant data, undermining accuracy. KARE improves upon this by integrating a multi-source knowledge graph (KG) with LLM reasoning. The KG is built from biomedical databases, clinical literature, and LLM-generated insights, structured using hierarchical community detection for precise information retrieval. Key innovations include dense medical knowledge structuring, dynamic retrieval of multi-faceted medical insights, and reasoning-enhanced predictions. KARE outperforms existing models in MIMIC-III and MIMIC-IV datasets, improving prediction accuracy by up to 15%, while also enhancing the interpretability and trustworthiness of clinical predictions.", "questions": "1. The authors used Claude 3.5 Sonnet as an expert model to generate training samples and augment knowledge graph. However, since Claude is a general-purpose model, could it lack some medical knowledge, potentially leading to biased training samples and cumulative errors? As mentioned in your summary: \"Yet LLMs still suffer from hallucinations and lack fine-grained contextual medical knowledge, limiting their high-stakes healthcare applications such as clinical diagnosis.\" In the experiment section, there are many related LLMs in the medical domain. It would be better if the researcher could compare KARE with more related LLM-based baselines referred in [4].\n2. I didn't see any examples of training samples in the code you provided. Can you provide us with some examples?\n3. This parameter design is very challenging. In KARE, there are many hyperparameters, including but not limited to those in graph generation, summarization, model training, and model testing. Any slight change can lead to significant deviations in the model's results. Could you elaborate on how you adjust the parameters in such a large hyper-parameter space?\n\n\n- [4] https://huggingface.co/blog/leaderboard-medicalllm.", "limitations": "", "rating": 3, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The research team has achieved impressive results, exceeding established benchmarks. This progress marks a significant step forward in improving predictive model accuracy in medical data analysis.\n2. The authors' workload is immense, and the experimental details are thoroughly outlined, which I greatly appreciate.", "weaknesses": "1. By reviewing your code and the details in the article, I can see that your workload is immense, however, the contribution of this article is incremental. My understanding is that it is essentially a combination of GraphRAG and GraphCare [1]. Furthermore, many key baselines were not cited. Since the authors mentioned that this paper focuses on RAG for EHR, some essential RAG algorithms should have been introduced, such as MedRetriever [2], and commonly used GraphRAG algorithms like KGRAG [3]. \n2. In the experiment or appendix section, I did not clearly see the formulas for Sensitivity and Specificity, nor were there any corresponding references, which is quite confusing to me. Moreover, using Accuracy as a metric in cases of highly imbalanced labels is unreasonable. For instance, in the MIMIC-III Mortality Prediction task, the positive rate is 5.42%. If I predict that all patients will survive, I can still achieve an accuracy of 94.58%. Previous works, such as GraphCare [1], have adopted AUROC and AUPRC as evaluation metrics.\n3. The article is overly long and filled with detailed content, making it easy for readers to miss important points.\n\n- [1] GraphCare: Enhancing Healthcare Predictions with Personalized Knowledge Graphs. ICLR 2024\n- [2] MedRetriever: Target-driven interpretable health risk prediction via retrieving unstructured medical text. CIKM 2021\n- [3] Biomedical knowledge graph-enhanced prompt generation for large language models. Arxiv 2023", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "aHbwGLd9Tl", "reviewer": "Reviewer_oKgq", "summary": "In this manuscript, the authors propose KARE, a knowledge graph-enhanced large language model designed for predicting patient mortality and readmission, with an added aim to mitigate LLM hallucinations. The paper introduces a novel approach that effectively combines knowledge graphs with LLMs through a clustering method. Additionally, the authors present a reasoning-chain mechanism to enhance the LLM's inference capabilities and provide interpretable prediction results. Experimental results demonstrate that KARE achieves a marginal improvement in mortality and readmission predictions on the MIMIC-III and MIMIC-IV datasets. However, several major and minor issues need to be addressed. If author could well address my concern, I would like to update my ratings,", "questions": "* In Section 3.2, how does the method handle cases where the same patient has more than two visits with different labels? Is there a specific design to distinguish similar symptoms in different patients and multiple visits from the same patient?\n* In Section 3.2, how is the effectiveness of the relevance score (i.e., Formula (3)) validated? Would an ablation study on the components of Relevance(C_k) help to clarify this?\n* How does the method utilize patient longitudinal visit information?\n* For the baseline ML methods, it’s mentioned that most are implemented using PyHealth. Could you clarify the backbone model for each ML method? Also, is there a fair configuration in place for implementing language model-based encoders, such as ClinicalBERT, across these methods? \n* The experimental results for mortality prediction show that the baseline ML methods, such as ConCare and TCN, perform closely to KARE, with some evaluation metrics even exceeding KARE’s. How should these results be interpreted?\n* In experiment, sensitivity and specificity are indeed important metrics, but F1 score and accuracy are also widely used for assessing diagnostic accuracy. Interestingly, in ablation study, excluding similar patients often yields better or comparable results than including them, especially for mortality prediction on the MIMIC-III dataset. Do these results support the methodological choices made in Section 3.2?\n* As you state that your base model is fine-tuned on Mistral-7B-Instruct-v0.3, I wonder about the performance of fine-tuning an LLM (Mistral or Llama) with RAG. Can your model still show substantial improvements compared to other LLM-based models beyond just zero-shot or few-shot settings? I am concerned that most of the current improvements may be attributed primarily to supervised fine-tuning.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper is well-written and structured.\n* The authors propose a novel knowledge-graph learning framework combined with LLMs.\n* The integration of knowledge augmentation and knowledge graphs is both novel and interesting.\n* The experimental evaluation is thorough, conducted on both the MIMIC-III and MIMIC-IV datasets.", "weaknesses": "* The excessive use of symbols makes certain sections difficult to follow.\n* The paper lacks a detailed description of longitudinal EHR processing with LLMs.\n* The improvement in mortality prediction performance achieved by KARE is marginal.\n* There is a lack of sensitivity analysis for hyperparameters, such as the number of top co-occurring concepts, top documents, and the criteria for selecting optimal values. Additional hyperparameters, such as maximum sequence length and maximum path count, should also be discussed, or at least identified explicitly for clarity.\n* The paper would benefit from a more in-depth interpretability analysis to clarify the relationship between the generated reasoning, labels, and the knowledge graph.\n* The paper lacks an evaluation of LLM hallucinations, despite the stated aim to address hallucination issues in clinical decision-making.\n* An analysis of model parameters and time consumption is missing, which could provide valuable insights into the model’s computational efficiency and practical applicability\n* There is a lack of references to related work.\n\n  [1] . Kang, M., Lee, S., Baek, J., Kawaguchi, K., & Hwang, S. J. (2024). Knowledge-augmented reasoning distillation for small language models in knowledge-intensive tasks. Advances in Neural Information Processing Systems, 36. \n\n  [2]. Niu, S., Ma, J., Bai, L., Wang, Z., Guo, L., & Yang, X. (2024). EHR-KnowGen: Knowledge-enhanced multimodal learning for disease diagnosis generation. Information Fusion, 102, 102069.\n\n  [3]. Kwon, T., Ong, K. T. I., Kang, D., Moon, S., Lee, J. R., Hwang, D., ... & Yeo, J. (2024, March). Large language models are clinical reasoners: Reasoning-aware diagnosis framework with prompt-generated rationales. In Proceedings of the AAAI Conference on Artificial Intelligence (Vol. 38, No. 16, pp. 18417-18425).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)