SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 1: Results of domain generalization on the eICU and MIMIC-IV datasets. An asterisk (*) indicates that SLDG achieves a significant improvement over the best baseline method, with a p-value smaller than 0.05. The experimental results demonstrate that SLDG exhibits robustness against spatial (eICU) and temporal (MIMIC-IV) domain shifts.
Method	eICU	MIMIC-IV
Readmission	Mortality	Readmission	Mortality
AUPRC	AUROC	AUPRC	AUROC	AUPRC	AUROC	AUPRC	AUROC
Oracle	0.219 (0.01)	0.677 (0.01)	0.271 (0.01)	0.839 (0.01)	0.282 (0.01)	0.693 (0.00)	0.428 (0.00)	0.898 (0.01)
Base	0.104 (0.02)	0.510 (0.01)	0.230 (0.01)	0.803 (0.01)	0.237 (0.01)	0.665 (0.01)	0.374 (0.01)	0.861 (0.00)
DANN	0.135 (0.01)	0.538 (0.01)	0.245 (0.01)	0.808 (0.01)	0.247 (0.01)	0.673 (0.01)	0.380 (0.02)	0.873 (0.02)
MLDG	0.104 (0.01)	0.525 (0.01)	0.224 (0.01)	0.797 (0.01)	0.205 (0.01)	0.637 (0.02)	0.360 (0.01)	0.857 (0.01)
ManyDG	0.150 (0.01)	0.549 (0.01)	0.259 (0.01)	0.814 (0.01)	0.249 (0.01)	0.676 (0.01)	0.388 (0.01)	0.880 (0.01)
IRM	0.136 (0.01)	0.538 (0.01)	0.252 (0.02)	0.811 (0.01)	0.242 (0.00)	0.668 (0.01)	0.387 (0.01)	0.876 (0.01)
MMLD	0.167 (0.01)	0.578 (0.00)	0.256 (0.01)	0.818 (0.01)	0.250 (0.02)	0.679 (0.01)	0.393 (0.01)	0.887 (0.01)
DRA	0.148 (0.01)	0.551 (0.01)	0.249 (0.01)	0.810 (0.01)	0.246 (0.01)	0.670 (0.01)	0.387 (0.01)	0.875 (0.01)
SLDG	0.186 (0.01)*	0.623 (0.01)*	0.268 (0.01)*	0.824 (0.01)*	0.274 (0.01)*	0.690 (0.01)*	0.416 (0.00)*	0.899 (0.01)*
1000 times. Additionally, we conduct independent two-sample t-tests to assess whether SLDG achieves a significant improvement over the baseline methods.
Implementation Details. For all baselines, we use the same Transformer [48] architecture as the backbone encoder. Patient demographics features (age, gender, and ethnicity) are embedded with an embedding look-up table. We also embed the timestamps with sinusoidal positional encoding. The medical, patient demographics, and temporal embeddings are added together to form the overall sequence embedding. All models are trained for 100 epochs, and the best model is selected based on the AUPRC score monitored on the source validation set. For SLDG, UMAP [30] from UMAP-learn [41] is used with 2 components, 10 neighbors, and 0 minimum distance; and k-Means from Scikit-learn [35] is used with the default hyper-parameter. Further information regarding the detailed implementations can be found in Appx. B.5.
Table 1 presents the domain generalization results on the eICU [38] and MIMIC-IV [18] datasets. Firstly, we observe a significant performance gap between the Oracle and Base methods, indicating the presence of substantial spatial and temporal domain gaps. This supports the use of the DG setting. Notably, the readmission tasks exhibit larger domain gaps, which is reasonable since hospitals across different locations and timestamps may have varying criteria for patient readmission. Secondly, we note that the two DG methods, DANN [16] and MLDG [22], utilizing coarse domain partitions such as region and timestamp, achieve minimal or no improvements. This outcome is expected because the domain partitions are too coarse, making it challenging to identify consistent domain features. In comparison, ManyDG [56] achieves better performance by considering each individual patient as a unique domain. Among the remaining three baseline methods that do not rely on domain IDs, IRM [4] demonstrates the slightest improvement. DRA [14] performs better due to the usage of multi-head networks, which share a similar intuition as SLDG. MMLD [29] attains the highest performance among all baselines, showcasing the advantages of explicit domain discovery. Lastly, SLDG outperforms baselines for all tasks. Specifically, in terms of the AUPRC score, SLDG achieves an  $11\%$  relative improvement in eICU readmission prediction,  $3\%$  in eICU mortality prediction,  $10\%$  in MIMIC-IV readmission prediction, and  $6\%$  in MIMIC-IV mortality prediction.
4.3 Quantitative Analysis
This section provides quantitative analyses to elucidate the performance enhancements achieved by SLDG. The analyses encompass the evaluation of clustering results, ablation studies on the clustering algorithm, the impact of the number of clusters and iterations, and a runtime comparison.
Evaluation of clustering results. First, we evaluate the domain recovery ability of DG methods that do not rely on domain IDs, namely MMLD [29], DRA [14], and the proposed SLDG. Since the actual latent domain categorizations are unavailable, we assess the separability of the learned clustering results with the Silhouette score [40]. Note that the reported Silhouette score is calculated on the testing set, while the hyper-parameters are chosen based on the Silhouette score on the training set. As depicted in Fig. 3, DRA achieves the lowest score, which aligns with expectations as it solely learns latent domain categorizations through multi-head networks without explicit clustering. In contrast, MMLD generates more distinct clusters due to its iterative clustering and training setup.
However, it still necessitates manual specification of the number of clusters. In comparison, SLDG obtains the highest score using an automated hierarchical clustering technique.
Figure 3: Performance of latent domain clustering. A higher Silhouette score indicates improved cluster separability.
Influence of the clustering algorithm. Next, we assess the influence of different clustering algorithms. Naive k-Means and GMM require manual specification of the number of clusters, which we set to the same value as SLDG's. The results can be found in Tab. 2. We observe that naive k-Means and GMM achieve similar performance similarly to the best baseline methods in Tab. 1. This outcome is reasonable since the success of SLDG relies on both the accurate discovery of latent domains and customized models for each domain. Naive clustering techniques often fail to identify subtle yet important latent domains. In contrast, SLDG, utilizing the automatic hierarchical clustering technique, achieves the highest score.
Table 2: Ablation study on the influence of the clustering algorithm.
Method	eICU	MIMIC-IV
Readmission	Mortality	Readmission	Mortality
AUPRC	AUROC	AUPRC	AUROC	AUPRC	AUROC	AUPRC	AUROC
SLDG + k-Means	0.148 (0.01)	0.553 (0.01)	0.249 (0.01)	0.814 (0.00)	0.250 (0.01)	0.670 (0.01)	0.388 (0.01)	0.886 (0.01)
SLDG + GMM	0.143 (0.01)	0.549 (0.01)	0.240 (0.01)	0.808 (0.01)	0.250 (0.01)	0.688 (0.01)	0.390 (0.00)	0.888 (0.00)
SLDG	0.186 (0.01)*	0.623 (0.01)*	0.268 (0.01)*	0.824 (0.01)*	0.274 (0.01)*	0.690 (0.01)*	0.416 (0.01)*	0.899 (0.01)*
Influence of the number of clusters and iterations.
Next, we analyze the impact of the number of clusters and iterations on the eICU readmission prediction task. We also compare our results with the best-performing baseline, MMLD [29]. The results can be found in Fig. 4. The upper panel of the figure shows that the model's performance initially improves with an increasing number of clusters. This improvement can be attributed to the finer granularity of clustering, which enables better identification of domains and customization of experts. However, as the number of clusters continues to increase, the model's performance starts to decline. This decline is caused by the growing number of model parameters, making training more challenging and leading to overfitting on suspicious samples. Similarly, the trend observed in the lower panel of the figure for the number of iterations aligns with the number of clusters. The performance initially improves as
Figure 4: Results on the influence of the number of clusters and iterations.
the number of iterations increases, allowing the model to learn more from the data. However, after a certain point, the model starts overfitting on specific clusters, leading to decreased performance.
Runtime comparison. Lastly, we compare the training time of SLDG with the naive Base baseline. All runtimes are measured on a single NVIDIA A6000 GPU. The results can be found in Tab. 3. The use of the UMAP [30] dimensionality reduction technique enables SLDG to perform clustering quickly. As a result, the training time overhead of SLDG is reasonably low overall (18% on eICU and 20% on MIMIC-IV) compared to the significant performance improvement achieved (up to 79% relative improvement on AUPRC score on eICU and up to 15% on MIMIC-IV).
Table 3: Runtime comparison.
Method	eICU	MIMIC-IV
Base	67 min	93 min
SLDG	79 min	112 min
Latent Space	Frequent Clinical Events from the Top-2 Clusters
Diagnosis	Sepsis, Infection, Kidney failure
Cardiac Arrest, Congestive Heart Failure
Treatment	Heart valve procedures, Cardiovascular monitoring
Ventilation, Respiratory intubation
Medication	Propofol, Lorazepam, Fentanyl
Amiodarone, Noradrenaline
Figure 5: Left: Common clinical events observed in two largest domains identified from each latent space. Right: Learned similarity among the identified domains within the diagnostic latent space.
For the case study, we first examine the recovered domains within each latent space. The results are presented in the left panel of Fig. 5. It is evident that the common events identified within the same domains are consistent. For instance, propofol, lorazepam, and fentanyl is frequently used together, serving the purpose of anesthesia (pre-surgery), sedation (in-surgery), and pain management (post-surgery). Furthermore, the right panel of Figure 5 illustrates a visualization of the learned domain similarity, i.e., the distances between the domain prototypical weights  $\{\mathbf{w}_{t,k}\}_{k=1}^{K_t}$ . Notably, our method (SLDG) successfully captures meaningful relationships among the latent domains. For instance, a strong relationship is learned between pneumonia and sepsis. In clinical practice, when pneumonia is severe or if the infection spreads beyond the lungs, it can enter the bloodstream and trigger a systemic response, leading to sepsis [8]. Another example is the observed strong association between cardiac arrest and infarction. In practice, if a significant portion of the heart muscle is damaged, it can disrupt the heart's electrical system, potentially leading to cardiac arrest [19].
Domain Generalization The goal of DG is to learn a model using data from multiple source domains in order to achieve effective generalization to a distinct target domain [7]. To achieve this, domain alignment approaches try to match the feature distributions among multiple source domains with techniques such as moments minimization [32, 25], contrastive learning [31], adversarial learning [25, 45], regularizers [4, 24], and augmentation [62, 50, 44, 60]. Meta-learning frameworks have also been utilized to simulate new domain scenarios during training [23, 27, 26, 5, 22, 29]. Additionally, domain-specific model ensemble techniques have been employed [9, 42, 43, 61]. However, these conventional DG methods assume the availability of domain IDs, which may not be feasible in healthcare settings where patients can belong to numerous unobserved domains.
Recent advancements in DG have attempted to alleviate the reliance on domain IDs [29, 62, 14, 10, 33]. MMLD [29] is the most relevant prior work to ours. It simultaneously discovers latent domains and learns domain-invariant features through adversarial learning. However, it focuses on training a
single model that generalizes across all domains. Given that patients from different domains exhibit distinct characteristics and require different treatment approaches [2, 58], training a single model for all domains poses challenges and can result in sub-optimal performance.
Clinical Predictive Modeling. 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, based on existing patient information. Deep learning models have been widely used in clinical predictive modeling with EHR data [55, 57]. These models are designed to capture temporal patterns in patient data [11, 36, 6, 28], model structural information in medical codes [13, 52], augment the model using pre-training [39], or leverage patient similarities for better decision making [58, 2]. However, these models typically assume an unchanged test domain and may suffer from degraded performance with domain shift. To address this issue, AutoMap [53] solves the feature space shift issue by learning an auto-mapping function without considering any distribution shift. MedLink [54] aggregates de-identified patient data from different sites to enable joint training. ManyDG [56] tackles patient covariate shift by treating each patient as a unique domain and disentangling domain variant and invariant features. However, maintaining a large number of domains is unnecessary, as similar patients often exhibit similar clinical behavior and can share a common domain.
Clinical predictive models often exhibit degraded performance when applied to data from new regions or future periods due to distribution shifts. To address this, we propose SLDG, a self-learning framework that iteratively identifies decoupled domains and trains customized classifiers for each domain. We evaluate SLDG on two medical datasets, and our results show that it outperforms all baseline methods. In addition, we provide detailed qualitative analyses and case studies to support our findings.
This work was supported by NSF award SCH-2205289, SCH-2014438, and IIS-2034479. This project has been funded by the Jump ARCHES endowment through the Health Care Engineering Systems Center.
Proceedings of the Conference on Health, Inference, and Learning, pages 279-290, Virtual Event USA, April 2021. ACM. ISBN 978-1-4503-8359-2. doi: 10.1145/3450439.3451878. URL https://dl.acm.org/doi/10.1145/3450439.3451878.
[60] Kaiyang Zhou, Yongxin Yang, Timothy Hospedales, and Tao Xiang. Learning to generate novel domains for domain generalization. In Andrea Vedaldi, Horst Bischof, Thomas Brox, and Jan-Michael Frahm, editors, Computer Vision – ECCV 2020, pages 561–578, Cham, 2020. Springer International Publishing. ISBN 978-3-030-58517-4.
[61] Kaiyang Zhou, Yongxin Yang, Yu Qiao, and Tao Xiang. Domain adaptive ensemble learning. IEEE Transactions on Image Processing, 30:8008-8018, 2021. doi: 10.1109/TIP.2021.3112012.
[62] Kaiyang Zhou, Yongxin Yang, Yu Qiao, and Tao Xiang. Domain generalization with mixstyle. In International Conference on Learning Representations, 2021. URL https://openreview.net/forum?id=6xHJ37MVxxp.
Contents of Appendix
A Pseudocode of SLDG 18
B Additional Experimental Setup 18
B.1 Datasets 18
B.2 Clinical Predictive Tasks 18
B.3 Data Split 19
B.4Baselines 20
B.5 Implementation Details 20
C Limitations and Broader Impacts 20
D Notations 21
E Additional Illustrations 21
A Pseudocode of SLDG
Algorithm 1: Training and Inference for SLDG.
1: // Training  
Require: Source training data from  $P_{tr}$   
2: Pre-train patient encoder  $E(\cdot)$  on the same task with binary cross-entropy loss for 40 epochs  
3: for iteration ranging from 1 to 3 do  
4: Perform decoupled domain discovery with the encoder  $E(\cdot)$  by Eq. (2), (3)  
5: Initialize gating and classifier weights  $\mathbf{w}_{t,k}, \mathbf{w}_{t,k}^{+}, \mathbf{w}_{t,k}^{-}$ by Eq. (10), (8)  
6: for epoch ranging from 1 to 20 do  
7: for each patient  $(x,y) \sim P_{tr}$  do  
8: Obtain decoupled patient representations  $\{\mathbf{h}_t\}_{t\in \mathcal{T}}$  by Eq. (4), (5)  
9: Compute domain-specific predictions  $C_{t,k}(\mathbf{h}_t)$  by Eq. (7)  
10: Compute gating weights  $G_{t}(\mathbf{h}_{t})$  by Eq. (9)  
11:Obtain final prediction  $o$  by Eq. (6)  
12: Update model parameters with binary cross-entropy loss  
13: end for  
14: end for  
15: end for  
16: // Inference  
Require: Target testing data from  $P_{te}$   
17: for each patient  $(x,y) \sim P_{tE}$  do  
18: Obtain decoupled patient representations  $\{\mathbf{h}_t\}_{t\in \mathcal{T}}$  by Eq. (4), (5)  
19:Compute domain-specific predictions  $C_{t,k}(\mathbf{h}_t)$  by Eq. (7)  
20:Compute gating weights  $G_{t}(\mathbf{h}_{t})$  by Eq. (9)  
21:Obtain final prediction  $o$  by Eq. (6)  
22:end for
B Additional Experimental Setup
For both datasets, we select our cohorts by filtering out visits of patients younger than 18 or older than 89 years old, visits that last longer than 10 days, and visits with data from less than 3 or more than 256 timestamps. In the case of the eICU dataset, we additionally exclude visits lasting shorter than 12 hours, as the predictions are made 12 hours after admission. Similarly, for the MIMIC-IV dataset, we exclude visits where the patient ultimately passed away, as the predictions are made upon discharge. Tab. 4 provides detailed statistics of the two datasets.
B.2 Clinical Predictive Tasks
We focus on two common clinical predictive tasks: readmission prediction and mortality prediction.
In the case of the eICU dataset, the predictions are made 12 hours after admission. Readmission prediction aims to determine whether a patient will be readmitted within the next 15 days following discharge. Mortality prediction, on the other hand, aims to predict whether a patient will pass away upon discharge. The overall prevalence for these tasks is  $15\%$  for readmission and  $4\%$  for mortality.
For the MIMIC-IV dataset, the predictions are made at the time of discharge. Similar to the eICU dataset, the readmission prediction task is defined as predicting whether a patient will be readmitted within 15 days after discharge. To prevent information leakage, the mortality prediction task for MIMIC-IV is defined as predicting whether a patient will pass away within 90 days after discharge. The overall prevalence for these tasks is  $14\%$  for readmission and  $4\%$  for mortality.
Table 4: Dataset statistics.
Item	eICU	MIMIC-IV
#Patients	116075	156549
#Admissions	149227	353238
Readmission Rate	0.15	0.14
Mortality Rate	0.04	0.04
Region: Midwest	Year: 2008-2010
#Patients	29767	37328
#Admissions	35989	56433
Readmission Rate	0.10	0.14
Mortality Rate	0.03	0.04
Age	62	56
Gender	F: 0.46, M: 0.54	F: 0.53, M: 0.47
Race	African American: 0.09, Asian: 0.01, Caucasian: 0.83, Hispanic: 0.01, Native American: 0.01, Other: 0.04	African American: 0.15, Asian: 0.03, Caucasian: 0.71, Hispanic: 0.06, Native American: 0.00, Other: 0.04
Average #Events	90.01	31.87
Region: Northeast	Year: 2011-2013
#Patients	5886	39125
#Admissions	6958	62586
Readmission Rate	0.17	0.15
Mortality Rate	0.06	0.04
Age	62	57
Gender	F: 0.44, M: 0.56	F: 0.53, M: 0.47
Race	African American: 0.03, Asian: 0.01, Caucasian: 0.92, Hispanic: 0.01, Native American: 0.00, Other: 0.03	African American: 0.17, Asian: 0.03, Caucasian: 0.66, Hispanic: 0.07, Native American: 0.00, Other: 0.07
Average #Events	104.54	35.19
Region: South	Year: 2014-2016
#Patients	27584	41737
#Admissions	33033	64592
Readmission Rate	0.11	0.14
Mortality Rate	0.04	0.04
Age	62	57
Gender	F: 0.46, M: 0.54	F: 0.52, M: 0.48
Race	African American: 0.21, Asian: 0.01, Caucasian: 0.68, Hispanic: 0.05, Native American: 0.00, Other: 0.03	African American: 0.17, Asian: 0.04, Caucasian: 0.66, Hispanic: 0.06, Native American: 0.00, Other: 0.07
Average #Events	84.28	36.53
Region: West	Year: 2017-2019
#Patients	17670	40496
#Admissions	19803	63654
Readmission Rate	0.29	0.14
Mortality Rate	0.04	0.04
Age	63	58
Gender	F: 0.45, M: 0.55	F: 0.52, M: 0.48
Race	African American: 0.05, Asian: 0.03, Caucasian: 0.77, Hispanic: 0.05, Native American: 0.02, Other: 0.08	African American: 0.17, Asian: 0.04, Caucasian: 0.65, Hispanic: 0.06, Native American: 0.00, Other: 0.07
Average #Events	85.29	36.95
B.3 Data Split
The eICU dataset comprises data collected from hospitals across the United States, while the MIMIC-IV dataset spans a period of ten years. Therefore, we utilize the eICU dataset to evaluate the model's performance across spatial gaps, and the MIMIC-IV dataset to assess its performance across temporal gaps.
For the eICU dataset, we divide it into four spatial groups based on regions: Midwest, Northeast, West, and South. Each group is then split into  $70\%$  for training,  $10\%$  for validation, and  $20\%$  for testing. We evaluate the gap between groups by comparing the performance of the backbone model trained on data from within the same group and data from outside the group. The target testing data is selected as the group (Midwest) that exhibits the largest performance gap, while the remaining groups (Northeast, West, and South) are used as the source training data.
Regarding the MIMIC-IV dataset, we divide it into four temporal groups: 2008-2010, 2011-2013, 2014-2016, and 2017-2019. Each group is further split into training, validation, and testing sets with a ratio of  $70\%$ ,  $10\%$ , and  $20\%$  respectively. We consider patients admitted after 2014 as the target testing data, while all preceding patients are included in the source training data.
B.4Baselines
We first compare SLDG to two naive baselines.
- Oracle: We directly train a backbone model on the training set of the target domain, select the best model on the target validation set, and evaluate its performance on the target testing set. This model is trained with in-domain data and can be seen as a upper bound for all domain generalization method.
- Base: We train a backbone model on the training set of the source domain, select the best model on the source validation set, and evaluate its performance on the target testing set. This model is trained with out-domain data and should act as a performance lower bound.
We then compare SLDG to both classic and recent domain generalization methods. For a fair comparisons, all the methods below are trained on the source training set, selected on the source validation set, and tested on the target testing set.
- DANN [16]: Domain-Adversarial Neural Networks leverage a domain classifier and a gradient reversal layer to extract domain-invariant representations. This method uses the coarse regional and temporal groups as the domain definition.
- MLDG [22]: Meta-Learning for Domain Generalization adopts the Model-Agnostic Meta-Learning (MAML) [15] framework and simulates the new domain scenario during training. This method also uses the coarse regional and temporal groups as the domain definition.
- ManyDG [56]: Many-Domain Generalization disentangles domain-variant and invariant features through mutual reconstruction and orthogonal projection. This method treats each patient as a unique domain.
- IRM [4]: Invariant Risk Minimization learns domain-invariant representations by minimizing a bound on the expected generalization error under domain shifts. It acts as a regularizer and does not require domain IDs.
- MMLD [29]: Domain Generalization using a Mixture of Multiple Latent Domains iteratively assigns pseudo domain labels via clustering and trains the domain-invariant feature extractor through adversarial learning. This method does not rely on domain IDs.
- DRA [14]: Latent Domain Learning with Dynamic Residual Adapters uses layer-wise multi-head correction networks with a gating mechanism and residual connection to enhance model learning. This method does not rely on domain IDs.
B.5 Implementation Details