Hybrid AI Approach for Counterfactual Prediction over Knowledge Graphs for Personal Healthcare

Authors: hao  huang
Venue: KDD_org_2024_Workshop_AIDSH
Published: 2024-06-05 13:38:48

Abstract: Artificial intelligence (AI) has become an invaluable tool in healthcare for disease prediction and diagnosis. Despite their predictive accuracy, AI models may ignore causal relationships between patient characteristics (demographic or clinical). As a result, although AI models capture associative patterns, ignoring the causal relationships of predicted characteristics limits their ability to perform counterfactual reasoning about the predicted characteristics of a patient when dependent characteristics change. This limitation of AI models can affect the understanding of predicted outcomes. We have proposed hybrid AI methods that combine symbolic reasoning over knowledge graphs (KGs), large language models (LLMs), and causal reasoning techniques to infer causal relationships between patients' properties. As a result, a causal model is learned, enabling counterfactual prediction to support clinical decisions. We apply these AI methods to predict the counterfactuals of biomarker results in non-small cell lung cancer (NSCLC) patients under hypothetical treatments with different smoking habits. We have created synthetic datasets based on clinical records of NSCLC patients to evaluate the performance of the proposed methods. The observed results suggest that our methods are competitive with baseline methods in causal relationship discovery and counterfactual prediction.

Main paper: 
 introduction: 
 Explainable Artificial Intelligence (XAI) has emerged as a paradigm to facilitate the interpretability and transparency of traditional machine learning (ML) models [4, 17] . In healthcare, XAI facilitates clinical decisions with AI-driven insights [2] . Recent studies investigate the great possibilities of integrating AI in diagnostic and prognostic processes [5, 6, 10] , highlighting the potential of AI to transform medical practices. A key aspect of making AI insights actionable in healthcare is the understanding of causal relationships and counterfactual scenarios [24] . Causal inference and counterfactual prediction go beyond simple correlations, providing deep insights into the "why" and "what if " behind AI decisions, essential for clinical applications [22] . Despite advances, these techniques assume the availability of causal knowledge, such as relationships between variables, representing patient characteristics (e.g. age, sex, and biomarkers). Causal associations can be obtained from domain experts, learned from data [29] , or extracted from knowledge bases such as ontologies and knowledge graphs (KGs) [23] . KGs [13] integrate data and metadata (specified in ontologies), which encode the meaning of properties, such as their domain and range. Consequently, both ontologies and KGs provide contextual information crucial for discovering causal relationships and facilitating counterfactual reasoning in the medical domain.Motivating Example. Consider a non-small cell lung cancer (NSCLC) patient, Eva, who is a female non-smoker and tested positive for the anaplastic lymphoma kinase (ALK) biomarker. Figure 1 -(1) partially visualizes Eva's characteristics, including her gender, smoking habits, and biomarker results. The oncology researchers aim to investigate how smoking habits affect biomarker results in patients like Eva. Specifically, they want to examine the hypothetical outcome of Eva's biomarker results if she were a smoker, shown in Figure 1 - (2) . Current AI techniques in healthcare [5, 6, 10] rely on predictive models, such as the random forest [5] and deep learning models [10] , to make accurate predictions. However, "Can we simply modify Eva's smoking habits in the dataset and use a predictive model to predict her counterfactual biomarker outcome?" Proposed Solution. We propose hybrid AI methods that combine:(1) symbolic reasoning on the metadata modeled in the ontologies of KGs; (2) numerical approach implemented in LLMs and ML methods; and (3) causal discovery and reasoning. As a proof of concept, we implement the hybrid method over a medical KG that integrates clinical data from lung cancer patients for counterfactual reasoning. This KG was built following the methods proposed by Vidal et al. [27] and Aisopos et al. [1] . Due to privacy restrictions, the results presented in this paper were derived from synthetic KGs created from the original KG reported in [1, 27] . The metadata captured in the medical KG is used to improve the accuracy of causal discovery and the performance of counterfactual prediction.Evaluation. We evaluate the effectiveness of our method in a use case of predicting the counterfactual biomarker (outcome) of NSCLC patients if the patient's smoking habits (treatment) were changed. In each dataset, the factual and counterfactual biomarker outcomes of patients are generated using an Additive Noise Model (ANM) [14] learned from the original data and based on expertdesigned causal knowledge, including all causal relationships between the characteristics of NSCLC patients. Our empirical results show that our method outperforms other methods in discovering causal relationships and is competitive in counterfactual reasoning.Contributions. The contributions of this work are: (1) The hybrid AI methods for causal relationship discovery over KGs. (2) Empirical evaluation of the proposed method on synthetic KGs created from clinical data of patients with NSCLC. This paper is structured as follows: Section 2 presents the foundational concepts for the proposed approach. Section 3 defines the problem statement and introduces the proposed hybrid AI method. Section 4 details the experimental settings, reports the used metrics, and discusses the results. Section 5 analyzes the state-of-the-art. Finally, section 6 summarizes our findings and explores future directions. 
