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

## Abstract

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).

## 1 Introduction

Time-series representation learning plays a crucial role in various domains, as it facilitates the extraction of generalizable temporal patterns from large-scale, unlabeled data, which can then be adapted for diverse tasks (Ma et al., ). However, a major challenge arises when dealing with irregularly-sampled time series, in which observations occur at uneven intervals (Li & Marlin, ). This irregularity poses challenges for classical time-series models that are restricted to regular sampling (Ayala Solares et al., ; Zhang et al., ). This issue is particularly significant in the healthcare domain, since longitudinal electronic health records (EHRs) are updated sporadically during outpatient visits or inpatient stays (Zhang et al., ). Moreover, individual medical histories often span a limited timeframe due to a lack of historical digitization, incomplete insurance coverage, and fragmented healthcare systems (Wornow et al., ). These challenges make it difficult for time-series models to capture the underlying trajectory dynamics (Amirahmadi et al., ). Addressing these challenges requires the development of novel representation learning techniques that can extract generalizable temporal patterns from irregularly-sampled data through next-token prediction pre-training. The pre-trained model is then applied to forecast trajectories based on the learned transferable patterns, even when patient data is only partially observed.

Recent advances in modeling irregularly-sampled time series have been achieved through specialized deep learning architectures (Che et al., ; Horn et al., ; Rubanova et al., ; Shukla & Marlin, ; Zhang et al., ). However, these models fall short in pre-training generalizable representations. While time-series Transformer models have gained attention, they are primarily designed for consecutive data and fail to account for irregular intervals between observations (Nie et al., ; Zhou et al., ; Wu et al., ). To handle both regular and irregular time series, TimelyGPT incorporates relative position embedding to capture positional information in varying time gaps (Song et al., ). BiTimelyGPT extends this by pre-training bidirectional representations for discriminative tasks (Song et al., ). Despite these improvements, both models rely on a data-independent decay, which is not content-aware and thus cannot fully capture complex temporal dependencies in healthcare data. The key challenge remains to develop an effective representation learning approach that extracts meaningful patterns from irregularly-sampled data.

In this study, we propose Trajectory Generative Pre-trained Transformer (TrajGPT) for irregular time-series representation learning. Our research offers four major contributions: First, it introduces a Selective Recurrent Attention (SRA) mechanism with a data-dependent decay, enabling the model to adaptively forget irrelevant past information based on contexts. Second, by interpreting TrajGPT as discretized ODEs, it effectively captures the continuous dynamics in irregularly-sampled data; This enables TrajGPT to perform interpolation and extrapolation in both directions, allowing for a novel time-specific inference for accurate forecasting. Third, TrajGPT demonstrates strong zero-shot performance across multiple tasks, including trajectory forecasting, drug usage prediction, and phenotype classification. Finally, TrajGPT offers interpretable health trajectory analysis, enabling clinicians to align the extrapolated disease progression trajectory with underlying patient conditions.

## 2 Related Works

### 2.1 Time-series Transformer Models

Time-series Transformer models have demonstrated strong performance in modeling temporal dependencies through attention mechanisms (Wen et al., ). Informer introduces ProbSparse self-attention to extract key information by halving cascading layer input (Zhou et al., ). Autoformer utilizes Auto-Correlation to capture series-wise temporal dependencies (Wu et al., ). FEDformer adopts Fourier-enhanced attention to capture frequency-domain relationships (Zhou et al., ). PatchTST compresses time series into patches and forecasts all timesteps using a linear layer (Nie et al., ). Despite their effectiveness, these methods fail to account for irregular time intervals. TimelyGPT and BiTimelyGPT address this limitation by encoding irregular time gaps with relative position embedding (Song et al., ; ). However, these models rely on a data-independent decay, whereas TrajGPT introduces a data-dependent decay to adaptively forget irrelevant information based on contexts. ContiFormer integrates ODEs into attention’s key and value matrices to model continuous dynamics (Chen et al., ). However, it demands significantly more computing resources than a standard Transformer with quadratic complexity, due to the slow process of solving ODEs. In contrast, TrajGPT models continuous dynamics by pre-training on irregularly-sampled data with efficient linear training complexity and constant inference complexity.

### 2.2 Algorithms designed for Irregularly-sampled Time Series

Various techniques have been developed to model irregular temporal dependencies through specialized architectures. GRU-D captures temporal dependencies by applying exponential decay to hidden states (Che et al., ). SeFT adopts a set function based approach, where each observation is modeled individually and then pooled together (Horn et al., ). RAINDROP captures irregular temporal dependencies by representing data as separate sensor graphs (Zhang et al., ). mTAND employs a multi-time attention mechanism to learn irregular temporal dependencies (Shukla & Marlin, ). In continuous-time approaches, neural ODEs use neural networks to model complex ODEs, offering promising interpolation and extrapolation solutions (Chen et al., ). ODE-RNN further enhances this by updating RNN hidden states with new observations (Rubanova et al., ). However, these methods lack a representation learning paradigm and often struggle to capture evolving dynamics in partially-observed data. In contrast, our TrajGPT can be interpreted as discretized ODEs, allowing it to learn continuous dynamics via large-scale pre-training. Moreover, TrajGPT utilizes interpolation and extrapolation techniques from the neural ODE family to predict accurate trajectories.

Figure: 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.
Refer to caption: x1.png

## 3 Methodology

We denote an irregularly-sampled time series as $x=\{(x_{1},t_{1}),\ldots,(x_{N},t_{N})\}$, where $N$ is the total number of samples. Each sample $(x_{n},t_{n})$ consists of an observation $x_{n}$ (e.g., a structured diagnosis code) and its associated timestamp $t_{n}$. The notations of variables are defined in Appendix. A.

### 3.1 TrajGPT Methodology

In the TrajGPT architecture illustrated in Fig. 1.a, each input sequence $x$ is first projected onto a token embedding ${\bm{X}}\in\mathbb{R}^{N\times d}$, where $N$ and $d$ denote the number of tokens and embedding size, respectively. A Rotary Position Embedding (RoPE) is then added to the token embedding, encoding relative positional information between tokens $n$ and $m$ (Su et al., ). Specifically, RoPE handles varying time intervals by encoding its relative distance $t_{n}-t_{m}$:

$$ $\displaystyle{\bm{Q}}_{n}={\bm{X}}_{n}{\bm{W}}_{Q}e^{i\theta t_{n}},\>{\bm{K}} _{m}={\bm{X}}_{m}{\bm{W}}_{K}e^{-i\theta t_{m}},\>{\bm{V}}_{m}={\bm{X}}_{m}{ \bm{W}}_{V}.$ (1) $$

The resulting input embedding is then passed through $L$ SRA layers, each comprising an SRA module and a feed-forward layer. SRA module operates in either parallel or recurrent forms. In the *recurrent* forward pass, SRA computes the output representation ${\bm{O}}_{n}$ based on a state variable ${\bm{S}}$:

$$ $\displaystyle{\bm{S}}_{n}=\gamma_{n}{\bm{S}}_{n-1}+{\bm{K}}_{n}^{\top}{\bm{V}} _{n},\>{\bm{O}}_{n}={\bm{Q}}_{n}{\bm{S}}_{n},\>\text{where }\gamma_{n}=\text{ Sigmoid}({\bm{X}}_{n}\mathbf{w}_{\gamma}^{\top})^{\frac{1}{\tau}}.$ (2) $$

The *data-dependent* decay $\gamma_{n}\in(0,1]$ and learnable decay vector $\mathbf{w}_{\gamma}\in R^{1\times d}$ enable SRA to selectively forget irrelevant past information based on contexts. For chronic diseases, TrajGPT assigns higher $\gamma_{n}$ values to slow down forgetting and capture long-term dependencies. Conversely, lower $\gamma_{n}$ values accelerate decay and prioritize recent events, making it more responsive to acute conditions. To avoid rapid decay from small $\gamma_{n}$ values, we introduce a temperature parameter $\tau=20$ to help preserve information over long sequences. Given an initial state ${\bm{S}}_{0}$ = 0, we can rewrite the recurrent updates of ${\bm{S}}_{n}$ in Eq. 2 as a sum of contributions from all previous timesteps:

