Multi-task Heterogeneous Graph Learning on Electronic Health Records

TL;DR

MulT-EHR addresses the challenges of heterogeneous, noisy electronic health records by modeling them as a heterogeneous graph and applying a causal denoising module to mitigate confounding. The framework uses a Transformer-based heterogeneous GNN backbone with self-supervised TransE pretraining and a multi-task environment-invariant objective to share knowledge across four clinical prediction tasks. Empirical results on MIMIC-III and MIMIC-IV show consistent improvements over state-of-the-art baselines across mortality, readmission, length of stay, and drug recommendation, with ablations confirming the contribution of each component. The approach demonstrates strong potential for interpretable, generalizable EHR representations and could extend to other domains that involve heterogeneous graphs and multi-task learning.

Abstract

Learning electronic health records (EHRs) has received emerging attention because of its capability to facilitate accurate medical diagnosis. Since the EHRs contain enriched information specifying complex interactions between entities, modeling EHRs with graphs is shown to be effective in practice. The EHRs, however, present a great degree of heterogeneity, sparsity, and complexity, which hamper the performance of most of the models applied to them. Moreover, existing approaches modeling EHRs often focus on learning the representations for a single task, overlooking the multi-task nature of EHR analysis problems and resulting in limited generalizability across different tasks. In view of these limitations, we propose a novel framework for EHR modeling, namely MulT-EHR (Multi-Task EHR), which leverages a heterogeneous graph to mine the complex relations and model the heterogeneity in the EHRs. To mitigate the large degree of noise, we introduce a denoising module based on the causal inference framework to adjust for severe confounding effects and reduce noise in the EHR data. Additionally, since our model adopts a single graph neural network for simultaneous multi-task prediction, we design a multi-task learning module to leverage the inter-task knowledge to regularize the training process. Extensive empirical studies on MIMIC-III and MIMIC-IV datasets validate that the proposed method consistently outperforms the state-of-the-art designs in four popular EHR analysis tasks -- drug recommendation, and predictions of the length of stay, mortality, and readmission. Thorough ablation studies demonstrate the robustness of our method upon variations to key components and hyperparameters.

Figures (11)

• Figure 1: Illustration of EHR data and the challenges. EHRs contain enriched information on clinical visits of patients, including relevant medications and diagnoses. However, analysis of EHR data poses three challenges --- the sparsity in the data structure, noisiness, and heterogeneity of patients and their visits make it difficult to deliver accurate analysis.
• Figure 2: Four examples of meta-relations highlighted in the EHRs. These meta-relations involve structural connections between multiple nodes with different node types, highlighting the heterogeneity of the data. Consequently, a homogeneous graph model is insufficient to effectively capture and represent these complex meta-relations.
• Figure 3: An example of a heterogeneous graph constructed from EHR data, where $N$ represents the total number of patients, and $m_i$ represents the total number of visits by the $i$-th patient.
• Figure 4: Overview of our proposed framework. We first construct a heterogeneous graph from the raw EHR data, and then obtain node-level representation with heterogeneous GNNs. Causal and trivial representations are disentangled and the task-specific loss is obtained by combining the classification loss $\mathcal{L}_{\rm obj}$ and the uniform loss $\mathcal{L}_{\rm unif}$. We adopt a task-level aggregation module to obtain the multi-task learning loss. After training the GNN, we test the GNN with different downstream tasks (e.g., mortality prediction and drug recommendation).
• Figure 5: A causal diagram illustrating the effects of shortcut features. Without denoising, the model would make a prediction based on trivial features $S$ (i.e., the backdoor path $S$$\to$$R$$\to$$Y$)