Predictive Modeling with Temporal Graphical Representation on Electronic Health Records

1 Introduction

Electronic Health Records (EHR) play a crucial role in healthcare by enhancing the quality of care and improving the efficiency and accuracy of medical decisions. In recent years, with the advancement of artificial intelligence technologies, many deep learning models based on EHR data have been extensively explored and applied in predictive modeling, such as patient risk assessment [Bai et al., 2018; Luo et al., 2020; Jiang et al., 2023; Wang et al., 2024], prediction of drug interactions [Zhang et al., 2017; Yang et al., 2021], and support for clinical decision-making [Yin et al., 2022; Yang et al., 2023c; Chen et al., 2024].

∗The corresponding author.

Predictive modeling on EHR data is proposed to utilize clinical events from patients' historical visits, such as diagnoses, procedures and medications to forecast their future health status. These clinical events are typically represented as a set of medical codes. To enhance the quality and accuracy of medical predictive modeling, models need to learn three hierarchical representations in EHR: the patient-level, visit-level, and (medical) code-level. Specifically, the patient level representation should capture the temporal relationships between historical visits to build a comprehensive health profile of the patient. Then, at the visit level, the model analyzes the various medical events encountered in each visit to understand the specific circumstances of each visit. Finally, at the medical code level, the model needs to examine the characteristics of individual medical events and their higher-order patterns across historical visits, as well as the correlations between different medical events.

Patients' longitudinal EHR, as a series of successive visits, can be naturally conceptualized as a sequence, as illustrated in Figure 1a. Current literature extensively employs sequence-based models, such as RNN-based [Choi et al., 2016; Gao et al., 2020] and Transformer-based [Luo et al., 2020; Yang et al., 2023c], to track state changes across a patient's historical visits. These models excel in extracting temporal information at the visit level. However, they fall short in discerning structured relationships between medical events within visits due to weak temporal associations among these events. In contrast, graph-based models, another prominent category of EHR models, are adept at understanding relationships between medical codes in individual visits, especially regarding their structured connections. Some models [Choi et al., 2017; Ma et al., 2018; Lu et al., 2021; Jiang et al., 2023] enhance medical concept representations by integrating elements like medical ontology trees or subgraphs from medical knowledge graphs (KGs). As shown in Figure 1b, other approaches conceptualize EHR as assorted graphs, applying graph-based models to extract structured information. For instance, [Choi et al., 2020] has introduced a Graph Convolutional Transformer, aimed at deciphering the graphical structure of medical codes. While graph-based models are proficient in learning enhanced visit-level representations through structured features, they face challenges in integrating temporal data. DDHGNN [Wen et al., 2022]

tries to remedy this by embedding dynamic temporal edges

in heterogeneous patient graphs, achieving improved predic-

tions but lacking an explicit focus on patient state changes

over historical visits. This leads to a pivotal question: Is there

a model capable of concurrently capturing all hierarchical la-

tent relationships within EHRs, thereby achieving superior

patient, visit, and medical code level representations

To address this challenge, in this paper, we model EHR

as a novel temporal heterogeneous graph, as depicted in the

Figure 1c. The visit-level and medical code-level representa-

tions are learned through dedicated medical event nodes and

time-dependent visit nodes, respectively. The patient graph is

designed to model intrinsic graphical relationships via edges

connecting medical event nodes to visit nodes. Additionally,

it facilitates the propagation of temporal relationships reflec-

tive of changes in patient state through temporal edges in-

terlinking these visits. Building upon the patient graph rep-

resentation, we introduce the Temporal gRAph traNSformer

(TRANS), a novel model designed to deeply unravel these

complex relationships. In particular, our temporal convolu-

tion merges temporal embedding with heterogeneous mes-

sage passing. This fusion allows for the extraction of graph-

structured features from medical events within individual vis-

its, while also capturing the complex, higher-order tempo-

ral dependencies that span across historical visits. We en-

hance the learning of graphical features by integrating local

structural encoding and global positional encoding into the

medical event node representations, facilitated through meta-

paths. Additionally, a patient explainer has been developed

to ensure that our predictions are interpretable. We highlight

the contributions of this work as follows:

· We propose a novel EHR representation method, trans-

forming patient EHR into a heterogeneous graph with

