Multi-Distribution Robust Conformal Prediction

2026-01-06

arXiv.org

AI Summary

This work addresses conformal prediction under multi-source distributional shifts, where test samples may originate from any source distribution or an arbitrary mixture thereof, necessitating prediction sets that guarantee coverage uniformly across all distributions. The authors propose a max-p aggregation mechanism that combines conformity scores from individual source distributions to construct a single, universally valid prediction set. They further introduce an optimization algorithm to learn a compact conformity function, thereby reducing prediction set size. This approach is the first to achieve uniform finite-sample coverage guarantees across multiple distributions, with theoretical results establishing the optimality and tightness of the max-p aggregation. The framework naturally extends to settings involving fairness, subgroup shift, and multi-source learning. Experiments on both synthetic and real-world data demonstrate that the method significantly narrows prediction sets while maintaining worst-case coverage, achieving performance nearly on par with standard single-source conformal methods.

Technology Category

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

Application Category

• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Security and Privacy: Large-scale security measurements

Abstract

In many fairness and distribution robustness problems, one has access to labeled data from multiple source distributions yet the test data may come from an arbitrary member or a mixture of them. We study the problem of constructing a conformal prediction set that is uniformly valid across multiple, heterogeneous distributions, in the sense that no matter which distribution the test point is from, the coverage of the prediction set is guaranteed to exceed a pre-specified level. We first propose a max-p aggregation scheme that delivers finite-sample, multi-distribution coverage given any conformity scores associated with each distribution. Upon studying several efficiency optimization programs subject to uniform coverage, we prove the optimality and tightness of our aggregation scheme, and propose a general algorithm to learn conformity scores that lead to efficient prediction sets after the aggregation under standard conditions. We discuss how our framework relates to group-wise distributionally robust optimization, sub-population shift, fairness, and multi-source learning. In synthetic and real-data experiments, our method delivers valid worst-case coverage across multiple distributions while greatly reducing the set size compared with naively applying max-p aggregation to single-source conformity scores, and can be comparable in size to single-source prediction sets with popular, standard conformity scores.

Problem

• Conformal Prediction
• Distribution Robustness
• Multi-Source Learning
• Fairness
• Coverage Guarantee

Innovation

• multi-distribution robustness
• conformal prediction
• max-p aggregation
• uniform coverage
• conformity score learning