1 Introduction↩︎

Machine learning (ML) models are often trained on vast amounts of data, but will inevitably encounter test distributions that differ from the training set. Such distribution shift is one of the most common failure modes for ML in practice. When models do fail, model developers need to diagnose and correct the problem. In the simplest case, this may simply consist of gathering more data to retrain the model. However, in other cases, it may be necessary to fix issues in an underlying data pipeline, add new features to replace ones that have become uninformative, or undertake other more complex interventions. A necessary starting point for any such process is to understand what changed in the new dataset. Developing such understanding may even have scientific importance. For instance, a novel virus variant may emerge with new risk factors, lowering the performance of models that predict disease progression, or specific mutations in the genome could have differing relevance to disease between ancestries, weakening polygenic risk score models due to fundamentally different biology between populations [1], [2].

We propose methods for diagnosing distribution shift, focusing specifically on the case of concept shift, or when the conditional distribution of the label given the features, \(p(Y\mid X)\), differs between the source and target distribution. Concept shift represents the difficult case where the relationship between features and outcome has changed, as opposed to marginal covariate or label shifts impacting only \(p(X)\) or \(p(Y)\) by themselves. Indeed, [3] document concept shift as the primary contributor to performance degradation across a wide range of empirical examples of distribution shifts, and [4] describe it as affecting “most of the real world applications”. In this setting, our goal is to understand the features that led \(p(Y\mid X)\) to differ between the source and target domains.

Understanding distribution shift has been the subject of increasing interest. However, existing methods mostly operate relative to a structure for the data which is prespecified by the analyst, for example a known causal graph [5], [6], fixed decomposition of the variables [7], or particular assumed models for distribution shift in which out-of-distribution performance can be identified using only unlabeled data [8]. Methods that do not impose such structural conditions largely repurpose other tools to explain distribution shift as a secondary objective. For example [9] propose to look for changes in model explanations, while [10] fit a decision tree to explain differences in predictions from source and target domain models as part of a larger empirical investigation.

We introduce SGShift, a new method directly designed for diagnosing concept shift. SGShift offers robust statistical performance, particularly with limited target-domain samples and without requiring prespecified causal structure. Just as sparsity is an effective principle for learning predictive models in many settings due to sparse mechanism shift [11], we hypothesize that the update needed to adapt a source-domain model to the target domain may often be well approximated by a sparse function of the features (a fact that we empirically verify in several application domains). In this case, a useful explanation of concept shift is to identify a small set of features that drive the change between the two distributions, which could e.g.be the subject of potential modeling fixes. SGShift frames this problem as learning an update to a source distribution’s predictive model using a minimal set of features to recover the performance loss in the target distribution. We show how this formulation allows simple, principled, and easily implemented diagnoses of distribution shift, without requiring any prior knowledge, causal information, or parametric priors regarding the dataset.

We benchmark SGShift against several baselines on semi-synthetic datasets with known feature shifts, observing greatly superior performance at identifying concept shifted features (referred to as shifted features throughout). We then apply SGShift to two real-data settings and recover real-world concept shifts consistent with findings from medical and biological literature, such as respiratory failure’s reduced association with COVID-19 hospitalization after Omicron and ancestry-associated rare variants related to Lupus prediction. Together, these findings provide evidence that SGShift can recover accurate and interpretable descriptions of concept shift across a wide range of settings.

1.1 Additional Related Work↩︎

Covariate shift. Much of the existing work on distribution shift has focused on the setting where the marginal feature distribution \(P(X)\) changes while the conditional distribution \(P(Y\mid X)\) remains unchanged, i.e., covariate shift. For instance, [12] introduce statistical tests to identify which variables have shifted between source and target domains, while [13] propose explaining observed shifts via a learned transportation map between the source and target distributions, not distinguishing between features and labels. \(P(X)\) shift can be identified by methods like two-sample tests [14] or classifiers [15] and corrected by techniques such as importance sampling [16]. [17] further use these ideas to correct covariate shift by regarding the unexplained residual as a shift in \(P(Y\mid X)\), although they don’t correct or explain the concept shift. Although these methods can be effective for addressing covariate shift, they often do not delve into potential shifts in the conditional distribution. Explaining shifts in \(P(Y\mid X)\) typically involves performing feature-by-feature analyses of the conditional distribution \(P(Y\mid X_i)\) [18]. However, such univariate assessments risk detecting spurious shifts due to unadjusted confounding in the presence of collinearity among predictors [19]. [13] consider an unsupervised setting where the goal is to identify a set of features whose distribution differs (e.g., sensors that have been compromised by an adversary), as opposed to identifying features whose relationship with a supervised label has changed.

Conditional distribution shift. Recent efforts have begun to tackle shifts in the conditional distribution \(P(Y\mid X)\) more directly. For example, [5] consider changes in a causal parent set as a whole, relying on known causal structures. [9] propose a model-agnostic “explanation shift detector” that applies SHAP (Shapley additive explanations) to a source-trained model and covariates in both source and target domains, without including the outcomes in the target domain. They then use a two-sample test on the feature-attribution distributions from SHAP to detect whether the model’s decision logic has changed because of the changing of \(P(X)\) across domains. Despite its effectiveness in signaling shifts, this approach does not pinpoint which features are driving the changes in \(P(Y\mid X)\). [7] decompose the domain loss gap into predefined marginal and conditional segments, then allocate feature-level contributions, while [20] automatically discover subgroups within the data for which to produce feature-level explanations. [6] stress tests a source model before distribution shift, requiring a prespecified set of shifting variables. [8] focus on estimation of performance shift on an unlabeled dataset, but this require restrictive assumptions for identifiability, particularly that non-shifted features have no shifts at all when conditioned on the shifted features and label between datasets. WhyShift [10] compares two independently trained models - one from each domain - and analyze their difference to locate regions of covariate space with the largest predictive discrepancy. SGShift differs in that we aim to explicitly identify what the features contributing to conditional distribution shift are without requiring any prior knowledge of the dataset.

