A hierarchical decomposition for explaining ML performance discrepancies

The performance of an ML algorithm can differ across domains due to shifts in the data distribution. To understand what contributed to this performance gap, prior works have suggested decomposing the gap into that due to a shift in the marginal distribution of the input features $p(X)$ versus that due to a shift in the conditional distribution of the outcome $p(Y|X)$ (Cai et al., 2023; Zhang et al., 2023; Liu et al., 2023; Qiu et al., 2023; Firpo et al., 2018). Although such aggregate decompositions can be helpful, more detailed explanations that quantify the importance of each variable to each term in this decomposition can provide deeper insight to ML teams trying to understand and close the performance gap. However, there is currently no principled framework that provides both aggregate and detailed decompositions for explaining performance gaps of ML algorithms.

As a motivating example, suppose an ML deployment team has an algorithm that predicts the risk of patients being readmitted to the hospital given data from the Electronic Health Records (EHR), e.g. demographic variables and diagnosis codes. The algorithm was trained for a general patient population. The team intends to deploy it for heart failure (HF) patients and observes a large performance drop (e.g. in accuracy) in this subgroup. An aggregate decomposition of the performance drop into the marginal versus conditional components ($p(X)$ and $p(Y|X)$) provides only a high-level understanding of the underlying cause and limited suggestions on how the model can be fixed. A small conditional term in the aggregate decomposition but a large marginal term is typically addressed by retraining the model on a reweighted version of the original data that matches the target population (Quionero-Candela et al., 2009); otherwise, the standard suggestion is to collect more data from the target population to recalibrate/fine-tune the algorithm (Steyerberg, 2009). A detailed decomposition would help the deployment team conduct a root cause analysis and assess the utility of targeted variable-specific interventions. For instance, if the performance gap is primarily driven by a marginal shift in a few diagnoses, the team can investigate why the rates of these diagnoses differ between general and HF patients. The team may find that certain diagnoses differ due to variations in patient case mix, which may be better addressed through model retraining, whereas others are due to variations in documentation practices, which may be better addressed through operational interventions.

Existing approaches for providing a detailed understanding of why the performance of an ML algorithm differs across domains fall into two general categories. One category, which is also the most common in the applied literature, is to quantify how much the data distribution shifted (Liu et al., 2023; Budhathoki et al., 2021; Kulinski et al., 2020; Rabanser et al., 2019), e.g. one can compare how the distribution of each input variable differs across domains (Cummings et al., 2023). However, given the complex interactions and non-linearities in (black-box) ML algorithms, it is difficult to quantify how exactly such shifts contribute to changes in performance. For instance, a large shift with respect to a given variable does not necessarily translate to a large shift in model performance, as that variable may have low importance in the algorithm or the performance metric may be insensitive to that shift. Thus the other category of approaches—which is also the focus of this manuscript—is to directly quantify how a shift with respect to a variable (subset) influences model performance. Existing proposals only provide detailed variable-level breakdowns for either the marginal or conditional terms in the aggregate decomposition, but no unified framework exists. Moreover, existing methods either require knowledge of the causal graph relating all variables (Zhang et al., 2023) or strong parametric assumptions (Wu et al., 2021; Dodd & Pepe, 2003).

This work introduces a unified nonparametric framework for explaining the performance gap of an ML algorithm that first decomposes the gap into terms at the aggregate level and then provides detailed variable importance (VI) values for each aggregate term (Fig 1). Whereas prior works only provide point estimates for the decomposition, we derive debiased, asymptotically linear estimators for the terms in the decomposition which allow for the construction of confidence intervals (CIs) with asymptotically valid coverage rates. Uncertainty quantification is crucial in this setting, as one often has limited labeled data from the target domain. The estimation and statistical inference procedure is computationally efficient, despite the exponential number of terms in the Shapley value definition. We demonstrate the utility of our framework in real-world examples of prediction models for hospital readmission and insurance coverage.

Describing distribution shifts. This line of work focuses on detecting and localizing which distributions shift between datasets (Kulinski et al., 2020; Rabanser et al., 2019). Budhathoki et al. (2021) identify the main variables contributing to a distribution shift via a Shapley framework, Kulinski & Inouye (2023) fits interpretable optimal transport maps, and Liu et al. (2023) finds the region with the largest shift in the conditional outcome distribution. However, these works do not quantify how these shifts contribute to changes in performance, the metric of practical importance.

