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

Thank you for taking the time to review our responses and revised manuscript. We are grateful for your recognition of our efforts and the increased score.

A3 for Question3: In this response, we clarify the advantages of TrajGPT over Mamba-1 and Mamba-2. While Mamba-1 is an SSM without linear attention, TrajGPT and Mamba-2 [5] are linear attention framework but with different decay gating mechanisms. TrajGPT offers a more expressive data-dependent decay mechanism than Mamba-2, which helps prevent catastrophic forgetting and maintains long-term dependencies. Additionally, TrajGPT surpasses Mamba-1 with efficient matrix representation for parallel forward-pass, as well as larger hidden dimensions and head size.

The table below, adapted from Table 4 in [6], provides a comparison of these models:

Specifically, Mamba-2 defines an SSM with both recurrent and parallel forms, called State Space Duality in original paper [5]. The recurrent form is:

It also has a parallel form with a matrix representation:

where the matrix $L$ is defined as:

$

L = 

\begin{bmatrix} 1 & & & & & \\ a_1 & 1 & & & & \\ a_2 a_1 & a_2 & 1 & & & \\ \vdots & \vdots & \vdots & \ddots & & \\ a_{T-1} \dots a_1 & a_{T-1} \dots a_2 & \dots & a_{T-1} & 1 \end{bmatrix} \nonumber

$

By mapping $(a_t, L, C, B, X) \xrightarrow{} (\beta_t, D, Q, K, V)$, Mamba-2 can be identified as a Linear Attention model with a data-dependent decay, similar to our TrajGPT. As discussed in our work, however, the computation of decay $\beta_t$ in Mamba-2 can cause numerical instability if the decay values are consistently small. Mamba-2 addresses this by transitioning the computation to log space, but it still suffers from catastrophic forgetting if these values remain consistently small. In contrast, our TrajGPT introduces a temporal term $\tau$ to control the range of $\gamma_t$, defined as $\gamma_n = \text{Sigmoid} (X_n w_\gamma^\top)^{\frac{1}{\tau}}$. Consequently, TrajGPT offers a more expressive decay mechanism than Mamba-2, effectively avoiding catastrophic forgetting and preserving long-term dependencies.

[5] Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. ICML 2024.

[6] Parallelizing Linear Transformers with the Delta Rule over Sequence Length. Neurips 2024.

Thanks for your careful review and constructive suggestions. Please find our response to your questions and concerns.

Response for both Weakness 1 and 2: In the revised version, we have included a discussion of the suggested MGP-TCN in Section 2.2. We also incorporated Mamba-1, Mamba-2, and MGP-TCN as baseline models in the PopHR experiments discussed in Section 5.2. We utilized the public GitHub codes for these three models with default setups. For the mamba models, we set the model parameters to around 7.5 million to ensure comparability with other Transformer baselines. These models employ next-token prediction pre-training and follow the same setup as TrajGPT for zero-shot learning, few-shot learning, and fine-tuning, applied to both forecasting and classification tasks. For MGP, originally designed for classification, we directly applied its code for insulin and CHF prediction. Since MGP-TCN does not support forecasting, we modified its classification output layer to generate predictions for all timesteps within the forecasting window.

The table below, referenced from Table 1, compares the performance of TrajGPT with Mamba and MGP-TCN. The results demonstrate that TrajGPT with time-specific inference achieves the best performance across forecasting, insulin prediction, and CHF classification. Mamba-2 performs well on these two classification tasks with fine-tuning, while Mamba-1 and MGP-TCN fall behind.

R3 for Weakness 3: To clarify, our model is designed to handle multivariate time series but in the context of discrete data. Specifically, we use a multinomial distribution for discrete tokens (e.g., one-hot-encoded diagnoses at any given time), where one token is 1 and all other tokens are 0 for each timestamp. In future work, we plan to extend the TrajGPT model to continuous multivariate time series (e.g., physiological signals such as ECG, clinical measurements such as MIMIC data, etc).

A1 for Question1: As mentioned in R3, our model is designed to handle discrete multivariate time series. It is straightforward to extend it to continuous multivariate time series by projecting the values from multiple variables into a token embedding. Hence, we only need to modify the input projection layer and keep the main architecture unchanged (e.g., [1, 2]).

