Representation Learning of Lab Values via Masked AutoEncoders

TL;DR

This paper introduces Lab-MAE, a transformer-based masked autoencoder that improves the imputation of missing lab values in EHR data, outperforming existing methods and enhancing fairness across patient groups.

Contribution

We propose Lab-MAE, a novel self-supervised transformer framework that explicitly models temporal dependencies in lab data for improved imputation accuracy.

Findings

Lab-MAE outperforms baselines on MIMIC-IV across RMSE, R2, WD.

It achieves equitable performance across demographic groups.

Lab-MAE demonstrates robustness without follow-up lab data.

Abstract

Accurate imputation of missing laboratory values in electronic health records (EHRs) is critical to enable robust clinical predictions and reduce biases in AI systems in healthcare. Existing methods, such as XGBoost, softimpute, GAIN, Expectation Maximization (EM), and MICE, struggle to model the complex temporal and contextual dependencies in EHR data, particularly in underrepresented groups. In this work, we propose Lab-MAE, a novel transformer-based masked autoencoder framework that leverages self-supervised learning for the imputation of continuous sequential lab values. Lab-MAE introduces a structured encoding scheme that jointly models laboratory test values and their corresponding timestamps, enabling explicit capturing temporal dependencies. Empirical evaluation on the MIMIC-IV dataset demonstrates that Lab-MAE significantly outperforms state-of-the-art baselines such as XGBoost, softimpute, GAIN, EM, and MICE across multiple metrics, including root mean square error (RMSE), R-squared (R2), and Wasserstein distance (WD). Notably, Lab-MAE achieves equitable performance across demographic groups of patients, advancing fairness in clinical predictions. We further investigate the role of follow-up laboratory values as potential shortcut features, revealing Lab-MAE's robustness in scenarios where such data is unavailable. The findings suggest that our transformer-based architecture, adapted to the characteristics of EHR data, offers a foundation model for more accurate and fair clinical imputation. In addition, we measure and compare the carbon footprint of Lab-MAE with the a XGBoost model, highlighting its environmental requirements.