NEST: Nested Event Stream Transformer for Sequences of Multisets

TL;DR

The paper addresses modeling hierarchical event streams that are naturally sequences of multisets, such as clinical encounters in EHRs. It introduces NEST, a nested event stream Transformer with Set-Wise Encoder (SWE) and Cross-Set Encoder (CSE) to preserve multiset boundaries and enable dual-level contextualization. A dedicated Masked Set Modeling (MSM) objective complements standard MLM to produce ready-to-use set representations without post-hoc pooling. Across Instacart, MIMIC-IV hosp, and proprietary EHR data, NEST delivers improved pretraining efficiency and superior downstream performance over flattened-sequence baselines, validating the proposed architecture and MSM approach for real-world SeqSet tasks.

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.

Figures (7)

• Figure 1: NEST models event stream data as sequences of multisets, where temporal order is preserved across sets and intra-set tokens are unordered. It is a domain-agnostic FM architecture that explicitly models set-level representations.
• Figure 2: Dense attention is applied both within and across sets, yet their composition yields a sparse pattern in each NEST layer. The theoretical complexity of the composite pattern is $O(nN + m^2) = O(nN)$, where $N=mn$ is the context window size.
• Figure 3: NEST is a simple stack of hierarchical Transformer layers. Each layer is a composite Transformer block that integrates a Set-Wise Encoder (SWE) with a Cross-Set Encoder (CSE). T2V is used for absolute time encoding, and RoPE is used in each CSE for position encoding at the multiset-level.
• Figure 4: NEST validation NBR performance during the training. Set prediction is challenging in this dataset, where the validation MSM loss converges at 1.5.
• Figure D.1: Distributions for # multisets and # tokens per user.

Theorems & Definitions (2)

• Definition 1
• Proposition 1: Fast Soft Support Matching under KL Gradient Update