Learning Representations from Incomplete EHR Data with Dual-Masked Autoencoding
Authors: Xiao Xiang, David Restrepo, Hyewon Jeong, Yugang Jia, Leo Anthony Celi
Organizations: Massachusetts Institute of Technology · École Polytechnique Fédérale de Lausanne · Harvard University · Beth Israel Deaconess Medical Center
Electronic health records (EHR) arrive masked. Clinicians order measurements selectively, and any patient table thus contains only a subset of the values that characterize the underlying physiological state. Prior masked modeling approaches on EHR data either impute the table before learning, represent missingness through a dedicated placeholder signal, or optimize solely for imputation, which limits the representations they learn for downstream clinical tasks and carries every unobserved entry through the encoder. We introduce AID-MAE, an Augmented-Intrinsic Dual-Masked Autoencoder that learns directly from incomplete tables by combining the intrinsic mask the record already carries with an augmented mask that hides a subset of observed values for reconstruction during pretraining. Neither type of masked entry enters the encoder, so attention operates only over what was observed. AID-MAE achieves consistent improvements over strong baselines across multiple clinical tasks on two datasets. Across experiments, we discuss that recovering the missing entries is not a prerequisite for learning and show that the representations learned carry clinical structure without supervision.
Figures & tables
Figure 1: AID-MAE (Augmented-Intrinsic Dual-Masked AutoEncoder) A subset of observed tokens of the original data with inherent missingness MISSING is randomly masked ( [MASK] ). Each measurement (value and timestamp) is embedded with positional encodings. The encoder processes only unmasked tokens, and the decoder receives the encoded representations along with a learned padding token [PAD] in place of missing or masked entries. Training optimizes a dual loss of reconstructing unmasked values and predicting features under augmented masking, while intrinsically missing entries are excluded from the loss. Code: https://github.com/xiaoxiang0605/AID-MAE
Task
# Stays
Prevalence
Mortality
92,938
7.6%
LOS <72 h
92,938
66.1%
Table 1: Downstream tasks and cohort sizes for MIMIC-IV and PhysioNet Challenge 2012.
Model
MIMIC-IV
PhysioNet Challenge 2012
Mortality
LOS
Mortality
AKI
Logistic Regression
80.7 ± 0.0
68.8 ± 0.0
72.7 ± 0.0
74.8 ± 0.0
Supervised Transformer
83.9 ± 0.5
72.3 ± 0.5
76.7 ± 1.4
76.1 ± 0.9
XGBoost
86.7 ± 0.0
76.9 ± 0.0
76.9 ± 0.0
77.1 ± 0.0
DuETT
86.4 ± 0.2
77.2 ± 0.0
77.7 ± 0.4
76.7 ± 0.2
AID-MAE (ours)
87.7 ± 0.1
77.6 ± 0.1
78.2 ± 0.3
77.3 ± 0.1
Table 2: Fine-tuning results on MIMIC-IV and PhysioNet Challenge 2012. We report AUROC (mean ± SD over 5 seeds). Best scores per column are in bold . AID-MAE outperforms all baselines.
Figure 2: Linear probing results for mortality prediction and length of stay prediction tasks. We compare our model against Logistic Regression with raw data (median imputed) and DuETT with frozen encoder, across different training data percentages (1%, 5%, 10%, 50%, 100%). Error bars represent standard deviation across 5 seeds.
Figure 3: UMAP of feature embeddings. UMAP projection of N=100,000 randomly sampled 64-D embeddings for 50 lab features. Colors denote 13 highlighted lab types. Each point corresponds to one measurement embedding. We denote two important patterns: Bottom-left: neighboring pair of pink (Platelet Count) and green islands (WBC); Bottom-right: neighboring pair of purple (Hemoglobin) and green islands (Hematocrit). The geometrical neighboring is consistent with their clinical coupling.
Figure 4: UMAP visualization of first-day CLS embeddings for initial MICU and CVICU admissions. Colors represent clusters from K-means ( k=2 ) applied in the embedding space.
Training Regime
Method
Mortality
AKI
Linear Probing
Random Weights
49.8
52.7
DuETT
70.9
66.2
AID-MAE (ours)
73.1
69.7
Fine-Tuning
Random Weights
76.7 ± 1.4
76.1 ± 0.9
DuETT
77.8 ± 0.3
76.4 ± 0.1
AID-MAE (ours)
78.6 ± 0.2
77.0 ± 0.2
Table 3: Transfer learning performance from MIMIC-IV pre-trained weights to PhysioNet Challenge 2012. We report AUROC for mortality and acute kidney injury (AKI) prediction under linear probing and full fine-tuning.
a
b
Mortality
Length of Stay
AUROC
AUPRC
AUROC
AUPRC
Random masking ( a=0 )
0
0.125
87.4 ± 0.1
49.4 ± 0.2
76.8 ± 0.2
62.0 ± 0.1
0
0.25
87.7 ± 0.1
49.8 ± 0.1
77.6 ± 0.1
63.1 ± 0.1
0
0.5
87.3 ± 0.1
49.0 ± 0.1
77.2 ± 0.2
62.0 ± 0.3
0
0.75
87.5 ± 0.1
49.8 ± 0.2
76.8 ± 0.2
61.9 ± 0.1
Table 4: Ablation study on proportional masking hyperparameters. Performance is reported as mean ± standard deviation across 5 random seeds. Downstream performance is largely flat across the grid
Input type
Mortality
Length of Stay
AUROC
AUPRC
AUROC
AUPRC
Without 24-hour information
85.4 ± 0.1
45.7 ± 0.2
75.6 ± 0.1
60.4 ± 0.1
Zero-filling
87.2 ± 0.1
49.3 ± 0.2
77.3 ± 0.1
62.7 ± 0.1
AID-MAE
87.7 ± 0.1
49.8 ± 0.1
77.6 ± 0.1
63.1 ± 0.1
Table 5: Input ablation study results (mean ± standard deviation across 5 random seeds).
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
Mortality
Length of Stay
Model
AUROC
AUPRC
AUROC
AUPRC
Logistic Regression
80.7 ± 0.0
37.3 ± 0.0
68.8 ± 0.0
53.1 ± 0.0
XGBoost
86.7 ± 0.0
47.5 ± 0.0
76.9 ± 0.0
62.2 ± 0.0
Supervised Transformer
83.9 ± 0.5
42.2 ± 1.3
72.3 ± 0.5
56.7 ± 1.5
DuETT
86.4 ± 0.2
47.3 ± 0.2
77.2 ± 0.0
62.8 ± 0.2
AID-MAE (ours)
87.7 ± 0.1
49.8 ± 0.1
77.6 ± 0.1
63.1 ± 0.1
Appendix
Table 6: Downstream task performance on the MIMIC-IV ICU dataset. We compare our proposed model against established baselines across two clinical tasks. Results are reported as mean ± standard deviation across 5 random seeds. Our model (random masking 25%) achieves superior performance on both tasks, demonstrating the effectiveness of AID-MAE. Best results are shown in bold . AUROC and AUPRC are reported as percentages with one decimal place.
Mortality
Acute Kidney Injury
Model
AUROC
AUPRC
AUROC
AUPRC
Logistic Regression
72.7 ± 0.0
31.3 ± 0.0
74.8 ± 0.0
58.6 ± 0.0
Supervised Transformer
76.7 ± 1.4
34.1 ± 3.5
76.1 ± 0.9
59.6 ± 1.3
XGBoost
76.9 ± 0.0
36.9 ± 0.0
77.1 ± 0.0
61.8 ± 0.0
DuETT
77.7 ± 0.4
38.8 ± 0.9
76.7 ± 0.2
61.8 ± 0.5
AID-MAE (ours)
78.2 ± 0.3
39.3 ± 1.7
77.3 ± 0.1
62.5 ± 1.2
Appendix
Table 7: Downstream task performance on the PhysioNet 2012 Challenge dataset. We compare our proposed model against established baselines across two clinical tasks. Results are reported as mean ± standard deviation across 5 random seeds. Our model (random masking 25%) achieves superior performance on both tasks, demonstrating the effectiveness of AID-MAE. Best results are shown in bold . AUROC and AUPRC are reported as percentages with one decimal place.
Figure 5: Linear probing results for mortality prediction and length of stay prediction tasks. We compare our model against Logistic Regression with median imputation and DuETT across different training data percentages (1%, 5%, 10%, 50%, 100%). Results are shown for both AUROC (top row) and AUPRC (bottom row) metrics. Our model consistently outperforms baseline methods across all data regimes and tasks, with particularly strong performance in low-data scenarios. Error bars represent standard deviation across 5 random seeds. The x-axis uses logarithmic scaling to better visualize performance across the range of training data percentages.
Setting
Mortality
AKI
With MICE
73.9 ± 1.2
75.8 ± 0.1
No Imputation
78.2 ± 0.3
77.3 ± 0.1
Appendix
Table 8: Pretraining AID-MAE on MICE-imputed data ( Buuren and Groothuis-Oudshoorn, 2011 ) . Architecture, hyperparameters, and splits are held fixed. Mean ± SD over 5 seeds.
Feature
NRMSE
NMAE
R 2
Model
XGBoost
Model
XGBoost
Model
XGBoost
Glucose finger stick
0.2030
0.2033
0.1516
0.1523
0.4806
0.4787
Potassium (serum)
0.1736
0.2026
0.1328
0.1591
0.5883
0.4391
Sodium (serum)
0.0676
0.0846
0.0416
0.0626
0.9362
0.9000
Chloride (serum)
0.0556
0.0702
0.0348
0.0514
0.9567
0.9310
Hemoglobin
0.0571
0.0622
0.0391
0.0429
0.9602
0.9528
Appendix
Table 9: Performance comparison between our proposed model and XGBoost baseline on the top 20 frequently collected clinical features ranked by frequencies. Metrics include normalized root mean squared error (NRMSE), normalized mean absolute error (NMAE), and coefficient of determination (R 2 ). Normalization is based on the value range. Bold values indicate superior performance (lower NRMSE/NMAE or higher R 2 ). Our proposed model demonstrates consistent superiority across 19 out of 20 features across all metrics.
Learning with missing data develops methods to estimate, predict, and make decisions when variables, labels, modalities, or time points are unobserved.
Machine learning holds great promise for advancing the field of medicine, with electronic health records (EHRs) serving as a primary data source. However, EHRs are often sparse and contain missing data due to various challenges and limitations in data collection and sharing between healthcare providers. Existing techniques for imputing missing values predominantly focus on known unknowns, such as missing or unavailable values of lab test results; most do not explicitly address situations where it is difficult to distinguish what is missing. For instance, a missing diagnosis code in an EHR could signify either that the patient has not been diagnosed with the condition or that a diagnosis was made, but not shared by a provider. Such situations fall into the paradigm of unknown unknowns. To address this challenge, we develop a general purpose algorithm for denoising data to recover unknown missing values in binary EHRs. We design a transformer-based denoising neural network where the output is thresholded adaptively to recover values in cases where we predict data are missing. Our results demonstrate improved accuracy in denoising medical codes within a real EHR dataset compared to existing imputation approaches and leads to increased performance on downstream tasks using the denoised data. In particular, when applying our method to a real world application, predicting hospital readmission from EHRs, our method achieves statistically significant improvement over all existing baselines.
Masked language modelling (MLM) is the dominant pretraining objective for transformer encoders applied to longitudinal electronic health records (EHRs). However, reconstructing individual events may overemphasise noisy or weakly informative observations and does not explicitly encourage learning higher-level patient states. Joint-Embedding Predictive Architectures (JEPAs) instead learn by predicting representations in latent space and have shown strong performance in a variety of domains. In this paper, we investigate to what extent JEPA can be adapted to EHR sequence representation. Our empirical investigation using MIMIC-IV shows that although all three of our JEPA formulations train stably without representational collapse, they consistently underperform standard MLM. We find that target size, not target heterogeneity, increases prediction error, yet reducing the target size does not recover downstream performance. This suggests that stable latent prediction alone is insufficient for competitive EHR representation learning and highlights the challenge of defining clinically meaningful latent targets for EHR.
Clinical diagnosis often requires combining imaging with physiological measurements, yet deployed models typically operate on unimodal data. We present PaCX-MAE, a cross-modal distillation framework that injects physiological priors into chest X-ray (CXR) encoders while remaining strictly unimodal at inference. PaCX-MAE augments in-domain masked autoencoding with a dual contrastive-predictive objective, aligning CXR representations with paired ECG and laboratory embeddings. Extensive evaluation across nine benchmarks demonstrates consistent improvements over domain-specific MAE, particularly on physiology-dependent tasks (e.g., +2.7 AUROC on MedMod; +6.5 F1 on VinDr). The method proves highly label-efficient in the 1% regime and preserves anatomical fidelity, achieving parity with MAE on segmentation tasks. Zero-shot and attention analyses confirm that PaCX-MAE successfully learns to attend to physiological indicators, such as the cardiac silhouette, absent in standard visual pretraining.
Yancheng Liu, Kenichi Maeda, Manan Pancholy
Department of Computer Science, Brown University, Providence, RI, USA.