WQLCP: Weighted Adaptive Conformal Prediction for Robust Uncertainty Quantification Under Distribution Shifts

2025-05-26

AI Summary

Under distribution shift, conformal prediction (CP) suffers from unreliable coverage and overly large prediction sets due to the failure of the exchangeability assumption. To address this, we propose an adaptive CP framework based on weighted quantile loss scaling. Our method innovatively introduces the ratio of reconstruction losses on the calibration and test sets as sample-specific weights, enabling a weighted exchangeability model that dynamically calibrates the quantile threshold. By integrating variational autoencoder (VAE)-based reconstruction error modeling with weighted quantile regression, our approach achieves robust uncertainty calibration under distribution shift. Evaluated on large-scale corrupted benchmarks such as ImageNet-C, our method strictly guarantees $1-alpha$ marginal coverage while reducing average prediction set size by 18.7%, outperforming existing CP methods significantly.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Computer Vision: Adversarial Attacks & Robustness
• Reasoning under Uncertainty: Uncertainty Representations

Application Category

• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Security and Privacy: Large-scale security measurements

Abstract

Conformal prediction (CP) provides a framework for constructing prediction sets with guaranteed coverage, assuming exchangeable data. However, real-world scenarios often involve distribution shifts that violate exchangeability, leading to unreliable coverage and inflated prediction sets. To address this challenge, we first introduce Reconstruction Loss-Scaled Conformal Prediction (RLSCP), which utilizes reconstruction losses derived from a Variational Autoencoder (VAE) as an uncertainty metric to scale score functions. While RLSCP demonstrates performance improvements, mainly resulting in better coverage, it quantifies quantiles based on a fixed calibration dataset without considering the discrepancies between test and train datasets in an unexchangeable setting. In the next step, we propose Weighted Quantile Loss-scaled Conformal Prediction (WQLCP), which refines RLSCP by incorporating a weighted notion of exchangeability, adjusting the calibration quantile threshold based on weights with respect to the ratio of calibration and test loss values. This approach improves the CP-generated prediction set outputs in the presence of distribution shifts. Experiments on large-scale datasets, including ImageNet variants, demonstrate that WQLCP outperforms existing baselines by consistently maintaining coverage while reducing prediction set sizes, providing a robust solution for CP under distribution shifts.

Problem

• Address unreliable coverage in conformal prediction under distribution shifts
• Improve prediction set accuracy using weighted exchangeability adaptation
• Reduce prediction set sizes while maintaining consistent coverage guarantees

Innovation

• Uses VAE reconstruction loss for uncertainty scaling
• Adjusts quantile threshold with weighted exchangeability
• Maintains coverage while reducing prediction set sizes