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

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

Abstract Summary

The paper introduces ChronoSSM, an autoregressive State Space Model (SSM) that jointly models events and timestamps with a shared backbone trained using combined token and temporal generation objectives. The authors compare a Joint training regime, where temporal supervision updates the backbone, with a Two-Stage regime, where timing is learned only using frozen event representations. Across four domains spanning dense and partial timestamp supervision, Joint training consistently makes inter-arrival information more recoverable from frozen representations without any systematic degradation in content-generation quality overall. The results show that temporal supervision can produce more temporally informative representations without materially degrading autoregressive event modeling.

Introduction Summary

Modern autoregressive sequence models are used beyond natural language in domains such as business-process traces, clinical event sequences, network traffic, and temporal knowledge graphs. In these settings, events are associated with explicit temporal metadata, and downstream tasks depend not only on what happened but also on when it happened. The paper addresses the question: can autoregressive sequence models learn both event identity and event timing within a shared representation?

A common way to incorporate timing is to treat it as an auxiliary task: first train a sequence model for next-event prediction, then freeze its representations and fit a separate timing module on top. This two-stage strategy assumes that representations optimized only for token prediction already retain enough temporal information, which is questionable because the token generation objective only rewards features that improve next-event likelihood, not inter-arrival modeling.

ChronoSSM predicts both the next event and its timestamp from a shared backbone. A lightweight temporal head is trained jointly with the token prediction head, so temporal gradients can directly update the backbone. The paper conducts a controlled comparison between two training regimes that share the same architecture, data, and temporal objective: a Two-Stage regime and a Joint regime. The evaluation covers four key domains: business-process traces, clinical event sequences, network traffic, and temporal knowledge graphs, with additional experiments on symbolic-music generation.

The two main research questions are:

• RQ1 (Temporal Recoverability): Does Joint training yield representations from which inter-arrival information is more recoverable than under Two-Stage training? The paper finds that Joint training consistently produces representations from which timing is more recoverable across all four domains, with the magnitude of the effect varying by domain.

• RQ2 (Quality Degradation): Does injecting temporal supervision into the shared backbone systematically degrade downstream modeling quality? The paper finds no evidence of a systematic degradation under Joint training, with Joint generating higher quality data in 2 out of 4 domains.

Background Summary

The paper describes the State Space Architecture, noting that SSMs are sequence models with recurrent latent-state dynamics that support linear-time processing in sequence length. The core mechanism is a linear dynamical system that maps an input sequence to an output sequence through a hidden state. 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 parameters are conditioned on the current input, allowing the model to selectively propagate or discard information according to the sequence context.

The problem setting defines a dataset D of timestamped sequences, where each sequence is written as (x k, t k) k=0 T-1, with x k a discrete token and t k its associated timestamp. From successive timestamps, inter-arrival times Δt k = t k - t k-1 k=1 T-1 are derived. The objective is to model the joint distribution of the next token x k+1 and its associated inter-arrival time Δt k+1: p γ(x k+1, Δt k+1 X k). Using the chain rule, this is factored as p θ,α(x k+1 X k) p θ,φ(Δt k+1 x k+1, X k), where γ = θ, α, φ collects the parameters of the shared sequence model, token-prediction head, and temporal-prediction head.

The shared sequence model is a causal autoregressive model parameterized by θ that produces a token-wise hidden representation h k from the history X k. The token-prediction head is a linear map W α tok ∈ R V×D parameterized by α, modeling the next-token distribution as softmax(W α tok h k). The temporal-prediction head is a function f φ time parameterized by φ and conditioned on a transition representation r k+1, derived from both h k and h k+1.

ChronoSSM Model Design Summary

The temporal representation r k+1 supplied to the temporal-prediction head is defined as the change between consecutive hidden states: r k+1 = Δh k+1 = h k+1 - h k. This choice aligns the temporal target and the representation used to predict it, as both are defined as differences between consecutive positions: Δt k+1 = t k+1 - t k corresponds to Δh k+1 = h k+1 - h k.

The training objective decomposes the negative log-likelihood of each transition into a token-prediction term and a temporal-prediction term. The token loss is categorical cross-entropy, while the temporal loss is the negative log-likelihood of the observed inter-arrival times under a conditional density. The joint negative log-likelihood is L NLL = L tok + L time. In practice, the weighted objective L = L tok + λ τ L time is optimized, where λ τ ≥ 0 controls the relative contribution of temporal prediction.

Two training schedules are considered:

• Joint Training: The backbone and both prediction heads are optimized simultaneously. Gradients from both L tok and L time update the backbone parameters θ. The temporal loss backpropagates through Δh k+1 = h k+1 - h k, encouraging changes between consecutive hidden states to become informative about the corresponding inter-arrival times.

• Two-Stage Training: Token and temporal learning are separated into two phases. In the first phase, the backbone and token-prediction head are optimized using L tok alone. In the second phase, the temporal-prediction head is trained using L time on the hidden-state transitions produced by the token-trained sequence model. The backbone is not updated during the second phase.