$$ $\displaystyle{\bm{S}}_{n}=\sum_{m=1}^{n}\left(\prod_{t=m+1}^{n}\gamma_{t} \right){\bm{K}}_{m}^{\top}{\bm{V}}_{m}=\sum_{m=1}^{n}\left(\frac{b_{n}}{b_{m}} \right){\bm{K}}_{m}^{\top}{\bm{V}}_{m},\>\text{where }b_{n}=\prod_{t=1}^{n} \gamma_{t},$ (3) $$

$b_{n}$ is the cumulative decay term for token $n$, and $b_{n}/b_{m}$
captures the relative decay between tokens $n$ and $m$. To express Eq. 3 in matrix form, we introduce a casual decay matrix ${\bm{D}}$ as:

$$ $\displaystyle{\bm{D}}=\begin{bmatrix}\frac{b_{1}}{b_{1}}&0&\cdots&0\\ \frac{b_{2}}{b_{1}}&\frac{b_{2}}{b_{2}}&\cdots&0\\ \vdots&\vdots&\ddots&\vdots\\ \frac{b_{N}}{b_{1}}&\cdots&\cdots&\frac{b_{N}}{b_{N}}\end{bmatrix}=\begin{ bmatrix}1&0&\cdots&0\\ \gamma_{2}&1&\cdots&0\\ \vdots&\vdots&\ddots&\vdots\\ \prod_{t=2}^{N}\gamma_{t}&\cdots&\cdots&1\end{bmatrix}$ (4) $$

With the summation form in Eq. 3 and the decay matrix ${\bm{D}}$, we reformulate Eq. 2 into matrix form:

$$ $\displaystyle{\bm{O}}_{n}={\bm{Q}}_{n}\sum_{m=1}^{n}\left(\frac{b_{n}}{b_{m}} \right){\bm{K}}_{m}^{\top}{\bm{V}}_{m}=\left(({\bm{Q}}_{n}{\bm{K}}^{\top}) \odot{\bm{D}}_{n}\right){\bm{V}}$ (5) $$

By computing output representations for all tokens, the *parallel* form of SRA is given by:

$$ $\displaystyle{\bm{O}}=({\bm{Q}}{\bm{K}}^{\top}\odot{\bm{D}}){\bm{V}},\>{\bm{D} }_{nm}=\begin{cases}\frac{b_{n}}{b_{m}},\>&n\geq m\\ 0,&n<m\end{cases}$ (6) $$

To enhance computational efficiency, we directly calculate each element ${\bm{D}}_{nm}=\prod_{t=m+1}^{n}\gamma_{t}$. We detail the equivalence between recurrence and parallelism in Appendix B. To capture a broader range of contexts, We extend the single-head SRA in Eq. 2 to a multi-head SRA:

$$ $\displaystyle{\bm{O}}_{n}^{h}={\bm{Q}}_{n}^{h}{\bm{S}}_{n}^{h},\>{\bm{S}}_{n}^ {h}=\gamma_{n}^{h}{\bm{S}}_{n-1}^{h}+{\bm{K}}_{n}^{h\top}{\bm{V}}_{n}^{h},\> \text{where }\gamma_{n}^{h}=\text{Sigmoid}({\bm{X}}_{n}\mathbf{w}_{\gamma}^{h \top})^{\frac{1}{\tau}},$ (7) $$

Head-specific decay $\gamma_{n}^{h}$ adjusts the influence of past information based on contexts, with ${\bm{w}}_{\gamma}^{h}$ encoding different aspects of medical expertise as demonstrated in Section 5.1.

Figure: 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.
Refer to caption: x2.png

### 3.2 TrajGPT as Discretized ODEs

In this section, we establish theoretical connection between our proposed SRA module and ODE. The recurrent form of SRA in Eq. 2 is a discretization of continuous-time ODE with zero-order hold (ZOH) rule (Gu et al., ). We provide a detailed proof in Appendix C.

Given a first-order ODE, we can derive our recurrent SRA (Eq. 2) using a ZOH discretization with a discrete step size $\Delta$:

$$ $\displaystyle{\bm{S}}^{\prime}(t)={\bm{A}}{\bm{S}}(t)+{\bm{B}}{\bm{X}}(t),\>{ \bm{O}}(t)={\bm{C}}{\bm{S}}(t)$ $\displaystyle\text{where }{\bm{A}}=\frac{\ln(\bm{\Lambda}_{t})}{\Delta},\>{\bm {B}}={\bm{A}}(e^{\Delta{\bm{A}}}-{\bm{I}})^{-1}{\bm{K}}_{t}^{\top},\>{\bm{C}}= {\bm{Q}}_{t},\>\bm{\Lambda}_{t}=\text{diag}(\mathbf{1}\gamma_{t}).$ (8) $$

This ODE naturally models the continuous dynamics underlying irregularly-sampled data, with $\Delta$ corresponding to the varying time intervals between observations. Since the parameters $({\bm{A}},{\bm{B}},{\bm{C}})$ depend on the $t$-th observation ${\bm{X}}(t)$, this continuous-time model becomes a neural ODE, ${\bm{S}}^{\prime}(t)=f({\bm{S}}(t),t,\theta_{t})$, with a differentiable neural network $f$ and data-dependent parameters $\theta_{t}=({\bm{A}},{\bm{B}},{\bm{C}})$ (Chen et al., ). Consequently, a single-head SRA serves as a discretized ODE with data-dependent parameters (i.e., neural ODE). TrajGPT with multi-head SRA operates as discretized ODEs, where each head corresponds to its own ODE and captures distinct dynamics.

As illustrated in Fig. 2.a, TrajGPT functions as discretized ODEs, enabling both interpolation and extrapolation of irregular time-series data. By capturing the underlying continuous dynamics, TrajGPT handles irregular input through discretization with varying step sizes. For interpolation, it simply evolves the dynamics within the observed timeframe using a unit discretization step size. For extrapolation, it evolves the dynamics forward or backward in time beyond the observed timeframe. Additionally, TrajGPT estimates disease risk trajectories by computing token probabilities at specific timesteps and evolving the dynamics through interpolation and extrapolation. A detailed trajectory analysis is provided in Section 5.4.

At inference time, we explore two strategies for forecasting irregularly-sampled time series: auto-regressive and time-specific inference (Fig. 2.b). Auto-regressive inference, commonly used by standard Transformer models, makes sequential predictions at equal intervals and selects the target timesteps accordingly. Given that TrajGPT functions as discretized ODEs, we introduce a novel *time-specific inference* to predict at arbitrary timesteps. To forecast a target time point $(x_{n^{\prime}},t_{n^{\prime}})$, TrajGPT utilizes both the target timestep $t_{n^{\prime}}$ and the last observation $(x_{n},t_{n})$ to predict the corresponding observation $x_{n^{\prime}}$. It calculates the target output representation ${\bm{O}}_{n^{\prime}}={\bm{Q}}_{n^{\prime}}{\bm{S}}_{n^{\prime}}$, taking into account the discrete step size $\Delta t_{n^{\prime},n}=t_{n^{\prime}}-t_{n}$ and the updated state ${\bm{S}}_{n^{\prime}}={\bm{D}}_{\Delta{t}_{n^{\prime},n}}{\bm{S}}_{n}+{\bm{K}}
_{n}^{\top}{\bm{V}}_{n}$.

### 3.3 Computational complexity

TrajGPT with its efficient SRA mechanism achieves linear training complexity of $O(N)$ and constant inference complexity of $O(1)$ with respect to sequence length $N$. In contrast, standard Transformer models incur quadratic training complexity of $O(N^{2})$ and linear inference complexity of $O(N)$ (Katharopoulos et al., ). This computational bottleneck arises from the vanilla self-attention mechanism, where $\text{Attention}({\bm{X}})=\text{Softmax}({\bm{Q}}{\bm{K}}^{T}){\bm{V}}$, resulting in a training complexity of $O(N^{2}d)$. When dealing with long sequences (i.e., $N>>d$), the quadratic term $O(N^{2})$ becomes a bottleneck for standard Transformer models.

