Wasserstein-regularized Conformal Prediction under General Distribution Shift

Conformal prediction yields a prediction set with guaranteed $1-\alpha$ coverage of the true target under the i.i.d. assumption, which may not hold and lead to a gap between $1-\alpha$ 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 $\alpha$. Besides, existing methods are mostly limited to covariate shift,while general joint distribution shifts are more common in practice but lessthis http URLresponse, 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 errorthis http URL-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.

@article{xu2025_2501.13430,
  title={ Wasserstein-regularized Conformal Prediction under General Distribution Shift },
  author={ Rui Xu and Chao Chen and Yue Sun and Parvathinathan Venkitasubramaniam and Sihong Xie },
  journal={arXiv preprint arXiv:2501.13430},
  year={ 2025 }
}