[论文解读] CORE-BEHRT: A Carefully Optimized and Rigorously Evaluated BEHRT

Mikkel Odgaard, Kiril Vadimovic Klein+4 more|arXiv (Cornell University)|Apr 23, 2024

Advanced Neural Network Applications被引用 4

一句话总结

CORE-BEHRT 引入了一种系统性优化且经过严格评估的基于 BERT 的电子健康记录（EHR）模型，通过增强的数据表示（包括药物和时间戳）以及优化的训练协议，提升了性能。其在 25 项临床任务上的平均 AUROC 达到 0.801，其中 17 项任务取得显著提升，24 项任务性能提高，展现出强大的泛化能力与更高的临床部署可信度。

ABSTRACT

The widespread adoption of Electronic Health Records (EHR) has significantly increased the amount of available healthcare data. This has allowed models inspired by Natural Language Processing (NLP) and Computer Vision, which scale exceptionally well, to be used in EHR research. Particularly, BERT-based models have surged in popularity following the release of BEHRT and Med-BERT. Subsequent models have largely built on these foundations despite the fundamental design choices of these pioneering models remaining underexplored. Through incremental optimization, we study BERT-based EHR modeling and isolate the sources of improvement for key design choices, giving us insights into the effect of data representation, individual technical components, and training procedure. Evaluating this across a set of generic tasks (death, pain treatment, and general infection), we showed that improving data representation can increase the average downstream performance from 0.785 to 0.797 AUROC ($p<10^{-7}$), primarily when including medication and timestamps. Improving the architecture and training protocol on top of this increased average downstream performance to 0.801 AUROC ($p<10^{-7}$). We then demonstrated the consistency of our optimization through a rigorous evaluation across 25 diverse clinical prediction tasks. We observed significant performance increases in 17 out of 25 tasks and improvements in 24 tasks, highlighting the generalizability of our results. Our findings provide a strong foundation for future work and aim to increase the trustworthiness of BERT-based EHR models.

研究动机与目标

• 为解决基于 BERT 的 EHR 模型在评估与优化方面缺乏严谨性的问题，该问题阻碍了其在临床中的应用。
• 分离并量化关键设计选择（数据表示、模型架构与训练协议）对下游性能的影响。
• 通过逐步优化 BEHRT 并结合跨任务验证，系统性地提升模型的可信度。
• 在多样化临床预测任务中展示一致的性能提升，确保模型在狭窄基准之外也具备泛化能力。
• 为未来 EHR 建模研究提供可复现且经过优化的基础，并提供可操作的洞见。

提出的方法

• 通过调整数据表示（如包含药物、时间戳、ICD-10 编码）和训练配置，对 BEHRT 进行渐进式优化。
• 采用五折交叉验证，并进行多次预训练运行，以稳定性能指标并减少收敛不稳定性。
• 在三个初始任务（死亡、疼痛治疗、感染）上评估模型性能以指导优化，随后在 25 项多样化临床预测任务上进行验证。
• 应用不同的掩码策略与池化方法，评估其对下游 AUROC 的影响，重点关注与 EHR 任务目标的一致性。
• 采用严格的评估流程，结合显著性检验（p < 10⁻⁷），以验证性能提升的可靠性。
• 聚焦真实世界 EHR 数据，确保临床相关性，通过年龄窗口与模态约束模拟实际部署条件。

Figure 1: This figure depicts the selection process of the best configuration. Non-selected configurations are depicted in a lighter shade, indicating they were considered but not chosen based on their impact on the average Area Under the Receiver Operating Characteristic (AUROC), across three predi

实验结果

研究问题

• RQ1 提升数据表示——特别是包含药物和时间戳——如何影响基于 EHR 的 BERT 模型的下游性能？
• RQ2 在数据表示之外，架构与训练协议的优化在多大程度上提升了模型性能？
• RQ3 优化收益在广泛且多样的临床预测任务中具有多大程度的泛化能力？
• RQ4 模型规模与微调数据量对性能饱和与可扩展性有何影响？
• RQ5 为何不同的掩码比例与池化策略在不同任务中产生不一致的改进？如何缓解这一问题？

主要发现

• 优化数据表示使三项核心任务的平均 AUROC 从 0.785 提升至 0.797（p < 10⁻⁷），主要归因于药物与时间戳特征的引入。
• 进一步优化模型架构与训练协议后，25 项多样化临床任务的平均 AUROC 提升至 0.801，其中 17 项任务取得显著提升，24 项任务性能提高。
• 在中风预测任务中达到最先进性能（AUROC 70.7%），优于以往工作（66.9%），且在更受限制的数据条件下实现。
• 性能提升稳定且具有统计显著性，表明优化过程未偏向特定任务。
• 药物与时间戳的引入对性能具有主导性影响，表明其在捕捉临床背景中的关键作用。
• 尽管数据量增加，部分任务的性能早期即趋于饱和，表明仅靠数据量不足以提升性能——模型设计与表示方式更具决定性影响。

Figure 2: Overview of the iterative optimization procedure, highlighting adapted settings in bold. We augmented the input data with medication codes, patient sex, event timestamps, full International Classification of Diseases (ICD-10), and Anatomical Therapeutic Chemical (ATC) codes. Improvements t