🤖 machine learning

NEST: Nested Event Stream Transformer for Sequences of Multisets

The paper introduces NEST, a Nested Event Stream Transformer that preserves the hierarchical structure of event sequences (sequences of multisets) to overcome the computational inefficiencies and representation limitations of existing flattened models, utilizing a novel Masked Set Modeling paradigm to enhance both pretraining efficiency and downstream performance.

Original authors: Minghui Sun, Haoyu Gong, Xingyu You, Jillian Hurst, Benjamin Goldstein, Matthew Engelhard

Published 2026-05-15

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Minghui Sun, Haoyu Gong, Xingyu You, Jillian Hurst, Benjamin Goldstein, Matthew Engelhard

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 trying to understand a patient's medical history, or perhaps a customer's shopping habits. In the real world, these events don't happen one by one in a perfect line like beads on a string. Instead, they happen in groups.

In a hospital: A doctor sees a patient (an "encounter"). During that visit, the patient might get a blood test, a prescription, and a diagnosis all at once. The exact order of those three things might not matter, or the records might be messy. But the group of events belongs to that specific visit.
In a grocery store: A customer goes shopping. They put milk, bread, and eggs in their cart. That "basket" is a group. The order they picked up the items doesn't matter; the basket as a whole is what matters.

Most current AI models (called Transformers) try to flatten these groups into a single, long line of events. They treat the milk, bread, and eggs as if they happened one after another, ignoring that they were part of the same shopping trip. This is like trying to understand a movie by looking at a list of every single frame without knowing which scenes belong together. It's computationally expensive and often misses the bigger picture.

Enter NEST (Nested Event Stream Transformer).

The authors of this paper built a new AI model called NEST that respects these natural groups. Here is how it works, using simple analogies:

1. The Two-Step Dance (The Architecture)

Instead of flattening the data, NEST keeps the "groups" (like hospital visits or shopping baskets) intact. It uses a clever two-step process inside every layer of its brain:

Step A: The Group Huddle (Set-Wise Encoder): First, the model looks at one group (e.g., one hospital visit) and lets all the events in that group talk to each other. It figures out what happened inside that visit. Think of this as a team meeting where everyone discusses their specific tasks.
Step B: The Global Roundtable (Cross-Set Encoder): Next, the model looks at the "captains" of each group (special tokens called [CLS]). These captains meet to share what happened in their respective groups with the rest of the timeline. This helps the model understand how Visit A relates to Visit B.

The Magic Trick: Most models do Step A for all groups, then Step B for all groups. NEST does them interleaved. It does a little bit of Step A, then a little bit of Step B, then back to Step A.

Why this matters: Imagine a relay race. If you run the whole first leg before passing the baton, you might drop some information along the way. NEST passes the baton back and forth constantly. This ensures that no detail from a specific event is lost when the model tries to understand the big picture. The paper proves mathematically that this "bypass" keeps more information alive than the old way.

2. The "Captain's Report" (Masked Set Modeling)

In the old way, after the AI processed all the data, it had to guess how to summarize a whole visit into a single score. This was like asking a student to summarize a whole book chapter after reading it, but only giving them a vague hint.

NEST introduces a new training game called Masked Set Modeling (MSM).

How it works: The model is shown a group of events (like a shopping basket), but it has to guess the contents of that group based on the "Captain" token, while the actual items are hidden.
The Result: This forces the "Captain" token to become a perfect summary of the whole group. Instead of needing a messy, guesswork summary later, the model already has a high-quality summary ready to use for any future task.

3. Why It's Better (The Results)

The authors tested NEST on real-world data:

Medical Records (EHR): They used data from hospitals (MIMIC-IV) and a private pediatric database.
Shopping: They used grocery data (Instacart).

The Findings:

Faster and Leaner: Because NEST respects the groups, it doesn't have to calculate connections between every single event. It only calculates connections within a group and between group captains. This makes it faster and uses less computer memory than older models, even though it has more "brain power" (parameters).
Smarter Predictions: NEST was better at predicting things like:
- Will a patient die during their hospital stay?
- Will a patient be readmitted within 30 days?
- What items will a customer buy next?
No "Post-Processing" Needed: Because the model learned to summarize groups during training, it didn't need to use clumsy, heuristic tricks after training to understand the data. The summaries were ready to go.

Summary

Think of NEST as a model that finally understands that context matters. It knows that a blood test and a prescription given during the same hospital visit are a team, and that team is part of a larger story of the patient's life. By keeping these groups together and letting them communicate efficiently, NEST learns faster, uses less energy, and makes better predictions than models that try to force everything into a single, long line.

Technical Summary: NEST - Nested Event Stream Transformer for Sequences of Multisets

1. Problem Statement

Event stream data, such as Electronic Health Records (EHRs) or user purchase histories, often exhibit a hierarchical structure where multiple events co-occur within specific temporal windows (e.g., a clinical encounter or a shopping basket). This structure is formally defined as a Sequence of Multisets (SeqSet). In these sequences, the order of events within a set (multiset) may be unknown, unreliable, or irrelevant, while the temporal order between sets is critical.

Existing foundation models (FMs) for event streams typically flatten this hierarchy into a one-dimensional sequence. This approach introduces two primary limitations:

