Imputation of Unknown Missingness in Sparse Electronic Health Records

Jun Han; Josue Nassar; Sanjit Singh Batra; Aldo Cordova-Palomera; Vijay Nori; Robert E. Tillman

cs.LGFeb 24, 2026

Imputation of Unknown Missingness in Sparse Electronic Health Records

Authors: Jun Han, Josue Nassar, Sanjit Singh Batra, Aldo Cordova-Palomera, Vijay Nori, Robert E. Tillman

Organizations: Optum AI

Abstract

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.

Figures & tables

Figure 1: A) An example of imputing known unknowns. In the patient EHR, there is NaN in ICD 3. Thus, there are only two possible choices for ICD 3: 0 or 1. B) An example of imputing unknown unknowns. In the patient EHR, there is a 0 in ICD 2 and ICD 3. Due to reporting issues, it is unclear whether a 0 represents a negative diagnosis (0) or a missing diagnosis (1). Due to this ambiguity, there are 4 possibilities after imputing: i) ICD 2 and 3 are both negative diagnoses, thus they both are 0. ii) ICD 2 was missing, thus it is 1. ICD 3 was a negative diagnosis, thus it is a 0. iii) ICD 2 was a negative diagnosis, thus it is a 0. ICD 3 was missing, thus it is 1. iv) Both ICD 2 and 3 are missing, thus they both are 1. C) Our proposed approach, Denoise2Impute-T, comprises 2 components. A patient EHR is input into a Set Transformer based denoiser that outputs a denoised patient EHR. The EHR is also input into a Thresholding network that outputs thresholds for each ICD code. A greater than or equal to element-wise comparison is performed between the denoised EHR and the thresholds, that leads to the final imputed EHR.

Figure 2: Descriptive statistics for both datasets used in this paper. A) Comparison of the prevalence of ICD codes in $\mathcal{D}_{1}$ vs $\mathcal{D}_{2}$ . B) Comparison of the sparsity of patient EHRs in $\mathcal{D}_{1}$ vs $\mathcal{D}_{2}$ . C) The eigenvalue spectra of the covariance matrix of $\mathcal{D}_{1}$ (blue) with respect to 100 prevalence-matched random binary matrices (red). D) The eigenvalue spectra of the covariance matrix of $\mathcal{D}_{2}$ (blue) with respect to 100 prevalence-matched random binary matrices (red).

Methods	$\mathcal{D}_{1}$	Prevalence	MLP	DAE	CDAE	D2Impute	D2Impute-T
Avg. AUPRC	51.86	52.37	61.91	62.05	62.17	64.90	65.53

Table 1: Average of dimension-wise AUPRC (higher is better) on the test set where noisy vectors were obtained from

\mathcal{D}_{1}

. D2Impute is short for Denoise2Impute.

Figure 3: Dimension-wise AUPRC for

T=993

computed on the test data. A) Comparison between AUPRC of the Denoise2Impute and the column mean ( prevalence ) for each dimension of the noisy data

\mathcal{D}_{1}

. B) Comparison between AUPRC of the Denoise2Impute and that of the

\mathcal{D}_{1}

. C) Comparison between AUPRC of the Denoise2Impute and AUPRC of the MLP.

Figure 4: Results for ICD code prediction for six common chronic conditions (A) and eight randomly chosen ICD codes (B). Results for Hospital Readmission task (C). Difference in AUPRC relative to using

\mathcal{D}_{1}

and 95% confidence intervals are plotted. The common chronic conditions in (A) and the description of ICD codes in (B) are provided in Table 3 in Appendix D .

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Neural Network Architecture of the denoising model $g_{\bm{\theta}}$ , where one of the key modules is the set attention block (SAB). The tensor shape of outputs in each layer is provided. We apply $L$ layers of the SAB.

Methods	$\mathcal{D}_{1}$	Prevalence	MissForest	MLP	CDAE	D2Impute	D2Impute-T
Micro AUPRC	62.67	70.39	64.14	79.10	79.13	81.28	81.64
Macro AUPRC	51.86	52.37	52.23	61.91	62.16	64.89	65.53

Appendix

Table 2: Micro AUPRC and Macro AUPRC (higher is better) on the test set where noisy vectors were obtained from

\mathcal{D}_{1}

. D2Impute is short for Denoise2Impute.