2 Preliminaries and Problem Formulation↩︎

Problem formulation. We consider the problem of identifying a set of features that can represent an observed conditional distribution shift. We observe labeled i.i.d. samples \(\{(X_i^{(S)},Y_i^{(S)})\}_{i=1}^{n_S}\sim P_S\) and \(\{(X_i^{(T)},Y_i^{(T)})\}_{i=1}^{n_T}\sim P_T\) from source and target domains, respectively, where \(X_i\in\mathbb{R}^p\) and \(n_S\) and \(n_T\) are the number of samples in the source and target domain. A source model \(h_S:\mathbb{R}^p\to\mathbb{R}\) is trained on the source domain and evaluated in the target domain. Concept shift occurs when \(P_T(Y\mid X=x)\neq P_S(Y\mid X=x)\) for some covariate values \(x\). Let \(\mu_S(x)=\mathbb{E}_S[Y\mid X=x]\), \(\mu_T(x)=\mathbb{E}_T[Y\mid X=x]\). On a chosen modeling scale, for example, the link scale for a generalized linear model, define the conditional-response contrast \(\Delta_g(x)=g(\mu_T(x))-g(\mu_S(x))\). Our primary goal is to identify a small set of raw feature indices \(A\subseteq[p]\) such that \(\Delta_g(x)\) is well approximated by a function depending only on \(x_A\).

The difficulty is that the source and target samples are unpaired: we only observe labels from \(P_S(Y\mid X)\) for source-domain covariates and labels from \(P_T(Y\mid X)\) for target-domain covariates. Accordingly, it is not possible to directly apply existing methods for sparse regression. The most directly related work, the WhyShift framework introduced by [10] for diagnosing concept shift, takes a plugin approach. A plugin strategy first fits models on the two datasets separately to approximate \(\mathbb{E}_S[Y\mid X]\) and \(\mathbb{E}_T[Y\mid X]\). Second, it fits a second model regressing some difference metric of \(\widehat{\mathbb{E}}_S[Y\mid X]\) and \(\widehat{\mathbb{E}}_T[Y\mid X]\) on \(X\) to summarize the structure in \(\Delta\). However, this plugin approach risks an accumulation of errors, particularly when we are interested in recovering structure related to sparsity: given noisy approximations to the two conditional expectations, the difference between \(\widehat{\mathbb{E}}_S[Y\mid X]\) and \(\widehat{\mathbb{E}}_T[Y\mid X]\) will not necessarily display the same sparsity pattern as \(\Delta\) (as we observe experimentally). It is also potentially challenging when we have limited target-domain data, since separately fitting \(\mathbb{E}_T[Y\mid X]\) may be especially difficult in this setting.

3 Method↩︎

Our method, SGShift, circumvents these difficulties by reformulating the problem so that existing sparse regression and controlled-selection methods can be applied to a correction term. Instead of fitting separate models for \(\mathbb{E}_S[Y\mid X]\) and \(\mathbb{E}_T[Y\mid X]\) and then estimating their difference, SGShift starts with a source-domain model \(h_S(X)\) and learns a sparse target-domain correction \[\begin{align} \min_{\hat{\Delta}} \mathbb{E}_T[\ell(h_S(X) + \hat{\Delta}(X), Y)] \quad \text{s.t. }\hat{\Delta} \text{ is } k\text{-sparse}. \end{align}\] This formulation treats \(h_S(X)\) as a fixed offset and learns only the correction from target-domain data. Our theory gives estimation consistency for the sparse correction under standard high-dimensional M-estimation assumptions, while SGShift-K provides false-discovery control for selected features through knockoffs. Exact support recovery would require stronger beta-min and design assumptions, and is not claimed for the basic SGShift estimator.

3.1 SGShift: Instantiation with \(\ell_1\) regularization↩︎

Our suggested implementation of SGShift uses a generalized additive correction with \(\ell_1\) regularization. Let \(h_S(X)\) denote the fixed source-domain predictor on the same modeling scale as the link function \(g\). We model \[g\!\left(\mathbb{E}_{T}[Y \mid X]\right) = h_S(X) + \phi(X)^\top \delta \label{equation:glm}\tag{1}\] Here, \(g\) is a link function, \(\phi:\mathbb{R}^p\to\mathbb{R}^K\) is a fixed basis map. Raw-feature attributions are then obtained by grouping basis coordinates according to the raw features on which they depend. The default choice is the standardized identity basis, so \(K=p\) and \(\phi(X)=X\) after standardization. When richer but still interpretable corrections are desired, \(\phi\) can include low-order terms such as pairwise interactions. Appendix 16 suggests that SGShift is reasonably robust to both under- and over-specified interaction bases. In order to control the sparsity level of \(\delta\), SGShift solves the averaged target-domain penalized loss \[\label{equ:naive} \begin{align} \hat{\delta} = \arg\min_{\delta \in \mathbb{R}^K} \Bigl\{ L(\delta) + \lambda \|\delta\|_1 \Bigr\} \\ \text{where } \; L(\delta) := \frac{1}{n_T}\sum_{i=1}^{n_T} \ell\left(h_S(X_i^{(T)})+\phi(X_i^{(T)})^\top\delta, Y_i^{(T)}\right). \end{align}\tag{2}\]

