TrajGPT: Irregular Time-Series Representation Learning for Health Trajectory Analysis

TL;DR

Experimental results demonstrate that TrajGPT excels in trajectory forecasting, drug usage prediction, and phenotype classification without requiring task-specific fine-tuning.

Abstract

In many domains, such as healthcare, time-series data is often irregularly sampled with varying intervals between observations. This poses challenges for classical time-series models that require equally spaced data. To address this, we propose a novel time-series Transformer called Trajectory Generative Pre-trained Transformer (TrajGPT). TrajGPT employs a novel Selective Recurrent Attention (SRA) mechanism, which utilizes a data-dependent decay to adaptively filter out irrelevant past information based on contexts. By interpreting TrajGPT as discretized ordinary differential equations (ODEs), it effectively captures the underlying continuous dynamics and enables time-specific inference for forecasting arbitrary target timesteps. Experimental results demonstrate that TrajGPT excels in trajectory forecasting, drug usage prediction, and phenotype classification without requiring task-specific fine-tuning. By evolving the learned continuous dynamics, TrajGPT can interpolate and extrapolate disease risk trajectories from partially-observed time series. The visualization of predicted health trajectories shows that TrajGPT forecasts unseen diseases based on the history of clinically relevant phenotypes (i.e., contexts).

Figures (5)

• Figure 1: TrajGPT overview. (a). TrajGPT processes irregularly-sampled time series by embedding an input sequence with RoPE. (b). Each SRA layer comprises an SRA module and a feed-forward layer, with the SRA module capable of operating in both recurrent and parallel forms.
• Figure 2: TrajGPT as discretized ODEs. (a). TrajGPT performs interpolation and extrapolation by modeling continuous dynamics as discretized ODEs. (b). Time-specific inference directly predicts irregular samples using previous hidden states and target timesteps.
• Figure 3: Visualization of token embeddings and sequence representations. (a). Visualization of token embeddings across 15 disease categories, where token nodes are colored and clustered by categories. The head-specific decay vectors ${\bm{w}}_\gamma^h$ (marked with $\times$) indicate the alignment of heads with disease categories. (b). Visualization of sequence representations for diabetic patients, highlighting insulin usage within six months of diagnosis. The distinction of two classes enables zero-shot classification.
• Figure 4: Predicted health trajectories for a diabetic patient (top) and a CHF patient (bottom). Panels (a) and (b) show the inferred disease trajectories with look-up and forecast windows. Matched predictions (solid circles) occur when the top 10 predicted PheCodes match the ground-truth. Larger solid circles indicate correctly predicted diabetes or CHF. Panels (c) and (d) display the predicted risk trajectories, showing increasing risks with age. For a target disease, TrajGPT computes risk as the token probability at each timestep and calculates risk growth as the difference between consecutive values. We highlight key timesteps to indicate significant risk growth and the associated phenotypes.
• Figure 5: The distribution of top-10 recall performance for TrajGPT and baseline methods across three forecasting window sizes. The TrajGPT with time-specific inference achieves better and more stable performance compared with auto-regressive inference and other baselines.