[論文レビュー] BEDS-Bench: Behavior of EHR-models under Distributional Shift--A Benchmark

Anand Avati, Martin Seneviratne+4 more|arXiv (Cornell University)|Jul 17, 2021

Machine Learning in Healthcare被引用数 5

ひとこと要約

BEDS-Bench は、分布シフト下における EHR ベースの機械学習モデルの挙動を評価するベンチマークであり、2 つのオープン EHR データセットを用いて多様な分布外（OOD）設定を構築する。研究では、評価されたすべてのモデルが OOD 条件下で一般化性能およびキャリブレーション性能に劣っていることが明らかになった。これは、実臨床現場での導入において、より高い耐性を備えたモデル開発が急務であることを示唆している。

ABSTRACT

Machine learning has recently demonstrated impressive progress in predictive accuracy across a wide array of tasks. Most ML approaches focus on generalization performance on unseen data that are similar to the training data (In-Distribution, or IND). However, real world applications and deployments of ML rarely enjoy the comfort of encountering examples that are always IND. In such situations, most ML models commonly display erratic behavior on Out-of-Distribution (OOD) examples, such as assigning high confidence to wrong predictions, or vice-versa. Implications of such unusual model behavior are further exacerbated in the healthcare setting, where patient health can potentially be put at risk. It is crucial to study the behavior and robustness properties of models under distributional shift, understand common failure modes, and take mitigation steps before the model is deployed. Having a benchmark that shines light upon these aspects of a model is a first and necessary step in addressing the issue. Recent work and interest in increasing model robustness in OOD settings have focused more on image modality, while the Electronic Health Record (EHR) modality is still largely under-explored. We aim to bridge this gap by releasing BEDS-Bench, a benchmark for quantifying the behavior of ML models over EHR data under OOD settings. We use two open access, de-identified EHR datasets to construct several OOD data settings to run tests on, and measure relevant metrics that characterize crucial aspects of a model's OOD behavior. We evaluate several learning algorithms under BEDS-Bench and find that all of them show poor generalization performance under distributional shift in general. Our results highlight the need and the potential to improve robustness of EHR models under distributional shift, and BEDS-Bench provides one way to measure progress towards that goal.

研究の動機と目的

• 分布シフト下における EHR モデルの性能を評価するためのベンチマークが不足しているという問題に取り組むこと。特に、実臨床現場での導入状況を想定した文脈での評価を目的とする。
• トレーニングデータの分布とは顕著に異なるデータでテストされた場合に、EHR モデルがどのように振る舞うかを定量化すること。
• OOD 条件下におけるモデルのキャリブレーションと信頼性推定の一般的な失敗モードを同定すること。
• 多様な OOD 環境において、EHR モデルの耐性を測定・比較するための標準化された評価フレームワークを提供すること。
• 臨床 AI アプリケーションにおけるモデルの信頼性と安全性を高める研究を促進すること。

提案手法

• ベンチマークは、患者の人口統計、臨床的状態、および時間的ウィンドウを操作することで、2 つの匿名化済みでオープンアクセス可能な EHR データセットを用いて、複数の OOD データ設定を構築する。
• 時間経過に伴うイベントの定義を一貫させるために、疾患発症予測や死亡リスク推定といった複数の臨床予測タスクを評価する。
• OOD 検出は、OOD-AUC、ECE（期待キャリブレーション誤差）、および信頼性キャリブレーションといった指標を用いて評価され、単一予測モデルの場合は予測エントロピーを用いて信頼性を測定する。
• ベンチマークには、インディストリビューション（IND）および OOD データの分割が含まれており、OOD データは異なる患者集団、時間帯、または医療機関サイトから抽出される。
• モデルは、OOD 例に対して低い信頼度（高い不確実性）を割り当てる能力、特にキャリブレーションと識別性能の観点から評価される。
• 再現可能な評価を支援するフレームワークであり、コードは https://github.com/Google-Health/records-research/tree/master/beds-bench で公開されている。

実験結果

リサーチクエスチョン

• RQ1 EHR ベースの機械学習モデルは、トレーニングデータとは分布が異なるデータでテストされた場合、どのように性能を示すか？
• RQ2 実臨床 EHR 環境における OOD 例に対して、モデルはどの程度適切なキャリブレーションと不確実性推定を維持できるか？
• RQ3 どのモデルアーキテクチャーやトレーニング戦略が、EHR データにおける分布シフトに対してより高い耐性を示すか？
• RQ4 EHR におけるクラス不均衡および OOD 検出シナリオ下で、一般的な信頼性推定手法（例：エントロピー）はどのように機能するか？
• RQ5 OOD-AUC や ECE といった既存の OOD 検出指標は、臨床テーブルデータにおける分布シフト下でのモデル挙動を信頼性を持って捉えられるか？

主な発見

• EHR モデルの評価において、OOD 耐性を意図して設計されたモデルを含め、すべてのモデルが複数の OOD 環境下で一般化性能が著しく劣っていることが判明した。
• どのモデルアーキテクチャも、すべての OOD シナリオで一貫して優れた性能を示すことはなく、EHR モデリングにおける耐性の明確な解決策は存在しないことが示された。
• モデルは、誤った OOD 予測に対しても頻繁に高い信頼度を割り当てるため、不確実性のキャリブレーションが著しく劣っていることが判明した。
• OOD-AUC 指標だけでは不十分であり、クラス不均衡下では誤解を招く可能性がある。誤分類された OOD 例を「陽性」として高信頼度で予測するモデルが、誤って評価されてしまうことがある。
• 期待キャリブレーション誤差（ECE）と信頼性キャリブレーションは、モデルの信頼性を示す強力な指標のままであるが、テストされたすべてのモデルが OOD 条件下で顕著なキャリブレーション誤差を示している。
• 本研究では、現在の EHR モデルが OOD として信頼性を持って検出できないことが確認された。これは、分布シフトが生じる実臨床現場での導入において、臨床的リスクを引き起こす可能性がある。