distinct sequential visit nodes and medical event nodes.

This approach adeptly captures both structural relation-

ships among medical events and temporal dependencies across historical visits.

• A novel model TRANS with temporal graph convolution and the structural and positional encoding of heterogeneous graphs is proposed to extract the temporal graphical representation of a patient's graph, leading to interpretable predictions.
• We conduct extensive experiments on three real-world datasets to demonstrate that both the constructed graph and the TRANS model outperform baseline methods, demonstrating the effectiveness of our approach to healthcare predictive modeling. The implementation code can be found at Github 1 .

1 https://github.com/The-Real-JerryChen/TRANS

(c) Temporal Graphical Representation

Figure 1: Different EHR representation methods. (a) An example of sequential representation. (b) Examples of graphical representation, which include graphs at the visit level and the patient level. (c) The temporal graphical representation proposed in this paper.

2 Preliminary

In this section, we provide some background knowledge about EHR data and then formulate the diagnosis prediction task.

2.1 Medical Codes

Each medical event can be represented as a set of unique medical codes, denoted as C = {c1, c2, . . . , c|C|}, where |C| is the size of unique medical codes. Each code may represent a diagnosis, procedure, or medication. In this paper, we abstractly refer to these events as e.

2.2 EHRData

The EHR of each patient is represented as a time-sequenced array of medical events, encompassing diagnoses, procedures, and medications. Specifically, the historical visit information of patient p can be represented as Xp = [x1, x2, · · · , xT ], where the t-th visit is denoted by a multihot vector xt ∈ {0, 1}|C| , t = 1, . . . , T p indicates the index of visit, and T p denotes the total visit times of patient p. The i-th element of one visit vector is set to 1 if it contains the medical code ci. To simplify, we explain the algorithms for one patient, omitting the superscript p when clear.

2.3 Sequential Diagnosis Prediction

The goal of predictive modeling is to predict the label at each time stamp or at the end of the visit sequence. In this study, we focus on the task of next-visit diagnosis prediction. It involves predicting diagnosis codes Cdiag. ∈ {0, 1}|C| based on the medical events from historical visits as well as the medication codes and procedure codes for the current visit.

3 Methodology

In this section, we first introduce the method for temporal graph construction, followed by details of the proposed model TRANS. Figure 2 shows the overall architecture of TRANS, which is composed of three main modules: (1) heterogeneous temporal message passing; (2) spatial encoder; and (3) patient graph explainer.

3.1 Graph Construction

To capture the structured information of potential disease development trends, treatment pathways, and other temporal graphical relationships within the patient EHR, we construct a heterogeneous patient graph G = (N, E) for each patient, as illustrated in Figure 2, where N and E are sets of nodes and edges, respectively. The heterogeneous graph comprises four types of nodes: diagnoses, procedures, medications, and timestamped visit nodes: N = Ne ∪ Nv. where |Nv| = T, |Ne| = |e|. Here, e represents medical events that include diagnosis, medication, and procedure. The edge set contains two types of edges: (1) undirected edges between medical codes and visits when a visit includes medical events Eev = {(u1, u2) : ∀u1 ∈ Ne, ∀u2 ∈ Nv}; (2) directed edges between visits in chronological order Evv = {(ut, ut+1}) : ∀u ∈ Nv, 0 ≤ t < T }.

To the best of our knowledge, our graph construction differs from existing methods [Choi et al., 2020; Wen et al., 2022; Jiang et al., 2023] of EHR representation learning. We explicitly represent each visit with time-dependent visit node. This approach decouples the representations of medical codes and the historical visits, allowing the embedding of clinical

event nodes to focus on the characteristics of the codes themselves, while using the embedding of visit nodes to capture changes in the patient's state. Additionally, we incorporate temporal embedding on the edges between medical codes and visits, as well as between successive visits with different time indexes.

3.2 Temporal Heterogeneous Message Passing

To effectively gather temporal information and advanced connections between medical events across historical visits, we propose a tailored temporal-aware message-passing mechanism for heterogeneous graphs. This system not only propagates medical code information into patient visit representations but also subtly connects various medical events through the visit nodes. We represent the features of medical code node and visit node outputted by the (l)-th layer of the network as h(l) e ∈ Rd and h(l) v ∈ Rd , respectively. They also serve as the input for the subsequent (l + 1)-th layer. We denote linear mappings as L(·), with each instance representing a distinct mapping.