where \(\ell(\eta,y)\) is the per-sample negative log-likelihood evaluated at linear predictor \(\eta\). By default, we choose \(\lambda\) by cross-validation on target-domain loss when the goal is predictive correction. When shifted-feature selection is the main goal, we vary \(\lambda\) to obtain a ranking across sparsity levels or use the heuristic in Appendix 11. Empirically, performance was stable when varying \(\lambda\) from \(0.1\) to \(10\) times the selected value.

3.2 SGShift-A: Refined fitting considering source model misspecification↩︎

Prioritizing shifted features relies on an existing model trained on the source dataset. However, it may be that this model does not represent the data well due to difficulties in model fitting. To reduce the extent to which source-model misfit biases the selection of shifted features, we incorporate an additional absorption term. The main absorption idea is a modeling approximation: we assume that a component of source-model misspecification is sufficiently shared across source and target domains that it can be represented by a common correction term, while the remaining target-specific correction is attributed to conditional shift.

Let \(\phi_i^{(S)}=\phi(X_i^{(S)})\in\mathbb{R}^K\) and \(\phi_i^{(T)}=\phi(X_i^{(T)})\in\mathbb{R}^K\). SGShift-A solves: \[(\hat{\omega},\hat{\delta}) = \arg\min_{\omega,\delta \in \mathbb{R}^K} \Bigl\{L^A(\omega,\delta) + \lambda_\omega\|\omega\|_1 + \lambda_\delta\|\delta\|_1\Bigr\},\] where \[\begin{align} L^A(\omega,\delta) =&\; w_S\frac{1}{n_S}\sum_{i=1}^{n_S} \ell\left(h_S(X_i^{(S)})+(\phi_i^{(S)})^\top\omega,\;Y_i^{(S)}\right) \\ &\; + w_T\frac{1}{n_T}\sum_{i=1}^{n_T} \ell\left(h_S(X_i^{(T)})+(\phi_i^{(T)})^\top(\omega+\delta),\right. \\ &\qquad\qquad\qquad\qquad\left. Y_i^{(T)}\right). \end{align}\] Here, \(\omega\) is a shared correction and \(\delta\) is target-specific, \(w_S=w_T=1/2\) by default so that source and target losses contribute comparably even when the sample sizes differ. We induce hierarchical regularization \(\lambda_\omega < \lambda_\delta\) so that variation explainable by a shared correction is absorbed before introducing target-specific shifted features. In practice, we select \(\lambda_\delta\) using the same strategy as SGShift, either cross-validation on target-domain loss or the heuristic in Appendix 11. We set \(\lambda_\omega=c\lambda_\delta\) with \(c<1\); in our experiments, values of \(c\) between \(0.1\) and \(0.9\) gave similar behavior, and we use this range as a default sensitivity grid.

3.3 SGShift-K: Explicit false discovery control with knockoffs↩︎

