Multi-Distribution Robust Conformal Prediction

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.

Prediction models are often trained or calibrated using data from several different sources, such as different regions, institutions, or demographic groups. A key challenge is that a new data point may come from any of these sources, or from a mixture of them, and the system may not know which source it belongs to. This paper develops a way to report a set of possible answers with a reliability guarantee that new data points fall into the proposed set, and this guarantee holds across all the specified sources. The method combines information from the different sources and learns to make the set of possible answers as small as possible while maintaining the same guarantee. Experiments on simulated and real datasets show that the proposed approach maintains reliable coverage across heterogeneous populations while producing much smaller prediction sets than more naive methods.