For inference, generation proceeds autoregressively. Given the token history 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 Δh k+1. The temporal head then predicts the associated inter-arrival time Δt̂ k+1 from Δh k+1, and the timestamp is updated as t k+1 = t k + Δt̂ k+1.

Experimental Setup Summary

The experimental setup is designed to isolate the effect of allowing temporal supervision to update the shared backbone. The primary evaluation compares Joint and Two-Stage training, which use the same model components, data, temporal objective, and generation procedure within each domain, but differ in whether the temporal loss updates the shared backbone.

The controlled comparison is conducted on four domains:

• Business-process traces (BPI Challenge 2018): Each case is represented as a sequence of timestamped activities. The activity label defines the token, while its occurrence time defines the timestamp. Events are ordered chronologically within each trace.

• Clinical event sequences (MIMIC-IV): Uses the chartevents table, which records timestamped bedside observations. The event defines the token, and the recording time defines the timestamp.

• Network traffic: Uses captures of video-streaming traffic. Packet contents are serialized as sequences of byte tokens, with a dedicated token delimiting successive packets. Timestamps are attached only to these boundary tokens.

• Temporal knowledge graphs (GDELT): Facts are mapped to a fixed vocabulary of entity and relation identifiers and represented as structured token blocks. The timestamp of each fact is attached only to its marker.

For the dense domains (BPI2018 and MIMIC-IV), a log-transformed auxiliary target is used: z k = log(1 + Δt k/s), where s is a dataset-specific time scale. The temporal head produces a raw scalar output, mapped through a softplus nonlinearity to obtain a non-negative inter-arrival prediction. The temporal factor is instantiated as a fixed-variance lognormal density.

For the partially timestamped domains (network traffic and temporal knowledge graphs), timestamps are attached only to selected boundary tokens. The temporal loss is evaluated only for positions in I time, the set of positions carrying valid temporal targets. The same target scaling, logarithmic transformation, softplus-constrained prediction, and squared-error loss are applied.

All main experiments instantiate ChronoSSM with a Mamba2 backbone. The temporal-head architecture is selected separately for each domain.

Temporal Recoverability Results (RQ1)

Temporal recoverability is defined 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. Two metrics are used:

• Linear timing probe: A ridge-regression probe predicts the temporal target z k from the frozen temporal representation r k. Mean absolute error (MAE) and coefficient of determination (R2) are reported.

• Temporal Cohesion Score (TCS): Evaluates whether transitions associated with similar temporal targets occupy nearby regions of representation space. For each temporal representation, its K nearest neighbors under cosine similarity are identified, and the mean absolute difference between their associated temporal targets is computed. This is compared with the mean target difference obtained by assigning K randomly selected transitions.

The results show that in every domain, Joint achieves higher TCS and probe R2, together with lower probe MAE, than Two-Stage. The diagnostic profile differs by domain:

• BPI2018: TCS increases slightly from 0.794 to 0.803, while probe R2 increases from 0.461 to 0.631 and MAE decreases from 3.342 to 2.775.

• MIMIC-IV: TCS increases from 0.848 to 0.862, probe R2 from 0.588 to 0.658, and MAE decreases from 0.645 to 0.597.

• Network traffic: Large improvements in all three diagnostics, including TCS from 0.147 to 0.622, probe R2 from 0.190 to 0.394, and MAE from 0.382 to 0.308.

• Temporal KG: Changes are smaller across all three diagnostics, but gains remain modest and consistently favor Joint.

The takeaway is an affirmative answer to RQ1: compared with Two-Stage, Joint makes inter-arrival information more recoverable from the representations supplied to the temporal head.

Effect on Generation Quality Results (RQ2)

Content-generation quality is evaluated using domain-specific protocols:

• BPI2018: Damerau-Levenshtein similarity (DLS) following the suffix-prediction protocol.

• MIMIC-IV: Jensen-Shannon divergence and total variation distance between generated and reference event-token distributions, plus recall over the 50 most frequent reference event types.

• Network traffic: Jensen-Shannon divergence for source and destination IP addresses and TCP ports.

• Temporal knowledge graphs: Filtered link prediction with mean reciprocal rank, mean rank, and Hits@(1,3,10).

The results show:

• BPI2018: Joint and Two-Stage obtain identical mean and median DLS (0.487 and 0.500).

• MIMIC-IV: Two-Stage achieves slightly lower token JSD (0.302 vs 0.314) and total variation distance (0.451 vs 0.477), while both regimes obtain the same recall over the 50 most frequent clinical event types (0.780).

• Network traffic: Joint consistently lowers the JSD of all four evaluated packet-header fields, with particularly large improvements for source and destination TCP ports (0.0764 vs 0.1248 and 0.1120 vs 0.2183).

• Temporal knowledge graphs: Joint improves every filtered link-prediction metric, yielding substantially higher MRR (0.1461 vs 0.0576) and Hits@ 1, 3, 10, together with a markedly lower mean rank (269.06 vs 2155.60).

