Technical Summary: Beyond the Training Distribution: Evaluating Predictions Under Distribution Shift and Selection Bias

Problem Statement

Machine learning models frequently degrade when deployed in new environments due to two primary, often co-occurring, challenges: covariate shift and selective labels.

• Covariate Shift: The marginal distribution of input features ($P(X)$) differs between the training (source) and deployment (target) populations, while the conditional relationship between features and outcomes ($P(Y|X)$) remains stable. Standard evaluation on a held-out test set may obscure performance degradation if errors are concentrated in feature regions common in the target environment but rare in the source.
• Selective Labels: In high-stakes domains (e.g., healthcare, lending), outcome labels are observed only for individuals selected by a prior decision policy (e.g., patients who receive a diagnostic test). This induces selection bias, as the observed labels are drawn from a decision-dependent subset of the population, making the training data an inaccurate reflection of the full deployment population.

Existing literature typically addresses these issues in isolation. This paper studies the joint presence of covariate shift and selective labeling, a scenario common in clinical decision-making where models trained on specific demographic groups are deployed in diverse populations, yet outcomes are only observed for patients who undergo specific clinical protocols. The core problem is to evaluate the target risk (expected loss in the deployment environment) of an arbitrary black-box prediction model when target outcomes are unobserved and source outcomes are selectively observed.

Methodology

1. Problem Formalization and Identification

The authors define the target risk $\psi$ as the expected loss in the target population:
$$\psi := E_T [\ell(f(X), Y)]$$
where $f$ is a fixed black-box predictor and $\ell$ is a general loss function.

To identify $\psi$ from observable data, the paper relies on four standard assumptions:

1. Selection on Observables (Source): Conditional on covariates $X$, label observability $D$ is independent of the outcome $Y$ in the source population ($Y \perp\!\!\!\perp D \mid X, R=1$).
2. Covariate Shift: The conditional distribution of outcomes given features is invariant across domains ($P_S(Y|X) = P_T(Y|X)$).
3. Positivity: Every covariate profile has a positive probability of being observed in the source population.
4. Bounded Likelihood Ratio: The density ratio between target and source distributions is bounded.

Under these assumptions, the target risk is identified via two equivalent expressions (Proposition 3.5):

1. Target Expectation: $\psi = E_T [\mu(X)]$, where $\mu(X) = E[\ell(f(X), Y) \mid X, R=1, D=1]$ is the conditional expected loss in the source.
2. Source Expectation with Corrections: $\psi = E_S \left[ \frac{p_T(X)}{p_S(X)} \cdot \frac{D}{\pi_S(X)} \cdot \ell(f(X), Y) \right]$, where $\frac{p_T(X)}{p_S(X)}$ corrects for covariate shift and $\frac{D}{\pi_S(X)}$ corrects for selective labels via inverse propensity weighting.

2. Efficient Influence Function (EIF)

The authors derive the Efficient Influence Function (EIF) for the target risk estimand (Proposition 3.6). The EIF, $\phi(Z; P)$, combines signals from both the source and target distributions:
$$\phi(Z; P) = \frac{R D}{\pi(X)} \frac{g(X)}{\rho} (L - \mu(X)) + \frac{1-R}{\rho} (\mu(X) - \psi)$$
where:

• $R$ is the domain indicator (1 for source, 0 for target).
• $D$ is the label observability indicator.
• $\pi(X)$ is the pooled propensity score.
• $g(X)$ is the domain propensity score.
• $\rho$ is the proportion of the target population.
• $L$ is the observed loss.

The first term represents a residual correction for labeled source units, weighted to transport them to the target distribution. The second term is a plug-in estimate for unlabeled target units, centered by the target risk.

3. Double Machine Learning (DML) Estimator

