SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zhenbang Wu 1 Huaxiu Yao 2 David M Liebovitz 3 Jimeng Sun 1
<sup>1</sup> University of Illinois Urbana-Champaign, {zw12, jimeng}@illinois.edu
<sup>2</sup> University of North Carolina at Chapel Hill, huaxiu@cs.unc.edu
<sup>3</sup>Northwestern University, david.liebovitz@nm.org
Deep learning models have been widely used to assist doctors with clinical decision-making. However, these models often encounter a significant performance drop when applied to data that differs from the distribution they were trained on. This challenge is known as the domain shift problem. Existing domain generalization algorithms attempt to address this problem by assuming the availability of domain IDs and training a single model to handle all domains. However, in healthcare settings, patients can be classified into numerous latent domains, where the actual domain categorizations are unknown. Furthermore, each patient domain exhibits distinct clinical characteristics, making it sub-optimal to train a single model for all domains. To overcome these limitations, we propose SLDG, a self-learning framework that iteratively discovers decoupled domains and trains personalized classifiers for each decoupled domain. We evaluate the generalizability of SLDG across spatial and temporal data distribution shifts on two real-world public EHR datasets: eICU and MIMIC-IV. Our results show that SLDG achieves up to  $11\%$  improvement in the AUPRC score over the best baseline.
Deep learning techniques have been increasingly popular in clinical predictive modeling with electronic health records (EHRs) [12, 11, 47, 58, 2]. However, these models typically assume that the training (source) data and testing (target) data share the same underlying data distribution (i.e., domain). This assumption can become problematic when models are applied to new domains, such as data from different hospitals or future time points [17, 59, 20, 37]. In these situations, domain shifts caused by variations in patient cohorts, clinical standards, and terminology adoption can significantly degrade the model's performance.
This paper aims to develop a clinical predictive model on the source data that effectively handles potential domain shifts when applied to the target data. Domain generalization (DG) [7] methods have been widely utilized to address such problems, including techniques like domain alignment [32, 24, 25, 45, 31, 51, 62], meta-learning [23, 27, 26, 5, 22, 29], and ensemble learning [9, 42, 43, 61]. However, when applied in healthcare settings, these methods encounter the following limitations:
- Reliance on domain IDs. Most DG methods depend on the presence of domain IDs, which indicate the domain to which each sample belongs, to guide the model training [24, 16, 4, 9, 42]. However, as shown in Fig. 1, patients can be divided into numerous latent domains based on features such as age, medical history, treatment, and symptoms. The actual categorization of these latent domains can be difficult to obtain and vary across different tasks [1, 49]. Consequently, existing DG methods often resort to broad domain categorizations, such as hospital or timestamp, which provide limited information [59, 17].
Figure 1: Conventional domain generalization methods typically rely on domain IDs and shared characteristics across domains to train a single generalized model. However, in the medical field, patients can be classified into numerous latent domains that are not directly observable. Additionally, each patient domain exhibits unique clinical characteristics, making it sub-optimal to train a single model for all domains.
- **Attempt to train a single model.** While some recent DG methods have attempted to alleviate the reliance on domain labels [62, 29], they try to train a single model that generalizes across all domains. However, patients from different domains possess distinct characteristics and require different treatment approaches [2, 58]. For example, as shown in Fig. 1, fever is considered a normal symptom for patients with viral infections as it helps stimulate the immune system. On the other hand, it can be a bad signal for patients with cardiovascular disease, leading to complications. Thus, training a single model for all domains is challenging and can lead to sub-optimal performance.
To overcome these limitations, we propose SLDG, a self-learning framework for domain generalization that iteratively discovers decoupled domains and trains customized classifiers for each discovered domain. Specifically, SLDG consists of the following iterative steps:
- Decoupled domain discovery. While domain labels are not initially available, we posit that they can be recovered by clustering the learned latent representations. However, identifying all fine-grained domains across various clinical features (e.g., demographics, diagnosis, and treatments) can be challenging. Instead, we propose to decouple these clinical features and discover the clusters separately for each type of feature. To achieve this, we maintain a distinct latent space for each type of features using a feature-specific patient encoder. Within each latent space, we perform hierarchical clustering independently to discover the domain categorizations. By adopting this approach, we effectively reduce the number of domains from exponential to linear to the number of feature types.
- Domain-specific model customization. To account for the unique characteristics of patients in different domains, our approach involves training customized classifiers for each domain. To ensure parameter efficiency, we extract domain representations from the learned clusters and utilize them to parameterize the domain-specific classifiers. For a given patient, we determine the closest domain by comparing the patient's representations with the domain representations, and subsequently select the corresponding classifier for accurate inference.
To assess the generalizability of SLDG across spatial and temporal shifts, we conduct experiments on two publicly available EHR datasets: eICU [38] and MIMIC-IV [18]. Our results demonstrate that SLDG outperforms the best baseline by up to  $11\%$  in terms of AUPRC score. We also conduct detailed analyses and ablation studies to investigate the factors contributing to the performance gain achieved by SLDG.
Figure 2: An illustration of the SLDG framework. The feature-specific patient encoder maps each patient into multiple latent spaces, with each space capturing patient characteristics from a specific perspective. Next, SLDG iteratively performs decoupled clustering to identify latent domains and learns domain-specific classifiers customized for each domain.
In EHR data, a patient's hospital visit is represented by a sequence of events, denoted as  $x = [e_1, e_2, \ldots, e_m]$ , where  $m$  is the total number of events in the visit. Each event  $e$  characterizes features of a certain type  $t$ , such as diagnosis, prescription, and lab tests. This mapping is denoted by a function  $T(e): \mathcal{E} \to \mathcal{T}$ , where  $\mathcal{E}$  and  $\mathcal{T}$  denote the sets of all events and types, respectively. For example, a patient visit can be [Acute embolism (I82.40), Atrial fibrillation (I48.91), Ultrasound (76700), CT scan (G0296), ECG (93042), Heparin IV (5224)]. And each event corresponds to a specific type (e.g., diagnosis, procedure, medication). The main objective of clinical predictive modeling is to predict the occurrence of future events, such as 15-day hospital readmission and 90-day mortality, denoted as  $y \in \{+, -\}$ , based on the patient's current visit  $x$ .
Existing clinical prediction works typically train a model  $f_{\phi}(\cdot)$  with parameter  $\phi$  by minimizing a loss function  $l(\cdot)$  on source training data sampled from distribution  $P_{tr}$ , as in Eq. (1),
$$
\underset {\phi} {\arg \min } \mathbb {E} _ {(x, y) \sim P _ {t r}} [ l (f _ {\phi} (x), y) ], \tag {1}
$$
with the hope that the trained model can perform well on the target test data distributed according to  $P_{te}$ . However, in real-world settings, the source and target distributions can differ due to spatial and temporal shifts, i.e.,  $P_{tr} \neq P_{te}$ . Consequently, the model trained on source data may experience a drop in performance when applied to the target data.
3 The SLDG Approach
In this paper, our goal is to train a model  $f_{\phi}(\cdot)$  on source data  $P_{tr}$  that can generalize to target data  $P_{te}$  despite potential domain shifts. Existing DG algorithms face limitations due to their reliance on domain IDs and attempts to train a single model for all domains. To overcome these limitations, we propose to iteratively discover latent domains and train customized classifiers for each domain. However, we face the challenge of dealing with a large number of latent domains, which not only makes domain discovery difficult but also results in an exponential increase in the number of model parameters with respect to the number of feature types. In the following sections, we will describe how our method SLDG addresses this challenge through decoupled domain discovery and domain-specific model customization. Additionally, we will introduce the training and inference strategy. Fig. 2 illustrates the SLDG framework.
3.1 Decoupled Domain Discovery
Although domain labels are not initially available, we hypothesize that the domain information is encoded in the learned latent representations and can be recovered with the clustering technique. How-
ever, patients can be categorized into thousands of latent domains, determined by various features such as age, medical history, treatment, and symptoms. For instance, a patient can fall into the fine-grained domain of older male patients with a history of smoking and a diagnosis of type 2 diabetes. Identifying all such fine-grained domains can be challenging, as clustering methods may either overlook smaller domains or result in an excessive number of domains that would inflate the number of parameters in subsequent steps. To address this, we propose decoupling these clinical features and independently discovering clusters for each feature type. For example, the patient above can simultaneously belong to the decoupled domains of older, male, history of smoking, and diagnosis of type 2 diabetes. This approach effectively reduces the number of domains from exponential to linear with respect to the number of feature types.
Concretely, we maintain a distinct latent space for each type of feature. When given an input patient visit  $x$ , SLDG maps it to the latent space corresponding to the feature type  $t \in \mathcal{T}$  using a feature-specific patient encoder  $E_{t}(\cdot)$ , as in Eq. (2),
$$
\mathbf {h} _ {t} := E _ {t} (x), \quad \mathbf {h} _ {t} \in \mathbb {R} ^ {h}, \tag {2}
$$
where  $h$  denotes the hidden dimension. Next, within each latent space of type  $t$ , SLDG gathers all patient representations  $\{\mathbf{h}_t^{(i)}\}_{i=1}^{N_{tr}}$ , where  $N_{tr}$  is the number of source training data, and performs clustering to discover the domain categorizations, as in Eq. (3),
$$
\mathbf {M} _ {t} := \operatorname {C l u s t e r} \left(\left\{\mathbf {h} _ {t} ^ {(i)} \right\} _ {i = 1} ^ {N _ {t r}}\right), \quad \mathbf {M} _ {t} \in \{0, 1 \} ^ {N _ {t r} \times K _ {t}}, \tag {3}
$$
where  $K_{t}$  represents the number of discovered domains in the latent space of type  $t$ .  $\mathbf{M}_t$  denotes the learned domain assignment, where  $\mathbf{M}_t[i,k]$  is equal to one if and only if (i.f.f.) the patient  $x^{(i)}$  is assigned to the  $k$ -th domain. We will describe this procedure in detail in the following.
Feature-Specific Patient Encoding. This module is responsible for mapping each patient into multiple latent spaces, each capturing the patient's health status of a specific feature type. This enables subsequent modules to decouple the representations of different feature types. For a patient's hospital visit  $x$  with a list of events  $[e_1, \dots, e_m]$ , SLDG computes the contextualized representation for each event by applying the embedding function  $E(\cdot)$ , as in Eq. (4),
$$
[ \mathbf {e} _ {1}, \dots , \mathbf {e} _ {m} ] = E ([ e _ {1}, \dots , e _ {m} ]), \quad \mathbf {e} _ {j} \in \mathbb {R} ^ {h}, \tag {4}
$$
where  $\mathbf{e}_j$  is the contextualized representation for event  $e_j$  with dimension  $h$ . We model  $E(\cdot)$  using a three-layer Transformer [48] framework. To ensure that there are no unseen events in the target data, we initialize the event embedding look-up table with ClinicalBERT [3] embeddings of the event name and then project it down to our hidden dimension of size  $h$ . The embedding look-up table is fixed during training.
Next, SLDG aggregates the contextualized event representations  $[\mathbf{e}_1,\dots ,\mathbf{e}_m]$  based on their types, such as family history, diagnosis, and treatments. For each type  $t\in \mathcal{T}$ , the type-specific representation  $\mathbf{h}_t$  is computed by averaging the representations of all events of that type, as in Eq. (5),
$$
\mathbf {h} _ {t} = \operatorname {A v e r a g e} \left(\left\{\mathbf {e} _ {j} \mid T \left(e _ {j}\right) = t _ {j = 1} ^ {m}\right), \mathbf {h} _ {t} \in \mathbb {R} ^ {h}, \right. \tag {5}
$$
where  $T(e_i)$  indicates the type of event  $e_i$ . If no events belong to a certain type, the pooled sequence representation is used as a substitute. Consequently, each patient's hospital visit is represented by a set of vectors  $\{\mathbf{h}_t\}_{t \in \mathcal{T}}$ , with each vector capturing the patient's health status from a specific type of events. These decoupled patient representations are then utilized to perform per-feature-type domain clustering, described next.
Hierarchical Domain Clustering. This module is responsible for clustering patient representations in each latent space to discover latent domains, enabling subsequent modules to customize the classifier for each domain. In the previous step, we obtain a set of patient representations  $\{\mathbf{h}_t^{(i)}\}_{i = 1}^{N_{tr}}$  for each latent space of type  $t$ . To perform clustering, standard clustering techniques such as k-Means and Gaussian Mixture Model (GMM) require specifying the number of clusters, which is less ideal as the number of clusters can be difficult to choose and may vary across latent spaces. Inspired by GEORGE [46], we adopt a fully automated hierarchical clustering technique by monitoring the Silhouette score [40].
Specifically, in each latent space, SLDG first applies UMAP [30] for dimensionality reduction. Then, it runs k-Means with  $k \in \{2, \dots, 10\}$  to identify the optimal number of clusters based on the highest Silhouette score. Subsequently, SLDG further split each cluster into five sub-clusters. However, only sub-clusters surpassing the Silhouette score of the original cluster and containing at least 500 patients are retained. The final number of clusters in the latent space of type  $t$  is denoted as  $K_{t}$ . The cluster assignment is represented by a binary matrix  $\mathbf{M}_t$  of size  $N_{tr} \times K_t$ , where  $\mathbf{M}_t[i, k]$  is set to one i.f.f. the patient  $x^{(i)}$  is assigned to the  $k$ -th cluster. This automated approach allows us to effectively select the number of clusters in each latent space, balancing between discovering overly coarse or fine-grained clusters.
3.2 Domain-Specific Model Customization
To accommodate the unique characteristics of patients in different domains, we propose to train customized classifiers for each decoupled domain. Given an input patient visit  $x$  and its multi-vector representations  $\{\mathbf{h}_t\}_{t \in \mathcal{T}}$ , SLDG computes the predicted probability  $o$  of a specific event occurring by employing a weighted combination of domain-specific classifiers in each latent space  $t \in \mathcal{T}$ , as in Eq. (6),
$$
o := \frac {1}{| \mathcal {T} |} \sum_ {t \in \mathcal {T}} \sum_ {k = 1} ^ {K _ {t}} \underbrace {G _ {t , k} (\mathbf {h} _ {t})} _ {\text {g a t e}} \cdot \underbrace {C _ {t , k} (\mathbf {h} _ {t})} _ {\text {c l a s s i f i e r}}, o \in \mathbb {R}, \tag {6}
$$
where  $C_{t,k}(\cdot)$  refers to the customized classifier for the discovered domain  $k$  in the latent space of type  $t$ , while  $G_{t,k}(\cdot)$  corresponds to the gating function. In the following, we will elaborate on how SLDG leverages the clustering results to efficiently parameterize the domain-specific classifier and effectively determine the gating weights.
To efficiently parameterize the domain-specific classifier  $C_{t,k}(\cdot)$  for the  $k$ -th discovered domain in the latent space of type  $t$ , we define two learnable weight vectors of size  $h$ :  $\mathbf{w}_{t,k}^{+}$  and  $\mathbf{w}_{t,k}^{-}$ , which represent the prototypes of the positive and negative classes, respectively. The predicted probability of a specific event occurring is computed based on the relative distance between the patient representation  $\mathbf{h}_t$  and the positive and negative prototypical weights, as in Eq. (7),
$$
C _ {t, k} \left(\mathbf {h} _ {t}\right) = \frac {\exp \left(- d \left(\mathbf {w} _ {t , k} ^ {+} , \mathbf {h} _ {t}\right)\right)}{\sum_ {*} \in \{+, - \} \exp \left(- d \left(\mathbf {w} _ {t , k} ^ {*}, \mathbf {h} _ {t}\right)\right)}, C _ {t, k} \left(\mathbf {h} _ {t}\right) \in \mathbb {R}, \tag {7}
$$
where  $d(\cdot, \cdot)$  is the Euclidean distance. To facilitate efficient learning, we initialize the two prototypical weight vectors  $\mathbf{w}_{t,k}^{+}$  and  $\mathbf{w}_{t,k}^{-}$ , with the average representations of patients from the corresponding classes, as in Eq. (8),
$$
\operatorname {I n i t} \left(\mathbf {w} _ {t, k} ^ {*}\right) = \operatorname {A v e r a g e} \left(\left\{\mathbf {h} _ {t} ^ {(i)} \mid \left(M _ {t} [ i, k ] = 1\right) \wedge \left(y ^ {(i)} = *\right) \right\} _ {i = 1} ^ {N _ {t r}}\right), * \in \{+, - \}, \tag {8}
$$
where  $M_{t}[i,k] = 1$  includes only patients assigned to the  $k$ -th domain in the latent space of type  $t$ .
We adopt a similar approach to the gating function  $G_{t}(\cdot)$ . For each discovered domain  $k$  in the latent space of type  $t$ , we introduce a learnable prototypical weight vector  $\mathbf{w}_{t,k} \in \mathbb{R}^h$ . The gating weights are determined based on the distance between the patient representation  $\mathbf{h}_t$  and the corresponding prototypical weights, as in Eq. (9),
$$
G _ {t} \left(\mathbf {h} _ {t}\right) = \operatorname {S o f t m a x} \left(\left\{- d \left(\mathbf {w} _ {t, k}, \mathbf {h} _ {t}\right) \right\} _ {k = 1} ^ {K _ {t}}\right), \quad G _ {t} \left(\mathbf {h} _ {t}\right) \in \mathbb {R} ^ {K _ {t}}, \tag {9}
$$
where the prototypical weight vectors  $\{\mathbf{w}_{t,k}\}_{k = 1}^{K_t}$  are initialized as the average representations of patients in that domain, as in Eq. (10),
$$
\operatorname {I n i t} \left(\mathbf {w} _ {t, k}\right) = \operatorname {A v e r a g e} \left(\left\{\mathbf {h} _ {t} ^ {(i)} \mid M _ {t} [ i, k ] = 1 \right\} _ {i = 1} ^ {N _ {t r}}\right), \quad k = 1, \dots , K _ {t}. \tag {10}
$$
3.3 Training and Inference
To train SLDG, we begin by utilizing a pre-trained patient encoder  ${}^{1}$  for decoupled domain discovery. Then,we iteratively update the model weights and re-generate the clusters every 20 epochs. In each
<sup>1</sup>In practice, we pre-train the patient encoder on the same clinical predictive task for 40 epochs. This means that SLDG undergoes a total of 100 epochs of training, which is consistent with other baselines.
iteration, we re-initialize the classifier and gating parameters. This iterative process is repeated three times to enhance the model's performance. During training, we minimize the binary cross-entropy loss. For inferencing, given a target patient visit, SLDG first maps it to multiple decoupled latent spaces with the feature-specific patient encoders  $E_{t}(\cdot)$ . Subsequently, in each latent space of type  $t$ , the gating function  $G_{t}(\cdot)$  determines the weight combinations used to aggregate the predictions from domain-specific classifiers  $\{C_{t,k}(\cdot)\}_{k=1}^{K_t}$ . The final prediction is obtained by averaging the predictions from all latent domains. The pseudocode of SLDG can be found in Appx. A.
Datasets. We evaluate SLDG on two publicly available real-world EHR datasets: eICU [38] and MIMIC-IV [18], which are described as follows:
- eICU [38] covers over 200K visits for 139K patients admitted to the intensive care unit (ICU) in one of the 208 hospitals across the United States. The data was collected between 2014 and 2015. The 208 hospitals can be further categorized into four groups based on their location (Midwest, Northeast, West, and South). We use age, gender, and ethnicity as patient demographic information, and leverage the diagnosis, treatment, medication, and lab tables to gather patient visit information.
- MIMIC-IV [18] covers over 431K visits for 180K patients admitted to the ICU in the Beth Israel Deaconess Medical Center. The data was collected between 2008 to 2019. The approximate actual year of each admission is revealed as one of the four-year groups (2008-2010, 2011-2013, 2014-2016, and 2017-2019). We use age, gender, and ethnicity as patient demographic information, and leverage diagnoses, procedures, and prescriptions to gather patient admission information.
We elaborate on the cohort selection process and provide comprehensive dataset statistics in Appx. B.1. In the end, we extract 149227 visits from 116075 patients in the eICU dataset, and 353238 visits from 156549 patients in the MIMIC-IV dataset.
Clinical Predictive Tasks. We focus on two common clinical predictive tasks: (1) Readmission prediction, which aims to determine whether a patient will be readmitted within the next 15 days following discharge. (2) Mortality prediction, which aims to predict whether a patient will pass away upon discharge in the eICU setting, or within 90 days after discharge in the MIMIC-IV setting. A detailed explanation of this setting can be found in Appx. B.2.
Data Split. We evaluate the performance of our model across spatial gaps using the eICU dataset. For this purpose, we select the target testing data as the group (Midwest) that demonstrated the largest performance gap in a pilot study. The remaining groups (Northeast, West, and South) are used as the source training data. To assess the model's performance across temporal gaps, we utilize the MIMIC-IV dataset. Patients admitted after 2014 are used as the target testing data, while all preceding patients are included in the source training data. We elaborate more on the data split in Appx. B.3.
Baselines. We compare SLDG against three categories of baselines. (1) The first category consists of naive baselines, including Oracle, trained directly on the target data, and Base, trained solely on the source data. (2) The second category comprises DG methods that require domain IDs. These include DANN [16] and MLDG [22], which use coarse regional and temporal groups as the domain definition, and ManyDG [56], which treats each patient as a unique domain. (3) The last category consists of DG methods that do not rely on domain IDs, including IRM [4], MMLD [29], and DRA [14]. A detailed explanation of all the baselines can be found in Appx. B.4.
Evaluation Metrics. Both readmission prediction and mortality prediction are binary classification tasks. To evaluate the performance of the models, we calculate the Area Under the Precision-Recall Curve (AUPRC) and the Area Under the Receiver Operating Characteristic Curve (AUROC) scores. For each metric, we report the average scores and standard deviation by performing bootstrapping