Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

SMART: Towards Pre-trained Missing-Aware Model for Patient Health Status Prediction

Zhihao Yu, Xu Chu, Yujie Jin, Yasha Wang, Junfeng Zhao

TL;DR

SMART tackles the pervasive missing-data problem in electronic health records by introducing a self-supervised, missing-aware representation learning framework. It combines a Variable Independent Encoder with MART blocks that perform missingsensitive temporal and cross-variable attentions, and employs a two-stage training strategy that reconstructs missing data in latent space rather than input space. A self-supervised pre-training stage, using EMA targets, enhances imputation capabilities, followed by fine-tuning with a task-specific decoder and a CLS-based prediction head. Across six clinical tasks on Cardiology, Sepsis, and MIMIC-III, SMART achieves state-of-the-art performance, especially under high missingness, while remaining lightweight and efficient, highlighting its potential for robust clinical decision support.

Abstract

Electronic health record (EHR) data has emerged as a valuable resource for analyzing patient health status. However, the prevalence of missing data in EHR poses significant challenges to existing methods, leading to spurious correlations and suboptimal predictions. While various imputation techniques have been developed to address this issue, they often obsess unnecessary details and may introduce additional noise when making clinical predictions. To tackle this problem, we propose SMART, a Self-Supervised Missing-Aware RepresenTation Learning approach for patient health status prediction, which encodes missing information via elaborated attentions and learns to impute missing values through a novel self-supervised pre-training approach that reconstructs missing data representations in the latent space. By adopting missing-aware attentions and focusing on learning higher-order representations, SMART promotes better generalization and robustness to missing data. We validate the effectiveness of SMART through extensive experiments on six EHR tasks, demonstrating its superiority over state-of-the-art methods.

SMART: Towards Pre-trained Missing-Aware Model for Patient Health Status Prediction

TL;DR

SMART tackles the pervasive missing-data problem in electronic health records by introducing a self-supervised, missing-aware representation learning framework. It combines a Variable Independent Encoder with MART blocks that perform missingsensitive temporal and cross-variable attentions, and employs a two-stage training strategy that reconstructs missing data in latent space rather than input space. A self-supervised pre-training stage, using EMA targets, enhances imputation capabilities, followed by fine-tuning with a task-specific decoder and a CLS-based prediction head. Across six clinical tasks on Cardiology, Sepsis, and MIMIC-III, SMART achieves state-of-the-art performance, especially under high missingness, while remaining lightweight and efficient, highlighting its potential for robust clinical decision support.

Abstract

Electronic health record (EHR) data has emerged as a valuable resource for analyzing patient health status. However, the prevalence of missing data in EHR poses significant challenges to existing methods, leading to spurious correlations and suboptimal predictions. While various imputation techniques have been developed to address this issue, they often obsess unnecessary details and may introduce additional noise when making clinical predictions. To tackle this problem, we propose SMART, a Self-Supervised Missing-Aware RepresenTation Learning approach for patient health status prediction, which encodes missing information via elaborated attentions and learns to impute missing values through a novel self-supervised pre-training approach that reconstructs missing data representations in the latent space. By adopting missing-aware attentions and focusing on learning higher-order representations, SMART promotes better generalization and robustness to missing data. We validate the effectiveness of SMART through extensive experiments on six EHR tasks, demonstrating its superiority over state-of-the-art methods.
Paper Structure (23 sections, 4 equations, 3 figures, 7 tables, 1 algorithm)

This paper contains 23 sections, 4 equations, 3 figures, 7 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Methodology
  4. Variable Independent Encoder
  5. MART Block
  6. Two-Stage Training Strategy
  7. Experiments
  8. Experimental Settings
  9. Experimental Results
  10. Main Results
  11. Ablation Study
  12. Effect of Missingness
  13. Model Efficiency
  14. Conclusions
  15. Additional Experimental Results
  16. ...and 8 more sections

Figures (3)

  • Figure 1: Overview of SMART. Left: We randomly mask EHR data and conduct reconstruction in the latent space. The reconstruction targets are generated by EMA updated parameters. Right: We illustrate the detailed architecture of the input encoder and the MART block. The input encoder embeds each variable (which can also be referred to as a biomarker) and missing mask into a separate hidden space. The MART block employs various techniques to capture feature interactions in both the temporal and variable dimensions while further encoding missing information.
  • Figure 2: Performance on different observed ratio of EHR.
  • Figure 3: Training time, parameters, and AUPRC(%) of all models on the three datasets. The size of the circle represents the number of parameters. The GPU runtime is counted at the start of training.