Learnable Prompt as Pseudo-Imputation: Rethinking the Necessity of Traditional EHR Data Imputation in Downstream Clinical Prediction

TL;DR

This work tackles the challenge of missing data in Electronic Health Records (EHR) by questioning the necessity of traditional imputation for downstream clinical prediction. It introduces Learnable Prompt as Pseudo-Imputation (PAI), a plug-and-play protocol that replaces data imputation with a learnable input prompt trained end-to-end to guide predictions. Across four real-world datasets and multiple backbone models, PAI consistently improves mortality and length-of-stay predictions, and demonstrates robustness under data scarcity and higher missingness. The findings suggest that learning direct task-relevant information via prompts can outperform imputation-based strategies, offering a scalable and integration-friendly approach for EHR analysis.

Abstract

Analyzing the health status of patients based on Electronic Health Records (EHR) is a fundamental research problem in medical informatics. The presence of extensive missing values in EHR makes it challenging for deep neural networks (DNNs) to directly model the patient's health status. Existing DNNs training protocols, including Impute-then-Regress Procedure and Jointly Optimizing of Impute-n-Regress Procedure, require the additional imputation models to reconstruction missing values. However, Impute-then-Regress Procedure introduces the risk of injecting imputed, non-real data into downstream clinical prediction tasks, resulting in power loss, biased estimation, and poorly performing models, while Jointly Optimizing of Impute-n-Regress Procedure is also difficult to generalize due to the complex optimization space and demanding data requirements. Inspired by the recent advanced literature of learnable prompt in the fields of NLP and CV, in this work, we rethought the necessity of the imputation model in downstream clinical tasks, and proposed Learnable Prompt as Pseudo-Imputation (PAI) as a new training protocol to assist EHR analysis. PAI no longer introduces any imputed data but constructs a learnable prompt to model the implicit preferences of the downstream model for missing values, resulting in a significant performance improvement for all state-of-the-arts EHR analysis models on four real-world datasets across two clinical prediction tasks. Further experimental analysis indicates that PAI exhibits higher robustness in situations of data insufficiency and high missing rates. More importantly, as a plug-and-play protocol, PAI can be easily integrated into any existing or even imperceptible future EHR analysis models.

Figures (6)

• Figure 1: Imputation-based protocol vs. Prompt-based protocol. (a) illustrates typical examples of EHR with missing values and the inconsistency in performance between existing imputation models on the EHR imputation task and downstream outcome prediction tasks. (b) presents a typical framework of EHR analysis models and highlights the differences between the imputation-based protocol and the prompt-based protocol. (c) illustrates that on the MIMIC-IV dataset, compared to the imputation-based protocol, the prompt-based protocol improves the performance of all the baseline backbone models.
• Figure 2: The framework of PAI. The main idea is to model missing values using a learnable prompt vector, which can be optimized by minimizing the loss of downstream tasks.
• Figure 3: Performance of PAI with different learning rates on mortality outcome prediction task. The learning rate of the backbone&head modules is fixed at 1e-2. Please note that AUROC and AUPRC are presented using a dual-axis plotting approach for better visualization.
• Figure 4: Performance comparison of PAI and imputation-based protocol with different parameters on mortality outcome prediction task. The imputation-based protocol uses RNN/Transformer with 1/2/3 layers respectively, and PAI only use 1 layer. Please note that AUROC and AUPRC are presented using a dual-axis plotting approach for better visualization.
• Figure 5: Performance comparison of PAI and imputation-based protocol with different missing rates on mortality outcome prediction task. Please note that AUROC and AUPRC are presented using a dual-axis plotting approach for better visualization.