# Wasserstein-regularized Conformal Prediction under General Distribution Shift

arXiv: https://arxiv.org/abs/2501.13430 · HTML: https://mlsanity.com/paper/2501.13430/

## Key facts

- Cited by 15 papers in the map (data through 14 Oct 2026).
- Builds on 34 of its 57 references that are in the map.
- First posted on arXiv 23 Jan 2025.

## Builds on (most cited first, 30 of 34)

- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [A Gentle Introduction to Conformal Prediction and Distribution-Free Uncertainty Quantification](https://mlsanity.com/paper/2107.07511/) (2021)
- [Out-of-Distribution Generalization via Risk Extrapolation (REx)](https://mlsanity.com/paper/2003.00688/) (2020)
- [Conformalized Quantile Regression](https://mlsanity.com/paper/1905.03222/) (2019)
- [Conformal Prediction Under Covariate Shift](https://mlsanity.com/paper/1904.06019/) (2019)
- [Classification with Valid and Adaptive Coverage](https://mlsanity.com/paper/2006.02544/) (2020)
- [Adaptive Conformal Inference Under Distribution Shift](https://mlsanity.com/paper/2106.00170/) (2021)
- [Conformal prediction beyond exchangeability](https://mlsanity.com/paper/2202.13415/) (2022)
- [Conformal Risk Control](https://mlsanity.com/paper/2208.02814/) (2022)
- [Statistical Aspects of Wasserstein Distances](https://mlsanity.com/paper/1806.05500/) (2018)
- [Conformal Prediction With Conditional Guarantees](https://mlsanity.com/paper/2305.12616/) (2023)
- [Traffic Graph Convolutional Recurrent Neural Network: A Deep Learning Framework for Network-Scale Traffic Learning and Forecasting](https://mlsanity.com/paper/1802.07007/) (2018)
- [Conformal Inference for Online Prediction with Arbitrary Distribution Shifts](https://mlsanity.com/paper/2208.08401/) (2022)
- [Domain Adaptation by Using Causal Inference to Predict Invariant Conditional Distributions](https://mlsanity.com/paper/1707.06422/) (2017)
- [A Review and Comparative Study on Probabilistic Object Detection in Autonomous Driving](https://mlsanity.com/paper/2011.10671/) (2020)
- [Localized Conformal Prediction: A Generalized Inference Framework for Conformal Prediction](https://mlsanity.com/paper/2106.08460/) (2021)
- [An introduction to domain adaptation and transfer learning](https://mlsanity.com/paper/1812.11806/) (2018)
- [Learning Optimal Conformal Classifiers](https://mlsanity.com/paper/2110.09192/) (2021)
- [Robust Validation: Confident Predictions Even When Distributions Shift](https://mlsanity.com/paper/2008.04267/) (2020)
- [Batch Multivalid Conformal Prediction](https://mlsanity.com/paper/2209.15145/) (2022)
- [Training Uncertainty-Aware Classifiers with Conformalized Deep Learning](https://mlsanity.com/paper/2205.05878/) (2022)
- [Improving Conditional Coverage via Orthogonal Quantile Regression](https://mlsanity.com/paper/2106.00394/) (2021)
- [Federated Conformal Predictors for Distributed Uncertainty Quantification](https://mlsanity.com/paper/2305.17564/) (2023)
- [KL Guided Domain Adaptation](https://mlsanity.com/paper/2106.07780/) (2021)
- [Conformal Prediction for Federated Uncertainty Quantification Under Label Shift](https://mlsanity.com/paper/2306.05131/) (2023)
- [One-Shot Federated Conformal Prediction](https://mlsanity.com/paper/2302.06322/) (2023)
- [Provably Robust Conformal Prediction with Improved Efficiency](https://mlsanity.com/paper/2404.19651/) (2024)
- [Adaptive conformal classification with noisy labels](https://mlsanity.com/paper/2309.05092/) (2023)
- [Split Localized Conformal Prediction](https://mlsanity.com/paper/2206.13092/) (2022)
- [Adaptive Conformal Prediction by Reweighting Nonconformity Score](https://mlsanity.com/paper/2303.12695/) (2023)

## Built on by (most cited first, 15 of 15)

- [Conformal Prediction under Levy-Prokhorov Distribution Shifts: Robustness to Local and Global Perturbations](https://mlsanity.com/paper/2502.14105/) (2025)
- [Conformal Calibration: Ensuring the Reliability of Black-Box AI in Wireless Systems](https://mlsanity.com/paper/2504.09310/) (2025)
- [Non-exchangeable Conformal Prediction with Optimal Transport: Tackling Distribution Shifts with Unlabeled Data](https://mlsanity.com/paper/2507.10425/) (2025)
- [Conformal Prediction Meets Long-tail Classification](https://mlsanity.com/paper/2508.11345/) (2025)
- [Weight Clipping for Robust Conformal Inference under Unbounded Covariate Shifts](https://mlsanity.com/paper/2605.02072/) (2026)
- [Conformal Prediction Beyond the Horizon: Distribution-Free Inference for Policy Evaluation](https://mlsanity.com/paper/2510.26026/) (2025)
- [Conformal Prediction via Transported Beta Laws](https://mlsanity.com/paper/2605.19024/) (2026)
- [Coverage Guarantees for Pseudo-Calibrated Conformal Prediction under Distribution Shift](https://mlsanity.com/paper/2602.14913/) (2026)
- [Conformalized Large Language Models under Configuration Shift](https://mlsanity.com/paper/2608.01460/) (2026)
- [Power Interpretable Causal ODE Networks: A Unified Model for Explainable Anomaly Detection and Root Cause Analysis in Power Systems](https://mlsanity.com/paper/2602.12592/) (2026)
- [Robust Conditional Conformal Prediction via Branched Normalizing Flow](https://mlsanity.com/paper/2605.01868/) (2026)
- [DistMatch: Adaptive Binning via Distribution Matching for Robust Sequential Conformal Prediction](https://mlsanity.com/paper/2606.00690/) (2026)
- [Multi-source conformal prediction: leveraging heterogeneity via localization](https://mlsanity.com/paper/2609.14531/) (2026)
- [Conformal Bayes under Continuous Label Shift: Sensitivity Analysis and the Limits of Exact Validity](https://mlsanity.com/paper/2609.27976/) (2026)
- [Sketched Calibration for Conformal Prediction under Covariate Shift](https://mlsanity.com/paper/2610.09208/) (2026)

## Abstract (arXiv)

Conformal prediction yields a prediction set with guaranteed 1-α coverage of the true target under the i.i.d. assumption, which may not hold and lead to a gap between 1-α 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 α. 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.