While \(\ell_1\) regularization enables sparse correction learning, its selected support can include false discoveries, especially when features are correlated or the target sample is small. For shifted-feature selection, we adapt the Model-X knockoffs framework [21]. Knockoffs generate synthetic features that preserve the dependence structure of the original features while being conditionally unrelated to the response, enabling controlled variable selection. Let \(\tilde{X} = [\tilde{X}^{(1)}, \ldots, \tilde{X}^{(p)}] \in \mathbb{R}^{n_T \times p}\) denote a knockoff copy of \(X_T\) and let \([\phi_T\;\tilde{\phi}_T]=[\phi(X_T)\;\phi(\tilde{X})] \in\mathbb{R}^{n_T\times 2K}\). For the formal knockoff guarantee, the tested units should be interpreted as raw features or as pre-specified groups of basis functions associated with raw features. If \(\phi\) includes interaction or nonlinear basis terms, then either a valid knockoff construction for the transformed design must be used, or the knockoff statistics should be grouped at the raw-feature level. We fit the correction model using both original and knockoff basis functions: \[\hat{\delta'} = \arg\min_{\delta' \in \mathbb{R}^{2K}} \bigl\{ L^K(\delta') + \lambda \|\delta'\|_1 \bigr\},\] where \[\begin{align} L^K(\delta') =&\; \frac{1}{n_T}\sum_{i=1}^{n_T} \ell\!\left(h_S(X_i^{(T)}) \right. \\ &\left. \qquad\qquad + [\phi(X_i^{(T)})\;\phi(\tilde{X}_i^{(T)})]\delta', Y_i^{(T)}\right). \end{align}\] Write \(\delta'=(\delta,\tilde{\delta})\), where \(\delta\) corresponds to original basis functions and \(\tilde{\delta}\) to knockoff basis functions. We compute knockoff importance statistics \(W_j\) comparing each original feature to its knockoff copy. To reduce randomization variability across knockoff draws, we then apply a repeated-knockoff aggregation procedure for feature selection. For stability selection, we threshold empirical selection frequencies across knockoff draws; for nominal FDR control of the aggregated set, we use the e-value derandomized knockoff construction and e-BH procedure [22]. Notably, the objective of SGShift-K is shifted feature selection only with the generation of knockoff copies, while SGShift and SGShift-A can do simultaneous feature selection and target model correcting from the trained source model.

Appendix 9 gives the construction details, and Appendix 10 states the assumptions under which the stability-thresholded and e-BH aggregated procedures control errors.

3.4 Theoretical guarantees↩︎

We show that when the model in Equation 1 is well-specified, SGShift has desirable theoretical guarantees for estimating of the sparse correction coefficients \(\delta\) under proper choice of the regularization parameter \(\lambda\). Importantly, this only requires imposing assumptions on the form of the between-distribution difference \(\Delta_g\), rather than on the complete regression function \(\mathbb{E}_T[Y\mid X]\), which is allowed to be nonparametric (as opposed to the standard Lasso setting), while consistency still requires standard high-dimensional assumptions on the target design, loss curvature, and sparse correction class. In particular, we obtain the following:

The proof is in Appendix 7. Here, \(\lesssim\) means asymptotically bounded above up to a constant factor, and \(\asymp\) means asymptotically the same order up to constant factors.

An analogous result holds for SGShift-A under the corresponding augmented-design assumptions.

Further, the use of knockoffs in SGShift-K allows us to obtain error-control guarantees for false selections.

Theorem 3.2 guarantees per family error rate (PFER) control for the stability-thresholded set and false discovery rate (FDR) control for the e-BH aggregated knockoff set under their respective assumptions. The proof is in Appendix 10. We also provide a heuristic one-pass penalty-calibration rule for SGShift in Appendix 11.

4 Experiments↩︎

Evaluation setup We evaluate our method on three real-world healthcare datasets (details in Appendix 12) split in such a way to exhibit natural concept shifts, 30-day Diabetes Readmission [23] split by ER admission, COVID-19 Hospitalizations [24] split by pre and post-Omicron, and SUPPORT2 Hospital Expenses [25] split by death in hospital. For each of these 3 naturally shifted datasets, we construct semi-synthetic simulations, consistent with previous work [5], [20]. We fit a “generator” model” to the real labels in source domain, relabeling the source data, then simulate the target dataset’s labels with an induced conditional shift by perturbing \(g(E[Y\mid X])\) based on selected input features. A “base” model is then trained from the relabeled source domain. We vary base and generator models to be each combination of decision tree, logistic/linear regression, gradient boosting, and support-vector machines, for a total of 16 settings in each dataset and 48 total settings. We consider 3 scenarios in each setting, sparse shift, where a small set of features are shifted, dense shift, where \(>60\%\) of the features are shifted, and sample size, the sparse setting but with varying the amount of samples in the target domain. All features to shift are selected randomly. We construct additional synthetic simulations that allow us to isolate the impact of more complex shift cases by simulating 1000 samples in each domain alongside a specific distribution of features with the same base/generator setup as in the semi-synthetic case. This includes high dimensional shifts, where 100-500 features exist and 20% are shifted, feature correlation, varying the maximum feature correlation \(\rho\) from 0.1 to 0.9, with i-th and j-th predictors correlated as \(\rho^{|i-j|}\) with 500 features and 20% are shifted, and signal-to-noise, varying the signal-to-noise ratio (SNR) from 1 to 16 with maximum feature correlation 0.7, 200 features, and 20% are shifted. In all simulations, a feature ranking is obtained by varying the penalty parameter from 0.0001 to 100, and feature selection performance (a binary 0/1 label) measured on this ranking with AUC and recall at false positive rate 5%.

Baselines We consider 3 baseline models which also use both features and labels in source and target domain to identify shifted features. Diff, a method we construct where we simply compute the outcome discrepancies of two “base models” separately trained on source and target data, and apply sparse regression on held-out samples and the base models’ outcome probability differences to identify features contributing to the shifts. WhyShift [10] uses two “base models” separately trained on source and target domains and computes model outcome probability discrepancies, then trains a non-linear decision tree on these discrepancies to detect regions (paths in the tree) responsible for conditional shifts. We extract the features from any path in the learned tree with feature importance \(> 0\) and consider them as the shifted features. SHAP, a Shapley value-based method we adapt from [9] such that we can find individual features that differ between datasets. SHAP trains “base models” separately on source and target data, computes the Shapley value of each feature, and ranks the largest absolute differences between models.

4.1 Benchmarking↩︎

Accuracy in detecting shifted features. First, we examine the case of sparse shifts, in line with SGShift’s sparsity assumption (Table 1). Across model settings and datasets, SGShift-K achieves the strongest performance and each variant of SGShift strongly outperforms baselines Diff, WhyShift and SHAP, with AUC typically greater than 0.9, 0.1-0.2 higher than the nearest baseline. Despite the presence of model mismatch, SGShift variants still attain high performance in the mismatched setting, on average only 0.02 AUC below the matched setting.

We next examine the case of dense concept shift, violating SGShift’s sparsity assumption. Table 1 shows evaluation results. Despite the assumption of sparsity, SGShift-K still attains AUC greater than 0.8 and 0.9, and all variants of SGShift still broadly outperform baselines. Baseline methods generally experienced a substantial performance drop in the dense case, such as all methods in the Diabetes dataset, each with AUC around 0.6, down from around 0.75 previously. SGShift variants are robust towards dense shift and do not over-emphasize a few features when many may be shifted. While perhaps counterintuitive given the sparsity assumption, SGShift likely performs well as it effectively acts as a regularized feature-ranking procedure, and can still capture most of the signal even when most shift coefficients are nonzero. Additional experiments testing global and interaction shifts are available in Appendix 15 and 16, respectively.

Next, we vary the sample size available in the target domain, simulating an online learning setting where data is gradually streaming in. Results for SUPPORT2 are reported in Figure 1. SGShift-K is able to identify over half the shifting features given only 100 samples regardless the model setting, and over 85% given 500 samples. This indicates SGShift-K is indeed an effective diagnostic tool, not requiring many samples for accurate feature identification. Similar results are reported for Diabetes and COVID-19 in Appendix 14.

We analyze more complex cases of concept shift in Figure 2. In the high dimensional case with 500 features and just 1000 samples in each domain, all variants of SGShift can still strongly detect shifted features with AUC\(>0.9\) in the matched case and \(>0.8\). This is in contrast with baseline methods, where performance collapses at the 500 feature mark. When features are exceedingly correlated, all variants of SGShift again achieve strong performance relative to baseline methods with average AUC\(>0.8\) in both cases. As the signal-to-noise ratio decreases substantially, SGShift and SGShift-A are relatively unaffected and maintain strong feature selection capability with AUC\(>0.9\). SGShift-K loses power as the signal becomes exceedingly weak as the feature importance statistics calculated by knockoffs become noisier, but nevertheless still outperforms baselines by \(>0.15\) AUC when SNR=1.

4.2 Real data↩︎

Sparse predictive corrections in real data. One advantage of SGShift is that it offers simultaneous performance recovery alongside feature attribution. In real data, this lets us test whether target-domain performance loss can be well approximated by a sparse update, rather than verifying that the underlying data-generating concept shift is itself sparse, where in contrast all baselines considered provide shifted feature estimates only. Operationally, we say that a target-domain performance drop admits a sparse update if 90% of the performance loss can be recovered using only a small set of features. Applied to the original data with natural concept shift, SGShift is able to learn updates to the source model that recover 90% of the performance loss in the target domain (Figure 3), requiring less than 1/3 of the total features, and in some cases as little as one feature. We sought to validate if this was the result of complex concept shift or more simple covariate or label shift. We applied inverse propensity weighting (IPW) [26], a common method for correcting for covariate shift to each of our model settings. We find that IPW can only recover at most 15% of the performance drop, and often has little change or even reduces model performance, likely due to insufficient sample size for the IPW procedure. We then apply IPW to the target labels to account for label shift, and find this explains only one case of shift and less than 25% of the shift on average. We conclude that covariate or label shift alone cannot explain these performance drops. As an illustrative example, we plot the performance recovery curve of SGShift for an SVM model in the SUPPORT2 dataset, varying the penalty parameter to allow more features to contribute to SGShift’s source model update. With only 6 features, the performance loss can be completely recovered, whereas correcting for label or covariate shift offers little improvement. We next test for sparsity at scale, training models on cross-state splits of the ACS Income Datasets [26] and identifying 87 cases of concept shift, where source-trained models underperform in the target domain. Applying SGShift to these cases, we find that all 87 cases can be recovered, many with less than 5% of features, and all required less than 50%. These results suggest that, in these real-data settings, target-domain performance degradation can often be well approximated by a sparse update.

Case study in healthcare. We next evaluate whether the top features selected by SGShift-K for the COVID-19 pre-/post-Omicron split are consistent with known clinical patterns (data split in Table 3). The highest ranked feature across all models is “respiratory failure” with a negative sign, consistent with the broad observation of reduced severe lower-respiratory disease during Omicron compared to the previous Delta variant [27], partly due to Omicron’s decreased ability at infecting lung cells [28]. More severe cases may be taking place in other pathways, such as the upper respiratory tract [29]. Other selected features, including “abnormal breathing” and “other circulatory/respiratory signs,” also have negative corrections, plausibly reflecting the same shift away from lower-respiratory severity. Features such as “primary diagnosis,” “insurance claims,” “chronic ischemic heart disease,” and “age 45–64” may reflect that non-lung-related comorbidities contribute relatively more to hospitalization risk after the reduced role of severe lung involvement [30].

Case study in genetics. We consider a known case of concept shift in the difference in Lupus severity and prevalence between ancestries [31]. We use the gene expression from 149 healthy and Lupus-affected Europeans, and 107 healthy and Lupus-affected Asians [32], and aim to predict Lupus status using the top 1000 variable genes in B cells, a cell type commonly implicated in Lupus. We split by ancestry and apply SGShift-K to find genes contributing to concept shift. Expectedly, we first observe an XGBoost model trained on Europeans underperforms when applied to Asians (European AUC 1.0, Asian AUC 0.84). SGShift-K discovers 6 genes in B cells contributing to this shift: ERRFI1, RP11-666A1.5, CTD-2561B21.11, AC012309.5, AC074212.5, and AP001059.5, all with negative coefficients. ERRFI1 and RP11-666A1.5 are both differentially expressed in B cells between these ancestries [33]. A genetic basis of difference in Lupus between ancestries has been discovered, and CTD-2561B21.11, AC012309.5, AC074212.5, and AP001059.5 are underpinned by eQTLs or repeat variants common in Europeans but rare in East Asians [31], [34]. Interferon signatures commonly correlate with Lupus prevalence, and Asians have elevated background interferon levels compared to Europeans, such as RP11-666A1.5 [35]. These results suggest that SGShift-K selects biologically plausible features associated with ancestry-specific differences in Lupus prediction, although SGShift tests for predictive differences and is not designed to identify direct causal biological mechanisms.

5 Discussion↩︎

We have presented SGShift, a method for attributing concept shift between datasets to a sparse set of features. Our work contributes towards understanding what makes models fail between datasets. We prove statistical guarantees regarding SGShift’s false discovery control and demonstrate high power in detecting true shifted features, even when the assumption of sparsity is violated. Our real-data results suggest that target-domain performance degradation in tabular prediction tasks can often be well approximated by sparse update. Future work could include optimizing model performance by explicitly modeling the difference between datasets given the identified shifted features, disentangling various contributors to concept shift such as label or measurement drift, or extending SGShift to non-tabular data, e.g., images or graphs.

Acknowledgements↩︎

This material is based upon work supported by the United States of America Department of Health and Human Services, Centers for Disease Control and Prevention, under award numbers U01IP001121 and NU38FT000005; and contract number 75D30123C1590. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the United States of America Department of Health and Human Services, Centers for Disease Control and Prevention.

Impact Statement↩︎

This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.

6 Justification of Restricted Strong Convexity (RSC)↩︎

The RSC condition is used only locally around the true correction parameter \(\delta^*\) and only over the usual Lasso cone. We state it for the averaged loss in Equation 2 .

7 Proof of Theorem 1: Estimation Guarantee for SGShift↩︎

Remark. The fixed offset \(h_S\) does not affect the RSC argument because the optimization is over \(\delta\); the curvature is determined by the loss and the basis design. The result establishes estimation consistency, not exact support recovery.

8 Proof Sketch for Theorem 2: Estimation Guarantee for SGShift-A↩︎

Let \(d_\omega=\hat{\omega}-\omega^*\) and \(d_\delta=\hat{\delta}-\delta^*\), and let \(S_\omega=\operatorname{supp}(\omega^*)\), \(S_\delta=\operatorname{supp}(\delta^*)\). The optimality inequality for the weighted penalty gives \[L^A(\hat{\omega},\hat{\delta}) + \lambda_\omega\|\hat{\omega}\|_1 + \lambda_\delta\|\hat{\delta}\|_1 \leq L^A(\omega^*,\delta^*) + \lambda_\omega\|\omega^*\|_1 + \lambda_\delta\|\delta^*\|_1.\] If \(\lambda_\omega\geq 2\|\nabla_\omega L^A(\beta^*)\|_\infty\), \(\lambda_\delta\geq 2\|\nabla_\delta L^A(\beta^*)\|_\infty\), the same decomposability argument yields a weighted cone condition. Applying local RSC for the augmented design gives \[\|\hat{\omega}-\omega^*\|_2^2+\|\hat{\delta}-\delta^*\|_2^2 \lesssim \frac{s_\omega\lambda_\omega^2+s_\delta\lambda_\delta^2}{\kappa_A^2}.\]

9 Construction and selection with Knockoffs↩︎

Following [21], a Model-X knockoff copy \(\tilde{X}=[\tilde{X}^{(1)},\ldots,\tilde{X}^{(p)}] \in \mathbb{R}^{n \times p}\) of \(X=[X^{(1)},\ldots,X^{(p)}] \in \mathbb{R}^{n \times p}\) is constructed so that, for any subset \(S \subseteq [p]\), \[\begin{align} (X, \tilde{X})_{\text{swap}(S)} \overset{d}{=} (X, \tilde{X}). \end{align}\] where \((X,\tilde{X})_{\mathrm{swap}(S)}\) swaps \(X^{(j)}\) and \(\tilde{X}^{(j)}\) for each \(j\in S\). The knockoff variables are constructed without using \(Y\), so that \(\tilde{X} \perp\!\!\!\perp Y \mid X\).

We set variable importance measure as coefficients: \(Z_k = |\hat{\delta'}_{k}(\lambda)|\), \(\tilde{Z}_k = |\hat{\delta'}_{k+K}(\lambda)|\). Alternatively, we can also use \(Z_k = \sup\{\lambda \geq 0 : \hat{\delta'}_k(\lambda) \neq 0\}\), the lambda value where each feature/knockoff enters the lasso path (meaning becomes nonzero). The knockoff filter works by comparing the \(Z_k\)’s to the \(\tilde{Z}_k\)’s and selecting only variables that are clearly better than their knockoff copy. The reason why this can be done is that, by construction of the knockoffs, the correction-null statistics are pairwise exchangeable. This means that swapping the \(Z_k\) and \(\tilde{Z}_k\)’s corresponding to null variables leaves the koint distribution of \((Z_1, \ldots, Z_K, \tilde{Z}_1, \ldots, \tilde{Z}_K)\) unchanged. Once the \(Z_k\) and \(\tilde{Z}_k\)’s have been computed, different contrast functions can be used to compare them. In general, we must choose an anti-symmetric function \(a\) and we compute the symmetrized knockoff statistics \(W_k = a(Z_k, \tilde{Z}_k) = - a(\tilde{Z}_k, Z_k)\) such that \(W_k\) indicates that \(X_k\) appears to be more important than its own knockoff copy. We use difference of absolute values of coefficients by default, but many other alternatives (like signed maximum) are also possible.

For repeat \(b\), let \(W_k^{[b]}\) denote the knockoff statistic for coordinate or group \(k\). The knockoff+ threshold is \[T^{[b]}=\min\left\{t\in\mathcal{W}^{[b]}:\frac{1+\#\{k:W_k^{[b]}\leq -t\}}{\#\{k:W_k^{[b]}\geq t\}\vee 1} \leq q \right\},\] \[\mathcal{W}^{[b]}=\{|W_k^{[b]}|:|W_k^{[b]}|>0\}.\] If this set is empty, we set \(\hat{A}^{[b]}=\emptyset\); otherwise \(\hat{A}^{[b]}=\{k:W_k^{[b]}\geq T^{[b]}\}\).

10 Proof of Theorem 3: Stability Selection Control↩︎

We repeat the main experiments using this control, results are below.

11 Heuristic One-Pass Penalty Calibration for Naive SGShift↩︎

Given the assumption of i.i.d. observations and the exponential family distribution to generate the dependent variable \(y\), \(f(X_i)\) determines the \(\mathbb{E}[y_i \mid X_i]\) under domain \(S\), and \(\delta\) captures the shift.

The negative log-likelihood function of \(\delta\) can be written as \[L(\delta) = \frac{1}{n_T} \sum_{i=1}^{n_T} \Bigl\{ \psi\bigl(f(X_i) + \phi_i^\top \delta\bigr) - y_i \bigl(f(X_i) + \phi_i^\top \delta\bigr) \Bigr\} = \ell\bigl(f(X_T) + \phi_T^\top \delta, y_T\bigr)\] where \(\psi(\cdot)\) is uniquely determined by the link \(g(\cdot)\), \(\ell\bigl( \eta, y\bigr) = \psi\bigl(\eta\bigr) - y\eta\) where \(\eta = f(X) + \phi^\top \delta\).

We regularize the GAM loss with an \(\ell_1\)-penalty \[\begin{align} \hat{\delta}(\lambda) &= \arg\min_{\delta \in \mathbb{R}^K} \Bigl\{ \frac{1}{n_T}\sum_{i=1}^{n_T} \bigl(\psi(\eta_i) - y_i\eta_i\bigr) + \lambda \|\delta\|_1 \Bigr\} \end{align}\]

The score vector is \[\nabla L(\delta) = \frac{1}{n_T}\sum_{i=1}^{n_T} \left[ \psi'\bigl(f(X_i) + \phi_i^\top \delta\bigr) - y_i \right] \phi_i\]

Evaluated at \(\delta=0\), \(\gamma := \nabla L(\delta)\Big|_{\delta=0} = \frac{1}{n_T}\sum_{i=1}^{n_T} \left[ \psi'\bigl(f(X_i)\bigr) - y_i \right] \phi_i\), where \(\psi'(\cdot)\) is the canonical mean function.

For a general correlated design, the exact KKT condition is \[0\in \nabla L(\hat{\delta})+\lambda \partial\|\hat{\delta}\|_1,\] so selection depends on the gradient at \(\hat{\delta}\), not only on \(\nabla L(0)\). The following calculation should therefore be interpreted as a one-pass orthogonal-design heuristic for choosing a penalty scale.

Assume each row \(\phi_i\) is sub‑Gaussian with i.i.d. coordinates and that every coordinate of \(\phi_i\) and \(y_i\) has been centered and variance‑normalized.

Let \(\sigma_\gamma^2:=\mathbb{V}(y_i\mid X_i)=\psi''\!\bigl(f(X_i)\bigr)\), where \(\psi''(\cdot)\) is the variance function of the canonical exponential‑family model.

Let the true parameter be \(\delta^*\in\mathbb{R}^K\) with support \(A\subseteq[K]\), \(|A|=a\), so \(\delta^*_j=0\) for \(j\in A^c\).

11.0.0.1 Null coordinates.

Let \(j\) be a null coordinate \(j \in A^c\) among \(K-a\) null coordinates.

At the true parameter \(\delta^*\), null-score centering gives \[\mathbb{E}\left[ Y_i-\psi'\{f(X_i)+\phi_i^\top\delta^*\}\mid X_i \right]=0,\] and hence \[\mathbb{E}\left[ \left(Y_i-\psi'\{f(X_i)+\phi_i^\top\delta^*\}\right)\phi_{ij} \right]=0.\] The following one-pass calculation approximates this null score by the score at \(\delta=0\).

Under this one-pass approximation, define \(Z_{ij}=\bigl(y_i-\psi'(f(X_i))\bigr)\phi_{ij}\), we treat \(\{Z_{ij}\}_{i=1}^{n_T}\) as i.i.d., mean‑zero, and sub‑Gaussian for the purpose of the heuristic calculation. \[\mathbb{V}[Z_{ij}] = \mathbb{E}\bigl[\bigl(y_i-\psi'(f(X_i))\bigr)^2\phi_{ij}^2\bigr] = \mathbb{E}\left[\psi''(f(X_i))\phi_{ij}^2\right] = \mathbb{E}\left[\psi''(f(X_i))\right] = \sigma_\gamma^2\] where the second equality follows under the additional orthogonal-design and homoskedastic working approximation used for this heuristic.

Under mild moment conditions, the Lindeberg-Feller central limit theorem (CLT) implies \[\frac{1}{\sqrt{n_T}}\sum_{i}\left\{\bigl(y_i-\psi'(f(X_i))\bigr)\phi_{ij}\right\}\xrightarrow{d} N\bigl(0,\sigma_\gamma^2\bigr)\quad\text{if }j\in A^c\] Therefore, for null coordinates, we have \[\sqrt{n_T}\,\gamma_j=-\frac{1}{\sqrt{n_T}}\sum_{i=1}^{n_T} \bigl(y_i-\psi'(f(X_i))\bigr)\phi_{ij} \xrightarrow{d} N(0,\sigma_\gamma^2) \quad\text{if }j\in A^c.\]

11.0.0.2 False‑selection probability and plug‑in mFDR estimate.

Under the same orthogonal one-pass heuristic, we approximate selection by the event \(\lvert\gamma_j\rvert>\lambda\), so the null‑coordinate error rate is \[Pr(j \text{ selected} \mid j \in A^c) = Pr(|\gamma_j|>\lambda) = 2\Bigl\{1-\Phi\bigl(\frac{\lambda\sqrt{n_T}}{\sigma_\gamma}\bigr)\Bigr\}\] where \(\Phi\) is the standard normal CDF.

Following [36], the marginal FDR is \[\text{mFDR}(\lambda)=\frac{\mathbb{E}[\text{\#False Discoveries}]}{\mathbb{E}[\text{\#Selected}]}\]

Plugging in the null probability above yields \[\widehat{\text{FDR}}(\lambda) = \min\bigl\{\frac{2K\bigl(1-\Phi\bigl(\lambda\sqrt{n_T}/\sigma_\gamma\bigr)\bigr)}{|\hat{A}(\lambda)|\vee 1}, 1\bigr\},\] where we use \(K\) rather than the unknown number of null coordinates.

A practical one‑pass rule that targets an approximate marginal-FDR level \(\alpha\) under the orthogonal-design heuristic is \[\hat{\lambda}_\alpha = \min\Bigl\{\lambda: \widehat{\text{FDR}}(\lambda)\le\alpha\Bigr\}\]

12 Datasets↩︎

All datasets are listed as below, and the full preprocessing code from raw data, together with the preprocessed data, are available in the source code, except restricted access COVID-19 Hospitalization data where we provide detailed fetching code and data version information from NIH All of Us Research Program [24]. Standardization is performed within the pipeline to ensure that features with larger values don’t disproportionately influence the \(\ell_1\) regularization penalty.

12.0.0.1 Diabetes 30-Day Readmission

The Diabetes 130-US Hospitals dataset, available through the UCI Machine Learning Repository (https://archive.ics.uci.edu/dataset/296/diabetes+130-us+hospitals+for+years+1999-2008), comprises 101,766 encounters of diabetic patients across 130 U.S. hospitals between 1999-2008 [23]. We fetch the data following TableShift’s procedure [26].

We define the source domain as 49,213 non-ER admissions (elective or urgent) with 25,196 readmitted patients, and the target domain as 24,402 ER admissions with 10,684 readmitted patients, with the binary classification task being prediction of 30-day readmission risk.

12.0.0.2 COVID-19 Hospitalization

The COVID-19 cohort is part of the NIH All of Us Research Program [24], a (restricted access) dataset containing electronic health records for 16,187 patients diagnosed with COVID-19 between 2020-2022. Features include demographic variables (age, gender, race), temporal indicators (diagnosis date relative to Omicron variant emergence), comorbidity status for 13 chronic conditions (diabetes, COPD), and diagnostic context (EHR vs. claims-based). We partition the data into three temporal groups: a source domain of 11,268 patients diagnosed prior to the beginning of 2022 with 2,541 patients hospitalized, a target domain of 2,219 patients diagnosed in January 2022 (early Omicron era) with 359 patients hospitalized. The binary classification task predicts hospitalization status (inpatient vs. outpatient).

12.0.0.3 SUPPORT2 Hospital Charges

From the Study to Understand Prognoses Preferences Outcomes and Risks of Treatment (SUPPORT2), publicly available via the UCI repository (https://archive.ics.uci.edu/dataset/880/support2) containing 9,105 critically ill patients [25].

The source domain is specified as 5,453 patients who survived hospitalization and the target domain as 1,817 in-hospital deaths. The regression task is defined as a prediction of \(\log_{10}\)(total hospital costs per patient).

13 Model Hyperparameters↩︎

We used standard implementations of classical machine learning models from scikit-learn, with hyperparameters either set to commonly used defaults or manually tuned for stability and performance. Supplementary table 4 summarizes the key hyperparameters for each model. Unless otherwise stated, all models were trained using their default solver settings. Random seeds were fixed via random_state to ensure reproducibility.

14 Sample size↩︎

15 Global calibration shifts↩︎

Global shifts. It may be true that all features are shifted in the same direction in a new domain due to differences in sensor calibration. In this case identifying specific shifted features may be difficult as all are perturbed slightly. We simulate this global effect where only a few true features having a conditional shift, and the rest are perturbed by noise with absolute values from 0.01 to 0.3 while the absolute values of true shifts are 3. Results are reported in Figure 5. SGShift-K still strongly identifies true shifted features with AUC \(>0.9\), even when all features are shifted slightly, and individual features are not over or under prioritized. Interestingly, for all methods, performance is relatively unchanged as the scale of the background shift increases. This may be due to the intercept term accumulating the background shift, as opposed to attributing it to any individual feature.

16 Interaction shifts↩︎

Interaction shifts. We assess each method’s performance in identifying features which shift in the interaction space in the SUPPORT2 dataset, where features are continuous. The goal is to detect individual features contributing to shift through interactions with other features. We consider two cases of SGShift-K, underspecified, where the basis function does not include interaction terms, and overspecified, where SGShift contains both second and third tier interactions. Results are reported in Table 5. In both cases, regardless of how the basis function is specified, SGShift displays strong performance.

References↩︎