[1] TimelyGPT: Extrapolatable Transformer Pre-training for Long-term Time-Series Forecasting in Healthcare. ACM BCB 2024.

[2] Time-LLM: Time series forecasting by reprogramming large language models. ICLR 2024

A2 for Question 2: Thank you for this very interesting suggestion. In principal, TrajGPT can be combined with GNNs for spatio-temporal representations. Suppose we have a sequence of dynamic spatial graphs (i.e., one sequence per patient). Each patient- and time-specific graph may be obtained by processing the static biomedical knowledge graph via the patient's conditions and medications. At a given time point for a patient, the graph adjacency matrix can serve as a hard or soft constraint to compute the attention [3]. The resulting pooled embedding of the graph can then be fed to TrajGPT for temporal modeling.

However, end-to-end training such model may be computationally challenging (e.g., tracking all graph- and time-attentions for backpropagation). Some existing work [4] employs a two-stage approach by first employing a pre-trained Transformer to encode sequence embeddings, which are then utilized to train a GNN, facilitating spatio-temporal representations.

[3] Learning the graphical structure of electronic health records with graph convolutional transformer. AAAI 2020.

[4] Pre-training enhanced spatial-temporal graph neural network for multivariate time series forecasting. KDD 2022.

We are grateful for your constructive comments. Below we address each question and concern.