Temporal Embedding

The temporal embedding in TRANS consists of two parts: the time encoding of the visit nodes; and the temporal encoding of edge features. We employ Time2Vec [Kazemi et al., 2019]

to encode the absolute times of the visits, normalized for each patient, into vectors. This time-aware vector is then concatenated with the initial features of the visit node h(0) v , thereby incorporating information about the time intervals between different visits.

Time2Vec is suitable for handling real-world timestamps, such as specific visit times, while for the indexed time of edges between medical events and visits, we use functional time encoding. Inspired by [Xu et al., 2020], we use T (t) to encode visit index t into a continuous vector, which is then regulated by a linear layer to obtain the time factor αt.

where ω1, . . . , ωd are learnable parameters. And αt = L(T (t)) . We separate the temporal vector from node features, treating it as an edge type, to maintain focus on the nodes' semantic information while infusing temporal data during message propagation into visit node features. Details can be found in the appendix.

Heterogeneous Message Passing

The entire message-passing process can be formalized into two stages: aggregation and updating. The first step is to aggregate the information of neighboring nodes. Specifically, we use an attention mechanism to weigh and integrate the features of neighboring nodes and concatenate the output from multi-head attention to obtain the final message. Taking as an example the process of propagating the features of neighbor nodes et ∈ N(vt) to the visit node vt, the i-th head attention

where et represents the medical event node such as diagnosis, medication, and procedure contained in t-th visit as well as previous visit node vt-1. The target node for message propagation is t-th visit node vt. We use weight matrix Wµ(e) ∈ Rd×d to encode different types of medical events, where µ(e) represents the type of the node et, such as diagnosis, procedure, or medication. After obtaining the attention scores for different neighboring nodes, we combine them with the mapped neighbor features to complete the entire aggregation process. We formalize it as follows:

where N(vt) denotes neighbors of vt, and ˜hiv represents the message passed to the visit node by the i-th head attention. Following the acquisition of aggregated information, the next stage is to combine this aggregated information with the target node's ego information to update the representation of the target node, as follows:

where ∥ i∈[1,h] denotes concatenating the outputs of multiple

heads, γ is the coefficient for the skip connection and σ represents activation function. Similar to the network structure in [Chen et al., 2020] , it can be either a learnable parameter or a hyper-parameter. Leveraging heterogeneous message passing, we integrate the features of medical events in the visit, along with the graphical structure information and historical visit information, into the current visit node, thereby updating the patient's representation.

3.3 Spatial Encoding

We integrate global positional encoding and local structural encoding into the initial feature of medical event nodes by concatenating them. This approach effectively captures structured features from patients' historical visits, which is crucial for illuminating key patterns like diagnosis sequences and treatment regimes, essential for accurately modeling disease stages. Such an integration enables us to discern intricate relationships in medical event co-occurrences and disease trajectories, thereby enhancing our predictive accuracy for a patient's future health status.

Local structural encoding S(e) is designed to capture local structural information in the graph based on random walks. Positional encoding P(e) adds a global positional encoding based on the Laplacian eigenvectors of the patient graph. However, in heterogeneous graphs, unlike homogeneous graphs [Ramp´aˇsek et al., 2022], direct computation of certain properties is not straightforward, owing to the complexity introduced by diverse types of nodes and edges. To navigate this complexity, we employ medical event metapaths, offering a structured and meaningful method to traverse the heterogeneous graph. For instance, in constructing meta-paths for medication nodes, visit nodes are used as intermediaries to connect multi-hop neighbors, forming a sequence represented as:

where Nm, Nv, and Nd represent medication, visit, and diagnosis nodes, respectively, and - denotes the pseudo edge. This meta-path construction results in an adjacency matrix ˜A ∈ R|Nm|×|Nm| , where |Nm| is the number of unique medication nodes within graph.

Once the meta-paths are established, we can compute structural embedding based on the corresponding adjacency matrix ˜A. And for P(e), we construct a Laplacian matrix and compute its top k eigenvectors to represent the spatial information of nodes within the graph. P(e) ∈ Rk for node e is expressed:

where each v(i) e is the element corresponding to node e in the i-th eigenvector. For structural encoding S(e) ∈ Rk , we calculate k-step random walk matrices ˜Wk based on the adjacency matrix ˜A, and extract the cumulative probabilities of node e connecting to itself at each step to form S(e) as :

where d(i) e represents the diagonal element of node e after the i-th step of random walk. This approach not only enhances the representational ability of the graph but also provides a richer context for the deep analysis of complex interactions between different types of medical events.

3.4 Patient Explainer

Based on the patient graph, we are able to predict the future development of a patient's status using their historical diagnosis and medical information. However, the decisionmaking process of GNNs also requires intuitive and interpretable analysis to increase the credibility of decisions and to assist doctors in better understanding them. Therefore, we have incorporated a heterogeneous graph explainer module into our model. By masking certain medical event nodes, the explainer aims to identify a significant subgraph Gs crucial for the model's predictive results. Utilizing mutual information (MI) as the criterion, this process can be expressed as follows:

where ˆ Y is the prediction of the TRANS model. The obtained subgraph contains the most influential medical events in the EHR for predicting the patient's status. We will provide a detailed introduction to this module in the appendix

4 Experiments

4.1 Experimental Setting

Datasets

We conduct experiments on the MIMIC-III [Johnson et al., 2016], MIMIC-IV [Johnson et al., 2023], and MarketScan Commercial Claims and Encounters (CCAE) [Butler et al., 2021] datasets. MIMIC-III and MIMIC-IV are two publicly available medical databases released on PhysioNet. Following [Ma et al., 2018], we filtered the original datasets to include patients who made at least two visits. The CCAE 2 dataset is a healthcare database, which contains patients with longer historical visit information. The goal of the diagnosis prediction task is to predict the diagnostic labels for the next visit. Similar to [Ma et al., 2018], we chose to predict the category of the diagnosis, which is the CCS (Clinical Classifications Software) code corresponding to the ICD code. We use PyHealth [Yang et al., 2023a] for parsing and preprocessing the dataset. Table 1 presents some statistics of the processed datasets. Details on dataset processing can be found in the Appendix.

Baselines

We evaluate our model against the following baselines:

Sequence-based Models. We use four sequence models based on different backbones. The Transformer [Vaswani et al., 2017] consists of encoders stacked based on the selfattention mechanism. RETAIN [Choi et al., 2016] is an RNNbased model that integrates a Reverse Time Attention mechanism. StageNet [Gao et al., 2020] is enabled by a stage-aware LSTM module and a stage-adaptive convolutional module. HiTANet [Luo et al., 2020] proposes a hierarchical timeaware Transformer model.

Graph-based Models. We compared our model with three graph-based models: KAME [Ma et al., 2018], GCT [Choi et al., 2020], and DDHGNN [Wen et al., 2022]. KAME incorporates medical ontology knowledge of ICD codes into the sequence model. GCT utilizes graph convolutional networks to learn the graphical structure in EHR. We extend it by concatenating a transformer, making it suitable for sequential input tasks. DDHGNN constructs EHR as heterogeneous graphs of patients and employs dynamic graph convolution to learn patient's features. Additionally, the results of some models [Hu et al., 2020; Anonymous, 2023] can also be regarded as ablation studies for our proposed model.

2 Available at https://www.merative.com/real-world-evidence.

The experimental setup and implementation details can be found in the Appendix.

Evaluation Metrics

We use two evaluation metrics [Ma et al., 2018] for the diagnosis prediction task: visit-level precision@k and code-level accuracy@k. Visit-level precision@k is defined as the number of correct medical codes within the top k predictions, normalized by the minimum of k and the actual number of categories present in a visit. Code-level accuracy@k is the ratio of correct predictions to the total number of predictions made for individual codes. The value of k varies from 10 to 30, with higher values indicating better model performance, where visit-level precision assesses broader performance and code-level accuracy reflects more detailed accuracy.

4.2 Experimental Results

Table 2 demonstrates the performance of the proposed TRANS model and all the other baselines on three datasets, including MIMIC-III, MIMIC-IV, and CCAE. We ran 10 times for each model under random dataset splits and reported the mean and standard deviation of the results, where the best results are highlighted in bold and the second best scores are underlined.

