Tokenization Tradeoffs in Structured EHR Foundation Models

这是一份关于论文《TOKENIZATION TRADEOFFS IN STRUCTURED EHR FOUNDATION MODELS》（结构化电子健康记录基础模型中的分词权衡）的详细技术总结。

1. 研究背景与问题 (Problem)

背景：
基于 Transformer 的基础模型（Foundation Models）正在改变医疗人工智能，特别是针对结构化电子健康记录（EHR）的模型。这些模型通过在纵向时间序列的临床事件（如诊断、手术、药物、实验室结果）上进行预训练，学习可适应的患者表示，并在下游临床任务中展现出优越的性能。

核心问题：
在训练 EHR 基础模型时，分词（Tokenization）——即如何将患者的临床时间线转换为离散的模型输入序列——是一个关键但被低估的步骤。分词策略决定了：

1. 哪些信息被保留（如代码与属性的关联、时间间隔）。
2. 信息编码的效率（序列长度、上下文窗口利用率）。
3. 模型需要从数据中学习哪些关系，哪些关系可以在分词阶段预计算。

目前，现有的 EHR 基础模型在分词设计上存在多种选择（如联合编码 vs. 因子化编码、时间编码方式等），但缺乏受控的对比研究，导致最佳实践尚未确立。特别是在数据受限的医疗领域，如何在“在分词时硬编码结构”与“让模型从数据中学习结构”之间进行权衡，尚不明确。

2. 方法论 (Methodology)

本研究采用因子实验设计（Factorial Design），系统地评估了三个分词设计轴对模型性能和计算效率的影响。

实验设置：

• 数据集：
  • 主要开发： 多伦多病童医院（SickKids）的儿科 EHR 数据（MEDS 格式，OMOP CDM 标准化），包含约 200 万患者，1.69 亿个临床事件。
  • 外部验证： 贝斯以色列女执事医疗中心（BIDMC）的 MIMIC-IV 成人重症监护数据，用于评估跨机构泛化能力。
• 模型架构： 统一的 Decoder-only Transformer 架构（28 层，隐藏层维度 768，1.156 亿参数），在所有实验中保持骨干网络不变，仅改变分词策略导致的词汇表大小和序列长度。
• 训练任务： 自监督的下一个词预测（Next-token Prediction）。
• 评估任务： 在 SickKids 上评估了 74 个 临床预测任务（涵盖输血、手术、影像、实验室结果、用药、临床结局），在 MIMIC 上评估了 13 个任务。

三个分词设计轴（2x2x2 因子设计）：

1. 事件编码 (Event Encoding)：
   • 联合编码 (Joint)： 将临床代码及其属性（如数值、分类值）合并为单个 Token（例如：葡萄糖_高 作为一个 Token）。
   • 因子化编码 (Factorized)： 将事件分解为多个 Token（例如：葡萄糖 + 高 两个 Token），模型需学习它们之间的关联。
2. 时间编码 (Time Encoding)：
   • 位置编码 (Time-Positions)： 利用患者年龄（天数）作为位置索引，结合旋转位置编码（RoPE），不消耗额外的序列位置。
   • 时间 Token (Time-Tokens)： 在连续事件之间插入离散的时间间隔 Token（如 2-6 小时），消耗额外的序列位置。
3. 工作流标注 (Workflow Annotations)：
   • 无工作流： 仅记录最终事件（如实验室结果时间）。
   • 有工作流： 将多步骤临床操作（如：开单 -> 采样 -> 出结果）表示为具有不同时间戳的多个事件。

消融实验 (Ablation Studies)：

• 事件编码消融： 分离“序列长度效率”与“局部绑定效率（Local Binding Efficiency）”的影响。
• 时间编码消融： 比较仅顺序、位置编码、时间 Token 和标量特征补充的效果。

3. 关键贡献 (Key Contributions)

