Graph Transformers on EHRs: Better Representation Improves Downstream Performance
Authors: Raphael Poulain
Rahmatollah Beheshti
Venue: ICLR 2024 poster (2024)
Keywords: transformers
graph neural networks
electronic health records
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: Following the success of transformer-based methods across various machine learning applications, their adoption for healthcare predictive tasks using electronic health records (EHRs)  has also expanded extensively. Similarly, graph-based methods have been shown to be very effective in capturing inherent graph-type relationships in EHRs, leading to improved downstream performance. Although integrating these two families of approaches seems like a natural next step, in practice, creating such a design is challenging and has not been done. This is partly due to known EHR problems, such as high sparsity, making extracting meaningful temporal representations of medical visits challenging. In this study, we propose GT-BEHRT, a new approach that leverages temporal visit embeddings extracted from a graph transformer and uses a BERT-based model to obtain more robust patient representations, especially on longer EHR sequences. The graph-based approach allows GT-BEHRT to implicitly capture the intrinsic graphical relationships between medical observations, while the BERT model extracts the temporal relationships between visits, loosely mimicking the clinicians' decision-making process. As part of our method, we also present a two-step pre-training strategy for learning better graphical and temporal representations. Our proposed method achieves state-of-the-art performance in a variety of standard medical predictive tasks, demonstrating the versatility of our approach.

Published as a conference paper at ICLR 2024
GRAPH TRANSFORMERS ON EHRS: BETTER REPRE- SENTATION IMPROVES DOWNSTREAM PERFORMANCE
Raphael Poulain, Rahmatollah Beheshti University of Delaware
(rpoul ain, rbi] @udel edu
ABSTRACT
Following the success of transformer-based methods across various machine learning applications, their adoption to healthcare predictive tasks using electronic health records (EHRs) has also expanded extensively. Similarly, graph-based methods have been shown to be very effective in capturing inherent graph-type relationships in EHRs, leading to improved downstream performance. Although integrating these two families of approaches seems like natural next step, in prac- tice, creating such design is challenging and has not been done. This is partly due to known EHR problems, such as high sparsity, making extracting mean- ingful temporal representations of medical visits challenging. In this study, we propose GT-BEHRT, new approach that leverages temporal visit embeddings extracted from a graph transformer and uses a BERT- model to obtain more robust patient representations, especially on longer EHR sequences. The graph- based approach allows GT-BEART to implicitly capture the intrinsic graphical relationships between medical observations (concepts), while the BERT model ex- tracts the temporal relationships between visits, loosely mimicking the clinicians' decision-making process. As part of our method, also present a two-step pre- training strategy for learning better graphical and temporal representations. We show how such improved representations can ultimately achieve tfth--fftthe-at performance in variety of standard medical predictive tasks, demonstrating the versatility of our approach'.
INTRODUCTION
With the growing availability of electronic health records (EHRs), deep learning methods, capable of disentangling complex patterns EHRs, have been widely used in various healthcare predictive tasks (Choi et al., 2016; Ma al., 2017: Gupta al., 2022b; 2024; Lipton et al., 2017; Choi al., 2018; 2020). Among the numerous approaches presented in recent years in this domain, two families stand out: approaches based the transformer architecture (Vaswani et al., 2017) and those based on graph neural networks (GNNs). Fundamentally, these two families have much in common, as both aim learn better representations from the raw data (from geometric deep learning view (Dwivedi & Bresson, 2021; Bronstein et al., 2021), they have the same roots). However, in practical EHR use cases, the two families have distinguishable characteristics.
On the transformer-based side, most methods adopt an encoder-based architecture, while the decoder-based ones are also rapidly expanding (Lee et al., 2019 Alsentzer al., 2019; Singhal al., 2022). Many of the encoder-based methods adopt the BERT (Bidirectional Encoder Representations from Transformers) model (Devlin et al., 2019) and are the ones achieving ofate-of-the-a art perfor- mance in EHR predictive tasks, like those presented by Pang et al. (2021); Li et al. (2020); Poulain al. (2022); Rasmy al. (2021): Prakash al. (2021); Meng et al. (2021): Li et al. (2021); and Chen et al. (2022). These methods generally conceptualize the medical codes (i.e., the main pieces of EHRs) as analogs to words, patient visits as sentences, and complete patient medical records as documents. This approach generates medical code embeddings while accounting for the temporal structure of EHRs through self-supervised (pre-training) strategies. However, a critical limitation arises in the encoding of the relationships between the visits, as they only weakly represented
Our code is publicly available at httpt:tt////...... http://ggttbb...oc. thyl 1 a Fe/GT-BEHRR

---

