From development to deployment: dataset shift, causality, and shift-stable models in health AI

Abstract

The deployment of machine learning (ML) and statistical models is beginning to transform the practice of healthcare, with models now able to help clinicians diagnose conditions like pneumonia and skin cancer, and to predict which hospital patients are at risk of adverse events such as septic shock. A major concern, however, is that model performance is heavily tied to details particular to the dataset the model was developed on—even slight deviations from the training conditions can result in wildly different performance. For example, when researchers trained a model to diagnose pneumonia from chest X-rays using data from one health system, but evaluated on data from an external health system, they found the model performed significantly worse than it did internally (Zech and others, 2018). The model failed to generalize (i.e., predict accurately) due to the shifts between the training conditions (health system one) and the deployment/testing conditions (health system two). These shifts are very common when moving a model from the training phase to deployment and can take a variety of forms, including changes in patient demographics, disease prevalence, measurement timing, equipment, treatment patterns, and more. Beyond contributing to poor performance, failing to account for shifts can also lead to dangerous decisions in practice: the system can fail to diagnose severely ill patients or recommend harmful treatments. This problem of shifting conditions which prevent generalization is referred to as dataset shift (Quiñonero-Candela and others, 2009), and in this article, we explain what it is, why it occurs, give an overview of the types of existing solutions, and discuss open challenges that remain. Generalization is crucial for successfully deploying models since we want model predictions to be accurate when applied to new situations that were not in the training dataset. In order to ensure that a model will generalize, a core requirement is to check that the model or learning procedure satisfies the stability property: is model performance robust when the data are perturbed in various ways? In particular, when addressing dataset shift, we want models to be stable to perturbations in (i.e., shifts in) how the data were generated—such as changes to patient demographics or clinician treatment patterns—as opposed to leave-one-out sampling perturbations (Yu, 2013; Giordano and others, 2019; Yu and Kumbier, 2019) or adversarial perturbations of model inputs (Madry and others, 2018; Hendrycks and Dietterich, 2019). Dataset shift is of particular importance because it frequently occurs when deploying models, is difficult to test stability against (as we will discuss later), and is highly relevant to the ongoing discussion about the challenges of regulating ML-driven medical devices (U.S. Food and Drug Administration, 2019) (see accompanying Stern and Price (2020) for more on this). Consider the mortality risk prediction model trained by Caruana and others (2015) on a dataset of hospitalized pneumonia patients, using information such as lab measurements, vital signs, and comorbidities. While the model had high predictive accuracy on one dataset, it was unstable to shifts in the choices driving which patients get admitted to the ICU versus the floor. As a result, when they evaluated it for triaging pneumonia patients upon ED presentation, they found that their model incorrectly predicted lower risk for patients with pneumonia and asthma versus those with only pneumonia. This shift in ICU admission policy, while subtle, had big implications: had the model been deployed for triage, it would have greatly endangered asthmatic pneumonia patients by suggesting they should be sent home. In machine learning, a policy is defined as a distribution over the possible actions that can be taken in any given scenario. Even subtle shifts in when or whether or not a lab was ordered can impact predictions—these forms of shifts are called policy shifts (Schulam and Saria, 2017). More generally, dataset shifts can come in various forms (from shifts in patient demographics to policy shifts in measurement frequency), which makes dataset shift challenging to address. Given that dataset shift is highly varied and can lead to dangerous failures if left unaccounted for, what should we do to address it? Typically, learning methods are developed on a specific dataset. A common practice for moving the model to a different setting is to adapt or re-learn using a dataset obtained from this new setting. The challenge with this approach is that it is unrealistic to assume that data from all possible deployment settings are available upfront. Within a failure prevention paradigm, developers anticipate and guard against likely sources of shifts between environments during model learning without the need for data from each setting (Saria and Subbaswamy, 2019). This involves three main components: (i) determining likely shifts we want models to be rob