Overall results. Compared to baseline models, our proposed model TRANS demonstrated SOTA performance across all three datasets. Notably, the improvements on the MIMIC-IV and CCAE datasets were more pronounced than on MIMICIII. On the other hand, compared to MIMIC-III, almost all models performed better on the MIMIC-IV and CCAE datasets. This outcome suggests that more comprehensive datasets can enhance model training outcomes.

Sequence-based models. As mentioned before, sequence models treat patients' historical visits as text, where visits and medical events correspond to sentences and words, respectively. Table 2 shows that for the MIMIC-III and MIMIC-IV datasets, it's evident that the overall performance of sequence models is not as good as that of graph-based models. This may be attributed to the fact that in contrast to graph-based models like GCT and DDHGNN, sequence models are less effective at extracting structured information within a visit. Our proposed TRANS, which constructs a patient graph with temporally connected visit nodes, therefore achieves significant improvements. For the CCAE dataset, our model also outperforms those sequence model-based models, which further highlights the critical role of the graph module.

Graph-based models. In baseline models like KAME, treestructured information of ICD codes is incorporated into the sequence structure. However, this does not deeply relate to the patient's state, resulting in KAME's suboptimal performance on several datasets. On the other hand, GCT, DDHGNN, and our proposed TRANS model integrate a graph feature extraction module for visit data into the sequential network, leading to notable improvements on datasets with shorter historical sequences like MIMIC-III and MIMIC-IV. Yet, as Table 2 indicates, these graph-based models are outperformed by ours on the CCAE dataset. This disparity may stem from their limited focus on visit characteristics, essential for capturing state changes in longer sequences for precise diagnosis. HiTANe's good performance on longer sequence datasets also reflects its hierarchical approach to analyzing a patient's historical visit stages.

4.3 Model Analyses

Ablation Study

In this section, we conduct ablation studies to investigate the contributions of different components. Specifically, we removed three key parts of TRANS, namely Spatial Encoding (SE), Temporal Embedding (TE), and Sequence information in node feature initialization (Seq). Spatial Encoding includes both structural encoding and positional encoding. Temporal Embedding comprises the Time2Vec features of visit nodes and the temporal embedding of edge features in the graph. As shown in Table 3, the experimental results first reveal that each of these components contributes to improving the model's performance to some extent. Second, we observe that removing TE causes more performance degradation than removing SE, especially in the CCAE dataset where both Seq and TE play a more significant role. This suggests that graph convolution layers in TRANS are capable of capturing structured information to a certain degree, which overlaps with some features in SE. However, the inclusion of SE also results in better performance, indicating its important role.

Graph Construction

In this section, we delve into the distinctions among various graph construction methods. There are three notable approaches: (1) The first method, exemplified by GCT [Choi et al., 2020], prioritizes the structural relationships between medical events within individual visits, omitting temporal data. This method aggregates medical event nodes to create a patient representation. (2) The second approach, as introduced by DDHGNN [Wen et al., 2022], conceptualizes EHR as dynamic graphs with patient nodes at their core. It infuses temporal information into the dynamic edges linked to medical events. (3) Our proposed method, the third approach, models EHRs utilizing both time-dependent visit nodes and medical event nodes. The comparative performance of these methods is showcased in Table 2, where our method demonstrates superior efficacy.

DDHGNN's graph construction, while acknowledging historical visit temporal information, incorporates it solely into dynamic edge features. This leads to a conflation of medical event node representations with patient state data, diluting the focus on specific ICD code attributes. In our analysis of the CCAE dataset, we selected 4,000 codes that reflect a balanced

(a) DDHGNN

distribution across disease categories. We visualized these codes in a 2D space using t-SNE (Figure 3), where different colors represent different categories. The clustering patterns of ICD code representations in our model (Figure 3b) are notably distinct and category-specific, unlike the more chaotic representations observed in DDHGNN (Figure 3a). This observation underscores the efficacy and superiority of our approach in graph construction.

4.4 Graph Explainer