Explaining loss differences across subpopulations. Understanding differences in model performance across subpopulations in a single dataset is similar to understanding differences in model performance across datasets, but the focus is typically to find subpopulations with poor performance rather than to explain how distribution shifts contributed to the performance change. Existing approaches include slice discovery methods (Plumb et al., 2023; Jain et al., 2023; d’Eon et al., 2022; Eyuboglu et al., 2022) and structured representations of the subpopulation using e.g. Euclidean balls (Ali et al., 2022).

Attributing performance changes. Prior works have described similar aggregate decompositions of the performance change into covariate and conditional outcome shift components (Cai et al., 2023; Qiu et al., 2023). To provide more granular explanations of performance shifts, existing works quantify the importance of shifts in each variable assuming the true causal graph is known (Zhang et al., 2023); covariate shifts restricted to variable subsets assuming partial shifts follow a particular structure (Wu et al., 2021); and conditional shifts in each variable assuming a parametric model (Dodd & Pepe, 2003). However, the strong assumptions made by these methods make them difficult to apply in practice, and model misspecification can lead to unintuitive interpretations. In addition, there is no unifying framework for decomposing both covariate and outcome shifts, and many methods do not output CIs, which is important when the amount of labeled data from a given domain is limited. A summary of how the proposed framework compares against prior works is shown in Table 1 of the Appendix.

Decomposition methods in econometrics. Explaining performance differences is similar to the long-studied problem of explaining income and health disparities between groups in econometrics. There, researchers regularly use frameworks such as the Oaxaca-Blinder decomposition (Oaxaca, 1973; Blinder, 1973; Fortin et al., 2011), which decomposes disparities into components quantifying the impact of covariate shifts and outcome shifts. These frameworks commonly decompose the components further to describe the importance of each variable (Oaxaca, 1973; Kirby et al., 2006). Existing methods typically rely on strong parametric assumptions (Fairlie, 2005; Yun, 2004; Firpo et al., 2018), which is inappropriate for the complex data settings in ML.

In summary, the distinguishing contribution of this work is that it unifies aggregate and detailed decompositions under a nonparametric framework with uncertainty quantification.

Here we introduce how the aggregate and detailed decompositions are defined for explaining the performance gap of a risk prediction algorithm $f:\mathcal{X}\subseteq\mathbb{R}^{m}\to\mathcal{Y}$ for binary outcomes across source and target domains, denoted by $D=0$ and $D=1$, respectively. Let the performance of $f$ be quantified in terms of a loss function $\ell:\mathcal{X}\times\mathcal{Y}\to\mathbb{R}$, such as the 0-1 misclassification loss $\mathbbm{1}\{f(X)\neq Y\}$. Suppose variables $X$ can be partitioned into $W\in\mathbb{R}^{m_{1}}$ and $Z=X\setminus W\in\mathbb{R}^{m_{2}}$, where $m=m_{1}+m_{2}$. This partitioning allows for a factorization of the data distribution $p_{D}$ in the source domain $D=0$ and target domain $D=1$ into

(see Fig 2 top left). As such, we refer to $W$ as baseline variables and $Z$ as conditional covariates. For the readmission example, $W$ may refer to demographics and $Z$ to diagnosis codes. The total change in performance of $f$ can thus be written as $\Lambda=\mathbb{E}_{111}\left[\ell(W,Z,Y)\right]-\mathbb{E}_{000}\left[\ell(W,Z,Y)\right],$ where $\mathbb{E}_{D_{\texttt{W}}D_{\texttt{Z}}D_{\texttt{Y}}}$ denotes the expectation over the distribution $p_{D_{\texttt{W}}}(W)p_{D_{\texttt{Z}}}(Z|W)p_{D_{\texttt{Y}}}(Y|W,Z)$. A summary of notation used is provided in Table 2 of the Appendix.

Aggregate. At the aggregate level, the framework quantifies how replacing each factor in (1) from source to target contributes to the performance gap. We refer to such replacements as “aggregate shifts,” as the shift is not restricted to a particular subset of variables. This leads to the aggregate decomposition $\Lambda=\Lambda_{\texttt{W}}+\Lambda_{\texttt{Z}}+\Lambda_{\texttt{Y}}$, where $\Lambda_{\texttt{W}}$ quantifies the impact of a shift in the baseline distribution $p(W)$ (also known as marginal or covariate shifts (Quionero-Candela et al., 2009)), $\Lambda_{\texttt{Z}}$ quantifies the impact of a shift in the conditional covariate distribution $p(Z|W)$, and $\Lambda_{\texttt{Y}}$ quantifies the impact of a shift in the outcome distribution $p(Y|W,Z)$ (also known as concept shift). More concretely,

