Wasserstein-regularized Conformal Prediction under General Distribution Shift

TL;DR

This paper introduces a Wasserstein distance-based method to improve conformal prediction under general distribution shifts, effectively reducing coverage gaps and producing smaller, more reliable prediction sets.

Contribution

It proposes a novel Wasserstein distance-based upper bound for coverage gaps and develops an algorithm that reduces this gap under distribution shifts using importance weighting and regularized learning.

Findings

Reduces coverage gaps to 3.2% across datasets

Produces prediction sets 37% smaller than worst-case methods

Effectively handles general distribution shifts

Abstract

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 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.