🤖 machine learning

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

The paper introduces ChronoSSM, an autoregressive State Space Model that jointly trains on event and timestamp generation to produce more temporally informative representations without compromising content prediction quality, outperforming traditional two-stage approaches that treat timing as a secondary signal.

Original authors: Adrien Schoen, Nachiketa Ratnakar Patil, Arjun Bhagoji, Francesco Bronzino

Published 2026-08-12

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Adrien Schoen, Nachiketa Ratnakar Patil, Arjun Bhagoji, Francesco Bronzino

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot to tell a story. For a long time, scientists have been very good at teaching robots what happens next in a story—like predicting that after "The king died," the next line is likely "The queen died of grief." But usually, the robot doesn't care when these things happen. It treats time like a boring background detail, just a secondary note in the margin. This works fine for writing novels, but it falls apart when you need to understand real-world events, like a doctor tracking a patient's heartbeat or a security system spotting a network attack. In those cases, knowing that a heart skipped a beat three seconds ago is just as important as knowing the beat happened at all. If the robot can't learn the rhythm of the story, it can't truly understand the plot.

This is where a new idea called "ChronoSSM" comes in. The researchers behind it asked a simple but tricky question: Can we teach a robot to learn the story and the timing at the exact same time, using the same brain? Usually, scientists use a two-step method: first, they teach the robot the story perfectly, freeze its brain, and then attach a separate, smaller module just to guess the time. But the authors suspected this was like trying to teach someone to drive by first teaching them to steer, freezing their hands, and then hoping they could figure out the gas pedal later. They wondered if teaching the robot to handle the story and the clock together from the start would make it smarter, without making it forget how to tell the story in the first place.

The team built a new kind of AI model called ChronoSSM to test this. Think of the model as a musician learning to play a complex song. The "Two-Stage" method is like teaching the musician the notes first, locking their fingers in place, and then asking them to figure out the tempo later. The "Joint" method, which ChronoSSM uses, is like teaching the musician the notes and the rhythm simultaneously, letting the rhythm help shape how they press the keys. The researchers tested this on four very different types of data: business process logs (like a checklist of office tasks), hospital patient records (a stream of medical events), network traffic (data packets zipping across the internet), and a map of world events (facts about history and news).

The results were surprisingly clear. When they looked at the "frozen" brains of the models after training, the ones trained with the "Joint" method were much better at revealing the timing information. It was as if the Joint-trained models had a hidden rhythm section that the Two-Stage models simply didn't have. The researchers could easily pull the timing data out of the Joint models, whereas the Two-Stage models kept their timing information buried deep and hard to find. This happened across all four different types of data, whether the timing information was dense (every single step had a time) or sparse (only some steps had a time).

But there was a catch: could teaching the robot to care about time make it worse at telling the story? The researchers were worried that by splitting the robot's attention between "what" and "when," it might start making mistakes in the story itself. They checked this carefully. On the business process data, the story quality was exactly the same. On the hospital data, the story was almost identical, with only tiny, negligible differences in how well it matched the most common events. However, on the network traffic and world event data, the Joint method actually made the story better. It seems that learning the rhythm helped the model understand the structure of the data even more deeply.

In short, the paper suggests that we don't have to choose between a smart storyteller and a precise timekeeper. By training them together, we get a model that understands both the plot and the pace, without losing its ability to tell a good story. This opens the door for smarter AI that can handle real-world situations where time matters just as much as the events themselves.

Technical Summary: ChronoSSM

Problem Statement

Modern autoregressive sequence models, including Transformers and State Space Models (SSMs), are predominantly trained to predict "what" happens next in a sequence, often treating "when" it happens as a secondary concern or an auxiliary signal. In domains where events carry explicit temporal metadata—such as business-process traces, clinical records, network traffic, and temporal knowledge graphs—this separation limits the model's ability to perform temporal reasoning, anomaly detection, and faithful reconstruction of event chronology.

A common strategy to address this is a Two-Stage approach: first training a model for next-event prediction, freezing its representations, and then fitting a separate timing module. This method implicitly assumes that representations optimized solely for token prediction already contain sufficient temporal structure. The authors argue this assumption is questionable, as the token generation objective rewards features that improve next-event likelihood but does not explicitly encourage the encoding of inter-arrival times.

Methodology: ChronoSSM

The authors introduce ChronoSSM, an autoregressive State Space Model designed to jointly model events and timestamps within a shared backbone.

Model Architecture

