Multi-Source Conformal Inference Under Distribution Shift

Abstract

Recent years have experienced increasing utilization of complex machine learning models across multiple sources of data to inform more generalizable decision-making. However, distribution shifts across data sources and privacy concerns related to sharing individual-level data, coupled with a lack of uncertainty quantification from machine learning predictions, make it challenging to achieve valid inferences in multi-source environments. In this paper, we consider the problem of obtaining distribution-free prediction intervals for a target population, leveraging multiple potentially biased data sources. We derive the efficient influence functions for the quantiles of unobserved outcomes in the target and source populations, and show that one can incorporate machine learning prediction algorithms in the estimation of nuisance functions while still achieving parametric rates of convergence to nominal coverage probabilities. Moreover, when conditional outcome invariance is violated, we propose a data-adaptive strategy to upweight informative data sources for efficiency gain and downweight non-informative data sources for bias reduction. We highlight the robustness and efficiency of our proposals for a variety of conformal scores and data-generating mechanisms via extensive synthetic experiments. Hospital length of stay prediction intervals for pediatric patients undergoing a high-risk cardiac surgical procedure between 2016-2022 in the U.S. illustrate the utility of our methodology.

Figure 4:

Boxplots of prediction interval widths

Figure 5:

Boxplots of coverage probability, under homogeneous covariate distributions

Figure 6:

Boxplots of prediction interval width, under homogeneous covariate distributions

Figure 7:

Boxplots of coverage probability, under weakly heterogeneous covariate distributions

Figure 8:

Boxplots of prediction interval width, under weakly heterogeneous covariate distributions

Figure 9:

Boxplots of coverage probability, under strongly heterogeneous covariate distributions

Figure 10:

Boxplots of prediction interval width, under strongly heterogeneous covariate distributions

Figure 11:

Local coverages, under CCOD is strongly violated and strongly heterogeneous covariate distributions and $n_k=3000$

Figure 12:

Weights vs. ${\chi }_k^2$ values, using $n_k=3000$ data under heteroscedasticity. The green points are by Federated I, the orange points are by Federated II (ours), the blue points are by Federated III, and the red dashed lines are for a reference line weights = 0.2.

Figure 13:

Comparison of coverage probabilities and average interval width when modifying the propensity score of observing the outcome between (0.4,0.6) (panel (a)) and (0.1,0.9) (panel (b)).

Figure 1:

Illustration of the proposed robust algorithm for multi-source conformal prediction. Each $\widehat{\theta }$ represented by a different color is the estimated $(1-\alpha )$-quantile of the conformal score using data from the site with the same color. ${\widehat{m}}_0$ (in red) is the estimated CDF of the conformal score using only the target site data. The other ${\widehat{m}}_k(k\ge 1)$ are the estimated CDFs of the conformal scores from source sites, and ${\widehat{\omega }}_{k,0}(k\ge 1)$ is the density ratio of site $k$ versus the target site. The federated ${\widehat{r}}_{\text{fed},0}$ is a weighted average of the site-specific quantiles, with weights given by $\widehat{\mathit{w}}$. The prediction interval ${\widehat{C}}_{\alpha }(\mathit{X})$ is the set of outcomes $y$ such that the corresponding conformal scores $S(\mathit{x},y)$ in the target are below the threshold ${\widehat{r}}_{\text{fed,}0}$.

Figure 2:

A: Marginal coverage, B: Prediction interval width, C: Conditional coverage, and D: Weights for our proposed federated method compared to the pooled sample and target only methods, where sample size $n_k=3000$, $k=0,\dots ,4$ under strongly heterogeneous covariate distributions and strong violation of CCOD.

Figure 3:

Each panel represents the prediction intervals for hospital LOS for a randomly selected individual following a Norwood procedure across $\alpha =\{0.1,0.2,0.3,0.4,0.5\}$ and conformal score $\in \{\mathrm{A}\mathrm{S}\mathrm{R},\mathrm{l}\mathrm{o}\mathrm{c}\mathrm{a}\mathrm{l}\mathrm{A}\mathrm{S}\mathrm{R},\mathrm{C}\mathrm{Q}\mathrm{R}\}$ for A: a patient in South, B: a patient in Midwest, C: a patient in West, D: a patient in Northeast.