DES-JEPA: Learning a Shared Vocabulary for High-Dimensional Discrete Event Sequences

AI summary

Overview

Research area: Artificial intelligence for temporal systems — specifically self-supervised representation learning and in-context (prior-fitted network) foundation models for discrete event sequences.

Technical level: Advanced. The paper assumes familiarity with joint-embedding predictive architectures, vector quantization, temporal point processes, and in-context learning.

Scope: The paper proposes DES-JEPA, which learns a single shared discrete codebook across synthetic stochastic processes whose event vocabularies have no correspondence, and DES-PFN, which reuses that frozen codebook for zero-shot in-context next-event prediction on unseen real-world datasets with 3 to 700 event types.

What This Paper Is About

Continuous temporal foundation models such as Chronos and TimesFM can transfer across domains because numerical values live on a comparable scale, but discrete event sequences (operational logs, clinical event trajectories) carry categorical labels whose meaning changes completely from dataset to dataset, so no shared embedding table is possible. Existing approaches either refit expensive statistical models like Hawkes processes from scratch per dataset, or rely on foundation inference models such as FIM-PP, which supports at most 22 event types. The paper's goal is a representation that transfers across processes rather than a better fit to any single one, so that realistic event streams with hundreds of distinct event types become tractable without task-specific fine-tuning.

Key Contributions

1. DES-JEPA: A joint-embedding predictive architecture that combines predictive latent learning, vector quantization, and prior-fitted-network-style meta-pretraining over synthetic data-generating processes to learn a shared codebook from unaligned discrete event sequences.
2. DES-PFN: A discrete event sequence prior-fitted network that consumes the frozen teacher codebook from DES-JEPA and autoregressively predicts next events in the shared code space, conditioned on an in-context reference set, with no gradient updates.
3. A kernel-based decoder: A Nadaraya-Watson estimator with a cosine-based kernel that maps a predicted code back to a context-specific raw event type, since quantization is many-to-one and non-invertible and event labels are unaligned across processes.
4. An empirical demonstration of scale: Zero-shot evaluation on eight real-world datasets spanning 3 to 700 event types, reported as the only method producing non-trivial results once cardinality reaches hundreds of event types.

Main Findings

• Zero-shot results without any fine-tuning: DES-PFN achieves next-event top-1 accuracy on eight real datasets at n_context = 128, n_query = 20, L = 32: retweet 0.520 ± 0.023, StackOverflow 0.356 ± 0.030, Taobao 0.420 ± 0.057, Amazon 0.241 ± 0.035, HDFS 0.200 ± 0.028, BGL 0.136 ± 0.020, MIMIC-IV 0.029 ± 0.010, and Vehicle 0.162 ± 0.014 (500 types) / 0.094 ± 0.009 (700 types).

• Specialists still win on small vocabularies: Fully-trained gradient-trained baselines (RMTPP, NHP, SAHP) outperform DES-PFN on the four lower-cardinality datasets — for example Taobao 0.559 for RMTPP versus 0.420 for DES-PFN, and StackOverflow 0.420 for both NHP and SAHP versus 0.356 for DES-PFN — but each requires a real held-out split and per-dataset retraining from scratch.

• FIM-PP is competitive only at low cardinality: FIM-PP beats DES-PFN on three of the five datasets where it can be evaluated (StackOverflow 0.441 versus 0.356, Taobao 0.516 versus 0.420, HDFS 0.266 versus 0.200), and loses on retweet (0.462 versus 0.520) and Amazon (0.207 versus 0.241).

• The 22-event ceiling is the decisive gap: FIM-PP's fixed architecture prevents evaluation on BGL (120 types), MIMIC-IV (441) and Vehicle (500/700). On these datasets all gradient-trained baselines either lack a usable data split or exceed the prescribed 600 s training-time budget (marked with a dagger for Vehicle), leaving DES-PFN as the only method with reported results.

• Weakest case is MIMIC-IV: The authors attribute the 0.029 accuracy to shared blood-draw timestamps across 95 % of consecutive lab events, which may leave little time information to exploit.

• The codebook does not collapse: With K = 768 codes, aggregated perplexity reaches approximately 500 over the training mixture; measured on 800 fresh synthetic tasks per family with n_active in [3, 300] and n_context = 32, all six DGP families use 601 to 716 of the 768 codes. Round-trip decode accuracy from true code to raw id is uniformly modest at 2.5 to 12.7 %, which the authors read as a capacity/collision cost rather than a degeneracy.

• Deliberate train/inference asymmetry: With window W = 2, only the single most recent position of each tokenized sequence is coded from a repeated-embedding fallback instead of a genuine future window; the authors state this position's code is systematically less informative than every earlier position's.

Methodology in Plain English

The authors start from the observation that a raw event id is meaningless across datasets, so instead of learning what events are, they learn what kinds of patterns recur. They generate a large synthetic training corpus from six families of stochastic processes — higher-order Markov chains, Hawkes processes sampled by Ogata thinning, randomly initialised and never-trained LSTMs, hidden Markov models, bursty renewal processes, and template/noisy-template generators — each with a freshly drawn, mutually incompatible event alphabet of up to 1024 types. No real-world data is used in either pretraining stage.

