Wasserstein-regularized Conformal Prediction under General Distribution Shift

TL;DR

This work tackles the degradation of conformal prediction guarantees under joint distribution shift by deriving a Wasserstein-distance–based upper bound on the coverage gap between calibration and test conformal score distributions. By introducing pushforward measures, the bound is decomposed into covariate-shift and concept-shift components, enabling targeted minimization via importance weighting and regularized representation learning. The proposed WR-CP framework integrates an ERM objective with a Wasserstein regularization term across multiple sources, providing finite-sample guarantees and a tunable trade-off between accuracy and efficiency. Empirical results on six datasets show substantial reductions in coverage gaps (to about $3.2\%$) and notably smaller prediction sets (approximately $37\%$ smaller than worst-case baselines) while offering flexible adaptivity to unseen mixtures of source distributions.

Abstract

Conformal prediction yields a prediction set with guaranteed $1-α$ coverage of the true target under the i.i.d. assumption, which may not hold and lead to a gap between $1-α$ and the actual coverage. Prior studies bound the gap using total variation distance, which cannot identify the gap changes under distribution shift at a given $α$. Besides, existing methods are mostly limited to covariate shift,while general joint distribution shifts are more common in practice but less researched.In response, we first propose a Wasserstein distance-based upper bound of the coverage gap and analyze the bound using probability measure pushforwards between the shifted joint data and conformal score distributions, enabling a separation of the effect of covariate and concept shifts over the coverage gap. We exploit the separation to design an algorithm based on importance weighting and regularized representation learning (WR-CP) to reduce the Wasserstein bound with a finite-sample error bound.WR-CP achieves a controllable balance between conformal prediction accuracy and efficiency. Experiments on six datasets prove that WR-CP can reduce coverage gaps to $3.2\%$ across different confidence levels and outputs prediction sets 37$\%$ smaller than the worst-case approach on average.

Figures (20)

• Figure 1: (a) Joint distribution shift can include both covariate shift ($P_X\neq Q_X$) and concept shift ($f_P\neq f_Q$). Coverage gap (Eq. (\ref{['eq: coverage gap']})) is the absolute difference in cumulative probabilities of calibration and test conformal scores at the $1-\alpha$ quantile $\tau$. We address covariate-shift-induced Wasserstein distance by applying importance weighting tibshirani2019conformal to calibration samples, and further minimize concept-shift-induced Wasserstein distance to obtain accurate and efficient prediction sets; (b) $Q_V^{(1)}$ and $Q_V^{(2)}$ are two distinct test conformal score distributions. Wasserstein distance (Eq. (\ref{['eq: Wasserstein distance']})) integrates the vertical gap between two cumulative probability distributions overall all quantiles, and is sensitive to coverage gap changes at any quantile. Total variation distance fails to indicate coverage gap changes thoroughly as it is agnostic about where two distributions diverge.
• Figure 2: Pushforward measures.
• Figure 3: Comparison of vanilla CP, IW-CP, and WR-CP based on normalized Wasserstein distance between calibration and test conformal scores: IW-CP can only address the distance caused by covariate shift, while WR-CP reduces the distance from concept shift. The $\beta$ values for the WR-CP method are 9, 11, 9, 10, 13, and 13, respectively.
• Figure 4: Coverages and prediction set sizes of WR-CP and baselines with $\alpha=0.2$: WR-CP makes coverages on test data more concentrated around the $1-\alpha$ level compared to vanilla CP, IW-CP, and CQR. While WC-CP ensures coverage guarantees, it leads to inefficient predictions due to large set sizes, whereas WR-CP mitigates this inefficiency. The $\beta$ values for the WR-CP method are 4.5, 9, 9, 6, 8, and 20, respectively.
• Figure 5: Pareto fronts of coverage gap and prediction set size obtained from WR-CP with varying $\beta$: WR-CP effectively balances conformal prediction accuracy and efficiency, providing a flexible and customizable solution. When $\beta=0$, WR-CP returns to IW-CP.

Theorems & Definitions (16)

• Definition 1: Kolmogorov Distance
• Definition 2: $p$-Wasserstein Distance
• Proposition 1
• Definition 3: Pushforward Measure
• Theorem 1
• Theorem 2
• Definition 4: Upper Wasserstein Dimension
• Proposition 2
• Theorem 3
• Theorem 4