Computational Inefficiency: Flattening leads to dense attention mechanisms with quadratic complexity relative to the total number of tokens, ignoring the inherent sparsity of the data structure.
Representation Quality: Standard models learn representations only at the token level. Set-level representations (e.g., a patient's state after an encounter) are recovered only through heuristic post-training pooling (e.g., mean pooling or using a [CLS] token without explicit supervision), which prior work has shown to be suboptimal. Furthermore, flattening forces models to learn spurious within-set relationships that may not exist.

2. Methodology

2.1 Architecture: Nested Event Stream Transformer (NEST)

NEST is a hierarchical Transformer designed specifically for SeqSet data. Its core innovation is the interleaved composition of two distinct encoder types within each layer, rather than stacking them in separate phases.

Set-Wise Encoder (SWE): Operates within each multiset independently. It processes tokens (events) and a learnable [CLS] token specific to that set. The SWE enforces permutation invariance over intra-set tokens (unless fine-grained timestamps are available, in which case Rotary Positional Embeddings (RoPE) can be applied). It aggregates set-level features into the [CLS] token.
Cross-Set Encoder (CSE): Operates exclusively on the sequence of [CLS] tokens from all sets. It models the temporal dynamics between encounters using RoPE to encode inter-set time intervals. Non-[CLS] token hidden states bypass the CSE.
Interleaving and Bypass: The SWE and CSE are interleaved layer-by-layer ( $L$ layers total). Crucially, non-[CLS] token states are propagated through bypass paths around the CSE blocks. This ensures that token-level information is not irrecoverably discarded when passing from set-level processing to sequence-level processing.

Information-Theoretic Justification: The paper argues via Proposition 1 that sequential designs (stacking all SWEs before all CSEs) violate the Data Processing Inequality by discarding token-level information at the SWE-to-CSE boundary. The interleaved design preserves this information, allowing the model to revisit raw token states in light of cross-set context at every layer.

2.2 Pretraining Objective: Masked Set Modeling (MSM)

To complement the standard Masked Language Modeling (MLM) objective, the authors introduce Masked Set Modeling (MSM).

Mechanism: During pretraining, entire multisets (encounters) are masked with a probability of 40%. Within a masked set, all non-special tokens are replaced with [MASK].
Goal: The model must predict the token distribution of the masked set using the [CLS] token's representation.
Significance: This forces the [CLS] token to explicitly aggregate set-level information and learn a ready-to-use set representation, eliminating the need for heuristic post-hoc pooling. It complements MLM, which focuses on fine-grained token-level contextualization.

3. Key Contributions

Architectural Primitive for SeqSet: NEST is a hierarchical foundation model that respects the opposed symmetries of SeqSet data: intra-set exchangeability and inter-set temporal order.
Information Preservation via Bypass: The paper characterizes how interleaved encoders preserve token-level information that sequential designs discard, confirmed by a topology-isolating simulation study.
Set-Level Pretraining: The formulation of MSM, a permutation-invariant objective that explicitly trains the [CLS] token to predict multiset token distributions, a feature absent in prior EHR foundation models.
Efficiency and Performance: NEST achieves structured sparsity, reducing attention complexity from $O(N^2)$ to $O(Nn + m^2)$ (where $N$ is total tokens, $n$ is set size, $m$ is number of sets), leading to higher throughput and lower memory usage compared to dense attention baselines.

4. Experimental Results

The authors evaluated NEST on three real-world datasets: Instacart (Next Basket Recommendation), MIMIC-IV (EHR clinical prediction), and a Proprietary Institutional EHR (Pediatric risk prediction).

Computational Efficiency: On a single NVIDIA H200 GPU, NEST demonstrated significantly higher throughput (90.46 k tokens/s) and lower peak memory (34.9 GiB) compared to BERT-style and approximate attention models (e.g., Longformer, Nyströmformer), despite having a larger parameter count (~120M).
Downstream Performance:
- MIMIC-IV: NEST achieved state-of-the-art results across five tasks (inpatient mortality, 30-day readmission, length of stay, ICD chapter/category classification), outperforming strong baselines like Hi-BEHRT, CORE-BEHRT, and GT-BEHRT. For example, NEST achieved 95.57% AUROC for mortality prediction vs. 93.80% for GT-BEHRT.
- Instacart: While NEST alone was competitive with BERT4Rec, the combination of NEST's pretrained representations with a repeat-aware model (ReCANet) outperformed all baselines in NDCG@10, indicating superior capture of user-item dynamics.
- Proprietary EHR: NEST outperformed LightGBM and CORE-BEHRT on pediatric risk prediction tasks (rAOM, ASD, 30-day readmission) under both linear probing and full fine-tuning.
Ablation Studies:
- Removing the CSE or replacing SWE with mean pooling significantly degraded performance, confirming the necessity of both encoders.
- Shuffling the temporal order of encounters at fine-tuning time degraded performance to the level of removing CSE, proving CSE encodes critical temporal structure.
- The MSM objective was shown to improve set-level representation learning, with the benefits residing in the encoder representations rather than just the task head.

5. Significance and Claims

The paper claims that the correct inductive bias for event-stream foundation models is to reflect the data's nested set structure rather than flattening it. By preserving multiset boundaries and inducing structured sparsity, NEST achieves a dual benefit:

Efficiency: It reduces computational costs associated with dense attention.
Quality: It produces high-quality, directly usable set-level representations without post-hoc pooling.

The authors position NEST as a step toward foundation models that respect the structural realities of real-world event streams (like healthcare and commerce) rather than forcing them into a flat sequence paradigm. They explicitly note limitations, such as the focus on discriminative representation learning rather than generative trajectory modeling, and the need for further cross-domain validation beyond healthcare and retail.