Considerations for Distribution Shift Robustness in Health

Arno Blaas, Andrew Miller, Luca Zappella, Joern-Henrik Jacobsen, Christina Heinze-Deml

Keywords: Distribution shift robustness, health

TL;DR: We demonstrate that distribution shifts that are likely to occur in healthcare settings require inclusion of auxiliary covariates for distribution shift robustness of models.

Abstract: When analyzing robustness of predictive models under distribution shift, many works focus on tackling generalization in the presence of spurious correlations. In this case, one typically makes use of covariates or environment indicators to enforce independencies in learned models to guarantee generalization under various distribution shifts. In this work, we analyze a class of distribution shifts, where such independencies are not desirable, as there is a causal association between covariates and outcomes of interest. This case is common in the health space where covariates can be causally, as opposed to spuriously, related to outcomes of interest. We formalize this setting and relate it to common distribution shift settings from the literature. We theoretically show why standard supervised learning and invariant learning will not yield robust predictors in this case, while including the causal covariates into the prediction model can recover robustness. We demonstrate our theoretical findings in experiments on both synthetic and real data.