ChronoSSM: Training for Temporally Aware Representations in Autoregressive State Space Models

State Space Models (SSMs) are sequence models with recurrent latent-state dynamics that support linear-time processing in sequence length (Fichtl et al., 2025; Gu et al., 2022). This is useful in long-sequence settings, where computational efficiency becomes important as context length grows.

The core mechanism of an SSM is a linear dynamical system that maps an input sequence to an output sequence through a hidden state (Gu et al., 2022). At step $k$, the model updates a latent state $s_{k}$ from the previous state and the current input $x_{k}$, and then produces an output representation $h_{k}$:

In modern selective SSMs, some of these parameters are conditioned on the current input, allowing the model to selectively propagate or discard information according to the sequence context. This input-dependent selection mechanism enables the model to adapt its effective memory to the content of the sequence.

Stacking several SSM layers produces a causal sequence model that processes the input in a single forward pass and generates a hidden representation $h_{k}$ at each position. These token-wise hidden states can then be used by lightweight prediction heads for tasks such as next-token prediction or temporal prediction. This representation-level property is the one exploited in this paper.

Let $\mathcal{D}$ denote a dataset of timestamped sequences. Each sequence of length $T$ in $\mathcal{D}$ is written as $\{(x_{k},t_{k})\}_{k=0}^{T-1}$, where $x_{k}$ is a discrete token and $t_{k}$ is its associated timestamp. From successive timestamps, we derive the inter-arrival times $\{\Delta t_{k}=t_{k}-t_{k-1}\}_{k=1}^{T-1}$.

Denoting the token history up to step $k$ by $\mathcal{X}_{k}=\{x_{i}\}_{i=0}^{k}$, our objective is to model the joint distribution of the next token $x_{k+1}$ and its associated inter-arrival time $\Delta t_{k+1}$:

Using the chain rule, we factor this joint distribution as

where $\gamma=\{\theta,\alpha,\phi\}$ collects the parameters of the shared sequence model, token-prediction head, and temporal-prediction head, respectively. Let $V$ denote the vocabulary size and $D$ the hidden-state dimension.

The shared sequence model is a causal autoregressive model parameterized by $\theta$ that produces a token-wise hidden representation $h_{k}$ from the history $\mathcal{X}_{k}$. In our main instantiation, we use a State Space Model, although the framework is not tied to this particular architecture. Any causal autoregressive backbone that produces token-wise hidden representations can be used instead; Appendix appendix C illustrates this with an additional Transformer-based instantiation. Formally, the backbone computes

The token-prediction head is a linear map $W^{\mathrm{tok}}_{\alpha}\in\mathbb{R}^{V\times D}$ parameterized by $\alpha$, and it models the next-token distribution as

The temporal-prediction head is a function $f^{\mathrm{time}}_{\phi}(\cdot;r_{k+1})$ parameterized by $\phi$ and conditioned on a transition representation $r_{k+1}$. After $x_{k+1}$ is processed, the sequence model produces $h_{k+1}$ according to Eq. 5. Because $\Delta t_{k+1}$ describes the transition from position $k$ to position $k+1$, its prediction is based on a representation $r_{k+1}$ derived from both $h_{k}$ and $h_{k+1}$. The temporal-prediction head models the second factor as

The construction of $r_{k+1}$ is introduced in section 3.

An inter-arrival time measures the elapsed time between two consecutive tokens, rather than a property of either token considered in isolation. We therefore represent it using the change between their corresponding hidden states:

The temporal target and the representation used to predict it are thus both defined as differences between consecutive positions:

This choice for $r_{k+1}$ is a modeling decision rather than a structural requirement of the framework. An alternative temporal representation $r^{\prime}_{k+1}=[h_{k};h_{k+1}]$ is considered in Appendix appendix F.

Under the factorization in Eq. 4, maximum-likelihood training decomposes the negative log-likelihood of each transition into a token-prediction term and a temporal-prediction term:

We now instantiate these two terms using the token-prediction and temporal-prediction heads introduced in section 2.2.

For token prediction, the factor $p_{\theta,\alpha}(x_{k+1}\mid\mathcal{X}_{k})$ is categorical, so the sequence-level token loss is

where $\mathrm{CE}$ denotes categorical cross-entropy, which is appropriate because $\{x_{k}\}_{k=0}^{T-1}$ is a sequence of discrete tokens (Goodfellow et al., 2016) (see section 2.2).

