SeqRisk: Transformer-augmented latent variable model for robust survival prediction with longitudinal data

In healthcare, risk assessment of patient outcomes has been based on survival analysis for a long time, i.e. modeling time-to-event associations. However, conventional approaches rely on data from a single time-point, making them suboptimal for fully leveraging longitudinal patient history and capturing temporal regularities. Focusing on clinical real-world data and acknowledging its challenges, we utilize latent variable models to effectively handle irregular, noisy, and sparsely observed longitudinal data. We propose SeqRisk, a method that combines variational autoencoder (VAE) or longitudinal VAE (LVAE) with a transformer-based sequence aggregation and Cox proportional hazards module for risk prediction. SeqRisk captures long-range interactions, enhances predictive accuracy and generalizability, as well as provides partial explainability for sample population characteristics in attempts to identify high-risk patients. SeqRisk demonstrated robust performance under conditions of increasing sparsity, consistently surpassing existing approaches.

💡 Research Summary

SeqRisk introduces a unified generative‑discriminative framework for survival analysis on irregular, high‑dimensional longitudinal electronic health records (EHR). The method first learns compact latent representations of multivariate time‑series measurements using either a standard variational autoencoder (VAE) or a longitudinal VAE (LVAE). The LVAE extends the VAE by placing an additive multi‑output Gaussian‑process (GP) prior on each latent dimension, conditioning on covariates such as patient ID, age, and gender. This GP prior explicitly models within‑subject temporal correlations and handles irregular sampling without requiring imputation.

The latent vectors for each visit are then treated as tokens for a transformer encoder. Tokens are embedded into a common hidden dimension, augmented with positional or time‑gap encodings that preserve the order and irregular intervals of observations. Multi‑head self‑attention captures both short‑range and long‑range dependencies, and an attention‑pooling layer aggregates the entire visit sequence into a single subject‑level vector. A small multilayer perceptron (MLP) maps this pooled representation to a nonlinear risk function fω(Zp, Xp).

Risk prediction follows the Cox proportional‑hazards formulation: h(t|Zp, Xp) = h0(t) · exp