1. 系统性量化分词权衡： 首次通过大规模因子实验，量化了事件编码、时间编码和工作流标注对 EHR 基础模型下游性能和预训练计算成本（FLOPs）的独立影响。
2. 揭示“局部绑定效率”机制： 通过消融实验证明，联合编码的优势主要源于局部绑定效率（Local Binding Efficiency）。即，将代码 - 属性对预计算为单个 Token，避免了模型在数据受限的情况下学习复杂的上下文依赖（例如：同样的“高”数值，对于葡萄糖意味着高血糖，对于肌酐意味着肾损伤）。
3. 跨机构泛化性验证： 证明了联合编码的优势在词汇表严重不匹配（Out-of-Vocabulary 率近 70%）的跨机构场景（儿科到成人 ICU）中依然有效，而时间编码和工作流的效果则具有机构特异性。
4. 效率与性能的双重提升： 发现更优的分词策略（联合编码 + 位置时间编码）不仅提高了预测性能，还显著降低了预训练所需的计算量。

4. 主要结果 (Results)

1. 事件编码 (Event Encoding)：

• 性能： 联合编码在 73/74 个任务中优于因子化编码（平均 AUROC 提升约 0.008）。
• 效率： 联合编码所需的预训练 FLOPs 比因子化编码少 39.5%。
• 原因： 消融实验显示，当添加属性值时，联合编码性能显著提升，而因子化编码没有。这是因为因子化编码迫使模型学习“局部绑定”（Local Binding），而在数据有限的 EHR 场景中，模型难以可靠地学习这种上下文相关的组合关系。

2. 时间编码 (Time Encoding)：

• 性能： 基于年龄的位置编码（Time-Positions）优于离散时间 Token（Time-Tokens），在 71/74 个任务中表现更好。
• 效率： Time-Positions 比 Time-Tokens 少消耗 9.6% 的 FLOPs。
• 原因： 离散时间 Token 消耗了宝贵的上下文窗口位置，却未提供成比例的信息增益。位置编码利用 RoPE 隐式编码时间，效率更高。
• 泛化性： Time-Positions 的优势在跨机构（儿科到成人）时消失，因为基于特定人群年龄分布校准的位置编码在年龄分布不同的群体中表现不佳。

3. 工作流标注 (Workflow Annotations)：

• 性能： 包含工作流阶段（如开单、采样、结果）在本地（SickKids）提升了性能（64/74 任务）。
• 成本： 增加了 35.6% 的预训练计算成本（序列变长）。
• 泛化性： 在 MIMIC 外部验证中未显示显著优势。因为不同机构的工作流时间戳模式差异巨大，这种特定于机构的信息难以泛化。

4. 外部评估 (MIMIC)：

• 尽管 SickKids 模型在 MIMIC 上的词汇表不匹配率高达 69.8%，但联合编码的优势依然显著转移（效应量与本地评估一致）。
• 这证明了预计算的结构性关联（代码 - 属性绑定）比学习到的时间或上下文模式具有更强的鲁棒性和泛化性。

5. 意义与结论 (Significance)

• 设计原则： 对于结构化 EHR 基础模型，在分词阶段预计算有意义的临床结构（如代码与属性的绑定） 是优于让模型从数据中学习这些结构的。这解决了数据受限环境下的“局部绑定问题”。
• 计算效率： 更紧凑、语义更丰富的 Token 设计（联合编码 + 位置时间编码）不仅能提高预测准确性，还能大幅降低预训练成本，这对于资源受限的医疗 AI 开发至关重要。
• 跨机构部署： 在跨机构部署时，应优先考虑词汇表对齐和结构化的事件表示（如联合编码），而不应过度依赖特定机构的工作流细节或特定人群的时间统计特征。
• 未来方向： 建议未来的 EHR 基础模型研究将“联合事件编码”和“基于位置的时间编码”作为默认设置，并明确报告分词策略。

总结： 该论文确立了分词策略是优化 EHR 基础模型性能和效率的一个可操作的关键杠杆。通过预计算临床语义结构，可以在不增加模型复杂度的情况下，显著提升模型在数据稀缺和跨机构场景下的表现。