Reasoning-Enhanced Healthcare Predictions with Knowledge Graph Community Retrieval

INTRODUCTION

Large language models (LLMs) (Touvron et al., 2023a;b;OpenAI et al., 2024;Team et al., 2024) has revolutionized natural language processing, offering unprecedented capabilities in understanding and generating human-like text. In the healthcare domain, LLMs hold the potential to transform clinical decision-making 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 et al., 2016;2017;Shickel et al., 2018;Choi et al., 2018;Ma et al., 2020a;Gao et al., 2020;Zhang et al., 2021;Wu et al., 2023;Jiang et al., 2024a;Zhu et 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-augmented generation (RAG) techniques (Lewis et al., 2021), which aim to mitigate hallucinations by retrieving external knowledge, often fall short in healthcare applications. They tend to retrieve information that, while semantically similar in latent space, fails to provide meaningful clinical insights, leading to suboptimal outcomes for precise healthcare predictions (Shi et al., 2024;Magesh et al., 2024;Li et al., 2024). For instance, when dealing with the diagnosis of heart failure, a 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 or specific biomarkers like NT-proBNP levels.

Knowledge graphs (KGs) offer a promising solution by providing structured representations of medical knowledge, capturing complex relationships between clinical entities (Liu et al., 2019;Yasunaga et al., 2022;Zhang et al., 2022). Integrating KGs with LLMs can enhance the models' reasoning capabilities and provide domain-specific 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 as 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 naïve RAG in the query-focused summarization task. However, the application of graph retrieval for LLM-based healthcare prediction remains largely unexplored.

In this paper, we introduce KARE (Knowledge Aware Reasoning-Enhanced HealthCare Prediction), a new framework that combines KG community-level retrieval (e.g., retrieving relevant subgraphs) 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 a novel method to construct and index multi-source medical concept KGs by integrating concept-specific knowledge 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.

MEDICAL CONCEPT-SPECIFIC KNOWLEDGE GRAPH EXTRACTION

For each medical concept c i in the EHR coding system, we extract a concept-specific knowledge graph G ci = (V ci , E ci ) globally tailored to the EHR datasets from three sources:

(a) Biomedical KG (e.g., UMLS (Bodenreider, 2004)): For each medical concept c i in EHR data, we extract a subgraph G KG ci by first iterating through the patient EHR dataset to collect the top X co-existing concepts appeared in each patient's data, forming the set of related concepts R ci . We then find the shortest path p ij in the KG for each pair (c i , c j ) in R ci , with a specified maximum path length. G KG ci = (V KG ci , E KG ci ) is constructed by combining all these shortest paths, where V KG ci and E KG ci are the union of nodes and edges in all p ij , respectively. See more details in Appendix B.1. (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 n documents from the corpus based on these medical concepts. For each retrieved document, we perform entity extraction and relation extraction to extract KG triples. The extracted triples are then added to the KG of the medical concepts mentioned in the document. By doing so, G BC ci is built for each concept c i . 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-specific KG G ci is the union of the subgraphs from each source:

Finally, we integrate all concept-specific KGs for the medical concepts in our EHR coding system. The resulting knowledge graph G ′ = (V ′ , R ′ , E ′ ) is defined as G ′ = ci∈C G ci where C is the set of all medical concepts in the specified EHR coding system.

Note: Different from the KG construction method introduced by GraphCare (Jiang et 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.

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 community structures, and generate multi-theme summaries for each community, providing a more comprehensive understanding of the knowledge contained withtin the KG.

To keep computation manageable, we limit the maximum number of triples per community (Z c ) and the maximum number of triples per initial summary (Z s ).

For each community, we generate two types of summaries using an 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-specific summary: A summary that highlights how the knowledge in the community relates to a 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 ≤ Z s ), we directly summarize all triples.

• For large communities (Z s < size ≤ Z c ), we shuffle and split the triples into subsets, summarize each subset, and then iteratively aggregate the summaries until we get a single comprehensive summary (prompt shown in Figure 13).

• For extremely large communities (size > Z c ), we do 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 themespecific summaries. Running the Leiden algorithm multiple times with different randomness parameter gives us diverse communities, allowing entities to contribute to multiple summaries. This rich, multi-level representation of the KG is the foundation for later steps.

Objective of

Step 2: This step constructs patient's EHR context with the highly relevant and finegrained medical knowledge attached.

Base Context Construction. For a patient p, we construct a base context B p with their EHR data with a standardized template. This context focuses on (1) task description, (2) the patient's conditions, 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 p and the other has a different label (Cui et al., 2024). Figure 11 shows an 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, we first construct a patient-specific knowledge graph G p by aggregating the concept-specific graphs G ci (defined in Eq. 1) for all medical concepts c i in the patient's EHR, using the mappings ϕ e and ϕ r from §3.1.2:

From G p , we derive two sets of nodes: V direct p , representing medical concepts that directly appear in the patient's EHR, and V indirect p , containing the remaining nodes in G p .

We then introduce a combined relevance score for each community C k to select the most relevant summaries for context augmentation:

) calculate the normalized counts of direct and indirect node hits by comparing the nodes in community C k with the corresponding sets of direct and indirect nodes. The parameter α ∈ [0, 1) weights the importance of indirect hits relative to direct hits. The decay function Decay(C k , V direct p ) reduces the contribution of previously hit nodes in community C k by a factor β H(v) , where β ∈ (0, 1] is a decay constant and H(v) is the hit count of node v in previous selections, considering only the direct nodes V direct p . Additional factors are defined as:

Here, e(•) denotes a text embedding function, cos(•, •) is the cosine similarity between embeddings, visit(v) returns the visit index of node v, and λ 1 , λ 2 , λ 3 ∈ [0, 1] control the weights of the metrics. The set T τ contains representative terms for the theme τ (e.g., {end-stage, life-threatening, ...} for mortality prediction), same as 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 patient'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 maintains task-oriented selection. In addition, we propose a Dynamic Graph Retrieval and Augmentation (DGRA) method to iteratively select the most relevant summaries to augment the patient's context.

At each iteration, it performs as: Algorithm 1 Dynamic Graph Retrieval and Augmentation Input:

(1) Compute the relevance scores for all candidate communities C k ∈ C using Eq. 3.

(2) Identify the community C best with the highest relevance score and add its summary S Cbest to the set of selected summaries S p .

(3) Increment the hit count H(v) for each node v in V Cbest , which will impact the decay in future relevance calculations.

(4) Remove C best from the candidate communities C, ensuring it is not reconsidered in subsequent iterations.

The process continues until N summaries have been selected. The final augmented patient context A p is obtained by concatenating the base context B p with the selected summaries.

By dynamically updating the node hits and recalculating relevance scores at each iteration, we prioritize 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.

MULTITASK-BASED FINE-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 of the task description D τ and the augmented patient context A p , 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 ρ best p,τ obtained from the previous step. When instructed to predict the label (with the prefix [Label Prediction]), it aligns its output with the ground truth label y * p,τ . 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 a new patient p new and task τ , we provide A pnew , τ , and the appropriate instruction to the fine-tuned model. Based on the instruction, the model can either generate the reasoning chain ρ pnew,τ or predict the label y pnew,τ . This flexible approach allows us to obtain detailed reasoning when necessary or perform efficient label prediction, leveraging the multitask training to effectively handle both tasks during inference.

EXPERIMENTAL SETTING

Tasks. In this work, we focus on the following EHR-based prediction tasks:

• Mortality Prediction. This task estimates mortality outcome for next visit, defined as f :

(x 1 , x 2 , . . . , x t-1 ) → y[x t ], where y[x t ] ∈ {0, 1} is patient's survival status during visit x t .

• Readmission Prediction. This task predicts if patient will be readmitted within σ days, defined as