For temporal prediction, the factor $p_{\theta,\phi}(\Delta t_{k+1}\mid x_{k+1},\mathcal{X}_{k})$ is defined by a conditional density over the continuous-valued target $\Delta t_{k+1}$. This definition is well posed provided that the chosen conditional density has support on the observed inter-arrival targets, including zero when zero inter-arrival times are present. We therefore define the sequence-level temporal loss as the negative log-likelihood of the observed inter-arrival times under that density:

where $\Delta h_{k+1}$ is the hidden-state transition defined in Eq. 8. The concrete form of $f^{\mathrm{time}}_{\phi}$ is specified for each experimental regime in section 4.

Summing Eq. 10 over the sequence gives the joint negative log-likelihood

Minimizing this objective corresponds to maximum-likelihood training of the factorized distribution over tokens and inter-arrival times.

In practice, the token and temporal losses may have different numerical scales. We therefore optimize the weighted objective

where $\lambda_{\tau}\geq 0$ controls the relative contribution of temporal prediction.

We consider two approaches for optimizing the parameter set $\gamma=\{\theta,\alpha,\phi\}$. Both regimes use the same shared sequence backbone model, token-prediction head, temporal-prediction head, and loss components. They differ only in whether the temporal loss is allowed to update the shared backbone. The two regimes are illustrated in Figure 1: Figure 1(a) shows the Joint setup, whereas Figure 1(b) shows the Two-Stage schedule. The corresponding optimization procedures are summarized in Algorithms 1 and 2.

In particular, the temporal loss backpropagates through $\Delta h_{k+1}=h_{k+1}-h_{k}$, encouraging changes between consecutive hidden states to become informative about the corresponding inter-arrival times. Because the token-prediction head uses hidden states produced by the same backbone, temporal supervision can also affect the representations used for next-token prediction.

Because the backbone is not updated during the second phase, temporal gradients cannot modify its representations. The temporal-prediction head can therefore exploit only the temporal information already present in the token-trained hidden-state transitions. This regime provides a controlled baseline for isolating the effect of allowing temporal supervision to shape the backbone.

Generation proceeds autoregressively. Given the token history $\mathcal{X}_{k}$, the token head first predicts the next token $x_{k+1}$. After this token is processed by the backbone, the model obtains $h_{k+1}$ and computes $\Delta h_{k+1}$. The temporal head then predicts the associated inter-arrival time $\widehat{\Delta t}_{k+1}$ from $\Delta h_{k+1}$, and the timestamp is updated as $t_{k+1}=t_{k}+\widehat{\Delta t}_{k+1}$.

The token-decoding strategy and the point estimate used for temporal prediction are specified in section 4.

Unless stated otherwise, all experiments in the main paper instantiate ChronoSSM with a Mamba2 backbone (Dao and Gu, 2024). To assess the robustness of the controlled comparison to the choice of causal autoregressive backbone, and to illustrate that the framework can be instantiated beyond SSMs, Appendix appendix C repeats the comparison with a GPT-2 backbone. The temporal-head architecture is selected separately for each domain to account for differences in sequence structure and temporal distribution.

Complete domain-specific preprocessing settings, model architectures, temporal-head configurations, loss weights, and training hyperparameters are reported in Appendix appendix H. Appendix appendix G reports the corresponding training-time and peak-memory overhead. Appendix appendix E reports an auxiliary-loss-weight sensitivity study for network traffic and temporal knowledge graph.

In this subsection, each sequence in $\mathcal{D}$ follows the canonical formulation introduced in section 2.2: every token $x_{k}$ is associated with an observed timestamp $t_{k}$, so inter-arrival times are defined for every consecutive pair of events.

Events are ordered chronologically within each trace. When multiple activities share the same timestamp, they are retained as separate tokens and therefore produce valid zero-valued inter-arrival targets. The original event log also contains resource identifiers, lifecycle transitions, application metadata, permit information, and other event- and case-level attributes, but we discard these auxiliary attributes and retain only the ordered activity labels and their timestamps to keep the modeling simple.

Events are ordered chronologically within each stay. When multiple observations share the same timestamp, they are retained as separate tokens, yielding valid zero-valued inter-arrival times, as in BPI2018. We discard the remaining clinical and administrative attributes and retain only the ordered event identities and their timestamps.

Because both datasets contain simultaneous events, the temporal prediction mechanism must accommodate observed inter-arrival times that are exactly zero. In particular, a plain $\log(\Delta t_{k})$ would be undefined when $\Delta t_{k}=0$. We therefore define

where $s$ is a dataset-specific time scale. This transformed quantity $z_{k}$ defines the auxiliary temporal target used for supervision. The temporal head produces a raw scalar output, which is mapped through a softplus nonlinearity to obtain a non-negative inter-arrival prediction $\widehat{\Delta t}_{k}$. This prediction is then transformed in the same way:

We instantiate the temporal factor as the fixed-variance log-normal density

where $\sigma^{2}$ is fixed. Its negative log-likelihood is equivalent, up to terms independent of the model parameters, to squared error between $\widehat{z}_{k}$ and $z_{k}$. This is a modeling choice rather than a requirement of the framework. Appendix appendix B reports an auxiliary GiantMIDI experiment that instead applies squared-error regression to scaled raw inter-arrival times, testing whether the results depend on the logarithmic target transformation.

Many temporally structured sequences do not provide a timestamp for every token but only for selected structural boundaries. These settings are therefore important for evaluating whether temporal supervision can still shape useful representations when it is available only sparsely within the sequence.

In this subsection, we consider domains in which timestamps are attached only to selected boundary tokens rather than to every token in the sequence. Let $\mathcal{I}_{\mathrm{time}}$ denote the set of positions carrying valid temporal targets. For each $k\in\mathcal{I}_{\mathrm{time}}$, let $k^{-}$ denote the preceding position in $\mathcal{I}_{\mathrm{time}}$. The temporal target and its corresponding hidden-state representation are then defined as

The temporal loss is evaluated only for $k\in\mathcal{I}_{\mathrm{time}}$. Tokens outside $\mathcal{I}_{\mathrm{time}}$ continue to contribute to the token-generation objective but do not contribute directly to the temporal objective. This defines a different setting from the canonical formulation introduced in section 2.2.

Facts are mapped to a fixed vocabulary of entity and relation identifiers and represented as structured token blocks of the form <|STM|>, h, r, o. Consecutive facts are grouped into fixed-length training sequences. The timestamp of each fact is attached only to its <|STM|> marker, so $\mathcal{I}_{\mathrm{time}}$ contains the <|STM|> positions.

Token decoding is domain-specific: BPI2018, MIMIC-IV, and dynamic knowledge graphs use greedy decoding, whereas network traffic uses stochastic sampling, following Chu et al. (Chu et al., 2026). Complete decoding parameters are reported in Appendix appendix H. For temporal prediction, the raw output of the temporal head is passed through a softplus nonlinearity to obtain the non-negative point estimate $\widehat{\Delta t}_{k}$, which is added to the preceding generated timestamp.

For additional descriptive context, Appendix appendix I reports one retrained domain-adapted baseline per dataset, including its architecture, probing adaptation, results, and interpretation. These external models are not used to answer RQ1 or RQ2, because they differ from ChronoSSM in architecture, objective, temporal representation, and generation procedure.

In order to respond to RQ1, we define temporal recoverability as the extent to which the temporal target associated with an event transition can be extracted from frozen backbone representations using a simple readout or reflected in their geometry. For each valid temporal transition, we evaluate the representation supplied to the temporal head and the transformed inter-arrival target used during training.

We report mean absolute error (MAE) and the coefficient of determination ($R^{2}$). A lower MAE indicates that the temporal target can be reconstructed more accurately from the representation, while a higher $R^{2}$ indicates that the probe explains a larger fraction of its variability. Lower MAE and higher $R^{2}$ therefore indicate stronger temporal recoverability.

We compare this quantity with the mean target difference obtained by assigning $K$ randomly selected transitions to each $r_{k}$, denoted $d_{\mathrm{rand}}$, and report

where $\varepsilon$ prevents division by zero. In all experiments, we set $K=5$ and use cosine similarity. Sensitivity to the neighborhood size and similarity measure is reported in Appendix appendix J.

A high TCS indicates that transitions with similar inter-arrival times are substantially closer in representation space than expected under random pairing. Temporal information is therefore more strongly reflected in the geometry of the learned representations. Appendix appendix K provides a separate descriptive layer-wise alignment analysis of the same representation geometry across domains.

Table 1 reports transition-level temporal recoverability across the four evaluation domains. In every domain, Joint achieves higher TCS and probe $R^{2}$, together with lower probe MAE, than Two-Stage. Thus, joint training improves both local temporal coherence and the linear accessibility of inter-arrival information, under both dense and partial timestamp supervision.

The diagnostic profile nevertheless differs by domain. On BPI2018, TCS increases only slightly, from $0.794$ to $0.803$, while the probe improvements are clearer: $R^{2}$ increases from $0.461$ to $0.631$ and MAE decreases from $3.342$ to $2.775$. This pattern may suggest that joint training makes temporal information more linearly accessible without substantially changing local temporal neighborhoods. MIMIC-IV shows a similar pattern. Network traffic, by contrast, shows large improvements in all three diagnostics, including TCS from $0.147$ to $0.622$, probe $R^{2}$ from $0.190$ to $0.394$, and MAE from $0.382$ to $0.308$. The concurrent gains may suggest changes in both local temporal coherence and linear accessibility. For the temporal knowledge graph, the changes are smaller across all three diagnostics than on the other domains: unlike BPI2018 and MIMIC-IV, where the probe metrics improve more clearly than TCS, and network traffic, where all three diagnostics improve substantially, its gains remain modest but consistently favor Joint.