Prior works have proposed similar aggregate decompositions (Cai et al., 2023; Liu et al., 2023; Firpo et al., 2018).

Detailed. At the detailed level, the framework outputs Shapley-based variable attributions that describe how shifts with respect to each variable contribute to each term in the aggregate decomposition. Based on the VI values, an ML development team can better understand the underlying cause for a performance gap and brainstorm which mixture of operational interventions (e.g. changing data collection) and algorithmic interventions (e.g. retraining the model with respect to the variable(s)) would be most effective at closing the performance gap. For instance, a variable with high importance to the conditional covariate shift term $\Lambda_{\texttt{Z}}$ can be due to differences in missingness rates, prevalence, or selection bias; and a variable with high importance to the conditional outcome shift term $\Lambda_{\texttt{Y}}$ may indicate inherent differences in the conditional distribution (also known as effect modification), differences in measurement error or outcome definitions, or omission of variables predictive of the outcome. Note that the framework does not output a detailed decomposition of the baseline shift term $\Lambda_{\texttt{W}}$, for reasons we discuss later.

We leverage the Shapley attribution framework for its axiomatic properties, which result in VI values with intuitive interpretations (Shapley, 1953; Charnes et al., 1988). Recall that for a real-valued value function $v$ defined for all subsets $s\subseteq\{1,\cdots,m\}$, the attribution to element $j$ is defined as the average gain in value due to the inclusion of $j$ to every possible subset, i.e.

So to define a detailed decomposition, the key question is how to define the value of a partial distribution shift only with respect to a variable subset $s$; henceforth we refer to such shifts as $s$-partial shifts. It turns out that the answer is far from straightforward.

The key estimands in this hierarchical attribution framework are the aggregate terms $\Lambda_{\texttt{W}},\Lambda_{\texttt{Z}}$, and $\Lambda_{\texttt{Y}}$ and the Shapley-based detailed terms $\phi_{\texttt{Z},j}$ and $\phi_{\texttt{Y},j}$ for $j=0,\cdots,m_{2}$. We now provide nonparametric estimators for these quantities as well as statistical inference procedures for constructing CIs. Nearly all prior works for decomposing performance gaps have relied on plug-in estimators, which substitute in estimates of the conditional means (also known as outcome models) or density ratios (Sugiyama et al., 2007), which we collectively refer to as nuisance parameters. For instance, given ML estimators $\hat{\mu}_{\cdot 10}$ and $\hat{\mu}_{\cdot 00}$ for the conditional means $\mu_{\cdot 10}(w)=\mathbb{E}_{\cdot 10}[\ell|W]$ and $\mu_{\cdot 00}(w)=\mathbb{E}_{\cdot 00}[\ell|W]$, respectively, the empirical mean of $\hat{\mu}_{\cdot 10}-\hat{\mu}_{\cdot 00}$ with respect to the target domain is a plug-in estimator for $\Lambda_{\texttt{Z}}=\mathbb{E}_{1\cdot\cdot}[{\mu}_{\cdot 10}-{\mu}_{\cdot 00}]$. However, because ML estimators typically converge to the true nuisance parameters at a rate slower than $n^{-1/2}$, plug-in estimators generally fail to be consistent at a rate of $n^{-1/2}$ and cannot be used to construct CIs (Kennedy, 2022). Nevertheless, using tools from semiparametric inference theory (e.g. one-step correction and Neyman orthogonality), one can derive debiased ML estimators that facilitate statistical inference (Chernozhukov et al., 2018; Tsiatis, 2006). Here we derive debiased ML estimators for terms in the aggregate and detailed decompositions. The detailed conditional outcome decomposition is particularly interesting, as its unique structure is not amenable to standard techniques for debiasing ML estimators.

For ease of exposition, theoretical results are presented for split-sample estimators; nonetheless, the results readily extend to cross-fitted estimators under standard convergence criteria (Chernozhukov et al., 2018; Kennedy, 2022). Let $\{(W_{i}^{(d)},Z_{i}^{(d)},Y_{i}^{(d)}):i=1,\cdots,n_{d}\}$ denote $n_{d}$ independent and identically distributed (IID) observations from the source and target domains $d=0$ and $1$, respectively. Let a fixed fraction of the data be partitioned towards “training” (Tr) and the remaining to “evaluation” (Ev); let $n_{\texttt{Ev}}$ be the number of observations in the evaluation partition. Let $\mathbb{P}_{d}$ denote the expectation with respect to domain $d$ and $\mathbb{P}^{\texttt{Ev}}_{d,n}$ denote the empirical average over observations in partition Ev from domain $d$. All estimators are denoted using hat notation. Proofs, detailed theoretical results, and psuedocode are provided in the Appendix.

