CSAI: Conditional Self-Attention Imputation for Healthcare Time-series

TL;DR

CSAI addresses the challenge of missing data in healthcare time series by extending BRITS with three mechanisms: an attention-based conditional hidden-state initialisation to capture long-range dynamics, a domain-informed temporal decay that aligns with feature-specific clinical recording patterns, and a non-uniform masking strategy that reflects non-random missingness. The method achieves superior imputation accuracy and downstream predictive performance on multiple EHR benchmarks, with ablations confirming each component's contribution. Integrated into PyPOTS, CSAI provides a practical, open-source tool for reliable imputation and prognosis in clinical data, offering potential for improvements in tasks like sepsis detection and outcome prediction. Overall, CSAI advances neural imputation for EHRs by aligning algorithmic techniques with the realities of clinical data collection and missingness patterns.

Abstract

We introduce the Conditional Self-Attention Imputation (CSAI) model, a novel recurrent neural network architecture designed to address the challenges of complex missing data patterns in multivariate time series derived from hospital electronic health records (EHRs). CSAI extends state-of-the-art neural network-based imputation by introducing key modifications specific to EHR data: a) attention-based hidden state initialisation to capture both long- and short-range temporal dependencies prevalent in EHRs, b) domain-informed temporal decay to mimic clinical data recording patterns, and c) a non-uniform masking strategy that models non-random missingness by calibrating weights according to both temporal and cross-sectional data characteristics. Comprehensive evaluation across four EHR benchmark datasets demonstrates CSAI's effectiveness compared to state-of-the-art architectures in data restoration and downstream tasks. CSAI is integrated into PyPOTS, an open-source Python toolbox designed for machine learning tasks on partially observed time series. This work significantly advances the state of neural network imputation applied to EHRs by more closely aligning algorithmic imputation with clinical realities.

Figures (5)

• Figure 1: An example of multivariate time-series. Observations $\boldsymbol{x_{1-5}}$ in time-stamps $\boldsymbol{s_{1-5}}=0,\,4,\,5,\,7,\,9$. Feature $d_2$ was missing during $\boldsymbol{s_{2-4}}$, the last observation took place at $\boldsymbol{s_1}$. Hence, $\boldsymbol{\delta_5^2}=\boldsymbol{t_5-t_1} =9-0=9$.
• Figure 2: The BRITS backbone process
• Figure 3: The CSAI architecture, which equips BRITS with self-attention. CSAI begins with an input embedding layer, followed by positional embedding to capture time dependencies. These embeddings are processed through multi-head attention, normalization, and feed-forward layers. The output is used to initialize hidden states for subsequent recurrent layers, accounting for both temporal dynamics and domain-specific variability in recording patterns.
• Figure 4: Visualisation of the incremental increase in imputation performance (decrease in MAE) as the different components of CSAI are incrementally added. The figure shows CSAI's MAE using the three datasets and 10% masking ratio and the following variations: a) baseline model (i.e. a BRITS architecture), b) adding the temporal decay function, c) adding the attention-base initialisation mechanism, and finally d) using non-uniform masking.
• Figure 5: Impact of Adjustment Factor on Model Performance in the Physionet Dataset.

Theorems & Definitions (3)

• Example 1
• Example 2
• Example 3