preliminaries: 
 Knowledge Graph [13] . A KG is a directed edge-labeled graph 𝐾𝐺 = (𝑉 , 𝐿, 𝐸), where 𝑉 , 𝐿 are respectively a set of entities and property labels, and 𝐸 ⊆ 𝑉 × 𝐿 ×𝑉 is a set of triples. A 𝐾𝐺 metadata is part of 𝐾𝐺 and corresponds to triples 𝐸 * = {(𝑠, 𝑝, 𝑜) ∈ 𝐸|𝑝 ∈ 𝐿 * } expressing the meaning of data, where 𝐿 * ⊆ 𝐿 is a set of properties used to annotate entities, classes, and properties.Ego Network [11] . Given a knowledge graph 𝐾𝐺 = (𝑉 , 𝐿, 𝐸), an entity 𝑒 ∈ 𝑉 , the ego network of 𝑒 in 𝐾𝐺 based on a set of properties 𝑃 ⊆ 𝐿 is defined as 𝜙 (𝑒, 𝑃) = {(𝑒, 𝑝, 𝑜)|𝑝 ∈ 𝑃 ∧ (𝑒, 𝑝, 𝑜) ∈ 𝐸}. Causal Concepts. In a causal analysis [18] , each unit 𝑒 refers to an entity subjected to an intervention 𝑑𝑜 (𝑇 =𝑡), where 𝑇 is a binary variable indicating treatment status, which is either the active treatment 𝑇 =1 intended for the impact analysis, or the control treatment 𝑇 =0 as reference for 𝑇 =1. The treatment effect of 𝑇 =1 versus 𝑇 =0 on unit 𝑒 is evaluated by comparing the corresponding potential outcomes 𝑌 (1) and 𝑌 (0) of 𝑒. The factual outcome under the actual treatment 𝑡 is 𝑌 𝐹 = 𝑌 (𝑡), while the counterfactual outcome under the hypothetical treatment 1-𝑡 is where𝑌 𝐶𝐹 = 𝑌 (1-𝑡). Causal Graph. A causal graph [19] is a directed acyclic graph (DAG) 𝐺 = (𝑋, 𝐸 ′ ),𝑃𝑎 𝐺 (𝑋 𝑖 ) = {𝑋 𝑗 |(𝑋 𝑗 , 𝑋 𝑖 ) ∈ 𝐸 ′ }. Let 𝑃 𝑡 (𝑋 ) = 𝑃 (𝑋 |𝑑𝑜 (𝑇 =𝑡)) be the distribution under intervention 𝑑𝑜 (𝑇 =𝑡) (𝑇 ⊆ 𝑋 ). The tuple (𝐺, 𝑃 𝑡 ) is a Causal Bayesian Network (CBN) [21] iff: (1) the distri- bution 𝑃 𝑡 (𝑋 ) over 𝑋 follows the factorization according to 𝐺; (2) 𝑃 𝑡 (𝑇 𝑖 = 𝑡 𝑖 ) = 1 for all 𝑇 𝑖 ∈ 𝑇 , whenever 𝑇 𝑖 = 𝑡 𝑖 consists with the in- tervention 𝑑𝑜 (𝑇 =𝑡); and (3) 𝑃 𝑡 (𝑋 𝑖 |𝑃𝑎 𝐺 (𝑋 𝑖 )) = 𝑃 (𝑋 𝑖 |𝑃𝑎 𝐺 (𝑋 𝑖 )) for all 𝑋 𝑖 ∈ 𝑋 \ {𝑇 }.With conditions (1-3), the interventional distribution 𝑃 𝑡 (𝑋 ) under any intervention 𝑑𝑜 (𝑇 =𝑡) is computed with a truncated factorization [21] : 𝑃 𝑡 (𝑋 ) = 𝑋 𝑖 :𝑋 𝑖 ∈𝑋 \{𝑇 } 𝑃 (𝑋 𝑖 |𝑃𝑎 𝐺 (𝑋 𝑖 )). 