Backbone: The model utilizes a causal autoregressive backbone (instantiated primarily with Mamba2, a selective SSM) that processes a history of tokens $X_k$ to produce token-wise hidden representations $h_k$ .
Dual Heads:
1. Token Prediction Head: A linear map predicting the next token $x_{k+1}$ .
2. Temporal Prediction Head: A function predicting the inter-arrival time $\Delta t_{k+1}$ .
Temporal Representation: To predict the time between events, the model defines the transition representation $r_{k+1}$ as the difference between consecutive hidden states: $r_{k+1} = \Delta h_{k+1} = h_{k+1} - h_k$ . This aligns the representation of the time gap with the change in the latent state.
Loss Function: The training objective combines token prediction and temporal prediction via a weighted negative log-likelihood:
$\mathcal{L} = \mathcal{L}_{tok} + \lambda_\tau \mathcal{L}_{time}$
Where $\mathcal{L}_{tok}$ is the categorical cross-entropy for tokens, and $\mathcal{L}_{time}$ is the negative log-likelihood of the inter-arrival times (modeled as a log-normal density over log-transformed targets to handle zero values and heterogeneity).

Training Regimes

The paper conducts a controlled comparison between two optimization schedules to isolate the impact of temporal supervision on the backbone:

Joint Training: The backbone, token head, and temporal head are optimized simultaneously. Gradients from both the token and temporal losses update the shared backbone parameters ( $\theta$ ). This allows temporal supervision to directly shape the learned representations.
Two-Stage Training: The backbone and token head are first trained using only $\mathcal{L}_{tok}$ . The backbone is then frozen, and the temporal head is trained on the resulting frozen representations. This serves as the baseline where temporal gradients cannot modify the backbone.

Experimental Setup

The evaluation spans four distinct domains with varying sequence structures and timestamp densities:

Business-Process Traces (BPI2018): Dense timestamps for activity logs.
Clinical Event Sequences (MIMIC-IV): Dense, bursty timestamps for patient vitals and observations.
Network Traffic: Partially timestamped (timestamps only at packet boundaries).
Temporal Knowledge Graphs (GDELT): Partially timestamped (timestamps only at fact markers).

The study also includes symbolic music generation in the appendix.

Key Results

RQ1: Temporal Recoverability

The authors investigate whether Joint training yields representations where inter-arrival information is more recoverable than in Two-Stage training. They use two diagnostic metrics on frozen representations:

Linear Probe: A ridge-regression probe predicting the temporal target from the representation.
Temporal Cohesion Score (TCS): A geometric measure of whether transitions with similar time gaps cluster together in representation space.

Findings: Across all four domains, Joint training consistently produces representations with higher temporal recoverability.

Network Traffic showed the most dramatic improvement (TCS increased from 0.147 to 0.622).
BPI2018 and MIMIC-IV showed significant improvements in linear probe performance ( $R^2$ and MAE), though geometric cohesion (TCS) improved more modestly.
Temporal Knowledge Graphs showed consistent, albeit smaller, gains across all metrics.

RQ2: Quality Degradation

The authors assess whether injecting temporal supervision systematically degrades the model's primary event-generation capability.

Findings: There is no evidence of systematic degradation in generation quality under Joint training.

BPI2018: Generation quality (measured by Damerau-Levenshtein similarity) was identical between Joint and Two-Stage.
MIMIC-IV: Joint training showed a slight degradation in marginal event-token distribution metrics (JSD and TV distance) but maintained identical recall for frequent clinical events.
Network Traffic & Temporal KG: Joint training actually improved content-generation quality. In these domains, Joint training yielded lower Jensen-Shannon divergence for packet headers and significantly better link prediction metrics (MRR, Hits@k) compared to Two-Stage.

Significance and Claims

The paper claims that explicit temporal supervision can improve the temporal recoverability of learned representations without materially degrading autoregressive event modeling.

Key takeaways include:

Unified Representation: Joint training successfully learns event semantics and temporal dynamics within a unified representation, challenging the assumption that token-only training suffices for temporal tasks.
Domain Dependence: The interaction between temporal recoverability and generation quality is domain-dependent. In some cases (Network Traffic, Temporal KG), temporal supervision appears to enhance the structural regularities captured by the token generator, leading to better overall performance.
Practical Implication: This approach paves the way for training temporally aware generative models that do not require a decoupled, two-stage pipeline, potentially benefiting downstream tasks like anomaly detection and forecasting that rely on faithful chronology.

The authors remain modest, noting that while representation-level gains are clear, a rigorous evaluation of the distributional fidelity of timestamps during free autoregressive generation remains a direction for future work. They also emphasize that the benefits of Joint training depend on the specific domain structure and the density of temporal supervision.