where y ∈ {0, 1}, φ(x t ) is timestamp of visit x t , and y[φ(x t ) -φ(x t-1 )] = 1 if φ(x t ) -φ(x t-1 ) ≤ σ, else 0. σ is set to 15 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 PyHealth (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 1 (mortality), excluding 54 patients with more than 10 visits. We then randomly (seed=42) sample unique patients with a label of 0, each having no more than 10 visits, until reaching a sample size of 10,000. For MIMIC-IV readmission prediction, we randomly (seed=42) select 5,000 unique patients with a label of 1 (will be readmitted) and 5,000 with a label of 0. (Zhang et al., 2021) 93.7 49.9 1.9 98.9 61.3 59.5 74.9 44.8 StageNet (Gao et al., 2020) 90.5 50.5 5.6 95.4 60.5 60.0 65.1 54.9 KerPrint (Yang et al., 2023b) 92.4 52.2 9.8 94.7 63.5 62.1 68.0 56.1 LM+ML GraphCare (Jiang et al., 2024a) 94.9 58.3 17.2 97.1 65.4 64.1 70.3 57.8 RAM-EHR (Xu et al., 2024) 94.4 59.6 14.8 98.9 64.8 63.5 74.7 52.4 EMERGE (Zhu et al., 2024a) 94.1 57.7 13.2 98.4 63.7 62.0 68.0 55.9 LLM Zero-shot (LLM: Claude 3.5 Sonnet) w/ EHR context only 89.5 50.4 6.4 94.4 54.3 35.4 98.9 0.2 w/ Classic RAG [a] 89.9 51.2 10.2 92.8 53.2 34.6 91.2 1.4 w/ KARE-augmented context [b] 92.3 54.6 14.2 94.6 56.3 43.8 93.9 10.6 Few-Shot (LLM: Claude 3.5 Sonnet) w/ exemplar only (N=2) [c] 88.7 49.5 5.6 93.4 52.7 42.2 87.0 11.1 w/ exemplar only (N=4) 88.0 49.2 5.6 92.7 53.6 44.7 84.0 15.7 w/ EHR-CoAgent [d] (Cui et al., 2024) 87.4 51.7 13.0 91.8 55.2 46.1 78.2 20.1 w/ KARE-augmented context 91.5 53.5 13.7 94.0 57.1 49.3 75.5 27.2 Fine-tuned (LLM: Mistral-7B-Instruct-v0.3) Backbone 90.4 53.0 11.4 94.3 57.6 57.6 50.5 66.3 w/ Classic RAG 90.1 51.4 12.5 91.6 60.2 59.9 56.1 64.5 KARE (ours) 95.3 64.6 24.7 98.3 73.9 73.7 76.7 70.7 MIMIC-IV Mortality Prediction (pos = 19.16%) Readmission Prediction (pos = 46.50%) Type Models Accuracy Macro F1 * Sensitivity * Specificity Accuracy Macro F1 Sensitivity Specificity ML GRU (Chung et al., 2014) 88.7 76.4 42.9 99.6 62.4 62.2 68.3 56.2 Transformer (Vaswani et al., 2017) 83.7 71.3 47.1 92.3 61.3 61.3 63.0 59.5 RETAIN (Choi et al., 2016) 84.8 73.8 52.4 92.4 62.8 62.6 68.7 56.6 GRAM (Choi et al., 2017) 86.4 74.4 50.6 93.9 62.5 62.5 67.4 57.8 Deepr (Nguyen et al., 2016) 89.2 78.9 50.8 98.2 59.2 59.2 57.0 61.5 TCN (Bai et al., 2018) 89.9 79.2 47.6 99.9 63.6 63.5 72.2 56.1 ConCare (Ma et al., 2020b) 89.8 78.9 47.1 99.9 59.8 59.8 63.5 56.6 AdaCare (Ma et al., 2020a) 88.7 78.2 50.3 97.8 62.9 62.9 58.4 67.7 GRASP (Zhang et al., 2021) 89.9 79.1 47.6 99.8 59.7 59.6 53.1 66.7 StageNet (Gao et al., 2020) 88.1 77.8 51.9 96.7 62.8 62.7 62.6 62.9 KerPrint (Yang et al., 2023b) 88.7 79.8 53.1 98.0 63.5 63.3 67.0 60.1 LM+ML GraphCare (Jiang et al., 2024a) 91.5 80.3 57.8 96.6 65.7 65.5 66.2 65.0 RAM-EHR (Xu et al., 2024) 90.5 78.4 52.6 97.0 65.5 65.5 64.0 67.0 EMERGE (Zhu et al., 2024a) 90.7 78.3 53.4 96.6 63.3 63.2 61.5 64.9 LLM Zero-shot (LLM: Claude 3.5 Sonnet) w/ EHR context only 80.5 47.0 2.7 98.7 49.4 45.7 81.8 21.5 w/ Classic RAG [a] 81.0 49.9 8.1 94.6 49.0 44.2 83.2 18.8 w/ KARE-augmented context [b] 83.2 54.3 12.7 96.3 52.3 49.7 80.6 27.7 Few-Shot (LLM: Claude 3.5 Sonnet) w/ exemplar only (N=2) [c] 80.8 46.7 2.1 99.5 49.3 44.7 84.0 19.1 w/ exemplar only (N=4) 81.6 49.9 5.3 99.8 49.0 44.1 84.3 18.2 w/ EHR-CoAgent [d] (Cui et al., 2024) 81.0 55.5 13.8 97.0 51.2 46.3 78.4 24.0 w/ KARE-augmented context 84.5 57.4 15.8 97.6 54.1 51.9 75.2 34.1 Fine-tuned (LLM: Mistral-7B-Instruct-v0.3) Backbone 92.2 83.1 65.0 96.2 56.1 46.7 23.1 76.2 w/ Classic RAG 92.5 83.8 63.2 97.6 58.8 52.1 46.7 57.5 KARE (ours) 94.1 90.4 73.2 99.8 73.9 73.8 85.6 63.7

[a] We retrieve up to ten documents from 30 M PubMed abstracts that are most similar to the base context. Dense retrieval is applied with Nomic (Nussbaum et al., 2024) (dim=768).

[b] Context is augmented as described in Section 3.2 using our constructed KG communities. Up to ten community summaries are appended as supplementary information.

[c] N=2 means that we retrieve two most similar patients (one from positive and the other from negative samples), similar to the strategy introduced in (Cui et al., 2024).

[d] Since the code for EHR-CoAgent was not open-sourced at the time of our paper submission, we implemented it ourselves, as detailed in Appendix C.

datasets;

(3) Sensitivity, quantifying the model's ability to identify high-risk patients; and (4) Specificity, 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., 2016), GRAM (Choi et al., 2017), Deepr (Nguyen et al., 2016), TCN (Bai et al., 2018), StageNet (Gao et 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) LLM-based 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 et al., 2024). We showcase the details of baseline implementation in Appendix C.

Implementation Details. We utilize Scikit-learn (Pedregosa et al., 2018) for agglomerative clustering and Graspy (Chung et al., 2019) for the hierarchical Leiden algorithm. Semantic similarity cal-

A ETHICS, BROADER IMPACTS, AND LIMITATIONS

Ethics: 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 interactions between the language models and the EHR data through Amazon Bedrock 34 , which provides rigorous compliance standards and privacy protection measures. This allowed us to fully leverage the capabilities of the LLMs while maintaining strict confidentiality of the patient data.

Broader Impacts: Our framework, KARE, 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.

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 a 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 outputs, 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 performance of our models. Third, our framework relies on the outputs of large language models, which are known to have biases and can generate hallucinations. While we have taken steps to mitigate these issues, such as collaborating with medical experts to validate the extracted knowledge, there remains a risk of the models producing incorrect or biased predictions in some cases. Fourth, our work is based on English biomedical literature and EHR data, and we did not evaluate the performance of KARE on data in other languages. The applicability and effectiveness of our approach for non-English clinical settings requires further investigation. Ongoing research is needed to develop 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 language 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 real-world healthcare settings. We encourage future work to focus on enhancing the generalizability, interpretability, and robustness of these approaches, as well as engaging with healthcare stakeholders to develop guidelines for the ethical and effective use of AI in clinical decision support.

Entity/Relation Extraction

< l a t e x i t s h a 1 _ b a s e 6 4 = " g X 2 a F 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 a concept set. This results in a list of concept sets, where each set represents the concepts associated with a specific patient visit. 2. Filtering Similar Concept Sets: To reduce redundancy and computational overhead, we filter out highly similar concept sets based on their concept multi-hot vector representation. We use the CountVectorizer from scikit-learn (Pedregosa et al., 2018) to create a vocabulary of unique concepts and transform each concept set into a multi-hot vector. We then compare the symmetric difference between pairs of concept sets and filter out sets that differ by fewer than a specified similarity threshold (5 in our case). This step helps to eliminate nearly duplicate concept sets while retaining a diverse range of concepts for knowledge graph construction. This process results in 26,134 concept sets in our experiment. 3. Dense Retrieval of Relevant PubMed Abstracts: For each filtered concept set, we retrieve the top n ( n = 3 in 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 (seed=42) select 1/10 abstracts as the subset corpus to proceed. We employ the nomic-ai/nomic-embed-text-v1.5 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 is optimized for efficiency by processing the embeddings in chunks and utilizing GPU acceleration when available. 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 a combination of long-term control medications, such as inhaled corticosteroids, and quick-relief medications like short-acting beta-agonists. Recent research has focused on 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 a, chronic respiratory condition], [asthma, characterized by, inflammation of airways], [inflammation, causes, narrowing of airways], [narrowing of airways, leads to, breathing difficulties], [wheezing, is a symptom of, asthma], [coughing, is a 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: {text}

You can introduce intermediate relationships with other entities based on your knowledge.

There's no need to keep all the relationships connected. To extract concept-specific knowledge graphs from large language models (LLMs), we follow a process similar to the initial steps of the "KG extraction from corpus" pipeline (Appendix B.2):

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 a concept set. This results in a list of concept sets, where each set represents the concepts associated with a specific patient visit. 2. Filtering Similar Concept Sets: To reduce redundancy and computational overhead, we filter out highly similar concept sets based on their concept multi-hot vector representation, as described in the previous subsection. This step helps to eliminate nearly duplicate concept sets while retaining a diverse range of concepts for knowledge graph construction. 3. Prompting LLMs for Relationship Extraction: For each filtered concept set, we prompt a large language model (Claude 3.5 Sonnet in our case) to identify relationships among the medical concepts that can be potentially helpful for clinical predictions, such as mortality prediction and readmission prediction.

The LLM is encouraged to introduce intermediate relationships with other entities based on its knowledge, and there is no requirement to keep all the relationships connected. The prompt instructs the LLM to use the original names of the provided concepts in the output, which should follow the format [[ENTITY1, RELATIONSHIP 1, ENTITY2], [ENTITY2, RELATIONSHIP 2, ENTITY3], ...]. The prompt we used for the relationship 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. for each concept c i . For each concept c i in the concept set, we store the connected 3-hop subgraph sourced from c i to its corresponding concept-specific knowledge graph G LLM ci . This step ensures that only the triples directly or indirectly connected to the concept c i 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 or 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.

B.4 CASE STUDIES OF KG CONSTRUCTION

"spinal fusion": [ "opioid analgesics", "other beta-lactam antibacterials in atc", "antithrombotic agents", "potassium supplements", "drugs for constipation", "other analgesics and antipyretics in atc", "i.v. solution additives", "drugs for peptic ulcer and gastro-oesophageal reflux disease (gord)", "iv solutions used in parenteral administration of fluids, electrolytes and nutrients", "spondylosis; intervertebral disc disorders; other back problems", "antiemetics and antinauseants", "anxiolytics", "other mineral supplements in atc", "antiinfectives and antiseptics, excl. combinations with corticosteroids", "antiepileptics", "other nutrients in atc", "antivaricose therapy drugs", "beta blocking agents", "lipid modifying agents, plain", "decompression peripheral nerve" ]

(Top-20 Co-existing Concepts for Each Concept) ⋮ ⋮ ================================================== Target Concept: opioid analgesics - ------------------------------- ------------------------------------------------------ ------------------------------------ ----------------------------------- [ "pleurisy; pneumothorax; pulmonary collapse", "coronary atherosclerosis and other heart disease", "potassium supplements", "other fractures", "anxiolytics", "opioid analgesics", "other gastrointestinal disorders", "epilepsy; convulsions", "i.v. solution additives", "antibiotics for topical use", "e codes: fall", "adrenergics, inhalants", "antiepileptics", "other diagnostic agents in atc", "diabetes mellitus without complication", "drugs for peptic ulcer and gastrooesophageal reflux disease (gord)", "other mineral supplements in atc", "spinal fusion", "beta blocking agents", "other analgesics and antipyretics in atc", "laminectomy; excision intervertebral disc", "intracranial injury", "other beta-lactam antibacterials in atc" ] (Concept set in a patient visit)

Doc. Retrieval

-----Rank 1, Similarity: 0.7296074032783508 PMID: 310377 Title: Management of acute pain in trauma. Authors: C N Shealy Abstract: In acute pain, TENS, ice packs, and a calm, reassuring attitude and voice are useful in reducing pain. Narcotic requirements can be reduced, and chronic pain may be prevented. Complications such as paralytic ileus and atelactasis can also be reduced. These techniques can be used in the emergency ward, the recovery room, and the doctor's office. Publication Year: 1979 -----Rank 2, Similarity: 0.7247765064239502 PMID: 235156 Title: Interaction between local anesthetics and analeptic drugs. Authors: R G Thompson, J A Aldrete Abstract: Although clinicall undesirable, the fortuitous pharmacologic interactions between local anesthetic agents and analeptic drugs may be protective when large doses of both agents are used. Mice pretreated with procaine, lidocaine, and tetracaine had a lower incidence of seizures when convulsive doses of either nikethamide or doxapram hydrochloride were given intraperitoneally. Mortality was also decreased in the groups given nikethamide and was zero in the animals treated with doxapram. All animals treated with pentylenetetrazol convulsed and only lidocaine (also used as an anticonvulsant) was able to reduce mortality in this group. Of the local anesthetic agents, tetracaine afforded the least protection from death, whereas lidocaine seemed to be most effective.

Publication Year: 1975 -----⋮ Triple Extraction [opioid analgesics, can be reduced by, TENS] [opioid analgesics, can be reduced by, ice packs] Triple Extraction [epilepsy; convulsions, can be caused by, large doses of analeptic drugs] [local anesthetic agents, may reduce, epilepsy; convulsions] [lidocaine, used as, antiepileptics] [lidocaine, most effective in reducing, mortality] [pentylenetetrazol, causes, epilepsy; convulsions] [lidocaine, reduces mortality from, epilepsy; convulsions] [ "pleurisy; pneumothorax; pulmonary collapse", "coronary atherosclerosis and other heart disease", "potassium supplements", "other fractures", "anxiolytics", "opioid analgesics", "other gastrointestinal disorders", "epilepsy; convulsions", "i.v. solution additives", "antibiotics for topical use", "e codes: fall", "adrenergics, inhalants", "antiepileptics", "other diagnostic agents in atc", "diabetes mellitus without complication", "drugs for peptic ulcer and gastrooesophageal reflux disease (gord)", "other mineral supplements in atc", "spinal fusion", "beta blocking agents", "other analgesics and antipyretics in atc", "laminectomy; excision intervertebral disc", "intracranial injury", "other beta-lactam antibacterials in atc" ] [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 on conditions, procedures, and medications. Labels: 1 = mortality, 0 = survival). We introduce a Regulator to import prior knowledge of the dataset to avoid the LLM to over-focus on improving true positive or true negative. For example, we set "**Must to Notice:** Only the patients with extremely very high risk of mortality should be predicted as 1. " 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 5. This is because that the instruction-updating approach used by EHR-CoAgent tends to excessively penalize false positives to produce instructions that boost true positives, especially in cases of imbalanced data like mortality prediction in the two datasets. However, this can compromise true negatives and accuracy in such scenarios. Refined Guidelines are the instructions refined by the critic agent, while Initial Guidelines are the seed instructions, which is as same as the task descriptions we used for other methods.

The critic agent, on the other hand, plays a different role in the EHR-CoAgent framework. It observes the predictor agent's outputs alongside the ground truth labels and identifies error patterns and discrepancies in the predictor's reasoning process. Based on this analysis, the critic agent refine the instructions to improve the reasoning process of the predictor. The prompt used for the critic agent is as follows:

You are an 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 are all incorrect. You should always use the ground truth as the final basis to discover many unreasonable aspects in 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. 5. The instructions should be generalizable to multiple samples, rather than specific to individual samples. 6. Conduct detailed analysis and write criteria based on the input samples, rather than writing some criteria without foundation. For UMLS, we utilize the "Full Release" version under "2024AA Full UMLS Release Files"foot_9 . For dense retrieval from PubMed abstracts, we utilize the local embedding model Nomic (dimension = 768) (Nussbaum et al., 2024). We use Amazon Bedrockfoot_10 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: we (1) first use the textembedding-3-large model (dimension = 1024) from Azure OpenAI to retrieve the text embeddings of entities and relations; and (2) use Scikit-learn (Pedregosa et al., 2018) to perform agglomerative clustering based on those embeddings. The optimal cosine distance thresholds θ e and θ r are both found to be 0.14, resulting in 513,867 triples in total after clustering.

Step 1.3: Hierarchical KG Community Detection and Indexing

We employ Graspy (Chung et al., 2019) to implement the hierarchical Leiden algorithm, setting the maximum size for each top-level community (max cluster size) to 5.

To enhance community diversity, the algorithm is run 25 times with different randomness at 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 s ), and 57,247 communities with the size smller than 150 (Z c ).