The takeaway is a negative answer to RQ2: although Joint does not dominate Two-Stage on every individual metric, there is no evidence that the improved temporal recoverability is obtained at the cost of a systematic degradation in the model's primary generative capability.

Discussion and Future Work Summary

The paper introduces ChronoSSM, an autoregressive model that jointly predicts tokens and inter-event times from a shared backbone. Across four temporally structured domains, Joint training yields representations from which temporal information is more easily recoverable (RQ1). At the same time, it does not systematically degrade downstream generation quality relative to Two-Stage training (RQ2): generation quality is unchanged on BPI2018, modestly reduced on selected MIMIC-IV distributional metrics, and improved on network traffic and temporal knowledge graphs.

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.

Future work should test more systematically how the gains from Joint training depend on temporal-supervision strength, temporal-head capacity, and the structure of the target domain. The practical value of improved temporal recoverability warrants direct evaluation through downstream tasks such as anomaly detection, forecasting, or temporal retrieval. Future work should also assess the distributional fidelity of timestamps produced during free autoregressive generation, and develop standardized evaluation procedures for jointly generated event content and timestamps.

Improvements to AI Systems:

1. Unified Temporal-Event Modeling: Implement a shared-backbone architecture where token prediction and timestamp prediction are jointly optimized in a single training pass. This allows the model to learn representations that encode both what happened and when it happened, rather than treating timing as a post-hoc task. The improved system can generate coherent event sequences with realistic inter-arrival times, useful for business process simulation, clinical trial design, and network traffic synthesis.

2. Temporal Recoverability via Representation Geometry: Adopt the temporal representation as the difference between consecutive hidden states (Δh = h k+1 - h k), aligning the representation space with the temporal target (Δt). This enables downstream tasks to extract timing information via simple linear probes, improving performance in anomaly detection (e.g., identifying unusual gaps in network packets) and forecasting (e.g., predicting equipment failure from irregular maintenance intervals).

3. Joint Training with Weighted Temporal Loss: Use a weighted objective L = L tok + λ τ L time, where λ τ controls temporal supervision strength. This allows system designers to tune the trade-off between content fidelity and temporal accuracy per domain. The improved system can adapt to domains with sparse timestamps (e.g., knowledge graphs) by reducing λ τ, or to dense temporal domains (e.g., clinical monitoring) by increasing it.

4. Domain-Adaptive Temporal Head: Implement a lightweight temporal head with domain-specific output distributions (e.g., lognormal for dense timestamps, squared-error for partial supervision). This enables the system to handle both fully timestamped sequences (e.g., electronic health records) and partially timestamped ones (e.g., network packets with boundary markers), improving robustness across heterogeneous data sources.

5. Autoregressive Joint Generation Pipeline: During inference, generate the next token first, then predict its timestamp from the transition representation, and update the timestamp cumulatively. This enables the system to produce complete event logs with both content and timing, useful for synthetic data generation in privacy-preserving settings (e.g., generating realistic patient trajectories without exposing real records).

6. Temporal Cohesion Scoring for Representation Quality: Integrate the Temporal Cohesion Score (TCS) as a training-time diagnostic to monitor whether similar temporal targets cluster in representation space. This allows the system to self-assess temporal awareness during training and adjust λ τ dynamically, improving performance in time-sensitive applications like fraud detection (where transaction intervals are discriminative) and predictive maintenance.

7. Cross-Domain Transfer with Temporal Supervision: Leverage the finding that Joint training improves temporal recoverability without systematic content degradation across four domains. The improved system can transfer pre-trained temporal representations to new domains with limited labeled data, accelerating deployment in emerging applications like smart-grid load forecasting or social media event detection.

8. Two-Stage Fallback for Stability: Retain the Two-Stage training regime as a fallback for domains where content quality is paramount (e.g., MIMIC-IV marginal distributions). The system can automatically select between Joint and Two-Stage based on validation metrics, ensuring no degradation in critical content-generation tasks while maximizing temporal awareness where beneficial.

Modern sequence models, from Transformers to State Space Models, have enabled powerful generative modeling across diverse domains, yet they are typically trained to predict what happens while treating when it happens as a secondary concern. In data-mining settings where events are associated with explicit timing information, this separation can limit temporal reasoning, anomaly detection, and faithful reconstruction of event chronology. A common strategy is to treat timing as an auxiliary signal, training a separate timing model using representations learned solely for event prediction. However, this two-stage approach implicitly assumes that representations optimized for event prediction already contain sufficient temporal structure. We introduce ChronoSSM, an autoregressive State Space Model (SSM) that jointly models events and timestamps with a shared backbone trained using combined token and temporal generation objectives. We compare the joint regime, where temporal supervision updates the backbone, with the two-stage regime, where timing is learned only using the frozen event representations. Across four domains spanning dense and partial timestamp supervision, joint training consistently makes inter-arrival information more recoverable from frozen representations without any systematic degradation in content-generation quality overall. Our results show that temporal supervision can produce more temporally informative representations without materially degrading autoregressive event modeling.