R1 for Weakness 1: In the updated version, we have included a discussion of the suggested PrimeNet model in Section 2.1. We also incorporated it as a baseline in the experiments conducted on the PopHR and eICU datasets in Sections 5.2 and 5.4, respectively. We utilized the publicly available code from GitHub with default hyperparameters (https://github.com/ranakroychowdhury/PrimeNet). We follow the setups for pre-training, few-shot learning, and fine-tuning on the full dataset, while adopting the zero-shot learning from our work. Since PrimeNet does not support forecasting, we modified its interpolation output layer to directly generate predictions for all timesteps within the forecasting window. For classification tasks, we used the original code without modifications.

We present the results in the two tables below, corresponding to the experiments on the PopHR and eICU datasets, as referenced in Table 1 and Table 2 of the updated manuscript. TrajGPT with time-specific inference outperforms PrimeNet in forecasting tasks, attributable to its dynamic modeling and reduced error accumulation. Additionally, TrajGPT achieves superior performance in insulin usage prediction and early sepsis detection, while performing comparably to PrimeNet in CHF classification.

R2 for Weakness 2: Thanks you for your feedback regarding the evaluation and benchmarking. We experimented with the publicly available eICU dataset from PhysioNet because it provides irregularly-sampled time series with discrete diagnoses and drugs provided to patients during their ICU stay. We extracted a dataset containing 288 unique PheCodes and 228 unique drugs. The final dataset includes 139,367 patients, with each patient having an average of 19 drugs and 3 ICD codes recorded within 15-minute intervals. We conducted the forecasting tasks on irregular clinical events (diagnoses and drugs) and early detection of sepsis.

We have updated the paper to include a description of dataset and preprocessing in Section 4.1, as well as an explanation of the experimental design in Section 4.3. Additionally, the updated Section 5.4 includes a discussion of the eICU experimental results, with the outcomes presented in Table 2. Below, we show the table on eICU benchmarking, which evaluate TrajGPT and baselines on the top-K diagnosis/medication forecasting and sepsis prediction. The experimental results highlight that TrajGPT achieves the highest performance in both forecasting tasks (top-10 and top-20 recall rates) and zero-shot sepsis classification.

R4 for Weakness 4: To address the concern about robustness, we have added a bootstrap evaluation of standard deviation in the revised manuscript. By resampling the dataset multiple times with replacement, we provide a robust estimate of the variance in our metrics. The table below provides a summary of TrajGPT's results, with the complete results available in the updated paper.

The reviewer also raises an important point about intrinsic biases, such as the limit cycle commonly observed in ODEs. Our model addresses this by capturing trend patterns using an exponential decay mechanism, as discussed in existing studies [4, 5]. This mechanism enables efficient extrapolation by extending learned trend patterns to longer sequences beyond the training length. To demonstrate this, we include a risk trajectory analysis for diabetes and CHF (Fig. 4b and 4d). The results show that our model effectively captures clear trends of increasing risks with age, reflecting age-related vulnerability to chronic diseases. As a result, our model avoids the intrinsic biases of classical ODEs, such as limit cycles, which can lead to inaccurate predictions.

[4] A Length-Extrapolatable Transformer. ACL 2023.

[5] TimelyGPT: Extrapolatable Transformer Pre-training for Long-term Time-Series Forecasting in Healthcare. ACM BCB 2024.

A1 for Q1: We have clarified that our model is a discretized ODE in our response R1. Similar to ODEs, our model makes predictions in a sequential manner. For each timestamp $t_{n}$, it computes the current state $S_{n}$ using the previous state $S_{n-1}$, and then use $S_{n}$ to predict the corresponding token $x_{n}$.

However, the lack of future observations may cause biased prediction, leading to accumulated errors over time. TrajGPT mitigates this through time-specific inference, resulting in more stable and accurate predictions compared to other baselines, as discussed in Section 5.2 (Fig. 5). Specifically, it predicts directly at any timestep using prior hidden states and current time, leveraging the additional information from the target timestamp. Therefore, time-specific inference reduces computational steps and error accumulation.

A2 for Q2: We have addressed the concerns about the UMAP visualization in our response R2.

A3 for Q3: A3: We have addressed the concerns in our R3 by including additional discussion and relevant literature from the medical domain to support our result analysis.

We sincerely appreciate your thorough review and valuable feedback. We provide our responses to your insightful questions and concerns.

R1 for Weakness1: First, we would like to clarify the definition of ODE and PDE. The primary distinction is that ODEs involve only one independent variable time $t$, while PDEs involve two or more independent variables, such as time $t$ and multiple time-independent variables $\mathbf{z} = [z_1, z_2, \ldots, z_D]$.

According to [1], ODE is defined with the time $t$ and a time-dependent input $x(t)$, $\frac{d s}{dt}(t) = f(s(t), x(t))$, where the output $s$ is referred to as the ``state" in the SSM. On the other hand, PDEs involve multiple independent variables in its definition $\partial_t s = f(s(t), \mathbf{z}, \partial_\mathbf{z} s, \partial_{\mathbf{z}\mathbf{z}} s, \ldots)$ [2], where $\mathbf{z}$ are time-independent variables. Same as ODEs, our model processes a time-dependent input $\mathbf{X}(t)$ and does not introduce additional time-independent variables.

Furthermore, existing studies confirm that continuous-time SSMs are ODEs [3]. We have shown in Appendix C that our model is a discretization of a continuous-time SSM (i.e., an ODE). Therefore, our proposed SRA module is a discretized ODE.

For a continuous-time SSM, we have $S^\prime (t) = A S(t) + B X(t)$, where $X(t)$ is time-dependent input and $S(t)$ is the state variable with $S_0 = 0$. This continuous-time SSM (i.e., ODE) can be discretized by various methods including Euler's method, zero-order hold (ZOH), and bilinear transform. In Appendix D, we show that our model is a discretization of SSM via ZOH.

We revised our paper accordingly. In Section 3.2, we explicitly state that our SRA module involves only one independent variable $t$, making it a discretized ODE rather than a PDE. Additionally, we explicitly stated that Appendices C and D cover the theoretical connection to ODEs and the derivation of the ZOH discretization.

[1] Parallelizing non-linear sequential models over the sequence length. ICLR 2024.

[2] Message Passing Neural PDE Solvers. ICLR 2022.

[3] Modeling sequences with structured state spaces. Dissertation of Albert Gu.

R2 for Weakness 2: Although global head-specific decay vectors $\mathbf{w}^{(h)}_\gamma$ and token embeddings $\mathbf{x}_n$ are different types of information, we can still innovatively visualize them in the same manifold because of their dot product in the decay function : $\gamma^{(h)}_n = \text{Sigmoid} (\mathbf{x}_n \mathbf{w}^{(h)\top}_\gamma)$.

Although global head-specific decay vectors $\mathbf{w}^{(h)}_\gamma$ and token embeddings $\mathbf{x}_n$ are different types of information, we can still innovatively visualize them in the same manifold because of their dot product in the decay function $\gamma^{(h)}_n = \text{Sigmoid} (\mathbf{x}_n \mathbf{w}^{(h)\top}_γ)$. Therefore, our UMAP visualization provides meaningful interpretation of the attention head by comparing the positions of their corresponding decay vectors with the token embeddings of known identities. For instance, when a category of token embeddings (e.g., endocrine/metabolic diseases) and a decay vector (e.g., head 8) are positioned closely in the latent space, this suggests a higher similarity between them. This proximity indicates that certain heads may specialize in certain disease categories, allowing the model to slow the forgetting processing and retain long-term dependencies.

R3 for Weakness 3: The parallel form achieves linear training complexity, as discussed in Section 3.3. To improve clarity, however, we moved the derivation of the parallel form to the Appendix.

To expand the experiment section, we conducted experiments on forecasting clinical events and early sepsis detection using the eICU dataset. In the updated version, we provided a discussion of eICU experimental results in Section 5.4 and Table 2, which show that TrajGPT achieves the highest performance in both forecasting tasks (top-10 and top-20 recall rates) and zero-shot sepsis classification.

To address concerns about medical findings, we have updated Section 5.1 (Qualitative Analysis of Embeddings) and 5.3 (Trajectory Analysis) with additional references from the clinical domain. We have also expanded the discussion of our work's clinical impact in Section 5.3 (Trajectory Analysis) and Section 6 (Conclusion), incorporating relevant references.

We appreciate the reviewer’s comments. We update our paper and address your concerns here:

R1 for Weakness 1: Thanks you for your feedback regarding the evaluation and benchmarking. We experimented with the publicly available eICU dataset from PhysioNet because it provides irregularly-sampled time series with discrete diagnoses and drugs provided to patients during their ICU stay. We extracted a dataset containing 288 unique PheCodes and 228 unique drugs. The final dataset includes 139,367 patients, with each patient having an average of 19 drugs and 3 ICD codes recorded within 15-minute intervals. We conducted the forecasting tasks on irregular clinical events (diagnoses and drugs) and early detection of sepsis.

We have updated the paper to include a description of dataset and preprocessing in Section 4.1, as well as an explanation of the experimental design in Section 4.3. Additionally, the updated Section 5.4 includes a discussion of the eICU experimental results, with the outcomes presented in Table 2. Below, we show the table on eICU benchmarking, which evaluate TrajGPT and baselines on the top-K diagnosis/medication forecasting and sepsis prediction. The experimental results highlight that TrajGPT achieves the highest performance in both forecasting tasks (top-10 and top-20 recall rates) and zero-shot sepsis classification.

R2 for Weakness 2: In the updated version, we have included a discussion of the suggested HeTVAE in Section 2.2. We also incorporated it as a baseline in the PopHR experiments discussed in Section 5.2. We utilized the publicly available code from GitHub with default hyperparameters (https://github.com/reml-lab/hetvae). While HeTVAE was originally designed for interpolation tasks, we modified its output linear layer to accommodate our forecasting and classification tasks. For forecasting, the output layer directly generates predictions for all timesteps within the forecasting window. For classification, it directly outputs the predicted class. Since HeTVAE does not support a representation learning approach, we trained it from scratch on the entire dataset for the insulin and CHF classification tasks.

The table below, referenced from Table 1, compares the performance of TrajGPT and HeTVAE. While HeTVAE demonstrates strong forecasting performance for top-K recall with K=10 and K=15, TrajGPT with time-specific inference surpasses it by 1.6% and 0.9%, respectively. However, HeTVAE, originally designed for generative interpolation tasks, performs poorly on discriminative prediction tasks such as insulin usage and CHF phenotyping.

R3 for Weakness 3: Thank you for your question. We have revised Section 3.2 to clarify this. For auto-regressive method, it generates samples sequentially and selects from irregularly spaced timestamps, which can lead to error accumulation. Our proposed time-specific inference directly extrapolates to the target timestep using prior hidden states and the gap from the last observed timestep. Therefore, it reduces computational steps and error accumulation for better performance.