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

2026-01-26

AI Summary

This study addresses the challenge of learning generalizable temporal representations from longitudinal electronic health records of patients with chronic kidney disease (CKD). The authors propose a time-aware LSTM (T-LSTM) to learn task-agnostic clinical embeddings from the MIMIC-IV database. Compared to standard LSTM and attention-augmented LSTM variants, T-LSTM achieves improved generalization while preserving structured temporal representations. Experimental results demonstrate that the learned embeddings yield a CKD staging classification accuracy of 0.74 (Davies–Bouldin index: 9.91) and significantly enhance downstream ICU mortality prediction performance, achieving accuracies of 0.82–0.83—outperforming end-to-end trained baselines. These findings validate the effectiveness of T-LSTM as a robust framework for generating universal clinical time-series representations.

Technology Category

• Machine Learning: Time-Series/Data Streams
• Data Mining & Knowledge Management: Mining of Spatial, Temporal or Spatio-Temporal Data
• Knowledge Representation and Reasoning: Geometric, Spatial, and Temporal Reasoning

Application Category

• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• User Modeling, Personalization and Recommendation: Studies of user behavior, including longitudinal effects of personalized systems
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• temporal embeddings
• electronic health records
• chronic kidney disease
• representation learning
• disease dynamics

Innovation

• temporal embeddings
• time-aware LSTM
• representation learning
• chronic kidney disease
• electronic health records