Technical Summary: A Practical Upper Bound on Selection Bias Effects in Medical Prediction Models

Problem Statement

Selection bias is a pervasive issue in real-world data that compromises the generalizability of machine learning models, particularly in high-stakes domains like healthcare. When models are trained on biased data (e.g., specific hospital records or volunteer-based biobanks) and deployed on a broader target population, they may fail catastrophically, potentially reifying discrimination against underrepresented groups.

Existing methods for assessing model generalizability under selection bias suffer from unrealistic assumptions. Techniques such as density ratio estimation require full access to the target distribution, while inverse probability weighting (IPW) assumes full observability of the variables causing selection. In practice, practitioners often lack access to representative target samples (due to cost or privacy) and rarely observe the explicit selection mechanism. Consequently, there is a lack of principled, practical tools to audit models for selection bias prior to deployment.

Methodology

The authors propose a novel framework to estimate a practical upper bound on the worst-case performance (expected loss) of a prediction model on an unobserved target population. The method operates under a realistic setting of partial observability:

1. Observed Data: A biased dataset $D_P$ (where selection indicator $S=1$) containing covariates $X$, outcomes $Y$, and a subset of selection variables $\tilde{U}$.
2. Target Constraints: Access to the target distribution $Q$ is limited to:
   • A subset of selection variables $\tilde{U}$ (fully observed).
   • Marginal summary statistics (means and variances) for all candidate variables.
   • No direct samples of $(X, Y)$ from the target distribution $Q$.

Key Theoretical Components

The method relies on three core assumptions:

• Conditional Independence: The selection indicator $S$ is conditionally independent of $(X, Y)$ given the full set of selection variables $U$.
• Common Support: The biased distribution $P$ and target distribution $Q$ share common support.
• Non-negative Loss: The loss function $\ell$ is non-negative.

Under these assumptions, the authors derive an upper bound $\hat{R}_Q$ for the expected loss $R_Q$ on the target distribution:
$$R_Q \leq \hat{R}_Q := E_P \left[ w(\tilde{U}) \cdot \phi(\tilde{X}, Y, U^*) \cdot \ell(f_P(\tilde{X}), Y) \right]$$
Where:

• $w(\tilde{U})$ is a propensity weight derived from the ratio of marginal densities $p(\tilde{U}) / p(\tilde{U}|S=1)$.
• $\phi(\tilde{X}, Y, U^*)$ is a worst-case factor defined as the maximum ratio of conditional densities over the space of unobserved selection variables $U^* = U \setminus \tilde{U}$.

Practical Implementation

To make the bound tractable, the authors introduce two heuristics:

1. Heuristic Identification of $U^*$: Since $U^*$ is unknown, the authors propose a moment-matching algorithm. This algorithm treats the selection process as a logistic function and uses the observed summary statistics from $Q$ to solve for coefficients that identify which remaining variables in the biased dataset are likely drivers of selection. Variables with non-zero coefficients in the confidence interval are nominated as $U^*$.
2. Diagnostic Tools: To assess the validity of the bound in practice (where assumptions may be violated), the authors propose three diagnostics:
   • KS Test: Evaluates overlap between propensity distributions in $P$ and $Q$.
   • Weight Design Effect ($d_{eff}$): Measures the stability of the resulting weights to detect lack of common support.
   • Propensity Invariance: A novel test to approximate the conditional independence assumption by checking if the distribution of covariates remains invariant across different propensity score bins.

Key Contributions

1. Realistic Problem Formulation: The work addresses the gap between theoretical bounds and practical constraints by requiring only partial observability of the selection mechanism and target distribution (specifically, marginal statistics and a subset of selection variables).
2. Novel Upper Bound: The authors derive a theoretical upper bound on generalization error that explicitly models the selection process without requiring full knowledge of the underlying mechanism or the target distribution's joint density.
3. Tractable Estimation and Diagnostics: The paper provides a practical algorithm for nominating unobserved selection variables and a suite of diagnostic tools to help practitioners interpret the bound and detect assumption violations.
4. Empirical Validation: The method is validated across fully synthetic data, semi-synthetic data from the All of Us Research Program, and real-world selection bias in MIMIC-IV.

Results

The authors evaluate their method against baselines including naive Inverse Probability of Participation Weighting (IPPW), empirical calibration, entropy balancing, and raking.

• Validity and Tightness: In experiments with synthetic and All of Us data, the proposed bound (using the heuristic $U^*$) yielded valid upper bounds (i.e., $\hat{R}_Q \geq R_Q$) in 97% of tasks for All of Us and 90% for synthetic data. In contrast, the next best baseline (IPPW) achieved validity in only 82% of tasks and frequently underestimated the true error.
• Robustness: The method remained robust even when the number of observed selection variables $\tilde{U}$ was limited (e.g., $dim(\tilde{U})=1$).
• Real-World Application: In MIMIC-IV experiments, the method successfully identified generalization gaps for tasks like hypertension and Type 2 diabetes prediction. For the mortality prediction task (where the true target error was unknown), the method provided a non-vacuous bound consistent with prior literature suggesting single-site hospital data often generalizes poorly.
• Heuristic Performance: The moment-matching heuristic for identifying $U^*$ achieved F2 scores of 0.84 (synthetic) and 0.81 (All of Us), closely approximating the performance of an oracle that knows the true $U^*$.

Significance and Claims

The authors position this work as a practical tool for model auditing rather than a perfect solution. They claim that their method enables practitioners to:

• Estimate the impact of selection bias in settings where target data is otherwise intractable.
• Make informed deployment decisions: If the estimated upper bound is small and assumption diagnostics are satisfied, deployment may proceed with confidence. If the bound is large or diagnostics indicate violations, practitioners are guided to delay deployment, collect more representative data, or apply debiasing techniques.
• Avoid the "false confidence" often provided by existing reweighting methods that tend to underestimate generalization error.

The paper concludes by noting limitations, including a focus on low-dimensional categorical data (though extensions to continuous data are discussed) and the inherent difficulty of density estimation. However, the authors assert that the proposed framework offers a principled step toward safer and more generalizable medical prediction models.