Leveraging the EIF, the paper proposes a Double Machine Learning (DML) estimator (Equation 7):
$$\hat{\psi} = \frac{1}{n} \sum_{i=1}^n \left[ \frac{R_i D_i}{\hat{\pi}(X_i)\hat{g}(X_i)} (L_i - \hat{\mu}(X_i)) + \frac{1-R_i}{\hat{\rho}} \hat{\mu}(X_i) \right]$$
Key features of this estimator include:

• Bias Correction: It uses the EIF structure to correct for first-order biases arising from errors in estimating nuisance functions ($\mu, \pi, g$).
• Sample Splitting/Cross-fitting: Nuisance functions are estimated on an auxiliary sample independent of the evaluation sample to ensure asymptotic normality.
• Robustness: The estimator is semiparametrically efficient and $\sqrt{n}$-consistent, provided the nuisance estimators converge at a rate of $o_p(n^{-1/4})$. This allows the use of flexible, non-parametric machine learning methods (e.g., random forests, neural networks) for nuisance estimation without requiring parametric rates.

Key Results

Theoretical Guarantees

The paper proves that the proposed DML estimator is asymptotically linear and semiparametrically efficient (Theorem 3.7). This implies:

• The estimator converges to the true target risk at a $\sqrt{n}$ rate.
• The estimator is asymptotically normal, enabling the construction of valid Wald-type confidence intervals for pre-deployment risk assessment.
• The estimator remains consistent even if the nuisance functions are estimated with high-dimensional or complex models, as long as their product of errors is negligible.

Empirical Evaluation

The authors evaluate the method using the eICU Collaborative Research Database (over 200,000 ICU admissions) and semi-synthetic experiments.

• Semi-Synthetic Experiments: Using real covariates from eICU with simulated outcomes and label observability, the DML estimator demonstrated the least bias and lowest RMSE compared to:
  • A naive Plug-in estimator (reweighting for both shift and selection).
  • A Covariate Shift (CS)-only estimator (ignoring selection bias).
  • A Selective Labels (SL)-only estimator (ignoring covariate shift).
  • Results showed that addressing only one source of bias leads to systematic misrepresentation of deployment risk, especially as selection strength increases.
• Real-World Experiments: Applied to four clinical outcomes (BNP, Troponin-I, CK, Hypoxemia) where labels are observed only when specific lab tests are ordered. The DML estimator produced risk estimates that differed significantly from baselines in several hospitals, highlighting that ignoring the joint nature of shift and selection can lead to underestimation or overestimation of model risk in new clinical settings.
• Robustness Checks: The estimator showed reasonable robustness to violations of positivity and bounded likelihood ratio assumptions. However, as expected, performance degraded under violations of selection on observables (unobserved confounding) and conditional outcome shift, though it remained competitive with other benchmarks.

Significance and Claims

The paper claims to be the first to formalize and provide a solution for evaluating black-box predictors under the simultaneous presence of covariate shift and selective labeling.

• Practical Impact: The work provides a rigorous framework for pre-deployment evaluation in high-stakes domains like healthcare, where models are often trained on one population but deployed on another, and where outcome data is inherently incomplete due to clinical decision-making processes.
• Methodological Contribution: By establishing the identification of the target risk and deriving the efficient influence function, the paper bridges the gap between distribution shift literature and selective label literature. It demonstrates that valid inference is possible even when target labels are entirely unobserved and source labels are biased.
• Fairness and Safety: The authors emphasize that failing to account for these joint challenges can systematically misrepresent model performance for underrepresented subpopulations, potentially perpetuating inequities. Their method allows practitioners to assess fairness and safety properties more accurately before deployment.
• Broader Applicability: While focused on clinical settings, the framework is applicable to any domain involving foundation models or LLMs where feedback is selectively observed and deployment distributions differ from training distributions.

The authors maintain a modest stance, acknowledging that identification relies on substantive assumptions (e.g., selection on observables) that may fail in practice (e.g., if clinical decisions depend on unmeasured severity). They suggest that while their method offers tolerance to small violations, future work could focus on relaxing these assumptions or using adaptive data collection.