We employ a patient graph explainer to enhance the interpretability of our model's predictions, which is crucial for dissecting the predictive role of each node, including medical event nodes like medications and diagnoses, and visit nodes. Leveraging our capability to discern the significance of various nodes, we have performed a detailed statistical analysis of key historical diagnoses within the CCAE dataset, which specifically pertains to patients with Chronic Obstructive Pulmonary Disease (COPD). For each COPD patient (labeled with CCS code '127'), we calculate the importance of each diagnosis occurrence and then average these values across the entire dataset. Our analysis of the diagnostic nodes related to COPD allows us to pinpoint several potential comorbidities. As depicted in Table 4, the explainer has identified a range of likely comorbidities, systematically categorized by disease type. Notably, these findings, encompassing prevalent complications such as Asthma, Bronchitis, and Hypertension, resonate strongly with established clinical knowledge 3 , demonstrating the efficacy of the TRANS on predictive modeling.

5 RelatedWork

In this section, We review two main EHR representation methods: sequence-based and graph-based, along with their predictive modeling algorithms.

Sequence-based models. Sequence-based models utilize sequential models to capture temporal information in EHR. In sequence-based studies, patients' EHRs are represented a sequence of visit. Each visit contains a list of medical events (e.g., diagnosis codes and medications). These studies focus on extracting temporal or contextual features between visits. For example, the classic model RETAIN [Choi et al., 2016] introduces a attention-based combination of forward and backward visit-level representations to obtain more accurate predictions. Dipole [Ma et al., 2017] employs bidirectional recurrent neural networks combined with three dis-

3 https://www.mayoclinic.org/diseases-

conditions/copd/symptoms-causes/syc-20353679

Respiratory System Diseases

J06 - Acute upper respiratory infections

J20 - Acute bronchitis

J45 - Asthma

Mental and behavioral disorders

F32 - Major depressive disorder

F17 - Nicotine dependence

Circulatory System Diseases

I10 - Essential hypertension

E78 - Disorders of lipoprotein metabolism

Table 4: Potential COPD comorbidities grouped by disease categories, as identified by the Patient Explainer.

tinct attention mechanisms for patient visit information prediction. More recently, many works [Gao et al., 2020; Luo et al., 2020; Chen et al., 2022; Yang et al., 2023c] have applied the powerful Transformer-based model adapted from natural language processing like BERT [Devlin et al., 2018] to EHR, treating each visit or medical code as a token.

Graph-based models. Different from sequence-based models, graph-based models represent EHR in the form of graphs, which naturally model the inherent structural information among medical events. For instance, many models leverage external medical KGs to improve medical code representation learning. To name a few, GRAM [Choi et al., 2017] utilizes medical ontologies and graph-based attention to develop strong medical code representations, while CGL [Lu et al., 2021] utilizes disease domain knowledge and unstructured text. KerPrint [Yang et al., 2023b] proposed a time-aware KG attention method to understand knowledge decay over time. There is also considerable research focusing on the extraction of graph structural features within EHR. [Choi et al., 2018; Choi et al., 2020] explored the prior and hidden structured relationships between medical events. [Wen et al., 2022] expanded the graph from visit-level to patient-level, constructing a dynamic graph to represent EHR. However, existing graph-based models fail to simultaneously extract the temporal and structured relationships between visits.

6 Conclusion

In this paper, we investigate the three levels of relationships with EHR and the limitations of existing patient representation methods for predictive modeling. To capture both temporal and structural information about EHR, we propose a novel temporal heterogeneous graph to represent each patient's EHR. Moreover, we propose a TRANS model to learn the three-level relationships, which use temporal graph transformer to extract temporal and structured information from the graph. We validated our methods and the TRANS model through experiments on three EHR datasets, demonstrating superior performance in diagnosis prediction. Utilizing KGs in healthcare AI is also a crucial research area, particularly for improving the effectiveness of predictive models in the field. We will explore the expansion of our models through the incorporation of KGs in the future work.

Ethical Statement

The datasets used in the paper are from PhysioNet (MIMICIII and MIMIC-IV) and Merative (CCAE), which are fully HIPAA-compliant de-identified, and have very minimal risk of the potential for loss of privacy. Moreover, Per the DUAs, all users to access the data will need to take full research, ethics, and compliance training courses. Thus, potential privacy and security risks would be eliminated and/or mitigated.

Acknowledgments

This work was funded in part by the National Science Foundation under award number IIS-2145625, by the National Institutes of Health under award number R01GM141279, and by The Ohio State University President's Research Excellence Accelerator Grant.