BEDS-Bench: Behavior of EHR-models under Distributional Shift - A Benchmark

Anand Avati, Martin Seneviratne, Yuan Xue, Zhen Xu, Balaji Lakshminarayanan, Andrew M. Dai

Keywords: Electronic Health Records, Distributional Shift, Benchmark

TL;DR: A new benchmark to evaluate behavior of ML models on EHR data under distributional shift

Abstract:

Machine learning (ML) 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 (a.k.a. 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, both in terms of methods as well as benchmarks, 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. The benchmark exercises several clinical prediction tasks, OOD data settings, and measures 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 provides one way to measure progress towards that goal. Code to reproduce the results in this paper and evaluate new algorithms against \bedS is made available at \url{https://github.com/Google-Health/records-research/tree/master/beds-bench}.

A full-length version of this paper is available at \url{https://arxiv.org/abs/2107.08189}.

Community Implementations: [ 1 code implementation](https://www.catalyzex.com/paper/beds-bench-behavior-of-ehr-models-under/code)