Using Claude 3.5 Sonnet as the LLM, we generate 147,264 community summaries (including both general and theme-specific summaries) with the prompts shown in Figure 12 and 13.

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, all LLM interactions are conducted via the Amazon Bedrock platformfoot_11 , a cloud infrastructure that allows us to fully harness LLM capabilities while maintaining strict privacy measures. The maximum output length for Sonnet is 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. We use Mistral-7B-Instruct-v0.3 (Jiang et al., 2023) as our local LLM, full-parameter fine-tuned using DeepSpeed (Rasley et al., 2020) Hardware Information: The experiments were conducted on a system with an AMD EPYC 7513 32-Core Processor and 1.0 TB of RAM. The setup includes eight NVIDIA A100 80GB PCIe GPUs, each with 81920 MiB of memory, providing a total of 640 GB GPU memory. The system's root partition has 32 GB of storage.

Training of each model runs on eight NVIDIA A100 GPUs and typically completes within five hours. After the training process, we select the best performing model checkpoint based on validation loss to perform the prediction.

F CASE STUDY F.1 REASONING-ENHANCED PREDICTION BY OUR FINE-TUNED MODEL # Reasoning# 1. **Patient Overview**:

The patient has a history of polyhydramnios, other complications of pregnancy, thyroid disorders, and ob-related trauma to the perineum and vulva. The patient underwent repair of current obstetric laceration and other procedures to assist delivery. Postpartum medications include throat preparations, antiinflammatory and antirheumatic products, analgesics, antipyretics, agents for hemorrhoids and anal fissures, antipruritics, drugs for constipation, bacterial and viral vaccines, antacids, calcium supplements, and thyroid preparations.