As a variant of linear attention (Katharopoulos et al., ), the SRA mechanism in TrajGPT overcomes this quadratic bottleneck of standard Transformer, achieving linear training complexity for long sequences. In recurrent SRA, ${\bm{O}}_{n}={\bm{Q}}_{n}{\bm{S}}_{n},{\bm{S}}_{n}=\gamma_{n}{\bm{S}}_{n-1}+{
\bm{K}}_{n}^{\top}{\bm{V}}_{n}$, both ${\bm{Q}}_{n}{\bm{S}}_{n}$ and ${\bm{K}}_{n}^{\top}{\bm{V}}_{n}$ have $O(d^{2})$ complexity. By recursively updating over $N$ tokens, the total complexity becomes $O(Nd^{2})$. For inference, TrajGPT proposes auto-regressive and time-specific methods. The auto-regressive inference sequentially generates sequences with equally spaced time intervals like the GPT model, incurring linear complexity of $O(N)$. In contrast, time-specific inference directly predicts the target time point with constant complexity of $O(1)$. Thus, TrajGPT achieves $O(N)$ training complexity and $O(1)$ inference complexity, making it computationally efficient for long sequences.

## 4 Experimental Design

### 4.1 Dataset and Pre-processing

Population Health Record (PopHR) database hosts massive amounts of longitudinal claim data from the provincial government health insurer in Quebec, Canada on health service use (Shaban-Nejad et al., ; Yuan et al., ). In total, there are approximately 1.3 million participants in the PopHR database, representing a randomly sampled 25% of the population in the metropolitan area of Montreal between 1998 and 2014. Cohort memberships are maintained dynamically by removing deceased residents and actively enrolling newborns and immigrants. We extracted irregularly-sampled time series from the PopHR database. Specifically, we converted ICD-9 diagnostic codes to integer-level phenotype codes (PheCodes) using the [PheWAS catalog](https://phewascatalog.org/phecodes) (Denny et al., ; ). We selected 194 unique PheCodes, each with over 50,000 occurrences. We excluded patients with fewer than 50 PheCode records, resulting in a final dataset of 489,000 patients, with an average of 112 records per individual.
The dataset was then split into training (80%), validation (10%), and testing (10%) sets.

### 4.2 Quantitative Experiment Design

Forecast irregular diagnostic codes We evaluated the long-term forecasting task using irregularly-sampled time series. Using a look-up window of 50 time points (e.g., diagnosis codes), we predicted the remaining codes in the forecasting windows. To measure performance, we used the top-$K$ recall metric with $K=(5,10,15)$. This metric mimics the behavior of doctors conducting differential diagnosis, where they list the most probable diagnoses based on a patient’s symptoms Choi et al. (). Since next-token prediction is inherently a forecasting task, TrajGPT enables zero-shot forecasting without requiring fine-tuning.

Drug usage prediction In this application, we predict whether each diabetic patient started insulin treatment within 6 months of their initial diabetes diagnosis. Following the preprocessing from previous work (Song et al., ), we extracted 78,712 diabetic patients with PheCode 250, where 11,433 patients were labeled as positive. Due to class imbalance, we use the area under precision-recall curve (AUPRC) as the evaluation metric. To avoid information leakage, we applied average pooling on the token embeddings truncated at the first diabetes record. To assess the generalizability of our pre-trained model, we performed zero-shot classification, few-shot classification with 5 samples, and fine-tuning on the full dataset for insulin usage prediction.

Phenotype classification PopHR database provides rule-based labels for congestive heart failure (CHF), with 3.2% of the total population labeled as positive. Given the class imbalance, we utilize the AUPRC metric to evaluate performance on the rare positive class. To assess the generalizability of the pre-trained TrajGPT, we conducted zero-shot classification, few-shot classification with 5 samples, and fine-tuning on the entire dataset.

### 4.3 Baselines

We compared our model against several time-series transformer baselines, including TimelyGPT (Song et al., ), BiTimelyGPT (Song et al., ), Informer (Zhou et al., ), Fedformer (Zhou et al., ), AutoFormer (Wu et al., ), PatchTST (Nie et al., ), TimesNet (Wu et al., ), and ContiFormer (Chen et al., ). BiTimelyGPT and PatchTST are encoder-only models that require fine-tuning for forecasting tasks, while other Transformer models with decoders can forecast without additional fine-tuning. We also evaluated models designed for irregularly-sampled time series, including mTAND (Shukla & Marlin, ), GRU-D (Che et al., ), RAINDROP (Zhang et al., ), SeFT (Horn et al., ), and ODE-RNN (Rubanova et al., ). Since these models do not have a pre-training method, they were trained from scratch on the training set. We followed previous works to set Transformer parameters to about 7.5 million (Table 4).

Transformer Pre-training paradigm With a cross-entropy loss, TrajGPT employs a next-token prediction task to pre-train generalizable temporal representations from unlabeled data (Radford et al., ). Given a sequence with a [SOS] token, TrajGPT predicts subsequent tokens by shifting the sequence to the right. The output representation of each token is fed into a linear layer for next-token prediction. For other models without an established pre-training paradigm, we employed a masking-based method by randomly masking 40% of timesteps with zeros (Zerveas et al., ). All Transformer models performed 20 epochs of pre-training with cross-entropy loss. When fine-tuning was applicable, we performed 5 epochs of end-to-end fine-tuning on the entire model.

Figure: 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.
Refer to caption: x3.png

## 5 Results

### 5.1 Qualitative Analysis of Embeddings

In this section, we provided a qualitative analysis of the token embeddings and sequence representations learned by our TrajGPT on the PopHR database (Fig. 3). We applied Uniform Manifold Approximation and Projection (UMAP) to visualize the global token embedding, with nodes colored and clustered by disease categories. The results reveal 12 clearly separated clusters. Some nodes are projected into other categories but still reflect meaningful clinical relationships; for instance, the mental disorders cluster (in green color) includes a black dot representing adverse drug events and drug allergies, implying high risk of opioid usage among the psychiatric group. Related disease categories with clinical relevance tend to cluster near each other. For example, mental disorders are closely clustered with neurological diseases, and circulatory diseases are adjacent to endocrine/metabolic diseases. We also visualized the head-specific decay vectors ${\bm{w}}_{\gamma}^{h}$ in Eq. 7 (marked as $\times$), which align well with specific disease categories. For instance, heads 7 and 8 align with circulatory and endocrine/metabolic diseases, respectively. These alignments suggest that these heads capture specialized clinical knowledge, effectively modeling the unique dynamics within each category.

In Fig. 3.b, we visualize the sequence representations to demonstrate TrajGPT’s ability to perform zero-shot classification of initial insulin usage among diabetic patients. To prevent information leakage, the sequence representations were truncated at the first diabetes record. These sequence representations were projected onto the same scale as the token embeddings in Fig. 3.a, allowing for direct comparison with the disease clusters. Patients taking future insulin treatment have embeddings closely aligned with the endocrine/metabolic cluster, indicating a strong association with diabetes-related conditions. In contrast, non-insulin patients are dispersed across various clusters, suggesting less severe diabetes histories. The clear separation between these groups highlights TrajGPT’s ability to perform zero-shot classification, showcasing the generalizability of its learned representations.

### 5.2 Forecasting Irregular Diagnoses Codes

**Table 1: The quantitative results on the diagnosis forecasting, insulin usage, and CHF classification. performance on PopHR’s irregular-sampled time series dataset. TrajGPT achieved the highest recall at $K=10$ and $K=15$. TrajGPT also outperforms other baseline methods on zero-shot evaluation on the prediction of insulin and CHF. The bold and underline indicate the best and second best results, respectively. $S$ indicates the number of few-shot examples. $-$ indicates non-applicable.**
| Methods / Tasks (%) | Forecasting | Diabetes-Insulin | CHF |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| K = 5 | 10 | 15 | S = 0 | 5 | all | S = 0 | 5 | all |  |
| TrajGPT (Time-specific) | 57.4 | 71.7 | 84.1 | 67.2 | 70.2 | 75.5 | 72.8 | 75.9 | 83.9 |
| TrajGPT (Auto-regressive) | 53.3 | 65.5 | 77.2 | — | — | — | — | — | — |
| TimelyGPT | 58.2 | 70.3 | 82.1 | 58.2 | 64.4 | 70.7 | 66.9 | 71.0 | 81.2 |
| BiTimelyGPT | 48.2 | 63.3 | 70.5 | 65.3 | 70.8 | 75.8 | 70.4 | 74.5 | 83.8 |
| Informer | 46.4 | 60.1 | 71.2 | 62.1 | 66.2 | 71.5 | 62.9 | 67.4 | 80.8 |
| Autoformer | 42.9 | 57.4 | 68.6 | 63.5 | 66.8 | 72.7 | 65.3 | 69.6 | 81.6 |
| Fedformer | 43.3 | 58.3 | 69.6 | 64.2 | 68.4 | 73.1 | 68.2 | 69.8 | 81.9 |
| PatchTST | 48.2 | 65.5 | 73.3 | 66.8 | 69.7 | 75.1 | 72.2 | 76.3 | 84.2 |
| TimesNet | 46.5 | 64.3 | 71.5 | 64.2 | 67.9 | 72.8 | 67.8 | 72.5 | 82.6 |
| ContiFormer | 52.8 | 67.2 | 76.9 | 63.5 | 68.0 | 75.0 | 68.4 | 74.9 | 83.1 |
| MTand | 52.6 | 70.2 | 83.7 | — | — | 74.6 | — | — | 85.4 |
| GRU-D | 54.2 | 69.5 | 80.5 | — | — | 72.1 | — | — | 79.9 |
| RAINDROP | 46.5 | 67.2 | 72.2 | — | — | 70.5 | — | — | 82.4 |
| SeFT | 49.3 | 68.1 | 79.4 | — | — | 71.7 | — | — | 83.4 |
| ODE-RNN | 54.7 | 70.6 | 78.6 | — | — | 73.5 | — | — | 82.9 |

TrajGPT with time-specific inference achieves the highest recall at $K=10$ and $K=15$, with scores of 71.7% and 84.1%, respectively (Table 1). At $K=5$, TrajGPT achieves the second-highest recall with 57.4%. Notably, time-specific inference outperforms the auto-regressive inference approach, demonstrating its effectiveness in forecasting based on the learned continuous dynamics. These results highlight TrajGPT’s strength in pre-training on underlying dynamics from sparse and irregular time-series data, facilitating accurate trajectory forecasting over irregular time intervals.

We then examined the distributions of top-10 recall across three forecast windows, comparing the two inference methods of TrajGPT as well as TimelyGPT, PatchTST, and mTAND (Fig. 5). TrajGPT’s time-specific inference consistently outperforms auto-regressive inference as the forecasting window increases, as it accounts for evolving states and query timesteps over irregular intervals. As expected, all models experience a performance decline as the forecast window increases, reflecting the increased uncertainty in long-term trajectory prediction. Despite this, TrajGPT achieves superior and more stable performance within the first 100 steps. In comparison, PatchTST shows a drastic decline as the window size increases, reflecting its difficulty with extrapolation. Therefore, TrajGPT excels in forecasting health trajectories through its time-specific inference.

### 5.3 Trajectory Classification

In this section, we evaluated trajectory classification on two applications: insulin usage prediction and CHF phenotype classification. We evaluated the performance of Transformer models under three settings: zero-shot learning, few-shot learning with $S=5$ samples, and fine-tuning on the entire dataset. Notably, non-Transformer models designed for irregularly-sampled time series (i.e., the last five methods in Table. 1) were trained from scratch. The results are summarized in Table. 1.

For insulin usage and CHF classification, TrajGPT achieves the highest zero-shot results, with 67.2% for insulin and 72.8% for CHF. This success can be attributed to TrajGPT’s ability to learn distinct clusters of sequence representations, as discussed in Section 5.1. For 5-shot classification, TrajGPT achieves the second-best performance in both tasks. For fine-tuning, it obtains the second best performance of 83.9% in insulin prediction, only 0.3% behind the best-performing BiTimelyGPT. We also compared TrajGPT with algorithms specifically designed for irregularly-sampled time series. These methods generally perform worse in insulin usage prediction, likely due to their difficulty in capturing meaningful temporal dependencies from truncated sequences. However, mTAND outperforms all models in the CHF task, achieving the best result at 85.4%.

Figure: 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.
Refer to caption: x4.png

### 5.4 Trajectory Analysis

In this analysis, we aimed to demonstrate TrajGPT’s effectiveness in trajectory modeling and provide insights into its classification performance. To achieve this, we conducted case studies on two patients: one diagnosed with diabetes and another with CHF. We visualized the observed and predicted disease trajectories for both patients, along with estimated risk trajectories over their lifetimes. As discussed in Section 3.2, we interpolated risks within the observed timeframe and extrapolated beyond it in both directions, computing risk as the token probability at each timestep. We also calculated risk growth by comparing each timestep to the previous one, identifying the ages with high risk growth as well as the associated phenotypes. By comparing disease and risk trajectories, we evaluated phenotype progression, disease comorbidity, and long-term risk development.

In Fig. 4.a, TrajGPT with time-specific inference achieves a top-10 recall of 90.1% for this diabetic patient. TrajGPT accurately predicts most diseases in the endocrine/metabolic and circulatory systems. Although this patient has no prior diabetes diagnosis in the observed data, TrajGPT successfully forecasts diabetes onset by identifying related metabolic and circulatory symptoms. Fig. 4.b illustrates the predicted risk trajectory for this patient, indicating a gradual increase in diabetes risk with age. We highlight specific phenotypes that contribute to the noticeable high risk growth between ages 59 and 62, including chronic IHD, hypothyroidism, obesity, and arrhythmia. These conditions are common comorbidities of diabetes, substantially elevating the likelihood of diabetes onset over time. In Fig. 4.c, we visualize the disease trajectory of a CHF patient, for whom TrajGPT produces a top-10 recall of 84.7%. TrajGPT accurately predicts a broad range of circulatory, respiratory, and endocrine/metabolic diseases. Despite the absence of prior CHF diagnosis, TrajGPT successfully predicts the onset of CHF based on a series of related circulatory conditions. In Fig. 4.d, the predicted risk trajectory reveals two spikes in risk growth at ages 65 and 74, corresponding to successive occurrences of circulatory diseases. This analysis demonstrates TrajGPT’s ability to forecast unseen phenotypes based on disease comorbidity and the risking risk with age. As a result, TrajGPT’s ability to model health trajectories and capture disease progression enhances its classification performance.

**Table 2: Ablation results of TrajGPT by selectively removing components and comparing inference methods. Performance is evaluated on forecasting task with a the of top-10 recall.**
| Forecast irregular diagnosis codes ($K$=10) | Time-specific inference | Auto-regressive inference |
| --- | --- | --- |
| TrajGPT | 71.7 | 65.5 |
| w/o decay gating (i.e., fixed $\gamma$) | 70.3 | 64.0 |
| w/o RoPE (i.e., absolute PE) | 67.8 | 63.2 |
| w/o linear attention (i.e., GPT-2) | — | 61.2 |

### 5.5 Ablation study

To evaluate the contributions of key components in TrajGPT, we performed ablation studies by selectively removing components such as decay gating, RoPE, and the linear attention module. We compared the time-specific inference and auto-regressive inference under different ablation setups. Notably, removing all components results in a vanilla GPT-2, which can only perform auto-regressive inference. The ablation studies were assessed on the forecasting task using the top-10 recall metric.

As shown in Table 2, removing the data-dependent decay and RoPE results in performance declines of 1.4% and 2.5%, respectively. This highlights the critical role of these modules in handling irregular time intervals. Replacing time-specific inference with auto-regressive inference leads to performance drops ranging from 4.6% to 6.2%, with the most significantly drop in TrajGPT. Furthermore, vanilla GPT-2 with auto-regressive inference produces the lowest performance, falling behind TrajGPT with time-specific inference by 10.5%.

## 6 Conclusion and Further Work

Our research proposes a novel architecture, TrajGPT, designed for irregular time-series representation learning and health trajectory analysis. To achieve this, TrajGPT introduces an SRA mechanism with a data-dependent decay, allowing the model to selectively forget irrelevant past information based on contexts. By interpreting TrajGPT as discretized ODEs, it effectively captures the continuous dynamics underlying irregularly-sampled time series, enabling both interpolation and extrapolation. For the forecasting task, TrajGPT provides an effective time-specific inference by evolving the dynamics according to varying time intervals. TrajGPT demonstrates strong zero-shot performance across multiple tasks, including diagnosis forecasting, drug usage prediction, and phenotype classification. TrajGPT also provides interpretable trajectory analysis, aiding clinicians in understanding the extrapolated disease progression along with risk growth.

To further validate generalizability, we plan to extend experiments to public datasets like MIMIC (Johnson et al., ), EHRSHOT (Wornow et al., ), and INSPECT (Huang et al., ). Additionally, we will compare TrajGPT with foundation LLMs, such as GPT-based (Wang et al., ) and Llama-based (Rasul et al., ) models. While we focus on in-domain data, we will explore representation learning and trajectory analysis on out-of-distribution data.

## Appendix A Denotations of Variables

**Table 3: Notations in TrajGPT**
| Notations | Descriptions | Notations | Descriptions |
| --- | --- | --- | --- |
| $x=\{(x_{1},t_{1}),\ldots,(x_{N},t_{N})\}$ | An irregularly-sample time series | $N$ | Number of tokens |
| $x_{n}$ | An observation | $t_{n}$ | Corresponding timestamp |
| ${\bm{X}}\in\mathbb{R}^{N\times d}$ | A sequence of tokens | $d$ | Hidden dimension |
| $L$ | Number of layers | $H$ | Number of Heads |
| ${\bm{Q}},{\bm{K}},{\bm{V}}\in\mathbb{R}^{N\times d}$ | Query, key, value matrices | ${\bm{W}}_{Q},{\bm{W}}_{K},{\bm{W}}_{V}\in\mathbb{R}^{d\times d}$ | Projection matrices for ${\bm{Q}},{\bm{K}},{\bm{V}}$ |
| ${\bm{O}}\in\mathbb{R}^{N\times d}$ | Output embedding | ${\bm{S}}\in\mathbb{R}^{d\times d}$ | State variable |
| $\theta$ | Rotary angle hyperparameter | $\gamma\in(0,1]$ | Data-dependent decay |
| $\mathbf{w_{\gamma}}\in\mathbb{R}^{d\times 1}$ | Decay weight vector | $\tau=20$ | Temperature term |
| $b_{n}=\prod_{t=1}^{n}\gamma_{t}$ | Cumulative decay | ${\bm{D}}\in\mathbb{R}^{N\times N}$ | Decay matrix |

## Appendix B Derivation of SRA Layer

Starting from the recurrent form of the TrajGPT model in Eq. 2, we derive the state variable $S$ by assuming ${\bm{S}}_{0}=0$:

$$ $\displaystyle{\bm{S}}_{n}=\gamma_{n}{\bm{S}}_{n-1}+{\bm{K}}_{n}^{\top}{\bm{V}} _{n}$ $\displaystyle{\bm{S}}_{1}={\bm{K}}_{1}^{\top}{\bm{V}}_{1}$ $\displaystyle{\bm{S}}_{2}=\gamma_{2}{\bm{K}}_{1}^{\top}{\bm{V}}_{1}+{\bm{K}}_{ 2}^{\top}{\bm{V}}_{2}$ $\displaystyle{\bm{S}}_{3}=\gamma_{3}\gamma_{2}{\bm{K}}_{1}^{\top}{\bm{V}}_{1}+ \gamma_{2}{\bm{K}}_{2}^{\top}{\bm{V}}_{2}+{\bm{K}}_{3}^{T}{\bm{V}}_{3}$ $\displaystyle\vdots$ $\displaystyle{\bm{S}}_{n}=\sum_{m=1}^{n}\left(\prod_{t=m+1}^{n}\gamma_{t} \right){\bm{K}}_{m}^{\top}{\bm{V}}_{m}=\sum_{m=1}^{n}\left(\frac{b_{n}}{b_{m}} \right){\bm{K}}_{m}^{\top}{\bm{V}}_{m},\>\text{where }b_{n}=\prod_{t=1}^{n} \gamma_{t},$ (9) $$

where we get the generalized updates of ${\bm{S}}_{n}$ using the cumulative decay term $b_{n}=\prod_{t=1}^{n}\gamma_{t}$. We can compute the output representation ${\bm{O}}_{n}$ using ${\bm{Q}}_{n}$ and ${\bm{S}}_{n}$:

$$ $\displaystyle{\bm{O}}_{n}$ $\displaystyle={\bm{Q}}_{n}{\bm{S}}_{n}={\bm{Q}}_{n}\sum_{m=1}^{n}\left(\frac{b _{n}}{b_{m}}\right){\bm{K}}_{m}^{\top}{\bm{V}}_{m}.$ (10) $$

To represent Eq. 10 in matrix form,
we introduce a causal decay matrix ${\bm{D}}$, where each element ${\bm{D}}_{nm}=\prod_{t=m+1}^{n}\gamma_{t}$
represents the decay relationship between two tokens $n$ and $m$:

$$ $\displaystyle{\bm{D}}=\begin{bmatrix}\frac{b_{1}}{b_{1}}&0&\cdots&0\\ \frac{b_{2}}{b_{1}}&\frac{b_{2}}{b_{2}}&\cdots&0\\ \vdots&\vdots&\ddots&\vdots\\ \frac{b_{N}}{b_{1}}&\cdots&\cdots&\frac{b_{N}}{b_{N}}\end{bmatrix}=\begin{ bmatrix}1&0&\cdots&0\\ \gamma_{2}&1&\cdots&0\\ \vdots&\vdots&\ddots&\vdots\\ \prod_{t=2}^{N}\gamma_{t}&\cdots&\cdots&1\end{bmatrix}.$ (11) $$

Using this decay matrix ${\bm{D}}$, we give the matrix form of the recurrent updates of $O_{n}$ in Eq. 10:

$$ $\displaystyle{\bm{O}}_{n}$ $\displaystyle={\bm{Q}}_{n}\sum_{m=1}^{n}{\bm{D}}_{nm}{\bm{K}}_{m}^{\top}{\bm{V }}_{m}$ $\displaystyle={\bm{Q}}_{n}\left({\bm{D}}_{n1}{\bm{K}}_{1}^{\top}{\bm{V}}_{1}+ \dots+{\bm{D}}_{nn}{\bm{K}}_{n}^{\top}{\bm{V}}_{n}\right)$ $\displaystyle={\bm{Q}}_{n}({\bm{D}}_{n1}{\bm{K}}_{1}^{\top}{\bm{V}}_{1}+\dots+ {\bm{D}}_{nn}{\bm{K}}_{n}^{\top}{\bm{V}}_{n}+\underbrace{{\bm{D}}_{n,n+1}}_{0} {\bm{K}}_{n+1}^{\top}{\bm{V}}_{n+1}+\dots+\underbrace{{\bm{D}}_{nN}}_{0}{\bm{K }}_{N}^{\top}{\bm{V}}_{N})$ $\displaystyle=\left(({\bm{Q}}_{n}{\bm{K}}^{\top})\odot{\bm{D}}_{n}\right){\bm{ V}}.$ (12) $$

To express the computation of all tokens, we obtain the parallel form of SRA as follows:

$$ $\displaystyle{\bm{O}}=({\bm{Q}}{\bm{K}}^{\top}\odot{\bm{D}}){\bm{V}},\>{\bm{D} }_{nm}=\begin{cases}\frac{b_{n}}{b_{m}},\>&n\geq m\\ 0.&n<m\end{cases}.$ (13) $$

## Appendix C TrajGPT as SSM and Neural ODE

The continuous-time SSM defines a linear mapping from an $t$-step input signal ${\bm{X}}(t)$ to output ${\bm{O}}(t)$ via a state variable ${\bm{S}}(t)$. It is formulated as a first-order ODE:

$$ $\displaystyle{\bm{S}}^{\prime}(t)={\bm{A}}{\bm{S}}(t)+{\bm{B}}{\bm{X}}(t),\>{ \bm{O}}(t)={\bm{C}}{\bm{S}}(t),$ (14) $$

where ${\bm{A}},{\bm{B}},{\bm{C}}$ denote the state matrix, input matrix, and output matrix respectively. Since data in real-world is typically discrete instead of continuous, continuous-time SSMs require discretization process to align with the sample rate of the data. With the discretization via zero-order hold (ZOH) rule (Gu et al., ), this continuous-time SSM in Eq. 14 becomes a discrete-time model:

$$ $\displaystyle{\bm{S}}_{t}=\bar{{\bm{A}}}{\bm{S}}_{t-1}+\bar{{\bm{B}}}{\bm{X}}_ {t},\>{\bm{O}}_{t}={\bm{C}}{\bm{S}}_{t}$ $\displaystyle\bar{{\bm{A}}}=e^{\Delta{\bm{A}}},\>\bar{{\bm{B}}}=(e^{\Delta{\bm {A}}}-{\bm{I}}){\bm{A}}^{-1}{\bm{B}},$ (15) $$

where $\bar{{\bm{A}}}$ and $\bar{{\bm{B}}}$ are the discretized matrices and $\Delta$ is the discrete step size. We provide a detailed derivation of ZOH discretization in
Appendix D.

Here, we show that a single-head SRA module is a special case of the discrete-time SSM defined in Eq. C, and then we derive its corresponding continuous-time SSM.
To achieve it, we first rewrite the recurrent SRA (Eq. 2) as follows:

$$ $\displaystyle{\bm{S}}_{t}$ $\displaystyle=\bm{\Lambda}_{t}{\bm{S}}_{t-1}+{\bm{K}}_{t}^{\top}{\bm{V}}_{t},$ $\displaystyle{\bm{O}}_{t}$ $\displaystyle={\bm{Q}}_{t}{\bm{S}}_{t},$ (16) $$

where $\bm{\Lambda}_{t}=\text{diag}(\mathbf{1}\gamma_{t})$ is a diagonal matrix with all diagonal elements equal to $\gamma_{t}$. In this way, the recurrent form of SRA in Eq. C corresponds to the discrete-time SSM defined in Eq. C, with $(\bar{\bm{A}},\bar{\bm{B}},{\bm{C}})=({\bm{\Lambda}}_{t},{\bm{K}}_{t}^{\top},{
\bm{Q}}_{t})$.
Assuming ZOH discretization, the parameters for the corresponding continuous-time SSM defined in Eq. 14 can be expressed as follows:

$$ $\displaystyle\begin{cases}\bar{{\bm{A}}}=e^{\Delta{\bm{A}}}=\bm{\Lambda}_{t}, \\ \bar{{\bm{B}}}=(e^{\Delta{\bm{A}}}-{\bm{I}}){\bm{A}}^{-1}{\bm{B}}={\bm{K}}_{t} ^{\top},\\ {\bm{C}}={\bm{Q}}_{t}.\end{cases}\Longrightarrow\begin{cases}{\bm{A}}=\frac{ \ln(\bm{\Lambda}_{t})}{\Delta},\\ {\bm{B}}={\bm{A}}(e^{\Delta{\bm{A}}}-{\bm{I}})^{-1}{\bm{K}}_{t}^{\top},\\ {\bm{C}}={\bm{Q}}_{t}\end{cases}$ (17) $$

As a result, our recurrent SRA can be interpreted as a discretized ODE.
Note that the ODE parameters $({\bm{A}},{\bm{B}},{\bm{C}})$ in Eq. 17 are data-dependent with respect to the $t$-th observation ${\bm{X}}_{t}$. Therefore, this continuous-time ODE is actually a neural ODE, ${\bm{S}}^{\prime}(t)=f({\bm{S}}(t),t,\theta_{t})$, with a differentiable neural network $f$ and data-dependent parameters $\theta_{t}=({\bm{A}},{\bm{B}},{\bm{C}})$ (Chen et al., ). The continuous dynamics underlying the irregular sequences are models by a neural ODE as follows:

$$ $\displaystyle{\bm{S}}^{\prime}(t)={\bm{A}}{\bm{S}}(t)+{\bm{B}}{\bm{X}}(t)=f({ \bm{S}}(t),t,\theta),\>{\bm{O}}(t)={\bm{C}}{\bm{S}}(t)$ $\displaystyle\text{where }{\bm{A}}=\frac{\ln(\bm{\Lambda}_{t})}{\Delta},\>{\bm {B}}={\bm{A}}(e^{\Delta{\bm{A}}}-{\bm{I}})^{-1}{\bm{K}}_{t}^{\top},\>{\bm{C}}= {\bm{Q}}_{t},\>\bm{\Lambda}_{t}=\text{diag}(\mathbf{1}\gamma_{t}).$ (18) $$

Consequently, a single-head SRA serves as a discretized (neural) ODE model. When we generalize the multi-head scenario, TrajGPT can be considered as discretized ODEs, where each head of SRA corresponds to its own ODE and captures distinct dynamics.

## Appendix D Proof of SSM Discretization via ZOH rule

To discretize the continuous-time model SSM, it has to compute the cumulative updates of the state ${\bm{S}}(t)$ over a discrete step size. For the continuous ODE in Eq. 14, we have a continuous-time integral as follows:

$$ $\displaystyle{\bm{S}}^{\prime}(t)={\bm{A}}{\bm{S}}(t)+{\bm{B}}{\bm{X}}(t)$ $\displaystyle{\bm{S}}(t+1)={\bm{S}}(t)+\int_{t}^{t+1}\left({\bm{A}}{\bm{S}}( \tau)+{\bm{B}}{\bm{X}}(\tau)\right)d\tau$ (19) $$

In the discrete-time system, we need to rewrite the integral as we cannot obtain all values of ${\bm{X}}(\tau)$ over a continuous interval $t\rightarrow t+1$:

$$ $\displaystyle{\bm{S}}(t+1)={\bm{S}}(t)+\sum_{t}^{t+1}({\bm{A}}{\bm{S}}(\tau)+{ \bm{B}}{\bm{X}}(\tau)\Delta\tau$ (20) $$

We replace ${\bm{X}}(t)$ in the time derivative ${\bm{S}}^{\prime}(t)$ as follows:

$$ $\displaystyle{\bm{S}}^{\prime}(t)={\bm{A}}{\bm{S}}(t)+{\bm{B}}{\bm{X}}(t)$ $\displaystyle{\bm{S}}^{\prime}(t)-{\bm{A}}{\bm{S}}(t)={\bm{B}}{\bm{X}}(t)$ $\displaystyle e^{-{\bm{A}}t}{\bm{S}}^{\prime}(t)-e^{-{\bm{A}}t}{\bm{A}}{\bm{S} }(t)=e^{-{\bm{A}}t}{\bm{B}}{\bm{X}}(t)$ $\displaystyle\frac{d}{dt}\left(e^{-{\bm{A}}t}{\bm{S}}(t)\right)=e^{-{\bm{A}}t} {\bm{B}}{\bm{X}}(t)$ $\displaystyle e^{-{\bm{A}}t}{\bm{S}}(t)={\bm{S}}(0)+\int_{0}^{t}e^{-{\bm{A}} \tau}{\bm{B}}{\bm{X}}(\tau)d\tau$ $\displaystyle{\bm{S}}(t)=e^{{\bm{A}}t}{\bm{S}}(0)+\int_{0}^{t}e^{{\bm{A}}(t- \tau)}{\bm{B}}{\bm{X}}(\tau)d\tau$ (21) $$

By introducing a discrete step size $\Delta=t_{k+1}-t_{k}$, we transform the above equation to a discrete-time system as follows.

$$ $\displaystyle{\bm{S}}(t_{k+1})=e^{{\bm{A}}(t_{k+1}-t_{k})}{\bm{S}}(t_{k})+\int _{t_{k}}^{t_{k+1}}e^{{\bm{A}}(t_{k+1}-\tau)}{\bm{B}}{\bm{X}}(\tau)d\tau$ $\displaystyle{\bm{S}}(t_{k+1})=e^{{\bm{A}}(t_{k+1}-t_{k})}{\bm{S}}(t_{k})+ \left(\int_{t_{k}}^{t_{k+1}}e^{{\bm{A}}(t_{k+1}-\tau)}d\tau\right){\bm{B}}{\bm {X}}(t_{k})\>\text{(assuming $x(\tau)\approx x(t_{k})$ over the interval)}$ $\displaystyle{\bm{S}}(t_{k+1})=e^{\Delta{\bm{A}}}{\bm{S}}(t_{k})+{\bm{B}}{\bm{ X}}(t_{k})\int_{t_{k}}^{t_{k+1}}e^{{\bm{A}}(t_{k+1}-\tau)}d\tau$ $\displaystyle{\bm{S}}(t_{k+1})=e^{\Delta{\bm{A}}}{\bm{S}}(t_{k})+{\bm{B}}{\bm{ X}}(t_{k})\int_{0}^{\Delta}e^{{\bm{A}}\tau^{\prime}}d\tau^{\prime}\>\text{( letting $\tau^{\prime}=t_{k+1}-\tau$)}$ $\displaystyle{\bm{S}}(t_{k+1})=e^{\Delta{\bm{A}}}{\bm{S}}(t_{k})+{\bm{B}}{\bm{ X}}(t_{k})\int_{0}^{\Delta}e^{{\bm{A}}\tau}d\tau$ $\displaystyle{\bm{S}}(t_{k+1})=e^{\Delta{\bm{A}}}{\bm{S}}(t_{k})+{\bm{B}}{\bm{ X}}(t_{k})\left(e^{\Delta{\bm{A}}}-{\bm{I}}\right){\bm{A}}^{-1}\>\text{( integral of matrix exponential function)}$ $\displaystyle{\bm{S}}_{k+1}={\bm{A}}{\bm{S}}_{k}+{\bm{B}}{\bm{X}}_{k}$ (22) $$

where the discretized state matrices $\bar{{\bm{A}}}=e^{\Delta{\bm{A}}}$ and $\bar{{\bm{B}}}=(e^{\Delta{\bm{A}}}-{\bm{I}}){\bm{A}}^{-1}{\bm{B}}$. Note that we apply the ZOH approach considering that $x(\tau)$ is constant between $t_{k}$ and $t_{k+1}$, we can rewrite the Eq. D by assuming ${\bm{X}}(\tau)\approx{\bm{X}}(t_{k}+1)$:

$$ $\displaystyle{\bm{S}}(t_{k+1})=e^{{\bm{A}}(t_{k+1}-t_{k})}{\bm{S}}(t_{k})+\int _{t_{k}}^{t_{k+1}}e^{{\bm{A}}(t_{k+1}-\tau)}{\bm{B}}{\bm{X}}(\tau)d\tau$ $\displaystyle{\bm{S}}(t_{k+1})=e^{{\bm{A}}(t_{k+1}-t_{k})}{\bm{S}}(t_{k})+ \left(\int_{t_{k}}^{t_{k+1}}e^{{\bm{A}}(t_{k+1}-\tau)}d\tau\right){\bm{B}}{\bm {X}}(t_{k+1})$ $\displaystyle{\bm{S}}_{k+1}=\bar{{\bm{A}}}{\bm{S}}_{k}+\bar{{\bm{B}}}{\bm{X}}_ {k+1}$ (23) $$

The resulting equation is the discrete-time SSM using ZOH discretization in eq. C.

#### Derivation of 𝑩 ¯ ¯ 𝑩 \bar{\bm{B}} over¯ start_ARG bold_italic_B end_ARG .

We use the equation $e^{{\bm{A}}\tau}=I+{\bm{A}}\tau+\frac{1}{2!}{\bm{A}}^{2}\tau^{2}+\cdots$, we have this integral of exponential function of ${\bm{A}}$:

$$ $\displaystyle\bar{\bm{B}}=\int_{0}^{\Delta}e^{{\bm{A}}\tau}{\bm{B}}d\tau$ $\displaystyle=\int_{0}^{\Delta}\left({\bm{I}}+{\bm{A}}\tau+\frac{1}{2!}{\bm{A} }^{2}\tau^{2}+\cdots\right)d\tau{\bm{B}}$ $\displaystyle=\left({\bm{I}}\Delta+\frac{1}{2}{\bm{A}}\Delta^{2}+\frac{1}{3!}{ \bm{A}}^{2}\Delta^{3}+\cdots\right){\bm{B}}$ $\displaystyle=\left(e^{\Delta{\bm{A}}}-{\bm{I}}\right){\bm{A}}^{-1}{\bm{B}}$ (24) $$

## Appendix E Details of Experiments

**Table 4: Configurations of TrajGPT and other transformer baselines on the PopHR dataset. We set TrajGPT and all Transformer baseline to 7.5 million parameters.**
| TrajGPT |  |
| --- | --- |
| Decoder Layers | 8 |
| Heads | 4 |
| Dim (${\bm{Q}}$, ${\bm{K}}$, ${\bm{V}}$, FF) | 200,200,400,400 |
| Transformer baselines including Encoder-decoder and Encoder-only models |  |
| Enc-Dec Layers | 4 & 4 |
| Encoder Layers | 8 |
| Decoder Layers | 8 |
| Heads | 4 |
| Dim (${\bm{Q}}$, ${\bm{K}}$, ${\bm{V}}$, FF) | 200,200,200,400 |

Figure: 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.
Refer to caption: extracted/5897467/topk_dist.png

## Peer Reviews

confidence: 5
contribution: 3
initial_rating: 6
presentation: 3
questions: Listed them out in the weaknesses section.
soundness: 3
strengths: - The paper focuses on an important problem of learning from irregularly sampled time series which is important for several practical use cases.
- The paper introduces a novel architecture for modeling irregularly sampled time series- selective recurrent attention (SRA). SRA enables time-specific inference that leads to better imputation and forecasting. 
- Experiments show that the framework achieves strong zero-shot performance across multiple tasks and achieves comparable results with SOTA in the finetuning setup.
- The paper is well written and easy to follow.
summary: The paper proposes an approach for representation learning of irregularly sampled time-series data called Trajectory Generative Pre-trained Transformer (TrajGPT). The TrajGPT architecture employs a Selective Recurrent Attention (SRA) mechanism that allows it to selectively focus on relevant past information and ignore irrelevant data. The proposed approach is able to capture continuous dynamics and perform interpolation and extrapolation. The proposed approach achieves strong performance on several tasks, including forecasting disease trajectories, predicting drug usage, and classifying patient phenotypes. It also show good zero-shot performance on several tasks and provides interpretable results.
weaknesses: - The evaluation of TrajGPT is limited, as it does not include benchmark irregularly sampled time series datasets such as PhysioNet and MIMIC-III. This makes it difficult to compare the performance of TrajGPT with other state-of-the-art models in this space despite the authors doing a good job of comparing on PopHR dataset.
- The paper overlooks important relevant work [1] and fails to provide a comparison with [1] which is the state-of-the-art for imputation in irregularly sampled time series.
- The paper lacks insight into why time-specific inference leads to better results than auto-regressive inference, which is the method used to train the model.


### References
[1] Heteroscedastic Temporal Variational Autoencoder For Irregularly Sampled Time Series, International Conference on Learning Representations, 2022.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: /
soundness: 3
strengths: The authors provide clear motivations for their methods design. Especially,  the SRA mechanism design acts like a discretization of continuous-time ODE not only helping the model to learn the hidden continuous dynamics but also can achieve linear training complexity and constant inference complexity. 
The visualizations in the result part looks really insightful and back up authors' assumption.
summary: The paper propose a novel transformer-based model tailored for irregularly-sampled time-series data, especially in healthcare domain. 
TrajGPT features the Selective Recurrent Attention (SRA) mechanism, which incorporates a data-dependent decay to filter out irrelevant historical data, thus improving its adaptability to variable intervals in time-series data. The model is interpreted as discretized ODEs, enabling both interpolation and extrapolation, which is essential for accurate, time-specific predictions.
weaknesses: 1. As the paper targeted representation-learning,  current representation-learning, like self-supervised learning for irregular time series lacks of discussion[1].
2. The paper only conducts experiments on PopHR dataset. Benchmark datasets such as MIMIC-III and physionet2012 are not included. 

[1].Ranak Roy Chowdhury, Jiacheng Li, Xiyuan Zhang, Dezhi Hong, Rajesh K. Gupta, and Jingbo Shang. 2023. PrimeNet: pre-training for irregular multivariate time series. In Proceedings of the Thirty-Seventh AAAI Conference on Artificial Intelligence and Thirty-Fifth Conference on Innovative Applications of Artificial Intelligence and Thirteenth Symposium on Educational Advances in Artificial Intelligence (AAAI'23/IAAI'23/EAAI'23), Vol. 37. AAAI Press, Article 807, 7184–7192. https://doi.org/10.1609/aaai.v37i6.25876

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 2
questions: - Can you provide more insights on the discrepancy between the proposed model and the definition of ODEs? Since the model relies on the PDE form to make predictions in the future, the lack of future observation, i.e., inputs, will almost surely lead to biased or at least inaccurate predictions. Can you explain how the model handles this issue and how it can be improved?
- Can you justify for the visualization of the learned embeddings and decay vectors? Or seperate the visualization of the learned embeddings and decay vectors to avoid confusion. This will help readers to better understand the model's inner workings and the role of each component in the model and avoid misinterpretation of the results.
- Can you explain more on the results of the experiments and find more support from the literature to validate the model's performance? This will help to better understand the model's effectiveness and accuracy in real-world applications.
soundness: 3
strengths: - Efficiency: TrajGPT achieves linear training and constant inference complexities, making it more scalable compared to traditional Transformers for long sequences.
- Performance and Generalizability: TrajGPT demonstrates impressive zero-shot capability across diverse tasks, highlighting its potential as a general-purpose tool for healthcare applications.
- Interpretability: The model’s architecture allows for meaningful health trajectory predictions and risk analysis, supporting clinical applications by enabling interpretative insights into disease progression.
summary: The paper introduces TrajGPT, a generative pre-trained Transformer designed for irregular time-series in healthcare applications. TrajGPT incorporates a novel linear Selective Recurrent Attention (SRA) mechanism that uses data-dependent decay to filter past information contextually. By interpreting the architecture as discretized ordinary differential equations (ODEs), the model effectively handles continuous dynamics, enabling it to interpolate and extrapolate for irregularly sampled data, thus enabling making direct predictions into desired future time step. TrajGPT is evaluated across trajectory forecasting, drug usage prediction, and phenotype classification, achieving strong zero-shot performance without task-specific tuning. Additionally, its design supports interpretability in health trajectory analysis, allowing clinicians to link disease trajectories with patient histories.
weaknesses: - Ambiguity in the definition of ODEs: The proposed TrajGPT model has a form that the gradient of the trajectory at certain time is not only dependent on the state at that time, but also the input at that time. This data-dependent gradient no longer satisfies the definition of ODEs, but more towards the definition of PDEs. The authors should clarify this discrepancy and provide a more accurate interpretation of the model.
- Poorly justified visualization: In Fig 3a, the authors visualize the learned embeddings along with head specific decay vectors using UMAP. However, these decay vectors do not have same meaning as the embeddings, as one is intended to be a weight matrix of the input data while the other serves as a representation of the input data. Therefore, although these vectors share the same dimensionality, they should not be embedded in the same manifold space by using UMAP.
- Improvements in paper structure: The paper would benefit from a more structured presentation of the proposed model, for example, in section 3.1, the authors could avoid presenting the parallel formulation of the proposed model as it is not directly related to the main contributions of the paper. On the other hand, the experiment section could be expanded to provide more insights on the model's performance across different tasks, including more detailed explanation of the results and comparison with established medical findings.
- Absence of robustness analysis: As the paper proposes a novel model for healthcare applications, it is important to evaluate the robustness and potential biases of the model. For instance, as the model draws similarities to ODEs, it is crucial to investigate intrinsic biases in the model such as the limit cycle bahavior, which could lead to biased predictions in certain scenarios. The authors should provide a more comprehensive analysis on the robustness of the model to ensure its reliability in real-world applications.

---

confidence: 4
contribution: 4
initial_rating: 6
presentation: 4
questions: - Do you believe this method could be easily extended to the context of multivariate time series?

- If expanded to multivariate time series, could it be combined with methods like GNNs to extract improved spatio-temporal representations?

- Despite the relationship, it remains unclear what the main advantages or potential inductive bias would be in using this method over Mamba blocks. Could you clarify this further? I would be happy to increase my score if the authors could expand on this information to clarify this important issue.
soundness: 4
strengths: - Comprehensive literature review.

- Thorough evaluation across different scenarios, with comparisons to multiple models and diverse datasets. The inclusion of interpretability through added figures is beneficial, and the ablation study is particularly valuable for assessing the contribution of each model component.

- The method is highly compelling, supported by a robust and thorough mathematical explanation.

- Novelty of the Attention Mechanism: The introduction of the Selective Recurrent Attention (SRA) module, which uses data-dependent decay, is an innovative contribution.

- Interpretation of TrajGPT as a Discretized ODE: Modeling TrajGPT as a discretized ODE is a valuable approach, allowing the model to capture continuous dynamics in irregular time series and perform time-specific inference, facilitating both interpolation and extrapolation in temporal data.

- Throughout the paper, I noted the similarities to State Space Models, which led me to find Appendix C particularly valuable.
summary: This paper presents the Trajectory Generative Pre-trained Transformer (TrajGPT), a Transformer model designed to handle irregularly sampled time series, particularly within the healthcare domain. TrajGPT introduces a novel attention mechanism, Selective Recurrent Attention (SRA). Inspired by discretized ordinary differential equations (ODEs), the model achieves interpolation and extrapolation in irregular time series, excelling in trajectory prediction tasks. Experimental results show outstanding performance in zero-shot learning scenarios, underscoring TrajGPT’s capability to capture complex health dynamics from incomplete data.
weaknesses: There are few weaknesses in this paper:

- Since the relationship with State Space Models (SSMs) is discussed, it might have been interesting to include a method using Mamba to compare results with this type of model.

- Including a comparison with methods such as Multi-Gaussian Processes, which have been previously used for this type of approach, could strengthen the evaluation [1].

- Ultimately, as I understand it, this method has only been applied to univariate time series; it would be beneficial to explore its extension to multivariate time series in future work.

[1] Moor, M., Horn, M., Rieck, B., Roqueiro, D., & Borgwardt, K. (2019, October). Early recognition of sepsis with Gaussian process temporal convolutional networks and dynamic time warping. In Machine learning for healthcare conference (pp. 2-26). PMLR.

Meta-review: **Summary**: The paper proposes a GPT framework for modeling irregularly sampled sequential data incorporating a selective attention mechanism to enable the model to adaptively use historical context. Empirical evaluation demonstrates good zero-short performance on healthcare data and under missingness. 

**Strengths**: 
1. Comparison to baselines and prior work is comprehensive. 
2. Empirical evaluation includes ablations of specific benefits of the model (with and without pre-training)
3. Reviewers have found the model to be fairly effective at modeling patient trajectories

**Weaknesses**:
1. Initial version of the paper did not have benchmark evaluations with respect to MIMIC-IV and eICU data. Subsequent rebuttal has incorporated results for these standard benchmark datasets.
2. There are concerns raised about the framing and analogy to ODEs. I agree with some of these concerns and over-interpretation.
3. Initial version of the paper also did not benchmark against State-space model analogues like MAMBA and did not have uncertainty quantification. Which have since been provided as a response in the rebuttal.
4. Over-interpretation and issues linger in the UMAP visualization of patient trajectories.

**Recommendation**:
Significant concerns about the paper have been addressed in the rebuttal, the updated version of the paper is improved in presentation. However, some fundamental technical framing issues remain. First, I do agree that the analogy to ODEs is slightly problematic and should be significantly tempered. Second, I believe all additional empirical evaluation provided in author response should be incorporated rigorously in the final version with reasonable bootstrapping. Finally, ablations should be more systematic. Considering all these issues, I believe the paper will benefit from an additional round of peer review and therefore I am recommending a reject.

Decision: Reject