Published as a conference paper at ICLR 2024
via the commonly used positional encodings. Furthermore, the association between natural lan- guage words and medical codes is not straightforward. In conventional natural language, sentences inherently comprise ordered sequence of words, while medical codes diagnosed during a single patient visit lack any temporal relationship (as they have been recorded simultaneously). Although in these methods positional tokens remain constant within each patient visit (a weak signal for inter- preting all codes in one visit similarly), considering visits as words rather than sentences is a more conceptually aligned analogy. This visit-level representation views a patient's medical history as an ordered sequence of visits, akin to a sentence with visits as its constituent words.
To generate a format suitable for such a visit-level (versus a medical code-level) transformer ar- chitecture, one can generate visit embeddings using fully connected linear layer, as what Luo al. (2020) and Ren et al. (2021) have done. However, due to the known problems of EHRs, such as noise, missingness, and sparsity, deriving meaningful visit representation through linear layers difficult and often fails to capture robust connections between medical codes. This is precisely where the second prominent family of EHR methods (GNN-based methods) comes in.
On the GNN-based side, EHRs are conceptualized as graphs, by capturing the inherent graphical relationships between various medical codes. Examples of these rich graphical relationships include the (causal or association) relationships between diagnoses and treatments, or between patients and providers. By following graph representation learning approach, medical codes can be viewed as nodes and their relationship as edges. This graphical structure can then be used to derive meaningful patient visit embeddings. A common concern with many previous studies of this type, however, is assuming to have prior knowledge of the structure of medical ontology trees (Shang et al., 2019; Choi et al., 2017a) not considering the temporal patterns between the medical codes (Choi al., 2020; Zhu & Razavian, ;021; Cai et al., 2022; Xu et al., 2023).
This study aims to reconcile these two connected but separate families of approaches. Our study runs on the hypothesis that extracting effective visit-level, graphical, and temporal representations can yield better performance in downstream predictive tasks on EHRs. Going to a coarser represen- tation (from code-level to visit-level) also helps with the quadratic complexity of transformer- based methods Keles et al. (2022), allowing learning longer temporal patterns using our method. Our new method,
(after, graph transformer BERT for EHRs), extends graph transformer archi- tecture to create visit embeddings used as input to a BERT-based encoder. To make such a hybrid design work, propose customized steps for representation learning and pre-training. Specifically, as common pre- pre-training approaches (like Masked Language Modeling) cannot be directly applied to our graph-based input representation, we also propose new pre-training process. We also leverage additional contextual information for the visits (such as the patient's age and the day of the year the visit occurred) to derive more robust time-aware visit embeddings. Consequently, the contributions of our work are as follows:
We present a novel architecture leveraging graph-based time-aware visit embeddings with a BERT-based model to generate patients' representations by extracting both code and visit- level information.
We propose a tailored pre-traine pre-training strategy to better capture the implicit graphical structure of EHR data and the temporal relationships between visits.
We evaluate our method in a series of experiments and achieve state-of-the-att performance on a variety of medical predictive tasks, showing the versatility of the proposed method and its robustness to the choice of downstream predictive task.
2
METHOD
Our proposed method, visit embeddings through graph transformer-hased architecture; then, we use the visit embeddings in transformer encoder
consists of two distinct steps (Figure 1). First, we extract the
encoder to generate patient representations used in downstream tasks.
2.1 INPUT REPRESENTATION
While there exists a wide range of medical information that can be stored in EHRs, in this paper without the loss of generality, we focus three more common elements: diagnoses, medications,
2

---

Published as a conference paper at ICLR 2024
G O
Psalin1
Grach Transformer Layens
Powty
Typezt)
Ageitt
Deyf'
Centerer Scooll
AgA:
Lineer Layer
Day,
6 8.0.000
Ne
Figure 1: GT-BEHRT's architecture. The blue block describes extracting the graph visit embed- dings, where Co. Ci..., C₁ represent four exemplar medical codes associated with the visit & of a fictional patient. The medical codes color-coded to represent different code types (condition, medication, procedure). The graph is based on the medical codes associated to a given visit (visit t here). Each visit goes through this process individually before being fed to the Transformer Encoder.
and procedures (our proposed method can be generalized to other cases). Therefore, we first define D, M, P the set of all observable diagnosis, medication, and procedure codes, respectively. For simplicity, we refer to = D+M + P as the set of all observable medical codes. Let us now define the medical record of patient p sequence of T visits, where, for each visit t T, a list of med- ical codes C? S OO is assigned to p. We initially assume that all medical codes in CP are fully con- nected and implicitly encode the relationship type between two codes (like, medication-medication intiicattoondiagnose
etc.) as a spatial encoding. Because obtaining a single meaningful graph em- bedding directly from the graph node features (graph pooling) is notoriously challenging (Liu al., 2022), we also add a virtual visit node <VST> and connect it to all other nodes in each visit graph with separate edge type. This node will be used as a readout module to obtain graph representa- tion. This specific architecture is particularly useful in settings where the true graphical structure is unknown.
We can now note GP (C, E)) as the graph defining p's th visit, with (ex, Ct, e,) 6 E & where Cs Ct, and e, denote the source, target, and relationship type of the edge, respectively. Therefore, the graph Gf has |CP] + nodes and (1C7)+1) undirected edges. Additionally, we leverage other visit-level contextual information help derive better visit embeddings. Specifically, we extract the visit type (in-patient, outpatient, ICU, ...) as well as the patient's age and the visit day of the year (1-366; help differentiate seasonal discases like the flu from age-related conditions such as Parkinson's 's or Alzheimer's 's diseases). Contextual information is of particular importance when ana- lyzing patient medical records. This process allows us to better mimic the thought process used by medical professionals who would make informed decisions considering the contextual information and the uniqueness of the patient.
2
We can now define V (C. Post. Typel" Age,' Day?') the tuple containing the graph, position of the visit the sequence, visit type, patient's age, and day the year for p's ht visit. Therefore, the medical sequence for each patient p can be formally defined as S, = Vp 1 ,..... Dp. For simplicity, we will drop the subscript p and only consider one patient for the remainder of the paper. We provide summary of the notations used throughout the paper in Table 4.
2.2 BERT WITH GRAPH-BASER VISIT EMBEDDINGS
To learn the hidden graphical structure of EHRs and obtain meaningful visit embeddings, make use of our graph structure with a fully connected transformer encoder with layers. Following the message passing framework proposed by Shi al. (2021), modify the traditional attention mechanism used in the original transformer (Vaswani et al., 2017), and add our edge encoding E to the Value and Key vectors of the neighboring nodes:
3

---

Published as a conference paper at ICLR 2024
A₂ = (. (H,WK + E,
Vd
Φ softmax(A)
Attn, Σ Si (H,WV + E.)WE)
where Wo. WK. Wv, WE Rd are trainable parameters, H is the node features matrix contain- ing the feature vectors for each node in the graph, and Eij is a d-dimensional vector representing the relationship between the nodes and j. Note that, for simplicity, we do not show the bias terms in the linear layers throughout the paper. We repeat the process for each attention head following the original transformer paper (Vaswani al., 2017):
MHA(H, E) = [head]: he.d heady,,,WW,,
where Wo dxxd is a learnable parameter matrix and head, the 7th attention head following Eq 1. We can now use the multi-head attention mechanism to update the features of the nodes for each layer E L. As proposed by Xiong et al. (2020), we perform the layer normalization (Layer (Ba et al., 2016) before the multi-head attention (M AA) and the linear layers (LN), which alleviates the warm-up stage issue and thus reduces the overall training time. The new node features H!{ are calculated following our updated transformer encoder architecture with residual connections:
At' M M H'-1
(2)
H' = LN (Layer
H
with 7 and E heing the nodes and edges feature vectors, obtained from the learned embeddings that map the categorical nodes and edge types to a d-dimensional latent space. We use the representation of the <VST> node of the last Graph Transformer layer (layer L) as a readout module to extract graph-level embedding g HTV for each visit graph. This method allows the <VST> node attend the important nodes the graph and derive meaningful representations. We denote the entire process as GG))), which is a function that takes graph as input and outputs a graph embedding.
This way, for each visit V4 E S, we denote = GC, E)) as the graph embedding resulting from our Graph Transformer. We embed the position, visit type, age, and day of the year in d- dimensional vectors, following the temporal embedding process proposed in CEHR-BERT by Pang et al. (2021). We refer to the resulting embedding vectors as Pos, Type, Age, and Day. We can now concatenate these four vectors with St to create 5dddimensional visit representation and feed it into fully connected layer to map it to a d d-dimensional representation iy:
0t [g:; Post; T Type;; Age:; Day.] W.,
where W, E Rgddd is learnable parameter matrix. We denote this entire process as Graph Visit Embeddings as shown by the blue block on Figure 1.
Following this process, we add another virtual token <CLS> to the visit embedding sequence and generate a sequence of visits representing each patient We feed this resulting sequence to Trans- former Encoder and denote its output sequence as S E (T++)xd with = Transformerfnncoder/ < CLS >, U1, ....... .... UT). Lastly, we use the representation of the <CLS> token for various down- stream medical tasks with a fully connected linear layer and define the loss function of the classifier Le with the binary cross-entropy loss function:
ÿ @S<CLS>> Wc),
(3)
L. = -y yogg(() (1 - y) log(1 ;)
where o is the sigmoid function and W. E dx is learnable vector.
2.3 PRE-TRAINING
To help our model learn better representations of EHRs, we pre-train GT-BEHRT in a two-step self-supervised learning process. We first pre-train the Graph Transformer only with node-level
4

---

Published as a conference paper at ICLR 2024
task, called Node Attribute Masking (NAM) (Hu et al., 2020). We then train the entire model combination of two tasks on graph and temporal level: Missing Node Prediction (MNP) and Visit Type Prediction (VTP). While the first step focuses on learning node embeddings (the predictions are made using the node embeddings only, before any graph pooling operation), the second step allows to leverage the graphical and temporal structure of EHRs, its learning objectives trained using the transformer's s output, thus having access to the other visits in the sequence. We provide a visual description of our proposed pre-training process in Figure 4.
Node Attribute Masking (NAM): Inspired by Masked Language Modeling (Devlin et al.. 2019) and the pre-training strategy proposed Hu et al. (2020) for GNNs, we propose to randomly mask 15% of all nodes in each visit graph. The features representing a selected node would be replaced by those of /MASKJ node, and the Graph Transformer would trained to retrieve the original nodes as shown in Figure 4 (Step 1), where, given an original visit input graph, two nodes are randomly masked (yellow nodes). We connect the node features of the last layer of our Graph Transformer H! to a linear layer and define NAM's loss as:
LNAM = CE(YNAM: WN.NA
where WNAM E Rd1 is a learnable matrix with O being the set of all medical codes as defined earlier, CE(y, x) - > Yo log( is the Cross-Entropy loss function with C' being the number of classes (here, = O.. and yN is the original masked node, As this task is visit-dependent, the node representations Hf do not receive any information on the previous/future visits of the sequence. This allows us to specifically learn better node representation, solely from a node-level task, before further pre-training graph-level tasks, following the promising results in Hu et al. (2020).
Missing Node Prediction (MNP): Additionally, we pre-train our model the graph and temporal levels using MNP. To do this, we randomly remove one node and its corresponding edges from every graph with three more nodes and train our model to predict the removed node:
LANP CE(yMNr, Siitnnr,SWWM.. SWMNP),
where WMNP E [Rdx(O) is a learnable matrix and PMNP is the removed node. Note that, while NAM is node-level task, MNP is a graph-level task performed using the transformer's 's output, after a graph pooling process.
Visit Type Prediction (VTP): Similar to the method proposed in CEHR BERT (Pang et al., 2021), we leverage another learning objective where we randomly mask 50% of the visit types and replace them with a [MASK] token and train the model to retrieve the masked visit types:
LVTP CE(YVTP, SWVTP),
where WyTP E [Rdx(k] learnable matrix with being the number of different visit types, and VvTP is the original visit type. Note that both MNP and VTP (Figure 4, Step 2) are performed using the transformer's output. Therefore, we train the model to predict the removed nodes and the masked visits given the context provided by the other visits in the sequence. The combination of the three tasks a two-step process allows our model to be trained on both node and graph levels, as proved important in Hu et al. (2020), and temporal level, having access to the other visits.
EXPERIMENTS
Data and Predictive Tasks We have evaluated the proposed method on a variety of predictive medical tasks on multiple datasets: in-hospital mortality (Mortality) and prolonged length of stay prediction in ICU (Length of Stay) on the MIMIC-IV dataset (Johnson et al., 2021), a dataset from the Beth Israel Deaconess Medical Center containing hospitalization data, and Heart Failure pre- diction on the All of Us dataset (All of Us Research Program Investigators, 2019), a large publicly- available EHR dataset of adult patients across the US. For all tasks, we extracted relevant information throughout patient visits, including conditions, medications, and procedures, as well as contextual information from the visit. Additionally, we have removed patients with less than 2 visits to leverage the temporality of patients' medical records. For brevity, we present a detailed description of cohort extractions, inclusion criteria, and dataset statistics in Appendix A.2. For preprocessing, adopt
5

---

Published as a conference paper at ICLR 2024
the standard procedure presented by Gupta al. (2022a) for MIMIC, and by Poulain et al. (2022) for All of Us.
The mortality prediction task is binary classification task aimed at predicting the death of patients during their stay in the ICU.
Mortality:
Length of Stay: The Length of Stay task is a binary classification task in which the objective is to predict which patients will stay in the ICU for prolonged period (more than 72 hours). Heart Failure: The Heart Failure task is binary classification task aimed at predicting the occur- rence of heart failure within the next 365 days.
Baselines We present several popular and state-of-the-arr baseline models that use various deep- learning paradigms as reference points, providing comprehensive evaluation landscape and a robust benchmark for our proposed method. These basclines span over several styles of methods suited for EHRs, namely, sequence-based (RNN transformer) and GNN-based methods, allowing us to highlight the efficacy of the proposed method and the importance of combining both paradigms in downstream tasks.
VGNN by Zhu & Razavian (2021) is a variationally regularized encoder-decoder graph network that initially considers all medical codes be fully connected and regularizes the node representations.
НyрEнR by Xu et al. (2023) is GNN-based method that considers EHRs hypergraph where medical codes are represented as nodes and patients as hyperedges, connecting the nodes.
Dipole by Ma et al. (2017) uses bi-directional RNNs with attention mechanisms. is one of the most popular models for EHR predictive tasks, presented before the emergence of transformers.
HiNeN by Luo et al. (2020) is a hierarchical wransfomer-basee model that uses local time em- beddings to provide visit representations. Additionally, it reformulates the self-attention mechanism to associate timesteps with visits.
BEHRT by Li et al. (2020) is a BERT-based model that considers medical codes as words and visits as sentences. Similar BERT, it leverages MLM to pre-train the model.
CEER-BERT by Pang al. (2021) extends BEHRT by incorporating time-aware emehddings. Addi- tionally, it leverages the visit type in an encoder-decoder architecture and pre-trains the model using MLM (encoder) and VTP (decoder).
CEHR-CAN-BERT by Poulain et al. (2022) is semi-supereised model based on CEHR-BERT that incorporated unlabeled data in an adversarial auxiliary task.
For all experiments, we randomly split the dataset into 80/10/10 train/validationntss regime and repeated the process five times, each time with a different random seed. We report the average and standard deviation for each metric (AUROC, AUPRC, and Fl Score). We note that the FI Score was computed with threshold of 0.5, which can be suboptimal for a model not calibrated around 0.5. We provide more information on the computing environment, the implementation of the baselines and the proposed method, and some analysis of the size and training time of each model in Appendix A.3.
Performance on the standard downstream tasks We report the results our initial experiments comparing method and the baselines the three predictive tasks in Table 1. GT-BEHRT con- sistently outperformed all baselines tested on all metrics, with the highest improvements on the Score. Interestingly, HiTANet, visit- wise transformer model, narrowly outperformed BEHRT, code-wise model, showing the potential of such an approach. Furthermore, CEHR-BERE demon- strated the best results amongst all the baselines tested, highlighting the efficacy of the time-aware embeddings and the added contextual information. It is also worth noting that the transformer- methods outperformed other methods, including graph-based models such as VGNN and HypEHR, showing the importance of temporality in understanding medical records. Lastly, on dataset with longer medical sequences (All of Us), visit-wise transformer models (CT-BEHRT and HANeN showed promising results compared to code-wise transformers.
Computational efficiency One major advantage of visit-level transformer model compared to medical code-level one is the reduced length the input sequence. Since the number of codes ob- served per visit C? is greater than or equal to 1, the resulting sequence lengths of GT-BEHRT are ders of magnitude shorter compared to other SERT-basee methods (i.e., BEHRT and CEHR-BERT), shown in Table 5. This shorter length also translates to larger than 70% reduction in training

---

Published as a conference paper at ICLR 2024
Monsality
Langlh of Sty
Failure
Model
AUROC
AUPRC
F1 Scoee
AUROC
AUPRC
Fi Score
AUROCC
AUPRC
P1 Scree
VONN
91.31 0.73
61.87 2.71
55.79
2.26
80.49 1,30
65.95
59.54
2.55
56.38 1.04
Htppörk
91,08 0.333
99.50 a: 1,08
54.95
1.48
79.29
1.18
64.94 1.00
59.11
1.48
$2.41 1.67
55.06
1.21
50.95 2,01
Dipol:
92.24 0.68
90.56 1:30
56.96
1.77
81 + 1.19
70.55+11333
60.04
1.31
56.87
60:36 2.93
55.34
HETANet
92.67 0.3.2
62.46 à: 1.44
2.57
82.04 6: 0.59
0:300
60.46 =: 1.66
91.59 = 1.07
65.29 =: 1.66
59.19 =: 1.41
BERRT
92.07 0.19
60.54 1: 1.66
59.22 2..7
81.80 0.52
69:73 0:99
60.82
1.43
$7.07
61.92
1.64
CEHR BURT
93.31 A 0.23
64.79 à 2.37
82.61 1 0.51
71.36
63.02
0.42
87.37 8: 1.24
63.68 1: 2.82
57.28 2: 2.10
CEHR-GAN-BERT
0.90
62.94
1.34
1.67
GT BEHRT
94.11
0.311
65.55 1.45
14.07:22.113
83.74
73.36
94,37 0.20
73.96
0.83 64.70 0.85
Table I: Comparison of CT-BEERT with the baselines on two standard prediction tasks on MIMIC- IV and one task on All of Us datasets. The boldface indicates the best. The star symbol (*) indicates significant improvements over all other models using two-tailed t-test at p=0.01.
time (and also reduces inference time), which can, again, be explained by the quadratic nature of the attention mechanism of the transformer Keles et al. (2022).
Performance on longer medical histories Due to the improvement in computational efficiency, we expected that, on longer sequences, the ability to access visit embeddings rather than code em- beddings would allow the models to derive better patient representations as it was the case on the All of Us dataset. To evaluate this further, we compare the performance of the baselines and GE-RE by isolating patients who had 10 more visits in the test set (n R 500).
Monality AUPRC
Longth Sty
Modal
AUROC
FI Scure
AUROC
AUPRC
FI Score
VONN
90.87 :: (-0.4)
62.41 2.20
2..1 (-004)
80.03 2: 1.34 (-044)
66.54
3.61 (+0.56) 56)
50:44 :: 2:83 (+0.80)
11y[URR
0.71 (0.57)
54.78 (.78
79.17 1.01
(+0.30)
59.30
(+0.19)
Dipole
89.95 + 0.98 (-2.29)
57.30
(-2.17)
54.11
62.17+ +431+
(40.35)
73,07
(42.5))
63.51 + 3:33 (-2.57)
HITANee
89.51 1 1.57 (-306)
62.26
3.91 (-0.20)
$1.16 1.000 (3347)
50.75
1.51 (-122)
68.08
(:2.99
60.56 1. 2.99 (+0.09)
BEHRT
86,32
1,32 (:.75)
55.87
(.6.7)
49.29 2:49 (-6.5)
79.15
1.7 (:2.5)
60.61
(9911)
56.14 4: 1 (-4.88)
CEHR-DERT
88.88 1: 0.84 (-4.77)
57,31 1: 4.55 (-7.48)
53.22 (-600)
81.23
(-1.38
69.14
5.06 (-2,22)
60.65 1. 5.39 (-2.57)
GAN BERT
88.21 1: 0.58 (5.26)
56.35
4.61 (8.11)
53.79
3:333 (:5288
50.02
0.94 (:2.11)
69.92
2.66 (1132)
60.20 à 3:09 (-2.4)
GT-BEHKT
91.64 0.74 (-2.47)
64.27 3.45 (-1.28)
62.70 2.1 (-1.37)
85.16 1.20 (+1.42)
76.36 (42.800
68.37 2.17 (+2.46)
Table 2: Results on patients with 10+ visits. Mean + Std (difference from Table 1 results).
Table 2 shows that the temporal models leveraging visits as inputs outperform code-wise methods in terms of performance drop or gain from the results on the entire testing set (the numbers in parentheses in the table). Notably, for the Mortality task, -ERR, code-wise model that outperforms Hi TANet on the entire dataset, suffers from much greater performance decrease. A similar pattern can be seen both Dipole and the proposed method, which shows, on average, the lowest performance degradation. While sode-wise methods also mostly decreased on the Lenghh of Stay task, i po le and GT-BEHRT on the other hand show increased performance, highlighting the importance of capturing temporality in complex downstream tasks. We note that, idcally, the number inside the parentheses should be close to O, meaning that the model's performance is agnostic to the length of the patient's 's medical history. Expectedly, VGNN and HyрEHR show an unchanged performance, regardless of the length of the medical history, due to the fact that the models do not consider temporality. Overall, these results show that GT-BEHRE can serve as well-rounded method, capable of achieving state-of-the-art performance on long medical sequences while still outperforming all baselines tested on the entire dataset.
To investigate the performance of GT-BEHRT on different subgroups (group fairness), we report the results of the experiments on the MIMIC-IV dataset for various genders and races/ethnicities. Table 6 shows that most results discrepancies remain within a single standard deviation. We believe fairness mitigation techniques (Wan al., 2021) could improve these results.
Subgroup analysis
To study the effectiveness of each component of GT GT-BEHRT, we conducted an ablation study by removing the key components of our method and comparing the results of each variant. In particular, we evaluated the efficacy of our pre- pre-training strategy and the use of the Graph Transformers to produce the visit embeddings. For the pre- pre-training strategy, we alternatively remove one of the two stages (NAM and VTP+MNP) and compare the results, while we replace the Graph
Ablation study
7

---

Published as a conference paper at ICLR 2024
Transformer (GT) with simple linear layer where the input becomes a simple multi-hot encoding. We present the results of our ablation study in Table 3. Note that, because NAM is designed to work on graphical data and cannot be applied to visit embeddings obtained through linear layers, do not report the variants with the NAM pre-training objective when removing the Graph Transformer component.
Monslity
LengBh
Fillure
AUBOC
AUBRCC
AURRI
AUPRC
V5.03 0.49
56.00
00.55
X2.14 0.45
71.50 1.95
3.67
1.41
68.59 1.29
01.12 1.17
90.63 0.55
56.30 1,51
61.54 1.92
92.70 0.65
71.39 1.00
03.92 1.45
92.83 0.55
70,80 1.92
61,82 d:
X
92.37 = 0.56
61.58 = 1.52
60.39 4. 1.62
81.79 4: 0.90
70:98 1.67
61.26 1.95
92.46 8: 1,30
69.03 1.55
61.28 4: 1.05
03:65
63.72
83.09 0.55
73.01 0.77
63.59
03.30 0.87
1.24
1.79
00.55
08.54 0.30
1:32
0.95
2:333
Table 3: Ablation study. v: Component is present, x: Component has been removed. The boldface and highlighted box indicate the best and the second-best performing variants, respectively.
Several takeaways can be drawn from the ablation study. First, both pre-training strategies seem to have a great impact on the overall results, as showcased by the second, fourth, and fifth rows. However, the impact of pre-training on the model with linear visit embeddings, while showing some improvements, remains mostly within a standard deviation. Second, NAM obtained better results than VTP+MNP on the Morality task which is rather unexpected, as only the Graph Transformer is trained during NAM as opposed to the entire network. However, VTP+MNP outperformed NAM on the Length of Stay task. Lastly, the use of the Graph Transformer without pre-training can impede the model's I performance. However, it benefits greatly from the knowledge acquired during pre-training. more so on average than its linear layer counterpart. It is also worth noting that both non-pre-trained models (first and third row) remain competitive compared to the basclines tested in Table 1, even the ones relying on a pre-training process.
Graph-level interpretation Amongst the many benefits of using graph-based visit embeddings is the possibility to interpret and visualize the model's 's output as well as analyze its mapping of the node to latent space. To understand how the model interprets different medical codes, we plot a 2D) visualization of the latent space of 1,000 most frequent medical codes using t-SNE (van der Maaten & Hinton, 2008) in Figure medical code, it visibly grouped codes of the same type (diagnoses, medications, or procedures) close to each other in the latent space, signaling the model's loose understanding of the medical codes and their purpose. Additionally, Figure 5 b shows heatmap of the attention weights from the graph transformer's last layer for one random visit. One can observe that the attention weights are not sparse, demonstrating that the model has indeed learned meaningful representations for code relationships. The sparsity of learned representations has been shown to be problematic in complex predictive medical tasks using GNNs (Zhu & Razavian, 2021; Choi et al., 2020). Lastly, 5 shows the graph inferred by the model from the attention heatmap, where we removed the edges with an attention score < 0.1. Using graphical structures of EHRs and being able show the connections between the medical codes can be very useful for the end-users (such as providers and practitioners) in interpreting the results of the models and increasing trust in the systems.
Although the model was never explicitly told the nature of each
4 RELATED WORK
Here we discuss a non-exhaustive list of studies that are most relevant to the proposed method.
Transformer-based models for EHRs Following the advances made by the Transformer archi- tecture (Vaswani et al., 2017) in natural language processing (NLP), researchers quickly started adopting it to longitudinal healthcare data. Most studies of this type fall into two subcategories by either considering the visits or the medical codes equivalents to NLP's tokens. Intuitively, the former better leverages the temporal relationships of visits while the latter is better suited to cap- earlier study
ture the relationships between the medical codes. SAnD (Song et al., 2017) that used a ID convolutional layer to obtain visit embeddings and kept the original transformer ar- chitecture. Similarly, Hi TANet (Luo et al., 2020) and RAPT (Ren et al., 2021) both leverage the
8

---

Published as a conference paper at ICLR 2024
transformer architecture and consider the visits the tokens by creating visit embeddings with fully-connected linear layer. RAPT proposed three pre-training tasks to learn robust patient represen- tations, while HiTANet proposed time-aware transformer to model time information in local and global stages without pre-training- One of the first studies that considered medical codes as tokens was the BEHRT (BERT EHR) model (Li et al., 2020). Building on the success of BEHRT, many studies extended this architecture. To name a few, model on prolonged length of stay prediction task, while RareBERT (Prakash et al., 2021) was trained to identify patients with rare diseases in highly imbalanced datasets. CEHR-BERT (Pang et al., 2021) introduced new embedding process using fully-connected linear layer to create temporal embeddings, instead of simply summing all embeddings as in previous BERT-based EHR applications and CEHR-GAN-BERE (Poulain et al., 2022) extended the method for few-shot learn- ing with a generative adversarial network (Goodfellow et al., 2014). BRITM (Meng et al., 2021) added gender embeddings, and BEHRT- CVD Poulain et al. (2021) included race and ethnicity. Re- cently, encoder-based models have also demonstrated very promising performances on both clinical notes Lee et al. (2019); Alsentzer et al. (2019) and structured EHRs Wornow et al. (2023). To the best of our knowledge, only DL-BERT (Chen et al., 2022) has proposed an architecture to capture both the visit-level and code-level relationships by stacking two BERT-based models. However, the method implies a temporal ordering between the medical codes of a given visit.
(Rasmy et al., 2021) pre-trained the
Graph-based methods for EHRs Given the recent advances in graph-based machine learning, many researchers have started investigating the graphical structure of EHRs aim to extract robust relationships within EHRs to allow for better learning of medical representations Li et al. (2022); Fouladvand et al. (2023). A common approach for such studies is to leverage prior domain knowl- edge through ontology trees. Notably, GRAM (Choi et al., 2017a) proposed to learn embeddings for a medical node through a convex combination of the embeddings of the code and its ancestors in an ontology graph. Similarly, G-BERT (Shang et al., 2019) used GNNs by constructing graphs from ontology trees to create code embeddings and fed them to a BERT-based model for each visit indi- vidually, preventing the model from truly learning the temporal information between visits. Also, Sun et al. (2021) combined the medical concept graph with the patient graph for disease prediction, and MiME (Choi et al., 2018) studied the hierarchical structure between diagnoses and prescribed medications. While these methods yielded promising results, they assume prior knowledge of the structure of the ontology tree which can be unrealistic in some real-world scenarios. Consequently, Choi et al. (2020) proposed to use Graph Convolution Transformer by representing a visit as fully- connected graph between all the codes associated with the given visit, while later Zhu & Razavian (2021) used Variationally Regularized GNN (Kipf & Welling, 2016) with similar representation on a petient-level instead of a single visit. Additionally, GNN-based methods can also be used with additional EHR data. For example, ME2Vec (Wu et al., 2021) expanded node2vec (Grover & Leskovec, 2016) to extract embeddings for patients, medical services, and doctors and Park et al. (2022) proposed multimodal approach with GNN to extract the graphical structure of EHR and a language model to handle clinical notes. Though these methods have been effective at capturing the underlying structure of EHRs, none of them studied their temporal dependencies. More closely related to our work, Liu et al. (2020) and Lee et al. (2020) combined GNNs with RNNs to repre- sent static and time-dependent nodes, respectively. While these studies extract temporal information alongside graph-related information, they do not study the structure of a visit to enhance the quality of the visit embeddings, to be used later in a sequence model (such as a transformer or RNN).
CONCLUSION
In this work, we presented GT-PEHRT, a new BERT-based model that leverages Graph Transform- ers that can: 1) capture better graphical structure of EHRs, 2) derive better visit-level representations, 3) learn temporal patterns. Through extensive experiments on a series of standard tasks, we showed that our model outperforms other -fate-of-the-art baselines. Specifically, due to the ability to learn longer temporal patterns, it outperforms other baselines by a wide margin when dealing with longer medical sequences. As a general method agnostic to the choice of downstream tasks, one area for future work would be evaluating our method across more diverse tasks and datasets.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
Our study was supported by NIH awards, P20GM103446 and P20GM113125.
REFERENCES
All Us Research Program Investigators. The All of Us Research Program. New England Journal of Medicine, (81(7)6668-677 2019. PMID: 31412182.
Emily Alsentzer, John Murphy, Willie Boag. Wei-Hung Weng, Jin, Tristan Naumann, and Matthew B. A. McDermott. Publicly available clinical bert embeddings, 2019.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization, 2016.
Michael M. Bronstein, Joan Bruna, Taco Cohen, and Petar Veličkovié. Geometric deep learning: Grids, groups, graphs, geodesies, and gauges, 2021.
Derun Cai, Chenxi Sun, Moxian Song, Baofeng Zhang, Shenda Hong, and Hongyan Li. Hy- pergraph Contrastive Learning for Electronic Health Records, pp. 127-135. 2022. doi: URL https://epult siam. sia...rg//eee...oo.moo www.oorg///o/eeeeeeo 113.713//
14.1177/1.7711111999 1.9781611977172.1
Xianlai Chen, Jiamiao Lin, and Ying An. DL-BERT: time-aware double-level -tRl-style model with pre-training for disease prediction. 2022. doi: 10.1 10.11999iiggaa55556622222.0000.0
Edward Choi, Mohammad Taha Bahadori, Jimeng Sun, Joshua Kulas, Andy Schuetz, and Walter Stewart. Retain: An interpretable predictive model for healthcare using reverse time attention mechanism. In D. Lee, M. Sugiyama, U. Luxburg, Guyon, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 29, Curran Associates, Inc., 2016.
Edward Choi, Mohammad Taha Bahadori. Le Song, Walter F. Stewart, and Jimeng Sun. Gram: braed attention model for healthcare representation learning. In Proceedings of the 23rd ACM SIGKDD International Conference Knowledge Discovery and Data Mining, KDD '17, pp. 787-795, New York, NY, USA, 2017a. Association for Computing Machinery. ISBN 9781450348874, doi: 10.144573099983333088666 URL https: /ttp////cii.//// http:/....rr/011111111.1 1097983.3098122
Edward Choi, Andy Schuetz, Walter F Stewart, and Jimeng Sun. Using recurrent neural network models for early detection of heart failure onset. Journal of the American Medical Informat- ics Association, 24(2):361-3) -300,,1--,,--0 2017b. ISSN 1067-5027. doi: 10.1093/jamia/ocw/ 112. URL
Edward Choi, Cao Xiao, Walter Stewart, and Jimeng Sun. Mime: Multilevel med- ical embedding of electronic health records for predictive healthcare. In S. Ben- gio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Gamett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Asso- ciates, Inc., 2018. URL https: htttps://proceedingg ings neurips. 934b535800bl cballcaafffee wtt.occa6f.f57722111111111111111 Papee .dff.
(paper/2018/Er/
Edward Choi, Zhen Xu, Yujia Li, Michael W. Dusenberry, Gerardo Flores, Yuan Xue, and An- drew M. Dai. Learning the graphical structure of electronic health records with graph convolu- tional transformer, 2020.
Jacob Devlin, Ming-We Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proc. of the Conference of the North American Chapter of the Association for Computational Linguistics, Vol I, pp. 4171-4186, 2019.
Vijay Prakash Dwivedi and Xavier Bresson. A generalization of transformer networks to graphs, 2021.
Matthias Fey and Jan Lenssen. Fast graph representation learning with PyTorch Geometric. In
ICLR Workshop Representation Learning Graphs and Manifolds, 2019.
10

---

Published as a conference paper at ICLR 2024
Sajjad Fouladvand, Federico Reyes Gomez, Hamed Nilforoshan, Matthew Schwede, Morteza Noshad, Olivia Jec, Jiaxuan You, Rok Sosic, Jure Leskovec, and Jonathan Chen. Graph-base clinical recommender: Predicting specialists procedure orders using graph representation learn ing. Journal of Biomedical Informatics, 143:104407, 2023. ISSN 1532-0464. doi: https://doi. 104407. URL https: //wp://wwwww///./// sci encedi rect. wwwccooc//cccoooo com/sciencee attic/e/pp1//1133222223300
lan Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K. Q- Weinberger (eds.), Advances in Neural Information Processing Systems, volume 27. Curran Associates, Inc., 2014.
Aditya Grover and Jure Leskovec. node2vec: Scalable feature learning for networks, 2016.
Mchak Gupta, Brennan Gallamoza, Nicolas Cutrona, Pranjal Dhakal, Raphael Poulain, and Rahma- tollah Beheshti. An extensive data processing pipeline for MIMIC-IV. In Proceedings of the 2nd Machine Learning for Health symposium, volume 193 of Proceedings of Machine Learning Re- search, pp. 311-325. PMLR, 28 Nov 2022a. URL https: https////oooceddeeeeee amm.iirr.s......m..mmm.mmmrmmm.m mlr.. press/ v193 /193//upeee hptp:/pppaaaaaaaaam .tml.
Mchak Gupta, Raphael Poulain, Thao-Ly T. Phan, H. Timothy Bunnell, and Rahmatollah Beheshti. Flexible-window predictions on electronic health records. Proceedings of the AAAI Conference on Artificial Intelligence, 36(11):12500-125
Jun. 2022b. doi: 10. 1009/aai/ w36i11.211110 URL attpps://jjs aaai. org/ index.nix....ii. wwtpnddwwppppp///a//mmppee phtp/AAT/article../ee//e witee//viw///ee view/21520.
Mehak Gupta, Thao-Ly T. Phan, Daniel Eckrich, H. Timothy Bunnell, and Rahmatollah Be- heshti. Reliable prediction of childhood obesity using only routinely collected EHRs is possible. medRxiv, 2024. doi: 10.1101/2024.01 content/ear"
2ttp://2022..11..224990.00 29............. 29.24301945. URL https: /www. www.miidddvvoooooommm org/ 01 29. 43001945.
Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Percy Liang, Vijay Pande, and Jure Leskovec. Strategies for pre-training graph neural networks, 2020.
Alistair Johnson, Lucas Bulgarelli, Tom Pollard, Steven Horng, Leo Anthony Celi, and Roger Mark. Mimic-iv, 2021.
Feyza Duman Keles, Pruthuvi Mahesakya Wijewardena, and Chinmay Hegde. On the computational complexity of self-attention, 2022.
Thomas N. Kipf and Max Welling. Variational graph auto-encoders, 2016.
Dongha Lee, Xiaoqian Jiang, and Hwanjo Yu. Harmonized representation learning dynamic ehr graphs. Journal of Biomedical Informatics, 106:103426, 2020. ISSN 1532-0464. doi: https://doi. 2it.::0000/...i.220.......
URL https: /www rect com/ scinnccence/ arriclee/s1//1333246600003
Jinhyuk Lee, Wonjin Yoon, Sungdong Kim, Donghyeon Kim, Sunkyu Kim, Chan Ho So, and Jae-
Kang. BioBERT: pre-trained biomedical language representation model for biomedical text mining. Bioinformatics, (664)):22444222) sep 2019. doi: 10. 10.00993hoinnfommatessrtz232
URL https: /tt:://ddo..oo org http:///.........ooooooo.aaaaaa.
Michelle M. Li, Kexin Huang, and Marinka Zitnik. Graph representation learning in biomedicine and healthcare. Nature Biomedical Engineering, 6(12): 2022. ISSN 2157-846X. doi:
10.10038/41151--22-00999
Yikuan Li, Shishir Rao, José Roberto Ayala Solares, Abdelaali Hassaine, Rema Ramakrishnan, Dexter Canoy, Yajie Zhu, Kazem Rahimi, and Gholamreza Aalimi-Khorshidi Behrt: Transformer for electronic health records. Scientific Reports, 10(1):7155, 2020. ISSN 2045-2322.
Yikuan Li, Mohammad Mamouei, Gholamreza Salimi-Khorshide Shishir Rao, Abdelaali Hassaine, Dexter Canoy, Thomas Lukasiewicz, and Kazem Rahimi. Hi-behrt: Hierarchical transformer- based model for accurate prediction of clinical events using multimodal longitudinal electronic health records. IEEE Journal of Biomedical and Health Informatics, 27:1 101-1117,7 2021. URL
https: //t:://app //ppi.ppp..//p semanticschol ar. org/ httg://orruss::23559000

---

Published as a conference paper at ICLR 2024
Zachary C. Lipton, David C. Kale, Charles Elkan, and Randall Wetzel. Learning to diagnose with 1stm recurrent neural networks, 2017.
Chuang Liu, Yibing Zhan, Chang Li, Bo Du, Jia Wu, Wenbin Hu, Tongliang Liu, and Dacheng Tao. Graph pooling for graph neural networks: Progress, challenges, and opportunities, 2022.
Sicen Liu, Tao Li, Haoyang Ding, Buzhou Tang, Xiaolong Wang, Qingcai Chen, Jun Yan, and Yi Zhou. hybrid method of recurrent neural network and graph neural network for next-period prescription prediction. International Journal of Machine Learning and Cybernetics, 11(12): 2849-2856, 2020. ISSN 1868-8071. doi: 10.1007/s 13011/10222-000015555 155-xx
Junyu Luo, Muchao Ye, Cao Xiao, and Fenglong Ma. Hitanet: Hierarchical time-aware attention networks for risk prediction on electronic health records. 2020. doi: 10. 144443344468666303007
Fenglong Ma, Radha Chitta, Jing Zhou, Quanzeng You, Tong Sun, and Jing Gao. Dipole: Di- agnosis prediction in healthcare via attention- based bidirectional recurrent neural networks. In Proceedings of the 23rd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, KDD '17, pp. 1903-1911, New York, NY, USA, 2017. Association for Computing Machinery. ISBN 978448884887
Yiwen Meng, William Speier, Michael K. Ong, and Corey W. Arnold. Bidirectional representation learning from transformers using multimodal electronic health record data to predict depression. IEFF Journal of Biomedical and Health Informatics, 25(8)33121-3)
2021.
Chao Pang, Xinzhuo Jiang, Krishna S. Kalluri, Matthew Spotnitz, RuiJun Chen, Adler Perotte, and Karthik Natarajan. TEHR-BERT Incorporating temporal information from structured ehr data to improve prediction tasks. In Proceedings of Machine Learning for Health, volume 158 of Proceedings of Machine Learning Research, pp. 239-260. PMLR, 04 Dec 2021.
Sungjin Park, Seongsu Bae, Jiho Kim, Tackeun Kim, and Edward Choi. Graph-text multi-modal pre-training for medical representation learning. In Gerardo Flores, George H Chen, Tom Pollard, Joyce C Ho, and Tristan Naumann (eds.), Proceedings of the Conference on Health, Inference, and Learning, volume 174 of Proceedings of Machine Learning Research, pp. 261-281. PMLR, / v 117414444
07-08 Apr 2022. URL https: / http:s::/prcccceeeeeeeee pwwppedinceeneeee/oee .omggs...... mlr. www.ss.ppsss..e /
html.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vetito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems 32, pp. 8024 8035.
Curran Associates, Inc., 2019.
Raphael Poulain, Mehak Gupta, Randi Foraker, and Rahmatollah Beheshti. -ransformer-basedd multi-target regression on electronic health records for primordial prevention of cardiovascular disease. In 2021 IEEE International Conference on Bioinformatics and Biomedicine (BIBM), pp- 726-731, 2021.
Raphael Poulain, Mehak Gupta, and Rahmatollah Beheshti. Few-shot learning with semi-supervised transformers for electronic health records. In Proceedings of the 7th Muchine Learning for Healthcare Conference, volume 182 of Proceedings of Machine Learning Research, pp. 853- 873. PMLR, 05-06 Aug 2022.
P. K. S. Prakash, Srinivas Chilukuri, Nikhil Ranade, and Shankar Viswanathan. Rarebert: Trans- former architecture for rare disease patient identification using administrative claims. In AAAI, 2021.
Laila Rasmy, Yang Xiang, Ziqian Xie, Cui Tao, and Degui Zhi. Med-bert: pretrained contextualized embeddings on large-scale structured electromic health records for disease prediction. npj Digital Medicine, 4(1), 2021. ISSN 2398-6352.
Houxing Ren, Jingyuan Wang, Wayne Xin Zhao, and Ning Wu. Rapt: Pre-training of time-aware transformer for learning robust healthcare representation. 2021. doi: 10. 40.1144444454444677766
12

---

Published as a conference paper at ICLR 2024
Junyuan Shang. Tengfei Ma, Cao Xiao, and Jimeng Sun. Pre-training of graph augmented trans- formers for medication recommendation. arXiv preprint rrXiv:1966000344 2019.
Yunsheng Shi, Zhengjie Huang, Shikun Feng, Hui Zhong, Wenjing Wang, and Yu Sun. Masked label prediction: Unified message passing model for semi-supervised classification. In Zhi-Hua Zhou (ed.), Proceedings of the Thirtieth International Joint Conference on Arlificial Intelligence, IJCAI-21, Pp. 1548-1552 International Joint Conferences on Artificial Intelligence Organiza- tion, 8 2021. doi: 10. 2244963/jcaii22222222 URL https: :/tp::://dci.o /10 24 24963/1jeaaa 2021 2/1 4. Main Track.
Karan Singhal, Shekoofeh Azizi, Tao Tu, S. Sara Mahdavi, Jason Wei, Hyung Won Chung, Nathan Scales, Ajay Tanwani, Heather Cole-Lewis, Stephen Pfohl, Perry Payne, Martin Seneviratne, Paul Gamble, Chris Kelly, Nathaneal Scharli, Aakanksha Chowdhery, Philip Mansfield, Blaise Aguera Arcas, Dale Webster, Greg S. Corrado, Yossi Matias, Katherine Chou, Juraj Gottweis, Nenad Tomasev, Yun Liu, Alvin Rajkomar, Joelle Barral, Christopher Semturs, Alan Karthikesalingam, and Vivek Natarajan. Large language models encode clinical knowledge, 2022.
Huan Song, Deepta Rajan, Jayaraman J. Thiagarajan, and Andreas Spanias, Attend and diagnose: Clinical time series analysis using attention models, 2017.
Zhenchao Sun, Hongzhi Yin, Hongxu Chen, Tong Chen, Lizhen Cui, and Fan Yang. Disease pre- diction via graph neural networks. IEEE Journal of Biomedical and Health Informatics, 25(3): 818 826, 2021. ISSN 2168-2194. doi: 10. 10.1 1099//hhi.2000.30441444
Laurens van der Maaten and Geoffrey Hinton. chine Learning Research, 986)):557--665,, 2008.
Visualizing data using t-sne. Journal of Ma- URL http: jmlr. wwt.org////o.o.o. wwggpapprs/vv//
mandermaatenio html.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E. ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In L. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Asso- ciates, Inc., 2017. URL https: http://:////.ooccigggg..so neurips cc wttpriiss..ppppeeeeeeee pppppp/pp1///iilll. pdf.
3f5ee243547dee 1 1
Mingyang Wan, Daochen Zha, Ninghao Liu, and Na Zou. Modeling techniques for machine learning fairness: survey, 2021. URL httLtt
/211 03015.
Michael Wornow, Rahul Thapa, Ethan Steinberg, Jason Fries, and Nigam Shah. Ehrshot: An ehr benchmark for few-shot evaluation of foundation models, 2023.
Tong Wu, Yunlong Wang, Yue Wang, Emily Zhao, and Yilian Yuan. Leveraging graph-based based hi- erarchical medical entity embedding for healthcare applications. Scientific Reports, 11(1), 2021. ISSN 2045-2322. doi: 10. (038/s415988 1088-02---525--
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In Hal Daumé III and Aarti Singh (eds.). Proceedings of the 37th International Conference on Machine Learning, volume 19 of Proceedings of Machine Learning Research, pp. 10534- PMLR, 13-18 Jul 2020. URL htwps:ttpt..... httpl:::///rooeeeini /pwppeddeneggssss//////// ml www.mmrrr.s.......m.mmmmmmmmm..e / 19 x ong html.
Ran Xu, Mohammed K Ali, Joyce C Ho, and Carl Yang. Hypergraph transformers for EHR-based clinical predictions. AMIA Summits Transl. Sci. Proc., 2023:582-591, June 2023.
Chaoqi Yang, Zhenbang Wu, Patrick Jiang, Zhen Lin, Junyi Gao, Benjamin Danek, and Jimeng Sun. PyHealth: A deep learning toolkit for healthcare predictive modeling. In Proceedings the 27th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining (KDD) 2023, 2023. URL https: /ttpttgtttttt////////////tc.hhtt com/ http////naauuuccyyeeeeeleee
Weicheng Zhu and Narges Razavian. Variationally regularized graph-based representation learn- ing for electronic health records. In Proceedings of the Conference on Health, Inference, and
13

---

Published as a conference paper at ICLR 2024
Learning. CHIL '21, pp. 1-13, New York, NY, USA, 2021. Association for Computing Machin- ery. ISBN 9781450383592. doi: 10.1145/2450439.44511 3451559533. URL htt.tt/s////// //doi. otg/ 1145/3450439 3451855.
A DETAILS FOR NOTATIONS, DATASET AND IMPLEMENTATIONS
A.1 NOTATIONS
For clarity, we include in Table 4 reference table to the notations used throughout the paper.
Table 4: Reference table to the notations used the paper.
Notation
Description
Set unique diagnoses codes D
M
Set of unique medication codes
P
Set of unique procedure codes
O
Set of unique medical codes O = D M++++
VP
Patient 1th visit
C²
Set of observed codes at patient p's tt visit Graph representations of C," Patient medical sequence
(C,,
8
Sp
d
Embedding size
L
Number of Graph Transformer layers
A.2 DATASET
The MIMIC-IV dataset Johnson et al. (2021) has been extensively used within the machine learning community. Among many possible downstream tasks Gupta et al. (2022a), the mortality in the ICU and ICU length of stay tasks are often used as benchmarks in the field. For that reason, we have extracted a cohort of patients who have been admitted the ICU using the MIMIC-IV dataset and evaluated our models in the two aforementioned tasks. Specifically, we define the mortality task as a binary classification task where our objective is to predict the patients who will die during their stay in the ICU, that is. between admission time and discharge time. Similarly, the length of stay task is designed as a binary classification task in which the goal is to predict patients who will have "prolonged" (> 3 days) length of stay in the ICU. Consequently, records of patients from the hosp module (non-ICU data) prior to admission to the ICU. That is, extracted the conditions, medications, and procedures codes, together with additional contextual visit information (patient's age, visit type, and day of the year) that appeared before the patient's admission the ICU, To leverage the temporality of patient data, we only include patients with two or more visits. Finally, this process results in cohort of 22,376 patients. Additionally, we have also extracted pre-training cohort of 76,253 patients. We provide further information on the cohort statistics in Figure 2.
have extracted all medical
Similar to Choi al. (2017b); Poulain et al. (2022), we evaluate our method on heart failure prediction task in patients between 18 and 85 years of age with the All of Us dataset. Specifically, our objective is to predict the occurrence of heart failure within the next 365 days using the past years of medical data as input. To prevent information leakage, we predict the first occurrence of heart failure for patients with multiple heart failure events. Furthermore, for each positive patient, randomly include up 10 control patients with similar ages and used the same time windows for all patients. This process results in cohort of 29.736 patients and a pre- vreitraining cohort of 95,358 patients. We provide further information on the cohort statistics in Figure 3.
14

---

Published as a conference paper at ICLR 2024
8000
500
1250
6000
400
Count
Count
Count
300
4000
200
2000
20
go
20
AP
60
80
100
100
200
Visit per patient
Age Admission
Unique per patient
(a)
(b)
(c)
Figure 2: Key statistics the MIMIC-IV dataset. (a): Distribution of the number of visits per
patient. (b): Distribution of the patient's s age at admission to the ICU, (c) Distribution of the number of unique medical codes per patient.
5000
800
2500
4000
2000
Count
Count
600
Count
1500
400
2000
1000
1000
500
0
AP
20
AP
80
100
200
# Visit per patient
Age at Prediction
# Unique codes patient
(b)
Figure 3: Key statistics of the All of Us dataset. (a): Distribution of the number of visits per patient. (b): Distribution of the patient's age at time of prediction. (c) Distribution the number unique medical codes per patient.
A.3 IMPLEMENTATIONS
We have used PyTorch Paszke et al. (2019) to implement all the baselines using their respective public codebases. Additionally, we have used PyHealth's 's Yang et al. (2023) implementation of the Dipole Ma et al. (2017) model. We implemented the graph component of GT BE-BBERTT using PyTorch Geometric Fey & Lenssen (2019) and all the experiments were run on an NVIDIA T4 GPU. For each baseline, as well as the proposed method, we have determined the optimal hyperparameters through grid search. We report a complexity analysis the models as well as the hyperparameters used throughout the experiments in Table 5. Note that, while VGNN (Zhu & Razavian, 2021) is graph-based method, the official code is not written using a Py Torch graph library, therefore having significantly slower runtimes.
VONN
HypllH
Dipolo
IRTANe
CEHR-BBET
CEHR-GANNBERY
GE-BEIK
Total number irantable porameterss Training Time per 1,000 patienss (s)
7.3M
270k
9.M
11.7M
11,00
12.74
13,99
1377M
40.31 1.55
1.8 à 0.25
2.59 a 0.7
4.09 0.56
25.23 8 0.93
2721 0.96
41 a 0.6
765 1.00
Inforencr Time 0000 pationss ts
1375 1.11
1.36 0.27
0.99 003
1,89
0.13
8.21 2.1:000 0.44
8.77 0.6
9.73 a 041
4.39 1998000 0.51
Sogsenee Lengh:
4.00 5.8
4.00
4.00 5.08
100.98 81.77
108.39 54.73
108.39 64.73
4.90 5,08
Number of Graph Layers
Number of Oraph Amerttion Head
Oraph
Transfomee Trarsfomer Transformer
512
Table 5: Running time statistics for each baseline and GT BEHRT with their respective hyperpa- rameters. "_" means that it is non-applicable the given method
15

---

Published as a conference paper at ICLR 2024
B PRE-TRAINING
To ensure the model acquires meaningful and transferable knowledge through pre-training, we first pre-train our graph
node-level, that is, the task based on the nodes themselves opposed to the entire graph. Then, we complete the pre-training process on a graph and temporal level. The tasks, realized after the transformer network, now involve the graph embeddings (as opposed the nodes) and have full access of the context surrounding them (future/past visits). This enables us to teach the model derive meaningful visits embeddings through the graph, and the temporality of EHRs with the transformer encoder.
?
Graph Pooling
Visit Type Prediction
Node Attribute Masking
Missing Node Prediction
Input Graph
Step
Step 2
Figure 4: Our two steps pre-training process with a given input graph (left). The code notation and color coding the same as the ones used in Figure The first step randomly masks nodes (yellow nodes) to be retrieved by the model. The second step randomly removed single node and the model trained retrieve the missing node while predicting the type of visit. Step solely pre-trains the graph transformer through Node Attribute Masking. The second step pre-trains the entire model with Missing Node Prediction and Visit Type Prediction.
C ADDITIONNAL EXPERIMENTS
Mortality
Length Stay
Model
AUROC
AUPRC
Fl Score
AUROC
AUPRC
Fl Score
Male
94.35 1 0.63
66.34
2.99
65.55 1: 2.89
84.34 d: 1.34
75.43 2.44
67.09 + 1.32
Fermaee
93.75 + 0.82
64.59 2.74
63.09 4: 2.55
82.86 0.35
71.71
0.86
63.69 +: 1.11
White
93.85 +: 0.32
64.91 +: 2.69
63.74 1: 1.69
83.55 +: 0.56
74.09 1: 0.70
65.39 4: 1.44
Black/African American
0.75
65.37 2.90
3.43
84.96
2.42
75.54
67.43 2.73
Hispanic/Latinn
95.89 + 0.87
68.45 4: 3.64
68.41 +: 2.36
84.08 +: 2.83
73.18 :: 2.00
63.74 +: 2.77
0.85
69.25 2.44
68.67 4.03
81.58
70.81
+ 3.03
Other
95.19 + 0.53
65.24 + 3.32
64,32 + 3.62
83.73 + 2.49
73.12 + 2.34
66,25 + 1.50
Aggregate
94.11 0.31
65.55 1.65
64.07 2.13
83.74 0.51
73.36 1.19
65.91 0.79
Table 6: Comparison of GT-BEHRT's results on the MIMIC-IV dataset across different demo- graphic subgroups.
16

---

Published as a conference paper at ICLR 2024
Karnatutt caldaa
Disbetas
Hyper
02
Diagnoses
Medications
Lovedly sodium
Procedures
0.1
2
4
Sbacody!
3.
(a)
(b)
Nodel
Hypathyroictoo
#sacody
Lovotay oon sod
(c)
Figure 5: (a): 2D visualization of the medical codes embeddings space using t-SNE. (b): Heatmap of the attention weights from the graph transformer's last layer for randomly selected visit. (c): Inferred graph from the attention weights displayed in (b).
17


Reviews:
[{"review_id": "W2JxvHeHgM", "reviewer": "Reviewer_74P1", "summary": "This paper presents GT-BEHRT, an innovative approach that integrates graph-based and transformer-based methodologies to enhance the analysis and predictive accuracy of electronic health records (EHRs). The technique specifically addresses the sparsity in EHR data and the need for capturing complex, graph-type relationships. GT-BEHRT combines graph transformer-derived embeddings for individual visits with a BERT-based framework, facilitating richer patient representations over extended EHR sequences. This method also features a novel two-step pre-training process that further refines the model’s capacity to decode both graphical and temporal patterns in the data. As a result, GT-BEHRT achieves leading performance across diverse medical predictive tasks, indicating its robustness and versatility.\n\nMain Contributions:\n\nGT-BEHRT Design: A new hybrid model that integrates graph transformer and BERT-based architectures to better handle the temporal and graphical nature of EHR data.\n\nTwo-Step Pre-training Strategy: A unique pre-training process that enhances the model's ability to understand complex relationships in EHRs, improving performance on predictive tasks.\n\nSuperior Predictive Performance: Through its innovative approach, GT-BEHRT sets a new benchmark for state-of-the-art results in various standard medical predictive tasks, demonstrating its potential to significantly impact healthcare analytics.", "questions": "Why did you only used one dataset to test you model?\nHave you looked at Michele Moore works?\nWhy you didn't compare your results with some of the state of the art works such as GMAI?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The authors present a commendable effort in intertwining GNN and Transformer methodologies, showcasing an innovative approach to a topic gaining traction in the field. Their literature review is generally thorough, but it's surprising to note the omission of recent contributions from Jure Leskovec's lab, which bear resemblance to this work. While the essence of the paper is intriguing, a deeper exploration of their model's specifics would have provided more comprehensive insights. Overall, the strength of this paper lies in its novel perspective on a burgeoning issue, even if there are areas left to be further elaborated.", "weaknesses": "The paper, while advancing an intriguing new architecture, does not fully address the breadth of state-of-the-art works within the Graph Neural Network (GNN) sphere. A more exhaustive acknowledgment and discussion of leading GNN research would have provided a richer context for the authors' contributions. Additionally, while the authors assert that their proposed architecture outperforms existing models, the paper falls short in offering visual depictions of the architecture. Such illustrations are crucial for readers to fully grasp the design and the innovations it purports to bring. Furthermore, the rationale behind the selection of certain models as baselines is not sufficiently elucidated. This lack of detailed justification and visual support may leave the reader questioning the thoroughness of the comparative analysis and the foundation of the authors' claims. The paper would benefit from more comprehensive visual materials and a deeper discussion on model selection to solidify its standing within the current scientific discourse.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qhMqbizG52", "reviewer": "Reviewer_omW6", "summary": "This paper focused on patient representations and developed a tailor transformer architecture leveraging both graph transformer and BERT-like encoder only transformer. Evaluation on MIMIC-III and eICU showed improved performance on two tasks, mortality and length of stay.", "questions": "What do authors would like to convey in the title \"BETTER REPRESENTATION IMPROVES DOWNSTREAM PERFORMANCE\"\nWhat's the definition of better and how to obtain it?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is well written and the developed method is easy to follow.\n\nThe architecture sounds reasonable to me.", "weaknesses": "Baselines are quite old. The most recent baseline is 2021.More literature research is necessary. For example,  [1,2,3] have dome similar things. Consequently, I am not convinced by the advantage of bringing GNN into BERT, \n\n[1] Hypergraph Transformers for EHR-based Clinical Predictions\n[2] EHRSHOT: An EHR Benchmark for Few-Shot Evaluation of Foundation Models\n[3] Unsupervised pre-training of graph transformers on patient population graphs", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ppqmwdilni", "reviewer": "Reviewer_efXj", "summary": "The authors presented an approach towards extending graphical neural networks to temporal domain using transformers for modeling EHRs. Past research has shown the importance of modeling EHR modalities as GNN that can better capture the inherent correlation better than a flat data structure. However, nuance of EHR data, especially around longitudinal aspects, makes it important to also capture the temporal dynamics. The authors compared their proposed method with several baselines and reported strong results.", "questions": "See above", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Some of the key strengths of the paper are as below\n- the authors have proposed a hybrid architecture combing GNN and transformer architectures to capture both the spatial and temporal dynamics of EHR data. Key contributions around this proposed architecture are around identifying some of the issues for extending GNN to temporal domain and propose a multi stage pre-training method to capture the dynamics accurately\n- The authors compared their methods against strong baselines and reported significant performance improvements\n- they key insights around the model being able to capture longer medical histories is very interesting and adds makes the model more applicable\n- The authors have also added ablation studies to capture the importance of their training strategy", "weaknesses": "The paper can be improved upon by addressing the following aspects\n - The method description can be improved upon. Consider adding an illustration of the training strategy and describing the methods using the illustrations\n- the authors can also consider adding sub-group analysis to further strengthen the claims around model performance\n- While being cognizant of the page limits, it would have been interesting to analyze some of the inferred graphical networks at individual example level", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QuXagndXKc", "reviewer": "Reviewer_CybP", "summary": "The paper proposes a BERT-based model to represent EHR with a graph-based time-aware visit embedding to better capture the implicit graphical structure of EHR data and the temporal relationships between visits. The paper shows better performances in multiple tasks on the MIMIC-IV dataset.", "questions": "- How is the F1 score calculated? Which threshold is picked? Some experiments with similar AUPRC result in substantially different F1 scores.\n\n- What is the difference between NAM and MNP? It seems they can be merged into one loss.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper demonstrates the importance of temporal effects in the EHR. This is insightful for research on EHR.\n\n- Comprehensive studies are conducted including various tasks and ablation settings. The confidence intervals are provided.", "weaknesses": "- The proposed method is a combination of building blocks originated from several previous papers, e.g. graph transformer, BEHRT. It is a good application paper that utilizes all these aspects, but the methodological impact is limited.\n\n- In the ablation study section, the author claims that \"both pre-training strategies seem to have a great impact on the overall results, as showcased by the second, fourth, and fifth rows.\". However, in fact, Table 3 shows the performance differences between experiments with and without NAM, MNP, and GT are marginal. The confidence interval of the last row and the first row (the simplest version) even overlap. The simple linear model seems to already perform very well compared to GT. A better way of justifying the impact of pretraining can be directly using linear probing on the frozen pre-trained representation.\n\n- The power of pretraining seems to be limited by the size of the dataset (around 20K patients). Under this condition, pretraining a large model may not significantly outperform a simple linear model. It will be more exciting to apply the proposed method to a larger dataset.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)