# Task #

Readmission Prediction Task: Objective: Predict if the patient will be readmitted to the hospital within 15 days of discharge. Labels: 1 = readmission within 15 days, 0 = no readmission within 15 days

Note: Analyze the information comprehensively to determine the likelihood of readmission. The goal is to accurately distinguish between patients who are likely to be readmitted and those who are not.

======================================== # Patient EHR Context # Patient ID: xxxxxxx Visit 0: Conditions: 1. polyhydramnios and other problems of amniotic cavity 2. other complications of pregnancy 3. thyroid disorders 4. ob-related trauma to perineum and vulva 5. other pregnancy and delivery including normal Procedures: 1. repair of current obstetric laceration 2. other procedures to assist delivery Medications: 1. throat preparations 2. antiinflammatory and antirheumatic products, non-steroids 3. other analgesics and antipyretics in atc 4. agents for treatment of hemorrhoids and anal fissures for topical use 5. antipruritics, incl. antihistamines, anesthetics, etc. 6. drugs for constipation 7. bacterial and viral vaccines, combined 8. antacids 9. calcium supplements 10. thyroid preparations ======================================== # Similar Patients # Patient ID: yyyyyyy Visit 0: Conditions: 1. polyhydramnios and other problems of amniotic cavity 2. ob-related trauma to perineum and vulva 3. other complications of pregnancy 4. thyroid disorders 5. other pregnancy and delivery including normal Procedures: 1. repair of current obstetric laceration 2. other procedures to assist delivery 3. artificial rupture of membranes to assist delivery Medications: 1. thyroid preparations 2. agents for treatment of hemorrhoids and anal fissures for topical use 3. antipruritics, incl. antihistamines, anesthetics, etc. 4. antacids 5. calcium supplements 6. other analgesics and antipyretics in atc 7. drugs for constipation 8. throat preparations 9. antiinflammatory and antirheumatic products, non-steroids 10. cough suppressants, excl. combinations with expectorants 11. opioid analgesics 12. expectorants, excl. combinations with cough suppressants Label: 0 Patient ID: zzzzzzz Visit 0: Conditions: 1. polyhydramnios and other problems of amniotic cavity 2. ob-related trauma to perineum and vulva 3. other complications of pregnancy 4. thyroid disorders 5. other complications of birth; puerperium affecting management of mother 6. other pregnancy and delivery including normal Procedures: 1. repair of current obstetric laceration 2. other procedures to assist delivery Medications: 1. thyroid preparations 2. cough suppressants, excl. combinations with expectorants 3. drugs for constipation 4. bacterial and viral vaccines, combined 5. opioid analgesics 6. other analgesics and antipyretics in atc 7. agents for treatment of hemorrhoids and anal fissures for topical use 8. antipruritics, incl. antihistamines, anesthetics, etc. 9. antacids 10. calcium supplements 11. throat preparations 12. antiinflammatory and antirheumatic products, non-steroids 13. urologicals Label: 1 ======================================== # Retrieved Medical Knowledge # [ 1. "Polyhydramnios, a condition of excess amniotic fluid, can increase the risk of preterm labor and premature rupture of membranes. This may lead to complications after delivery that require close monitoring and potential readmission." 2. "Thyroid disorders during pregnancy can impact both maternal and fetal health. If not properly managed, thyroid conditions may result in adverse outcomes and increase the likelihood of postpartum complications necessitating readmission." 3. "The combination of obstetric perineal trauma repair and the use of topical hemorrhoid treatments suggests uncomplicated healing, reducing the likelihood of readmission." 4. "The use of antiinflammatory and analgesic medications postpartum is typical for managing pain and discomfort. However, careful monitoring is necessary to ensure proper pain control and avoid adverse effects that could prolong recovery or require readmission." 5. "Calcium supplements are often prescribed during pregnancy and postpartum to support maternal bone health. Adequate calcium intake and monitoring can help prevent complications and reduce the risk of readmission related to calcium deficiency or imbalance." 6. " The use of bacterial and viral vaccines combined with thyroid preparations indicates proactive management of the patient's thyroid condition and immunization status, supporting a lower risk of readmission. 7. "Hemorrhoids and anal fissures are common during pregnancy and after vaginal delivery. The use of topical agents for treatment suggests active management of these conditions, which may help prevent potential complications and reduce the risk of readmission." 8. "Individuals taking thyroid preparations are likely managing a pre-existing or pregnancy-related thyroid condition. Proper thyroid hormone regulation is crucial in the postpartum period. Imbalances in thyroid hormones may lead to complications requiring hospital readmission if not closely monitored and adjusted" 9. "Constipation is a frequent issue postpartum, and the use of drugs for constipation indicates active management. Adequate bowel function and prevention of severe constipation can help avoid complications and reduce the likelihood of readmission." 10. "Extensive lacerations or complications from assisted delivery may require additional care and monitoring, potentially leading to readmission." ] INPUT (Reasoning Task) OUTPUT Figure 10: Case Study of the Fine-tuned KARE Model. F.2 ZERO-SHOT & FEW-SHOT COMPARISON