ICD-10 codes	Description of Diagnosis ICD-10 codes
Z37	Outcome of delivery
Z00	Encounter for general examination without CSRD
Z12	Encounter for screening for malignant neoplasms
M54	Dorsalgia
Z23	Encounter for immunization
Z01	Encounter for other special examination without CSRD

Appendix

Table 3: Description of additional ICD codes from Fig. 4 . CSRD is short for without complaint, suspected or reported diagnosis.

	T2D	AFib	Asthma	COPD	CKD	HTN	Osteoarthritis
Ground Truth	90.77	58.40	73.63	71.59	80.03	44.77	48.76
	$\pm$ 0.14	$\pm$ 0.50	$\pm$ 0.29	$\pm$ 0.46	$\pm$ 0.44	$\pm$ 0.64	$\pm$ 0.32
$\mathcal{D}_{1}$	78.47	40.85	47.26	58.08	61.90	32.29	38.13
	$\pm$ 0.14	$\pm$ 0.50	$\pm$ 0.29	$\pm$ 0.46	$\pm$ 0.47	$\pm$ 0.58	$\pm$ 0.31
softImpute	77.28	40.06	46.43	58.15	61.78	32.46	37.89
	$\pm$ 0.16	$\pm$ 0.48	$\pm$ 0.31	$\pm$ 0.48	$\pm$ 0.51	$\pm$ 0.63	$\pm$ 0.36

Appendix

Table 4: Results for 7 ICD codes corresponding to common chronic conditions. Mean AUPRC and

95\%

confidence interval (CI) are provided. We bootstrap 80% of the test data and repeat 50 times to compute CI. ”T2D”: type-II diabetes. ”AFib”: atrial fibrillation. ”COPD”: chronic obstructive pulmonary disease. ”CKD”: chronic kidney disease. ”HTN”: hypertensive heart disease. D2Impute is short for Denoise2Impute.

	Z37	Z00	Z12	M54	Z23	Z01	V22	Z34
Ground Truth	93.86	90.29	89.15	82.5	78.06	79.38	94.81	93.62
	$\pm$ 0.15	$\pm$ 0.07	$\pm$ 0.06	$\pm$ 0.12	$\pm$ 0.09	$\pm$ 0.13	$\pm$ 0.09	$\pm$ 0.16
$\mathcal{D}_{1}$	80.95	81.70	79.81	82.00	72.08	72.28	75.16	86.17
	$\pm$ 0.31	$\pm$ 0.09	$\pm$ 0.10	$\pm$ 0.09	$\pm$ 0.1	$\pm$ 0.09	$\pm$ 0.31	$\pm$ 0.21
softImpute	80.41	80.56	79.32	81.35	71.87	72.41	75.69	86.58
	$\pm$ 0.33	$\pm$ 0.10	$\pm$ 0.11	$\pm$ 0.10	$\pm$ 0.14	$\pm$ 0.11	$\pm$ 0.32	$\pm$ 0.26

Appendix

Table 5: Results for eight randomly chose ICD code predictions, where each ICD code has a prevalence larger than 6%. Mean AUPRC and

95\%

confidence intervals (CI) are provided. We bootstrap 80% of the test data and repeat 50 times to compute CI. D2Impute is short for Denoise2Impute.

	T2D	AFib	Asthma	COPD	CKD	HTN	Osteoarthritis
Ground Truth	90.77	58.40	73.63	71.59	80.03	44.77	48.76
	$\pm$ 0.14	$\pm$ 0.50	$\pm$ 0.29	$\pm$ 0.46	$\pm$ 0.42	$\pm$ 0.54	$\pm$ 0.32
$\mathcal{D}_{2}$	88.71	55.38	71.77	67.24	76.02	41.87	45.07
	$\pm$ 0.07	$\pm$ 0.32	$\pm$ 0.18	$\pm$ 0.38	$\pm$ 0.27	$\pm$ 0.68	$\pm$ 0.23
softImpute	86.98	54.62	70.06	66.13	74.06	40.14	43.29
	$\pm$ 0.14	$\pm$ 0.39	$\pm$ 0.21	$\pm$ 0.41	$\pm$ 0.29	$\pm$ 0.41	$\pm$ 0.24

Appendix

Table 6: Results for 7 ICD codes corresponding to common chronic conditions. Mean AUPRC and

95\%

