🤖 machine learning

NEST: Nested Event Stream Transformer for Sequences of Multisets

本文介绍了 NEST，一种嵌套事件流 Transformer，它通过保留事件序列（多重集序列）的层次结构，克服了现有扁平化模型的计算低效和表示局限性，并利用一种新颖的掩码集合建模范式来提升预训练效率和下游性能。

原作者： Minghui Sun, Haoyu Gong, Xingyu You, Jillian Hurst, Benjamin Goldstein, Matthew Engelhard

发布于 2026-05-15

📖 1 分钟阅读☕ 轻松阅读

原作者： Minghui Sun, Haoyu Gong, Xingyu You, Jillian Hurst, Benjamin Goldstein, Matthew Engelhard

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在试图理解一位患者的病史，或者某位顾客的购物习惯。在现实世界中，这些事件并非像串在绳子上的珠子那样，一个接一个地完美线性发生。相反，它们是以群组的形式发生的。

在医院里： 医生接诊一位患者（一次“就诊”）。在这次就诊期间，患者可能同时接受血液检查、获得处方并得到诊断。这三件事的确切顺序可能并不重要，或者记录可能有些混乱。但这组事件属于那次特定的就诊。
在杂货店里： 顾客去购物。他们将牛奶、面包和鸡蛋放入购物车。那个“购物篮”就是一个群组。他们拿起这些物品的顺序并不重要；重要的是整个购物篮。

大多数当前的 AI 模型（称为 Transformer）试图将这些群组压平为一条单一、漫长的线性事件序列。它们将牛奶、面包和鸡蛋视为依次发生的事件，忽略了它们属于同一次购物之旅的事实。这就像试图通过查看每一帧画面的列表来理解一部电影，却不知道哪些场景属于同一幕。这不仅计算成本高昂，而且往往忽略了大局。

引入 NEST（嵌套事件流 Transformer）。

本文的作者构建了一种名为 NEST 的新 AI 模型，它尊重这些自然的群组。以下是其工作原理，使用简单的类比：

1. 两步舞（架构）

NEST 没有将数据压平，而是保持“群组”（如医院就诊或购物篮）的完整性。它在模型的每一层内部使用了一个巧妙的两步过程：

步骤 A：群组内部讨论（集合编码器）： 首先，模型查看一个群组（例如一次医院就诊），并让该群组内的所有事件相互“交流”。它弄清楚在该次就诊内部发生了什么。这就像一次团队会议，每个人讨论各自的具体任务。
步骤 B：全球圆桌会议（跨集合编码器）： 接下来，模型查看每个群组的“队长”（称为 [CLS] 的特殊标记）。这些队长会面，向时间线的其余部分分享各自群组中发生的事情。这有助于模型理解就诊 A 与就诊 B 之间的关系。

魔法技巧： 大多数模型先对所有群组执行步骤 A，然后对所有群组执行步骤 B。而 NEST 则交错执行这两步。它执行一点步骤 A，然后执行一点步骤 B，然后再回到步骤 A。

为什么这很重要： 想象一场接力赛。如果你跑完第一整段才传递接力棒，你可能会在途中丢失一些信息。NEST 则不断地来回传递接力棒。这确保了当模型试图理解大局时，特定事件的任何细节都不会丢失。论文从数学上证明，这种“旁路”比旧方法保留了更多的信息。

2. “队长的报告”（掩码集合建模）

在旧方法中，AI 处理完所有数据后，必须猜测如何将整个就诊总结为单一分数。这就像要求学生在阅读完整个章节后总结该章，却只给他们一个模糊的提示。

NEST 引入了一种新的训练游戏，称为掩码集合建模（MSM）。

工作原理： 模型看到一个事件群组（如购物篮），但它必须根据“队长”标记来猜测该群组的内容，而实际物品是隐藏的。
结果： 这迫使“队长”标记成为整个群组的完美总结。因此，模型不再需要在稍后进行混乱的、靠猜测的总结，而是已经准备好高质量的总结，可用于任何未来任务。