DES-JEPA itself has two encoders. A student encoder reads only the past (causal attention with a 32-token sliding window and RoPE) and predicts a code; a teacher encoder is a shallow bidirectional prefix (2 of the student's 6 layers) that looks at the next two events and snaps their representation to the nearest entry in a 768-entry codebook. The teacher is never trained by gradients — it follows the student by an exponential moving average with τ = 0.999, and stop-gradients prevent both encoders from collapsing to a constant. Training minimises a cross-entropy loss on predicting the future code plus standard VQ-VAE codebook commitment terms with β = 0.25. Because naive vector quantization tends to collapse onto a handful of indices, the authors monitor codebook perplexity and revive dead codes using a usage-EMA criterion with a batched reset trigger and a threshold relative to each code's fair share.

Events are embedded purely from statistics of their own occurrence history — surprisal, running and local-window frequency, whether the type has been seen, and recency and dispersion of inter-arrival gaps, plus step-count analogues that need no timestamps — never from a table keyed on the raw id.

DES-PFN then freezes the codebook and treats it as a vocabulary. A query sequence and a reference context drawn from the same process are both tokenized into codes, and a transformer alternates two attention axes: ordinary causal attention along time and a column-wise attention that pools information across the reference sequences through 32 learnable inducing vectors, in the style of TabICLv2 and the Set Transformer. Inter-arrival times use a Chronos-style discretised embedding, turning time prediction into classification. At inference the model forecasts codes autoregressively, and the Nadaraya-Watson decoder converts the predicted code into an actual event type by letting each context occurrence cast a similarity-weighted vote for its own identity.

Why This Matters

Impact on research: The paper reframes cross-domain transfer for event data. Rather than trying to align symbols that have no shared meaning, it learns a vocabulary of recurring dynamics — the missing invariant is claimed to be a set of shared patterns, not the event labels themselves. This directly attacks the ceiling that has kept foundation inference models like FIM-PP at a few dozen event types, and offers an alternative to refitting classical point processes per dataset.

Real-world applications:

• Vehicle diagnostic and fault-code logs, where the authors' own industrial setting involves large event alphabets.
• Electronic health records and clinical-event trajectories such as MIMIC-IV lab-test sequences.
• System and infrastructure logs (HDFS, BGL) for anomaly-adjacent next-event modelling.
• User-behaviour streams such as Taobao and Amazon actions, retweet cascades and StackOverflow badge awards.

Industry relevance: The work comes from BMW Group and the University of Glasgow, and the evaluation covers exactly the setting industrial practitioners face — hundreds of distinct event codes, no labelled target domain, no budget for per-dataset retraining. Both models were pretrained on a single NVIDIA RTX 2080 Ti in about one day each, and all reported numbers were measured at inference on an Apple M2 Pro CPU with no GPU acceleration, which suggests the downstream cost profile is modest.

Future Directions

• Adaptive codebook capacity: The codebook size is fixed at pretraining time at K = 768 and is not adapted to a downstream dataset's true cardinality; the authors note that rate-distortion experiments are currently lacking.
• Long-horizon rollout: Every result is one-step-ahead. Genuine multi-step autoregressive rollout is untested, particularly error accumulation and representing multiple plausible futures beyond a single point forecast.
• Multi-head, multiple-choice decoding: The authors propose replacing the single-code predictor with a multi-head decoder trained under a winner-takes-all multiple-choice learning objective, so several plausible future codes can be emitted per query instead of one.
• Better handling of information-poor timing: The MIMIC-IV result points at settings where consecutive events share timestamps and timing carries little signal, an open question for the current embedding scheme.

Target Audience

Researchers and practitioners in temporal foundation models, event-sequence modelling and self-supervised representation learning, particularly those working with high-cardinality categorical streams in healthcare, automotive diagnostics, log analysis or user-behaviour analytics. It will also interest engineers evaluating whether an in-context, zero-shot model can replace per-dataset temporal point process fitting, and readers who want a concrete example of applying JEPA-style predictive representation learning outside vision and audio.

Authors’ abstract

Temporal foundation models such as Chronos and TimesFM generalize across domains through a shared numerical observation space. In contrast, discrete event sequences (DESs), including operational event logs or patient clinical-event trajectories, have domain-specific event vocabularies whose labels lack a consistent correspondence across domains. Beyond this lack of consistency, the number of distinct events in real-world settings is high enough to challenge both classical temporal point process modeling and current in-context foundation inference models such as FIM-PP. Critically, current foundation models for DESs are limited to a small, fixed number of event types, restricting their applicability to realistic, high-dimensional event streams. To address this, we propose two complementary modules: (i) DES-JEPA, a joint-embedding predictive architecture that learns a shared discrete vocabulary from synthetic stochastic processes with unaligned event vocabularies; and (ii) DES-PFN, a discrete event sequence prior-fitted network that reuses the learned vocabulary for in-context next-event prediction. Without task-specific fine-tuning, this framework achieves competitive zero-shot performance on eight unseen real-world datasets and scales up to 700 distinct event types.