NEST: Nested Event Stream Transformer for Sequences of Multisets

TL;DR

NEST is a novel transformer architecture designed to handle hierarchical event stream data as sequences of multisets, improving efficiency and representation quality over traditional flattening methods.

Contribution

The paper introduces NEST, a transformer model that preserves hierarchy in event streams and a new Masked Set Modeling paradigm for better set-level representations.

Findings

NEST outperforms existing models in real-world multiset sequence tasks.

NEST improves pretraining efficiency and downstream task performance.

Preserving hierarchy enhances both computational efficiency and representation quality.

Abstract

Event stream data often exhibit hierarchical structure in which multiple events co-occur, resulting in a sequence of multisets (i.e., bags of events). In electronic health records (EHRs), for example, medical events are grouped into a sequence of clinical encounters with well-defined temporal structure, but the order and timing of events within each encounter may be unknown or unreliable. Most existing foundation models (FMs) for event stream data flatten this hierarchy into a one-dimensional sequence, leading to (i) computational inefficiency associated with dense attention and learning spurious within-set relationships, and (ii) lower-quality set-level representations from heuristic post-training pooling for downstream tasks. Here, we show that preserving the original hierarchy in the FM architecture provides a useful inductive bias that improves both computational efficiency and representation quality. We then introduce Nested Event Stream Transformer (NEST), a FM for event streams comprised of sequences of multisets. Building on this architecture, we formulate Masked Set Modeling (MSM), an efficient paradigm that promotes improved set-level representation learning. Experiments on real-world multiset sequence data show that NEST captures real-world dynamics while improving both pretraining efficiency and downstream performance.