Table 7: Case Study of zero-shot and few-shot EHR prediction with LLM (Sonnet-3.5). Ground Truth:

The patient will die in the next visit (Prediction = 1). For the ethic concern, the patients involved are all from synthetic MIMIC-III by HALO (Theodorou et al., 2023).

Case Description

Input Prompt with EHR Given the following task description and patient context, please make a prediction with reasoning based on the patient's context. Task: Mortality Prediction Task Objective: Predict the mortality outcome for a patient's subsequent hospital visit based solely on conditions, procedures, and medications.

Labels: 1 = mortality, 0 = survival Key Considerations: Conditions:

• Severity of diagnosed conditions (e.g., advanced cancer, severe heart failure, sepsis)

• Presence of multiple comorbidities

• Acute vs. chronic nature of conditions Procedures:

• Invasiveness and complexity of recent procedures

• Emergency vs. elective procedures • Frequency of life-sustaining procedures (e.g., dialysis, mechanical ventilation) Medications: • Use of high-risk medications (e.g., chemotherapy drugs, immunosuppressants) • 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 a high risk of 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 Visit 0: Conditions: -Deficiency and other anemia -Essential hypertension -Complication of device; implant or graft -Congestive heart failure; nonhypertensive -Cancer of prostate -Anxiety disorders -Thyroid disorders -Disorders of lipid metabolism -Conduction disorders -Mycoses -Other diseases of kidney and ureters -Cancer of esophagus Procedures: -Diagnostic cardiac catheterization; coronary arteriography -Other or procedures on vessels other than head and neck -Colorectal resection Drugs: -Other drugs for obstructive airway diseases, inhalants in ATC -Lipid modifying agents, plain -Antithrombotic agents -Angiotensin II receptor blockers (ARBs), plain Visit 1: Conditions: -Congestive heart failure; nonhypertensive -Cardiac dysrhythmias -Shock -Alcohol-related disorders -Diabetes mellitus with complications -E codes: adverse effects of medical care -Disorders of teeth and jaw Procedures: -Other vascular catheterization; not heart -Blood transfusion Drugs: -Other drugs for obstructive airway diseases, inhalants in ATC -Other analgesics and antipyretics in ATC -Drugs for peptic ulcer and gastro-oesophageal reflux disease (GORD) -Beta blocking agents -Potassium supplements Continued on next page