3. 为什么它更好（结果）

作者在真实世界数据上测试了 NEST：

医疗记录（EHR）： 他们使用了来自医院（MIMIC-IV）和一家私立儿科数据库的数据。
购物： 他们使用了杂货数据（Instacart）。

发现：

更快更精简： 由于 NEST 尊重群组，它不必计算每一个单独事件之间的连接。它只计算群组内部以及群组队长之间的连接。这使得它比旧模型更快，且使用的计算机内存更少，尽管它拥有更多的“脑力”（参数）。
更智能的预测： NEST 在预测以下事项方面表现更好：
- 患者是否会在住院期间死亡？
- 患者是否会在 30 天内再次入院？
- 顾客接下来会购买什么物品？
无需“后处理”： 由于模型在训练期间学会了总结群组，因此它不需要在训练后使用笨拙的启发式技巧来理解数据。总结已经准备就绪。

总结

将 NEST 视为一个终于理解上下文至关重要的模型。它知道在同一次医院就诊中进行的血液检查和开具的处方是一个团队，而这个团队是患者生活更大故事的一部分。通过保持这些群组在一起并让它们高效沟通，NEST 比那些试图将所有内容强行塞入单一长线的模型学习得更快、使用能量更少，并能做出更好的预测。

技术摘要：NEST——面向多重集序列的嵌套事件流 Transformer

1. 问题陈述

事件流数据（如电子健康记录 EHR 或用户购买历史）通常表现出分层结构，即多个事件在特定的时间窗口内（例如一次临床就诊或一个购物篮）共同发生。这种结构被形式化定义为多重集序列（SeqSet）。在这些序列中，集合（多重集）内部事件的顺序可能是未知的、不可靠的或无关的，而集合之间的时间顺序则至关重要。

现有的事件流基础模型（FMs）通常将这种分层结构展平为一维序列。这种方法引入了两个主要局限性：

计算效率低下：展平导致注意力机制变得稠密，其复杂度相对于令牌总数呈二次方增长，忽略了数据结构固有的稀疏性。
表示质量：标准模型仅学习令牌级别的表示。集合级别的表示（例如就诊后患者的状态）仅通过启发式的后训练池化（例如平均池化或使用未进行显式监督的 [CLS] 令牌）来恢复，先前的工作已表明这并非最优。此外，展平迫使模型学习集合内部可能根本不存在的虚假关系。

2. 方法论

2.1 架构：嵌套事件流 Transformer（NEST）

NEST 是专为 SeqSet 数据设计的分层 Transformer。其核心创新在于在每个层内交错组合两种不同的编码器类型，而不是将它们堆叠在独立的阶段中。

集合级编码器（SWE）：独立地在每个多重集内运行。它处理该集合特有的令牌（事件）和一个可学习的 [CLS] 令牌。SWE 对集合内令牌强制执行置换不变性（除非存在细粒度的时间戳，此时可应用旋转位置嵌入 RoPE）。它将集合级特征聚合到 [CLS] 令牌中。
跨集合编码器（CSE）：仅在来自所有集合的 [CLS] 令牌序列上运行。它使用 RoPE 编码集合间的时间间隔，从而模拟就诊之间的时间动态。非 [CLS] 令牌的隐藏状态绕过 CSE。
交错与旁路：SWE 和 CSE 逐层交错（总共 $L$ 层）。关键在于，非 [CLS] 令牌状态通过旁路路径绕过 CSE 块进行传播。这确保了当从集合级处理过渡到序列级处理时，令牌级别的信息不会被不可恢复地丢弃。

信息论依据：本文通过命题 1论证，顺序设计（在所有 CSE 之前堆叠所有 SWE）违反了数据处理不等式，因为在 SWE 到 CSE 的边界处丢弃了令牌级别的信息。交错设计保留了这些信息，使模型能够在每一层结合跨集合上下文重新审视原始令牌状态。