We now present simulations that verify the theoretical results by showing that the proposed procedure achieves the desired coverage rates (Sec 5.1) and illustrate how the proposed method provides more intuitive explanations of performance gaps (Sec 5.2). In all empirical analyses, performance of the ML algorithm is quantified in terms of 0-1 accuracy. Below, we briefly describe the simulation settings; full details are provided in Sec E in the Appendix.

We now analyze two datasets with naturally-occurring shifts.

Hospital readmission. Using electronic health record data from a large safety-net hospital, we analyzed performance of a Gradient Boosted Tree (GBT) trained on the general patient population (source) to predict 30-day readmission risk but applied to patients diagnosed with HF (target). Features include 4 demographic variables ($W$) and 16 diagnosis codes ($Z$). Each domain supplied $n=3750$ observations.

Model accuracy drops from 70% to 53%. From the aggregate decompositions (Fig 5a), we observe that the drop is mainly due to covariate shift. If one performed the standard check to see which variables significantly changed in their mean value (MeanChange), then one would find a significant shift in nearly every variable. Little support is offered to identify main drivers of the performance drop. In contrast, the detailed decomposition from the proposed framework estimates diagnoses “Drug-induced or toxic-related condition” and “mental and substance use disorder in remission” as having the highest estimated contributions to the conditional covariate shift, and most other variables having little to no contribution. Upon discussion with clinicians from this hospital, differences in the top two diagnoses may be explained by (i) substance use being a major cause of HF at this hospital, with over eighty percent of its HF patients reporting current or prior substance use, and (ii) substance use and mental health disorders often occurring simultaneously in this HF patient population. Based on these findings, closing the performance gap may require a mixture of both operational and algorithmic interventions. Finally, CIs from the debiased ML procedure provide valuable information on the uncertainty of the estimates and highlight, for instance, that more data is necessary to determine the true ordering between the top two features. In contrast, existing methods do not provide (asymptotically valid) CIs.

ACS Public Coverage. We analyze a neural network trained to predict whether a person has public health insurance using data from Nebraska in the American Community Survey (source, $n=3000$), applied to data from Louisiana (target, $n=6000$). Baseline variables include 3 demographics (sex, age, race), and covariates $Z$ include 31 variables related to health conditions, employment, marital status, citizenship status, and education.

Model accuracy drops from 84% to 66% across the two states. The main driver is the shift in the outcome distribution per the aggregate decomposition (Fig 5b) and the most important contributor to the outcome shift is annual income, perhaps due to differences in cost of living across the two states. Income is significantly more important than all the other variables; the ranking between the remaining variables is unclear. In comparing the performance of targeted model revisions with respect to the top variables from each explanation method, we find that model revisions based on top variables identified by the proposed procedure lead to AUCs that are better or as good as those based on RandomForestAcc (Table 3 in the Appendix).

ML algorithms regularly encounter distribution shifts in practice, leading to drops in performance. We present a novel framework that helps ML developers and deployment teams build a more nuanced understanding of the shifts. Compared to past work, the approach is nonparametric, does not require fine-grained knowledge of the causal relationship between variables, and quantifies the uncertainty of the estimates by constructing confidence intervals. We present case studies on real-world datasets to demonstrate the use of the framework for understanding and guiding interventions to reduce performance drops.

Extensions of this work include relaxing the assumption of overlapping support of covariates such as by restricting to the common support (Cai et al., 2023), allowing for decompositions of more complex measures of model performance such as AUC, and analyzing other factorizations of the data distribution (e.g. label or prior shifts (Kouw & Loog, 2019)). For unstructured data (e.g. image and text), the current framework can be applied to low-dimensional embeddings or by extracting interpretable concepts (Kim et al., 2018); more work is needed to extend this framework to directly analyze unstructured data. Finally, the focus of this work is to interpret performance gaps. Future work may extend ideas in this work to design optimal interventions for closing the performance gap.

We would like to thank Lucas Zier, Avni Kothari, Berkman Sahiner, Nicholas Petrick, Gene Pennello, Mi-Ok Kim, and Romain Pirracchio for their invaluable feedback on the work. This work was funded through a Patient-Centered Outcomes Research Institute® (PCORI®) Award (ME-2022C1-25619). The views presented in this work are solely the responsibility of the author(s) and do not necessarily represent the views of the PCORI®, its Board of Governors or Methodology Committee, and the Food and Drug Administration.