We evaluate generated samples against held-out data using domain-specific protocols taken from the literature.

Table 2 reports content-generation quality across the four domains. The controlled comparison provides a negative answer to RQ2: Joint training does not systematically degrade generative performance relative to Two-Stage training. Its effect is nevertheless domain-dependent, but provides no evidence of a systematic degradation in generative performance.

On BPI2018, Joint and Two-Stage obtain identical mean and median DLS. The additional temporal supervision therefore changes neither the generated activity suffixes nor their similarity to the reference traces in this setting. On MIMIC-IV, Two-Stage achieves slightly lower token JSD and total variation distance, while both regimes obtain the same recall over the 50 most frequent clinical event types. Thus, Joint introduces a small degradation in matching the marginal event-token distribution, but does not reduce coverage of frequent clinical events. Appendix appendix D reports the complementary held-out validation-loss diagnostics: MIMIC-IV is also the only main domain in which Joint has a slightly higher token loss than Two-Stage (1.296 versus 1.293, a difference of about 0.2%). This small teacher-forced loss difference is consistent with the limited distributional change, but does not indicate a broad degradation in generation quality: frequent-event coverage is unchanged.

In contrast, Joint improves content-generation quality in both partially timestamped domains. On network traffic, Joint consistently lowers the JSD of all four evaluated packet-header fields, with particularly large improvements for source and destination TCP ports. On temporal knowledge graphs, Joint also improves every filtered link-prediction metric, yielding substantially higher MRR and Hits@$\{1,3,10\}$, together with a markedly lower mean rank. In these domains, temporal supervision not only preserves content generation but appears to improve the structural regularities captured by the token generator. The lower held-out $L_{\mathrm{tok}}$ under Joint in both domains is consistent with these output-level improvements (see Appendix D).

The relationship between temporal recoverability and content-generation quality is domain-dependent. On BPI2018, improved temporal recoverability is accompanied by unchanged activity-sequence quality. On MIMIC-IV, it coexists with a modest degradation in marginal event-token distribution metrics, but no reduction in frequent-event coverage. On network traffic and temporal knowledge graphs, Joint training is associated with better evaluated content-generation metrics alongside improved temporal recoverability. These results indicate that stronger temporal organization does not necessarily impose a uniform temporal–semantic trade-off. Instead, its effect depends on the domain structure, the density of temporal supervision, and the content properties captured by the evaluation protocol.

Appendix appendix K compares the layer-wise representations learned by the token-generation model with and without explicit temporal supervision. We use Centered Kernel Alignment (CKA), Centered Kernel Nearest Neighbor Alignment (CKNNA), and Mutual $k$-Nearest Neighbor (M-KNN). Across domains, temporal supervision changes the learned representations, but the form of this change varies: in some cases it affects both global and local geometry, whereas in others it is primarily local. Changes are also generally less pronounced in earlier layers, which may reflect their stronger focus on the input structure. The patterns nevertheless differ across the four domains, so they do not yet provide a single mechanistic account of how temporal supervision reshapes representation learning in SSMs.

Several questions remain about the conditions under which temporal supervision is beneficial. Building on the sensitivity analyses in Appendix appendix E, future work should test more systematically how the gains from Joint training depend on temporal-supervision strength across all domains, temporal-head capacity, and the structure of the target domain. This would clarify the conditions under which temporal gradients beneficially reshape the backbone.

The practical value of improved temporal recoverability also warrants direct evaluation. Future work should test whether it translates into better performance on downstream tasks that explicitly depend on temporal structure, such as anomaly detection, forecasting, or temporal retrieval. This would help establish whether the representation-level gains observed here are not only diagnostic but also practically useful across temporally structured domains.

Finally, future work should assess the distributional fidelity of timestamps produced during free autoregressive generation, testing whether the representation-level gains translate into faithful temporal synthesis beyond the content-generation metrics considered in this study. For RQ2, we prioritized domain-specific generation protocols established in the respective application literatures, allowing a controlled comparison of the two training regimes using metrics with existing task-specific meaning. These protocols primarily assess event content, however. Developing standardized evaluation procedures for jointly generated event content and timestamps is therefore an important next step, particularly because existing generative approaches often decouple temporal prediction from content generation (section 7).