2.2 预训练目标：掩码集合建模（MSM）

为了补充标准的掩码语言建模（MLM）目标，作者引入了掩码集合建模（MSM）。

机制：在预训练期间，以 40% 的概率掩码整个多重集（就诊）。在掩码的集合内，所有非特殊令牌均被替换为 [MASK]。
目标：模型必须利用 [CLS] 令牌的表示来预测掩码集合的令牌分布。
意义：这迫使 [CLS] 令牌显式地聚合集合级信息并学习即用的集合表示，从而消除了对启发式事后池化的需求。它补充了专注于细粒度令牌级上下文的 MLM。

3. 主要贡献

SeqSet 的架构原语：NEST 是一个分层基础模型，它尊重 SeqSet 数据的对立对称性：集合内可交换性和集合间时间顺序。
通过旁路实现信息保留：本文描述了交错编码器如何保留顺序设计所丢弃的令牌级别信息，这一点已通过拓扑隔离模拟研究得到证实。
集合级预训练：提出了 MSM，这是一种置换不变的目标，显式训练 [CLS] 令牌预测多重集令牌分布，这是先前 EHR 基础模型中缺失的功能。
效率与性能：NEST 实现了结构化稀疏性，将注意力复杂度从 $O(N^2)$ 降低到 $O(Nn + m^2)$ （其中 $N$ 是令牌总数， $n$ 是集合大小， $m$ 是集合数量），与稠密注意力基线相比，实现了更高的吞吐量和更低的内存使用。

4. 实验结果

作者在三个真实世界数据集上评估了 NEST：Instacart（下一个购物篮推荐）、MIMIC-IV（EHR 临床预测）以及一个专有机构 EHR（儿科风险预测）。

计算效率：在单张 NVIDIA H200 GPU 上，尽管参数量较大（约 1.2 亿），NEST 仍表现出显著更高的吞吐量（90.46k 令牌/秒）和更低的峰值内存（34.9 GiB），优于 BERT 风格和近似注意力模型（如 Longformer、Nyströmformer）。
下游性能：
- MIMIC-IV：NEST 在五项任务（住院死亡率、30 天再入院、住院时长、ICD 章节/类别分类）中取得了最先进（SOTA）的结果，优于 Hi-BEHRT、CORE-BEHRT 和 GT-BEHRT 等强基线。例如，在死亡率预测中，NEST 的 AUROC 达到 95.57%，而 GT-BEHRT 为 93.80%。
- Instacart：虽然 NEST 单独使用时与 BERT4Rec 具有竞争力，但将 NEST 的预训练表示与感知重复的模型（ReCANet）结合后，在 NDCG@10 指标上超越了所有基线，表明其能更好地捕捉用户 - 物品动态。
- 专有 EHR：在线性探测和全量微调两种设置下，NEST 在儿科风险预测任务（rAOM、ASD、30 天再入院）上均优于 LightGBM 和 CORE-BEHRT。
消融研究：
- 移除 CSE 或将 SWE 替换为平均池化会显著降低性能，证实了两种编码器的必要性。
- 在微调时打乱就诊的时间顺序会使性能下降至移除 CSE 的水平，证明 CSE 编码了关键的时间结构。
- MSM 目标被证明能改善集合级表示学习，其收益存在于编码器表示中，而不仅仅在于任务头。

5. 意义与主张

本文主张，事件流基础模型的正确归纳偏置应是反映数据的嵌套集合结构，而不是将其展平。通过保留多重集边界并诱导结构化稀疏性，NEST 实现了双重效益：

效率：降低了与稠密注意力相关的计算成本。
质量：无需事后池化即可生成高质量、可直接使用的集合级表示。

作者将 NEST 定位为迈向尊重现实世界事件流（如医疗和商业）结构现实的基础模型的一步，而不是强迫它们适应扁平序列范式。他们明确指出了局限性，例如侧重于判别式表示学习而非生成式轨迹建模，以及需要超越医疗和零售领域的进一步跨领域验证。