[论文解读] 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 基础机器学习模型在分布偏移条件下行为的基准，使用两个公开的 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 条件下的常见失败模式。
• 提供一个标准化的评估框架，用于衡量和比较 EHR 模型在多样化 OOD 场景下的鲁棒性。
• 推动研究以提升临床人工智能应用中模型的可靠性与安全性。

提出的方法

• 该基准使用两个去标识化、公开获取的 EHR 数据集，通过调整患者人口统计学特征、临床状况和时间窗口，构建多个 OOD 数据场景。
• 在多个临床预测任务上评估模型，包括疾病发作预测和死亡风险估计，使用一致的时间至事件定义。
• 使用 OOD-AUC、ECE（期望校准误差）和置信度校准等指标评估 OOD 检测性能，其中置信度通过单预测模型的预测熵进行衡量。
• 该基准包含分布内（IND）和 OOD 数据划分，其中 OOD 数据来自不同的患者群体、时间段或医疗系统机构。
• 评估模型对 OOD 样例分配较低置信度（更高不确定性）的能力，重点关注校准和区分性能。
• 该框架支持可复现的评估，并在 https://github.com/Google-Health/records-research/tree/master/beds-bench 发布代码。

实验结果

研究问题

• RQ1 当在与训练数据分布不同的数据上测试时，基于 EHR 的机器学习模型表现如何？
• RQ2 在真实世界的 EHR 场景中，模型在 OOD 样例上在多大程度上保持了适当的校准和不确定性估计？
• RQ3 哪些模型架构或训练策略在 EHR 数据的分布偏移下表现出更好的鲁棒性？
• RQ4 在类别不平衡和 OOD 检测场景下，常见的置信度估计方法（如熵）表现如何？
• RQ5 现有的 OOD 检测指标（如 OOD-AUC 和 ECE）能否可靠地捕捉临床表格数据在分布偏移下的模型行为？

主要发现

• 所有评估的 EHR 模型，包括专为 OOD 鲁棒性设计的模型，在多个 OOD 场景下均表现出较差的泛化性能。
• 没有单一模型架构在所有 OOD 场景中始终优于其他模型，表明在 EHR 建模中尚无明确的鲁棒性解决方案。
• 模型经常对 OOD 预测赋予高置信度，即使这些预测是错误的，表明其不确定性校准性能差。
• 仅依赖 OOD-AUC 指标在类别不平衡下可能不足且具有误导性，因为若 OOD 样例被错误分类为正例，该指标可能奖励那些对 OOD 样例赋予高置信度的模型。
• 期望校准误差（ECE）和置信度校准仍是模型可靠性的有力指标，但所有测试模型在 OOD 条件下均表现出显著的校准偏差。
• 本研究证实，当前的 EHR 模型无法可靠地被识别为 OOD，这在存在分布偏移的真实临床环境中部署时会带来临床风险。