Learning temporal embeddings from electronic health records of chronic kidney disease patients

TL;DR

The paper tackles the challenge of learning temporally rich, task-agnostic representations from longitudinal EHR data for model-guided medicine. It compares vanilla LSTM, attention-augmented LSTM, and time-aware LSTM (T-LSTM) on a CKD cohort from MIMIC-IV, training embeddings separately from end-to-end mortality predictors. The T-LSTM produced the most structured embedding space (DBI 9.91) and highest CKD-stage accuracy (0.74), while embedding-based downstream mortality prediction achieved higher accuracy (0.82–0.83) and AUROC/AUPRC (≈0.89–0.90) than direct end-to-end learning. The results support using temporal embeddings as a foundation for multi-task, clinically useful representation learning in digital patient models.

Abstract

We investigate whether temporal embedding models trained on longitudinal electronic health records can learn clinically meaningful representations without compromising predictive performance, and how architectural choices affect embedding quality. Model-guided medicine requires representations that capture disease dynamics while remaining transparent and task agnostic, whereas most clinical prediction models are optimised for a single task. Representation learning facilitates learning embeddings that generalise across downstream tasks, and recurrent architectures are well-suited for modelling temporal structure in observational clinical data. Using the MIMIC-IV dataset, we study patients with chronic kidney disease (CKD) and compare three recurrent architectures: a vanilla LSTM, an attention-augmented LSTM, and a time-aware LSTM (T-LSTM). All models are trained both as embedding models and as direct end-to-end predictors. Embedding quality is evaluated via CKD stage clustering and in-ICU mortality prediction. The T-LSTM produces more structured embeddings, achieving a lower Davies-Bouldin Index (DBI = 9.91) and higher CKD stage classification accuracy (0.74) than the vanilla LSTM (DBI = 15.85, accuracy = 0.63) and attention-augmented LSTM (DBI = 20.72, accuracy = 0.67). For in-ICU mortality prediction, embedding models consistently outperform end-to-end predictors, improving accuracy from 0.72-0.75 to 0.82-0.83, which indicates that learning embeddings as an intermediate step is more effective than direct end-to-end learning.

Figures (3)

• Figure 1: Overview of the approach. (a) Longitudinal EHR data are extracted and preprocessed to construct patient-specific time series. (b) Clinical trajectories are represented as irregularly sampled multivariate sequences with explicit time intervals. (c) Temporal models, including vanilla LSTM, attention-augmented LSTM, and T-LSTM, are trained as embedding models and end-to-end predictors. (d) The embedding models generate compact latent patient representations. (e) Embeddings are evaluated using intrinsic clustering analyses and extrinsic downstream prediction tasks.
• Figure 2: Recurrent model architectures evaluated in this study. (a) Vanilla LSTM, in which hidden and cell states are updated sequentially from the input time-series. (b) Attention-augmented LSTM, in which an attention mechanism over LSTM hidden states is applied to compute a context vector summarising the sequence. (c) Time-aware LSTM (T-LSTM), in which time elapsed between observations are explicitly incorporated into the cell state update to model irregularly sampled clinical data.
• Figure 3: t-SNE visualisation of embedding spaces separated by classes 0-7 based on ICD codes for CKD stages. (a) Embedding space of vanilla LSTM. (b) Embedding space of attention-augmented LSTM. (c) Embedding space of T-LSTM.