Table 7 -continued from previous page Case Description • Rank 3, Similarity: 0.7102002501487732 PMID: 42254 Title: [Non arrhythmogenic sudden death as complication of coronary heart disease]. Authors: H Zilcher, D Glogar Abstract: In a cohort of 417 patients admitted consecutively to the Coronary Care Unit for acute myocardial ischemia (unstable angina pectoris in 121, acute myocardial infarction in 296 patients) 21 cases of non arrhythmogenic 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 in 4 patients with acute myocardial infarction, whereas all patients with unstable angina pectoris died from pump failure. Patients with cardiac rupture within 24 hours after admission, had significantly higher systolic and diastolic blood pressure in comparison with the other groups and with patients dying from cardiac rupture on the third day, or 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 beta blocker therapy. ... Publication Year: 1979 • (...) • Rank 9, Similarity: 0.6921100616455078 PMID: 443259 Title: Chest pain, shock, arrhythmias and death in a young woman. Authors: Abstract: 9 figures form the core of this article describing and discussing a case of sudden death, 2 hours after a 30-year old woman presented at a hospital emergency with chest pains. She had taken no medications other than oral contraceptives (OCs) for 10 years. The patient was admitted to the coronary care unit where findings included a palpable blood pressure of 94 mm of Hg, a heart rate of 128/min, and a respiratory rate of 28/minute. Cyanosis was noted, jugular veins were distended, and there were rales over the lung bases bilaterally; cardiac sounds were soft and a third heart sound was audible. Arterial oxygen tension was 15 mm of Hg, and carbon dioxide tension was 42 mm of Hg; pH was 7.2. Ventricular tachycardia developed and ventricular fibrillation ensued. The patient was intubated and well oxygenated, external cardiac compression was performed, sodium bicarbonate, epinephrine, and calcium were administered, and electrical defibrillation was performed. After several attempts, the latter resulted in a slow idioventricular rhythm on the electrocardiogram, but neither the blood pressure nor pulse was detectable. ... Publication Year: 1979

• Rank 10, Similarity: 0.6920691132545471 PMID: 938220 Title: Left ventricular function during aortic surgery. Authors: R M Carroll, R B Laravuso, J F Schauble Abstract: Fourteen patients undergoing surgery for aneurysm or occlusive disease of the abdominal aorta were studied. Thirteen patients had a history of hypertension or myocardial infarction; two patients had chronic obstructive pulmonary disease. Tachycardia, hypertension, and elevated pulmonary artery occluded. (PAo) pressure occurred in response to laryngoscopy and intubation in two patients; elevation of PAo pressure in response to aortic cross-clamping occurred in two patients. In three of these four patients, electrocardiographic evidence of myocardial ischemia appeared. These events are important in a consideration of the occurence of myocardial infarction in patients undergoing abdominal aortic surgery. Satisfactory treatment of myocardial ischemia has been accomplished with the use of propranolol hydrochloride and sodium nitroprusside. Publication Year: 1976

Model Output: # Reasoning # Based on the patient's context and the given task description, I will analyze the key factors to predict mortality risk for the subsequent hospital visit: Conditions:

• The patient has multiple serious conditions, including congestive heart failure (nonhypertensive), cardiac dysrhythmias, and shock. These are severe cardiovascular issues that significantly increase mortality risk.

• The presence of diabetes mellitus with complications adds to the overall complexity of the patient's health status.

• Alcohol-related disorders may complicate treatment and recovery.

Procedures:

• The patient underwent blood transfusion, which might indicate severe blood loss or anemia, potentially linked to 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 surgeries.

⋮

Label: 1 -------------------------------------

