Wasserstein-regularized Conformal Prediction under General Distribution Shift

2025-01-23

AI Summary

Addressing the failure of coverage guarantees in conformal prediction under distribution shifts—particularly joint covariate and concept shift—this work relaxes the standard i.i.d. assumption and, for the first time, derives an upper bound on coverage deviation based on the Wasserstein distance, enabling decoupled analysis of both shift types. We propose Wasserstein-regularized representation learning with importance-weighted conformal prediction (WR-CP), integrating pushforward probability measure theory with finite-sample error analysis. Evaluated on six benchmark datasets, WR-CP reduces coverage deviation to 3.1% and shrinks average prediction set size by 38% relative to the worst-case baseline, thereby achieving a significantly improved and controllable trade-off between predictive accuracy and efficiency.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Constraint Satisfaction and Optimization: Distributed CSP/Optimization
• Reasoning under Uncertainty: Uncertainty Representations

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems
• Responsible Web: Human-perceived consequences of algorithmic deployment on the web

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.1%$ across different confidence levels and outputs prediction sets 38$%$ smaller than the worst-case approach on average.

Problem

• Consistent Prediction
• Data Variability
• Prediction Accuracy

Innovation

• WR-CP Algorithm
• Wasserstein Method
• Dynamic Error Assessment