Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Multi-Distribution Robust Conformal Prediction

Yuqi Yang, Ying Jin

TL;DR

This work tackles reliable uncertainty quantification when test data may come from any of several heterogeneous source distributions. It introduces MDCP, a max-p aggregation framework that combines per-source conformal p-values to produce a single prediction set with finite-sample uniform coverage across all sources. The authors establish population-level optimality and asymptotic equivalence to oracle sets when source scores converge to an optimal shared score, and they deploy an end-to-end algorithm that learns conformity scores via a dual objective for both classification and regression. Through extensive simulations and real-data experiments in satellite imagery, poverty mapping, and healthcare, MDCP demonstrates improved efficiency (smaller prediction sets) while preserving worst-case coverage across distributions. The approach offers practical and theoretically grounded guarantees for distribution-robust conformal prediction in multi-source settings, with potential extensions to broader conformity scores and test-time group identification.

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.

Multi-Distribution Robust Conformal Prediction

TL;DR

This work tackles reliable uncertainty quantification when test data may come from any of several heterogeneous source distributions. It introduces MDCP, a max-p aggregation framework that combines per-source conformal p-values to produce a single prediction set with finite-sample uniform coverage across all sources. The authors establish population-level optimality and asymptotic equivalence to oracle sets when source scores converge to an optimal shared score, and they deploy an end-to-end algorithm that learns conformity scores via a dual objective for both classification and regression. Through extensive simulations and real-data experiments in satellite imagery, poverty mapping, and healthcare, MDCP demonstrates improved efficiency (smaller prediction sets) while preserving worst-case coverage across distributions. The approach offers practical and theoretically grounded guarantees for distribution-robust conformal prediction in multi-source settings, with potential extensions to broader conformity scores and test-time group identification.

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.
Paper Structure (72 sections, 10 theorems, 78 equations, 23 figures, 2 algorithms)

This paper contains 72 sections, 10 theorems, 78 equations, 23 figures, 2 algorithms.

Table of Contents

  1. Introduction
  2. Preview of results
  3. Related work
  4. Multi-source/distribution conformal prediction.
  5. Distributionally-robust conformal prediction.
  6. Group-conditional conformal prediction.
  7. Distribution robustness and multi-source/group learning.
  8. Max-p conformal prediction
  9. Optimality of max-p aggregation
  10. Size optimization under uniform validity
  11. Optimal prediction set under conditional validity.
  12. Asymptotic optimality of max-p aggregation
  13. Practical algorithms
  14. Optimizing scores via an empirical dual objective
  15. Example: sieve estimation.
  16. ...and 57 more sections

Key Result

Theorem 1

Let $\{p^{(k)}(y)\}_{k=1}^K$ and $p(y)$ be defined above. Then, the aggregated set equals the union of the per-source conformal sets: For an independent test point $(X_{n+1},Y_{n+1})\sim P$ with any mixture distribution $P = \sum_k \pi_k P^{(k)}$ and arbitrary weights $\sum_{k=1}^K \pi_k=1$, $\pi_k\geq 0$, the prediction set achieves valid coverage

Figures (23)

  • Figure 1: Prediction sets with uniform coverage need to balance the coverage across multiple distributions. (a) When one distribution has heavier tails, a valid prediction set $\hat{C}(X)$ may coincide with the larger one $\hat{C}_1(X)$. (b) When two distributions partially overlap, a uniformly valid prediction set $\hat{C}(X)$ sits in between two distributions and is longer than single-source sets $\hat{C}_1(X)$ and $\hat{C}_2(X)$. (c) MDCP achieves uniform coverage by jointly training a conformity score and aggregating multiple prediction sets from the trained score.
  • Figure 2: Performance of MDCP and baselines in the classification Linear experiments, where the bars represent the result of each method averaged over $N=100$ runs, and the dots represent the result in each run. Left: coverage over all test data. Middle: worst-case coverage over single-source test data. Right: average set size over all test data.
  • Figure 3: Performance of MDCP and baselines in the classification Nonlinear experiments. The $x$-axis is the setting of the nonlinear term $g(x)$, with the linear setting presented for comparison. The connected dots are average results colored by method, with the colored, dimmed dots being the results in each of the $N=100$ runs. Left: coverage over all test data. Middle: worst-case coverage over single-source test data. Right: average set size over all test data.
  • Figure 4: Performance of MDCP and baselines in the classification Temperature experiments. The $x$-axis is the temperature parameter $\tau$. Each line shows the results of a method averaged over $N=100$ runs, with shaded $\pm1$ standard deviation across runs. Left: coverage over all test data. Middle: worst-case coverage over single-source test data. Right: average set size over all test data.
  • Figure 5: Evaluation with regression Linear suites; details are otherwise the same as Figure \ref{['fig:iter_eval_class_overall_vanilla']}.
  • ...and 18 more figures

Theorems & Definitions (15)

  • Theorem 1: Finite-sample uniform validity
  • Theorem 2: $X$-conditional optimality
  • Theorem 3
  • Theorem 4: Informal
  • Example 5: Polynomials
  • Example 6: Splines
  • Theorem 9
  • Theorem 10: Marginal optimality
  • Definition 11: Lévy distance
  • Definition 12: Generalized quantile
  • ...and 5 more