----Patient ID: zzzzzzzz Visit 0: Conditions: 1. pleurisy; pneumothorax; pulmonary collapse 2. congestive heart failure; nonhypertensive 3. complications of surgical procedures or medical care 4. screening and history of mental health and substance abuse codes 5. cardiac dysrhythmias ⋮ Procedures: 1. other diagnostic procedures of respiratory tract and mediastinum 2. other non-or therapeutic procedures on respiratory system 3. incision of pleura; thoracentesis; chest drainage Medications: 1. other agents acting on the reninangiotensin system in atc 2. antiadrenergic agents, peripherally acting KG Community Chronic Kidney Disease Diabetes Anemia Hypertension Iron Supplementation le a d s to c a u s e c a u s e treat "Chronic Kidney Disease (CKD), commonly caused by Diabetes and Hypertension, can lead to Anemia, characterized by a decreased number of red blood cells. Anemia may cause symptoms such as fatigue and shortness of breath, increasing the risk of hospital readmissions. To prevent readmissions in CKD patients, it is essential to manage the underlying causes, monitor hemoglobin levels, treat Anemia with Iron Supplementation, and ensure treatment adherence." Theme-Specific Summarization "Diabetes and Hypertension are major risk factors for Chronic Kidney Disease (CKD). As CKD progresses, it can lead to complications such as Anemia, which can further increase morbidity and mortality. Early detection and management of Diabetes and Hypertension are crucial in preventing the development and progression of CKD, ultimately reducing the risk of premature death in these patient populations." Theme: Mortality Theme: Readmission Figure 14: An example of theme-specific summarization of knowledge graph community.

Mortality Prediction Task: Objective: Predict the mortality outcome for a patient's subsequent hospital visit based solely on conditions, procedures, and medications. Labels: 1 = mortality, 0 = survival Key Considerations: 1. Conditions: -Severity of diagnosed conditions (e.g., advanced cancer, severe heart failure, sepsis) -Presence of multiple comorbidities -Acute vs. chronic nature of conditions 2. Procedures: -Invasiveness and complexity of recent procedures -Emergency vs. elective procedures -Frequency of life-sustaining procedures (e.g., dialysis, mechanical ventilation) 3. Medications: -Use of high-risk medications (e.g., chemotherapy drugs, immunosuppressants) -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 a high risk of 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. Readmission Prediction Task: Objective: Predict if the patient will be readmitted to the hospital within 15 days of discharge based solely on conditions, procedures, and medications. Labels: 1 = readmission within 15 days, 0 = no readmission within 15 days Key Considerations: 1. Conditions: -Chronic diseases with high risk of exacerbation -Conditions requiring close monitoring or frequent adjustments -Recent acute conditions with potential for complications 2. Procedures: -Recent major surgeries or interventions with high complication rates -Procedures that require extensive follow-up care -Incomplete or partially successful procedures 3. Medications: -New medication regimens that may require adjustment -Medications with narrow therapeutic windows or high risk of side effects -Complex medication schedules that may lead to adherence issues Note: Analyze the information comprehensively to determine the likelihood of readmission. The goal is to accurately distinguish between patients who are likely to be readmitted and those who are not. Task Description -Mortality Prediction Task Description -Readmission Prediction "Key Considerations" section is only for zero/few-shot prompting & training sample generation "Key Considerations" section is only for zero/few-shot prompting & training sample generation The reasoning should be comprehensive, medically sound, and clearly explain how the patient's information leads to the predicted outcome.

Important Notes:

-Do not mention the ground truth label in the reasoning process.

-Use the relevant knowledge as needed.

-Analyze the similarities and differences between the patient and similar patients to justify the prediction.

After generating the reasoning chain, please review it and indicate your confidence in the reasoning chain at the end. H HUMAN EVALUATION OF REASONING CHAINS

We hired three MD students and an MD to conduct expert evaluation of the reasoning chain generated by KARE. We instructed them to use the 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 were 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. It evaluates whether the reasoning contains contradictions.

-5 (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. -3 (Fair): Some inconsistencies exist that partially undermine the rationale's reliability.

-2 (Poor): Significant inconsistencies exist, making the rationale difficult to trust.

-1 (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.

-5 (Excellent): The rationale is entirely medically accurate and reflects evidence-based clinical knowledge. -4 (Good): Minor inaccuracies are present but do not affect the overall clinical validity.

-3 (Fair): The rationale contains some medically incorrect statements that could mislead clinicians. -2 (Poor): Multiple inaccuracies significantly reduce the credibility of the rationale.

-1 (Very Poor): The rationale is largely or entirely medically incorrect, making it unusable.

• SPECIFICITY Definition: Specificity evaluates how detailed and precise the reasoning is in addressing the clinical scenario.

-5 (Excellent): The rationale provides highly detailed and tailored insights specific to the patient case. -4 (Good): The rationale is detailed but occasionally includes generalities.

-3 (Fair): The rationale is moderately specific, with noticeable generalizations.

-2 (Poor): The rationale is vague and lacks sufficient detail to guide clinical decision-making.

-1 (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.

-5 (Excellent): The rationale strongly supports the correct prediction and adds valuable clinical insights. -4 (Good): The rationale is helpful overall but lacks some critical insights.

-3 (Fair): The rationale provides some useful guidance but is incomplete or not compelling.

-2 (Poor): The rationale adds minimal value to the prediction and lacks actionable insights.

-1 (Very Poor): The rationale is unhelpful and does not contribute meaningfully to the diagnosis. • HUMAN-LIKENESS Definition: Human-likeness measures how well the clinical rationale demonstrates insight and understanding of the patient description or diagnosis in a way that resembles human clinical reasoning.

-5 (Excellent): The rationale reflects deep clinical insight, contextual understanding, and reasoning that fully mimics human behavior.