our approach: 
 Let 𝐾𝐺 = (𝑉 , 𝐿, 𝐸) be a KG, 𝐶 be a class s. Our Solution. To solve the problem, we propose hybrid methods exploiting both data (all properties of patients, 𝑃 𝐶 ) and metadata (encoding the meaning of properties and classes) in 𝐾𝐺. Our approaches establish causal relationships among properties in 𝑃 𝐶 using data-driven and metadata-driven (e.g., LLM) causal discovery models, laying the foundation for a causal model 𝜗 to capture the causal mechanism, which generates the outcome property 𝑝 𝑌 of patients. The model 𝜗 is trained to maximize the likelihood of the observed data (in 𝐾𝐺), enabling it to learn a causal mechanism that is optimized to predict the most probable counterfactual outcomes of each patient under a hypothetical treatment.The design pattern of our proposal is illustrated in Figure 2 , following the design principles by Bekkhum et al. [26] , where the white rectangles represent outputs and inputs; the green rectangle represents symbol; the blue ovals represent processes; and the orange hexagons represent models. It comprises three components that tackle the tasks of causal relationship discovery and causal model learning towards solving the problem of counterfactual prediction.(1) Symbolic Reasoning. This component takes as input a medical KG 𝐾𝐺 = (𝑉 , 𝐿, 𝐸), a target class 𝐶 (e.g. NSCLC Patient), and all properties of 𝐶, i.e., 𝑃 𝐶 . It employs a semantic model, which defines the meaning of symbols in KG, to perform reasoning and query processing over KG. This process deduces data and metadata for the counterfactual prediction. Specifically, it queries the properties 𝑃 𝐶 (data) of all patients (units), and transforms the data into a dataset 𝐷 𝐾𝐺 (𝐶, 𝑃 𝐶 ). It also extracts the metadata that describes the specific domain of the medical knowledge graph 𝐾𝐺 (e.g., lung cancer), the meaning of the target class 𝐶, and the properties in 𝑃 𝐶 . The metadata is used to populate an LLM prompt [7] for querying causal relationships between properties in 𝑃 𝐶 . Although generic, we consider basic metadata about properties (e.g., domain, range, and notation) and classes. However, as shown by Hung and Vidal [15] , other reasoning methods can be used to improve causal reasoning.(2) Causal Discovery. This component aims to infer the causal graph or DAG that encodes the causal relationships between properties in 𝑃 𝐶 . It takes as input the prompt sections for querying causal relationships and the dataset 𝐷 𝐾𝐺 (𝐶, 𝑃 𝐶 ) derived from 𝐾𝐺; and outputs a causal graph 𝐺. To do this, it uses a data-driven model causal discovery and a metadata-driven model LLM. The data-driven method can be any traditional causal discovery method, such as PC [25] and FCI [25] , or GES [8] , which uses the dataset 𝐷 𝐾𝐺 (𝐶, 𝑃 𝐶 ) and produces a causal graph 𝐺 1 = (𝑃 𝐶 , 𝐸 1 ). The metadata-driven modal can be any LLM that takes the LLM prompt as input and outputs a causal model 𝐺 2 = (𝑃 𝐶 , 𝐸 2 ) represented by a set of causal relationships. We designed four sections for the LLM prompt, where the role section uses the metadata describing the domain information of 𝐾𝐺 for defining the persona or function the LLM should adopt; the context section uses the metadata describing the meaning of the target class 𝐶 and properties in 𝑃 𝐶 . For example, the domain (rdfs:domain) and range (rdfs:range) of properties, the humanreadable label (rdfs:label) and annotation (rdfs:comment); the objective section specifies the task of identifying causal relationships between properties in 𝑃 𝐶 ; the instruction section formats the output of the causal relationships. The final output of this component is a causal graph 𝐺 = (𝑃 𝐶 , 𝐸 ′ ) s.t. 𝐸 ′ = 𝐸 2 ∪ {(𝑐, 𝑒) ∈ 𝐸 1 |(𝑃 𝐶 , 𝐸 2 ∪ {(𝑐, 𝑒)}) 
experimental study: 
 In this study, we investigate the impact of smoking habits (treatment) on the biomarker results (outcome) of NSCLC patients. We evaluate the performance of HealthCareAI in the tasks of causal discovery and counterfactual prediction over synthetic KGs. Research Questions. Let 𝐾𝐺 be a KG of NSCLC patients. We raise two research questions. Q1: Given a causal graph 𝐺 * provided by an expert (a.k.a. the Expert DAG in Figure 3 ), which encodes the causal relationships between properties of patients. Can our method discover the same DAG? Q2: Given counterfactual biomarkers of patients in 𝐾𝐺, is our method able to correctly predict counterfactuals using the context of patients in 𝐾𝐺? Original NSCLC KG. The original NSCLC KG is reported in [1, 27] . In this study, we employed a fragment of the KG, containing 1,808 patients of class LCPatient. Each patient has the following properties: Biomarker that is either ALK or EGFR and other biomarkers, Age that is categorized as Young (≤ 50 years) or Old (> 50 years), Gender that is Male or Female, SmokerType that is Non-Smoker or Smoker, FamilyCancer which is OnlyMajor if a patient's family antecedents have only these cancers: Breast, Lung, Colorectal, Head and neck, Uterus/cervical, Esophagogastric, Prostate, otherwise has-Minor, FamilyGender that is Women if all family antecedents are women, Men if all of them are men, otherwise WomenorMen, and FamilyDegree that indicates the degree of relationship of family antecedents to the patient (i.e., first, second, or third degree). Synthetic NSCLC KGs. We use an ANM [14] , trained based on the Expert DAG 𝐺 * (in Figure 3 ), using the dataset derived from the KG fragment (including 1808 patients) to faithfully capture the causal mechanism in the dataset. Using this model, we generate synthetic KGs with various patient numbers 𝑁 ∈ {2𝑘, 5𝑘, 10𝑘 }. Specifically, Age, Gender, FamilyGender, and FamilyDegree are simulated from uniform distributions; while SmokerType, FamilyCancer, and Biomarker are simulated using logistic functions:EQUATIONwhere 𝑌 is the synthetic variable; 𝑋 ′ represents parent variables of 𝑌 in the DAG 𝐺 * (see Figure 3 ). The 𝛼 and 𝛽 are learned from the dataset using a logistic regression. A noise term N (0, 0. := 𝑥 ′ 1 , . . . , 𝑥 ′ 𝑘 s.t. (∀𝑖 ∈ [1, 𝑘]) [(𝑒, 𝑋 ′ 𝑖 , 𝑥 ′ 𝑖 ) ∈ 𝑆]where 𝑘 is the variable number of 𝑋 ′ . Metrics. For Q1, given the Expert DAG 𝐺 * and an estimated DAG 𝐺, we use the Jaccard Index [16] to evaluate the performance:𝐽 𝐼 (𝐺, 𝐺 * ) = |𝑒𝑑𝑔𝑒𝑠 (𝐺 )∩𝑒𝑑𝑔𝑒𝑠 (𝐺 * ) | |𝑒𝑑𝑔𝑒𝑠 (𝐺 )∪𝑒𝑑𝑔𝑒𝑠 (𝐺 * ) | .For Q2, given the ground truth counterfactual Biomarker of all patients 𝑌 𝐶𝐹 , we evaluate the performance of each CBN using the Pearson correlation coefficient [9] (PCC):𝑃𝐶𝐶 ( Ŷ𝐶𝐹 , 𝑌 𝐶𝐹 ) = 𝑐𝑜𝑣 ( Ŷ𝐶𝐹 ,𝑌 𝐶𝐹 ) 𝜎 Ŷ𝐶𝐹 𝜎 𝑌 𝐶𝐹, where Ŷ𝐶𝐹 represents the counterfactuals estimated by a CBN, 𝑐𝑜𝑣 (•) and 𝜎 represent covariance and standard deviation. Higher values of these criteria indicate better performance. Compared Methods. We compare HealthCareAI with all datadriven methods: PC [25] , FCI [25] , and GES [8] (implemented with the Causal-learn Python package [30] ) and metadata-driven method: LLM (via official website of GPT-4 [20] ). HealthCareAI is implemented using PC and GPT-4. CBNs of all methods are implemented using the pgmpy Python package [3] . Appendix A shows the LLM prompt for query causal relationships. The code and data for reproduction are available here 1 . Results of Causal Discovery. We evaluate the DAGS produced by all methods against the Expert DAG 𝐺 * (in Figure 3 ) over various synthetic KGs with different patient numbers 𝑁 . Table 1 presents the evaluation results (by Jaccard Index) of different methods. The results show that the HealthCareAI outperforms other methods with Jarccard Index of 88.9% (𝑁 =2𝑘), 100% (𝑁 =5𝑘), and 88.9% (𝑁 =10𝑘), followed by LLM with Jarccard Index of 66.7% in all settings of 𝑁 . The PC and FCI methods, perform better than the GES method. Surprisingly, when 𝑁 ≥ 5𝑘, increasing 𝑁 does not further improve the performance of causal discovery. These results answer Q1 that our method achieves excellent performance in causal discovery. The results also show that the metadata-driven model LLM can act as a complement to the data-driven model, highlighting the potential of combining data and metadata from KGs for causal discovery.Results of Counterfactual Prediction. Based on the DAGs estimated in the previous step, we train CBNs from the dataset over the synthetic KGs using the Maximum Likelihood Estimation. We denote the estimated counterfactual Biomarker by CBN as Ŷ𝐶𝐹 . For each patient 𝑖, Ŷ𝐶𝐹 𝑖 is 1 (represents ALK or EGFR) if 𝑃 𝐶𝐹 (𝑌 =1|𝑇 =𝑡, 𝑋 =𝑥) > 0.5, otherwise is 0 (represents other biomarkers). The evaluation results are presented in Table 2 ; the PCC metrics for each CBN are presented as the mean (± standard deviation), using 5-fold crossvalidation on all NSCLC patients and their counterfactuals. The results indicate that PC and FCI outperform others in scenarios with limited data (i.e., 𝑁 = 2𝑘). In contrast, the CBN trained based on the Expert DAG 𝐺 * exhibits the lowest performance, which may be explained by the overfitting issue [12] . This is likely due to the complex structure of the Expert DAG, which requires a large dataset to learn the conditional probability tables of the CBN. Conversely, the simpler structures of the DAGs derived from the PC and FCI methods allow for effective learning with less data.In the scenario of large datasets (i.e., 𝑁 ≥ 5𝑘), the Expert CBN outperforms all CBNs trained based on DAGs by other methods. Although enlarging the dataset generally enhances the generalization capability of CBNs, it is notable that the performance of all CBNs slightly declines as 𝑁 increases from 5k to 10k. The performance of the LLM CBN (using DAG by LLM), despite its high structural similarity to the Expert DAG, remains the worst across all datasets. This result underscores the limitations of DAGs estimated without considering the underlying data, emphasizing the crucial role of data-driven causal graph estimation for robust counterfactual reasoning By incorporating the causal relationships estimated by the constraint-based method PC and those inferred by LLM, the CBN produced by our HealthCareAI method achieves competitive performance across all settings of 𝑁 , compared to the CBN based on the Expert DAG. The results address Q2, confirming that our method can provide good counterfactual reasoning performance. 
related work: 
 Causal Discovery. Understanding causal relationships in data is crucial in healthcare [24] . Traditional approaches like the PC (Peter-Clark algorithm) and FCI (Fast Causal Inference) [25] rely on statistical tests to identify conditional independence and create the directed acyclic graph (DAG); the GES (Greedy Equivalence Search) [8] uses a score-based strategy to optimize the Bayesian Information Criterion (BIC) [8] to find the best DAG. Recent advanced methods have tried to use large language models for causal discovery [28] . Causal Models. Causal Bayesian Networks (CBNs) [21] and Structural Causal Models (SCMs) [21] are well-known causal models, where CBNs focus on causality within categorical data, enabling intervention analysis and interpretable inference; SCMs provide a general framework with structural equations for causal reasoning. Both models use DAGs to depict causal relationships, but CBNs are more suitable for categorical variables and easier to interpret. AI for healthcare. Recent advanced AI techniques [5, 6, 10] in healthcare apply various machined learning or deep learning models for prediction, offering great potential in healthcare, suggesting that AI can perform as well as or better than humans in some tasks, e.g., diagnosing disease. However, the potential of causality [22] is largely overlooked. To bridge the gap, we propose a hybrid method for causal discovery and reasoning over healthcare KGs. 
conclusion and future work: 
 In this paper, we propose hybrid AI methods, HealthCareAI, which incorporate advanced ML techniques with domain knowledge for causal relationship discovery and offer counterfactual reasoning over KGs. Through extensive experiments on synthetic KGs of nonsmall cell lung cancer (NSCLC) patients, we demonstrate the effectiveness of HealthCareAI in discovering causal relationships and predicting counterfactual outcomes. These capabilities underscore the potential of HealthCareAI to enhance clinical decision-making and deepen the understanding of personal clinical results. Currently, HealthCareAI is tailored to work with categorical data, which, while useful, limits its applicability across broader data types, such as numerical data. We plan to extend its capabilities to analysis using multiple data types and reasoning processes from biomedical ontologies. Additionally, we want to explore various prompting engineering techniques for improving causal discovery with large language models (LLMs). 
 
 Appendix: https://figshare.com/s/c50b45f2483987664c1c 



Reviews: review: title: title: Simple and effective method, Open-sourced code and data for easy reproduction
main_review: main_review: Summary Of Strengths:
- Clear motivating example and well-defined problem
- Simple and effective hybrid method for counterfactual prediction on personal healthcare
- Open-sourced code and data for easy reproduction of the result

Summary Of Weaknesses:
- Minor types need to be fixed, for example, Table 2 is too wide
- Can this method be effectively extended for other healthcare conditions and applied to real-life data
title: title: Hybrid AI Approach for Counterfactual Prediction over Knowledge Graphs for Personal Healthcare
main_review: main_review: **Summary**  
The paper designs a framework combining symbolic reasoning over knowledge graphs (KGs), large language models (LLMs), and causal reasoning techniques to infer causal relationships and predict counterfactual. The proposed hybrid method achieves best performance in causal discovery task.

**Strength and weakness**  
**Strengths**  
- The proposed method has strong ability for causal discovery.
- Assist the conventional causal discovery methods with LLM to enhance the performance.

**Weakness**  
- The evaluations primarily use synthetic datasets. While this is necessary for privacy reasons, the realism and complexity of real patient data can be substantially different. This could limit the applicability of the proposed methods to actual healthcare settings where data distributions and patient characteristics are more variable.
- While the paper aims at improving interpretability of AI in healthcare, the actual usability and interpretability of the outputs for clinical decision-making are not thoroughly assessed.
- There are some typos. For instance, in line 280, KG repeats twice.