[論文レビュー] 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 は、電子歴史記録（EHR）向けに体系的に最適化され、厳密に評価された BERT ベースのモデルを導入し、薬剤情報や時刻情報の包括的表現と洗練されたトレーニングプロトコルにより、性能を向上させた。25 の臨床的タスクで平均 0.801 の AUROC を達成し、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 に対する段階的最適化を実施。
• 性能指標の安定化と収束不安定性の低減を目的に、5 折り交差検証と複数回の事前学習ランを用いた。
• 最適化を導くために最初に 3 タスク（死亡、疼痛治療、感染症）でモデル性能を評価し、その後 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 異なるマスキング比やプーリング戦略が、タスクによって一貫性のない向上効果を示す理由は何か？また、これを緩和する方法は？

主な発見

• データ表現の改善により、3 つのコアタスクで平均 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