confidence intervals (CI) are provided. We bootstrap 80% of the test data and repeat 50 times to compute CI. ”T2D”: type-II diabetes. ”AFib”: atrial fibrillation. ”COPD”: chronic obstructive pulmonary disease. ”CKD”: chronic kidney disease. ”HTN”: hypertensive heart disease. D2Impute is short for Denoise2Impute.

	Z37	Z00	Z12	M54	Z23	Z01	V22	Z34
Ground Truth	93.86	90.29	89.15	82.5	78.06	79.38	94.81	93.62
	$\pm$ 0.15	$\pm$ 0.07	$\pm$ 0.06	$\pm$ 0.12	$\pm$ 0.09	$\pm$ 0.13	$\pm$ 0.09	$\pm$ 0.16
$\mathcal{D}_{2}$	91.98	89.36	87.11	78.28	75.92	75.20	93.95	90.91
	$\pm$ 0.14	$\pm$ 0.06	$\pm$ 0.07	$\pm$ 0.09	$\pm$ 0.10	$\pm$ 0.11	$\pm$ 0.08	$\pm$ 0.15
softImpute	89.74	87.25	86.33	76.44	73.18	74.46	90.16	88.17
	$\pm$ 0.21	$\pm$ 0.09	$\pm$ 0.10	$\pm$ 0.11	$\pm$ 0.14	$\pm$ 0.16	$\pm$ 0.11	$\pm$ 0.18

Appendix

Table 7: Results for eight randomly chose ICD code predictions, where each ICD code has a prevalence larger than 6%. Mean AUPRC and

95\%

confidence intervals (CI) are provided. We bootstrap 80% of the test data and repeat 50 times to compute CI. D2Impute is short for Denoise2Impute.

Methods	$\mathcal{D}_{1}$	softImpute	MLP	DAE	CDAE	D2Impute	D2Impute-T
Mean	20.74	18.58	20.81	20.87	20.89	21.47	21.51
CI	$\pm$ 0.15	$\pm$ 0.16	$\pm$ 0.13	$\pm$ 0.12	$\pm$ 0.14	$\pm$ 0.10	$\pm$ 0.10

Appendix

Table 8: Hospital readmission prediction on

\mathcal{D}_{1}

dataset. Mean AUPRC and

95\%

confidence intervals (CI). We bootstrap 80% of the test data and repeat 50 times to compute CI. D2Impute is short for Denoise2Impute.

Figure 6: Results on ICD code prediction where the denoising models are trained using Noisy $\mathcal{D}_{1}$ . We compare the per ICD code AUPRC of ImputEHR with those of A) imputing 0’s with prevalence B) using Noisy $\mathcal{D}_{1}$ and C) denoising using an MLP.

Incomplete Data Imputation

Incomplete data imputation estimates unobserved values while accounting for data structure, missingness patterns, and uncertainty.

See all 74 papers

Learning Representations from Incomplete EHR Data with Dual-Masked Autoencoding
Feb 16, 2026Xiao Xiang, David Restrepo, Hyewon Jeong +2Learning with Missing Data Masked Autoencoders
Massachusetts Institute of Technology · École Polytechnique Fédérale de Lausanne · Harvard University +1
FeatEHR-LLM: Leveraging Large Language Models for Feature Engineering in Electronic Health Records
Apr 24, 2026Hojjat Karami, David Atienza, Jean-Philippe Thiran +1Healthcare Electronic Health Records
1École Polytechnique Fédérale de Lausanne (EPFL), Lausanne, Switzerland
PAMF: Prior-Aware Multimodal Fusion for Incomplete Time Series Data
Jun 4, 2026Ziwen Kan, Wugeng Zheng, Tianlong Chen +1Missing-Modality Learning Flow Matching
Department of Computer Science, University of Central Florida · Department of Computer Science, University of North Carolina at Chapel Hill

Feb 16, 2026J/K move · Enter open · S save

Learning Representations from Incomplete EHR Data with Dual-Masked Autoencoding

Xiao Xiang, David Restrepo, Hyewon Jeong, Yugang Jia+1

Massachusetts Institute of Technology · École Polytechnique Fédérale de Lausanne · Harvard University+1

Learning with Missing Data Masked Autoencoders Tabular Representation Learning Electronic Health Records

Open paper AskPDF arXiv

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.

Abstract

Incomplete Data Imputation

Explore similar work

Learning Representations from Incomplete EHR Data with Dual-Masked Autoencoding

FeatEHR-LLM: Leveraging Large Language Models for Feature Engineering in Electronic Health Records

PAMF: Prior-Aware Multimodal Fusion for Incomplete Time Series Data