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

TL;DR

This work addresses uncertainty quantification under distribution shifts for conformal prediction (CP). It introduces Reconstruction-Loss-Scaled CP (RLSCP), which uses a variational autoencoder’s reconstruction loss as a global uncertainty cue to scale CP scores, and Weighted Quantile Loss-scaled CP (WQLCP), which further refines calibration by weighting calibration samples according to reconstruction-loss ratios to adapt to test distributions. Experiments on large-scale datasets with diverse distribution shifts show that WQLCP preserves nominal coverage while substantially reducing prediction-set sizes, outperforming state-of-the-art baselines. The findings highlight transformer-based VAEs and test-distribution-aware calibration as effective ingredients for robust CP in non-exchangeable settings, with practical impact for reliable, efficient uncertainty quantification in vision systems.

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.

Figures (4)

• Figure 1: (a) CP coverage violations under distribution shifts. While transformer-based networks (ViTs, DeiTs) perform better compared to CNNs (ResNet50, ResNet152), they still fail to maintain coverage in distribution-shifted scenarios. The desired coverage level is set to 0.90. (b) Prediction set sizes: The diagram illustrates larger prediction set sizes for ImageNet variant datasets.
• Figure 2: Comparison of reconstruction loss and softmax mean of the true label across ImageNet, ImageNetV2, ImageNetR, and ImageNetA.
• Figure 3: Overview of the WQLCP framework. The model leverages VAEs to compute reconstruction losses for uncertainty estimation, applying weighted quantile and scaling test scores to refine conformal prediction.
• Figure 4: $\beta$-VAE ablation study: (Left) Reconstruction MSE loss vs KL divergence trade-off. (Right) Coverage vs $\beta$ on ImageNetA, showing optimal $\beta=1.2$ achieves 0.7402 coverage.