📊 statistics

Multi-Distribution Robust Conformal Prediction

This paper proposes a max-p aggregation scheme and an associated learning algorithm for multi-distribution robust conformal prediction that guarantees uniform coverage across heterogeneous source distributions while significantly improving prediction set efficiency compared to naive approaches.

Original authors: Yuqi Yang, Ying Jin

Published 2026-07-10

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yuqi Yang, Ying Jin

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a detective trying to predict the future, but you have a weird problem: you have case files from three different cities, and each city has its own unique style of crime. One city has mostly pickpockets, another has mostly burglars, and the third has a mix of everything. You build a prediction tool to guess what crime will happen next.

The old way of doing this was to build a separate tool for each city. But what if a new case comes in, and you don't know which city it's from? If you use the "pickpocket tool" for a burglary, you'll be wrong. If you use the "burglar tool" for a pickpocket, you'll be wrong.

This paper introduces a new detective tool called MDCP (Multi-Distribution Conformal Prediction). Its superpower is that it builds one single prediction set that is guaranteed to be right, no matter which city the new case comes from. It doesn't need to know the city's name to do its job.

The "Too Big" Problem

Usually, to be safe, you might just take the prediction from the pickpocket city, the prediction from the burglar city, and the prediction from the mixed city, and smash them all together into one giant bag of possibilities. The paper calls this "naive aggregation."

The authors show that this giant bag is often way too big. It's like carrying a backpack full of every possible tool just in case you need one. It's safe, but it's clumsy and useless because it's so heavy. In their simulations, this "naive" method created prediction sets that were 34.39% larger than necessary in classification tasks and 22% wider in regression tasks.

The Magic Trick: The "Max-p" Score

The paper's main finding is a clever way to shrink that giant backpack without losing any safety. They propose a method called max-p aggregation.

Think of it like a group of judges. Each judge (representing a different city) gives a score on how likely a prediction is. Instead of averaging their scores, the new method says: "We only care about the judge who is most skeptical." If even the most skeptical judge thinks a prediction is possible, then we include it. If the most skeptical judge says "No way," we leave it out.

This sounds simple, but the paper proves mathematically that this specific way of combining the judges' opinions guarantees that your prediction set will be valid for every single city, even if the new case comes from a city you've never seen before.

Making it Efficient: The "Smart Score"

Just using the "most skeptical judge" rule still leaves the backpack a bit heavy. The authors realized that if you teach the judges to speak a common language (a shared "conformity score"), they can agree on a much smaller, tighter bag of possibilities.

They developed an algorithm that learns this shared language. In their experiments, this smart learning made the prediction sets almost as small as the best possible single-city tool, but with the safety of covering all cities.

What They Ruled Out

The paper explicitly argues against two common ideas:

Ignoring the source: You cannot just train one model on all the data mixed together and expect it to work for every specific city. The paper shows that if the test data comes from a specific source (like a specific hospital or region), a standard model often fails to cover the true outcome.
Needing to know the source at the end: Many fairness tools require you to know the group identity (like race or location) after the prediction is made to adjust the result. The paper argues this is impossible in sensitive scenarios where that information is hidden or protected. Their method works without ever seeing the group label for the new test point.

How Sure Are They?

The authors are very confident in the safety of their method. They have a mathematical proof that guarantees the prediction set will cover the true answer at least 90% of the time (for a significance level of $\alpha = 0.1$ ) across all sources, even with a small amount of data. This is a hard guarantee, not just a guess.

However, their claims about efficiency (how small the bag is) are based on simulations and real-world data tests, not a mathematical proof that it is the absolute smallest possible bag.

In their simulations with 3 sources and 2,000 samples per source, their method reduced the size of the prediction set significantly compared to the "naive" giant bag.
In real-world tests using satellite images (FMoW dataset), poverty maps (PovertyMap dataset), and medical records (MEPS dataset), their method consistently delivered tight coverage while keeping the prediction sets small.
They found that in some cases, their method produced sets even smaller than the single-source baselines, which is a surprising and helpful result.

The Bottom Line

The paper suggests that by using a "max-p" strategy combined with a smart learning algorithm, we can build prediction tools that are universally safe for mixed environments but efficient enough to be actually useful. It solves the problem of "which city is this from?" by building a tool that works for every city at once, without needing to know the answer beforehand.

Technical Summary: Multi-Distribution Robust Conformal Prediction

Problem Statement

The paper addresses the challenge of constructing prediction sets that maintain valid coverage guarantees when test data may originate from an arbitrary member or mixture of $K$ heterogeneous source distributions. In many real-world scenarios—such as clinical risk prediction across different hospitals, fairness across protected groups without observing group labels, or multi-source learning—the test distribution is unknown and may differ significantly from the training distribution.

Standard conformal prediction, when calibrated on pooled data, only guarantees coverage for the specific mixture distribution induced by the training sample. Conversely, calibrating on a single source fails to guarantee coverage when the test point comes from a different source. The authors formalize the goal as constructing a single prediction set $\hat{C}(X_{n+1})$ that achieves uniform coverage:
$\min_{k \in [K]} \mathbb{P}_{D \times P^{(k)}} \left( Y_{n+1} \in \hat{C}(X_{n+1}) \right) \ge 1 - \alpha$
This ensures the coverage guarantee holds simultaneously for all possible source distributions, regardless of which one generated the test point.

Methodology: Multi-Distribution Conformal Prediction (MDCP)

The proposed framework, MDCP, consists of two main components: a validity mechanism and an efficiency optimization strategy.

1. Max-p Aggregation for Validity

The foundation of the method is a distribution-free aggregation scheme.

Single-Source P-values: For each source $k$ , a conformity score $s_k$ is computed. A conformal p-value $p^{(k)}(y)$ is derived for a candidate label $y$ based on the calibration fold of source $k$ .
Aggregation: The method aggregates these p-values by taking their maximum: $p(y) = \max_{k \in [K]} p^{(k)}(y)$ .
Prediction Set: The final prediction set is the set of labels where the aggregated p-value exceeds the significance level $\alpha$ : $\hat{C}(X) = \{ y : p(y) \ge \alpha \}$ .
Theoretical Guarantee: The authors prove that this aggregated set is exactly the union of the single-source conformal sets. Consequently, it provides finite-sample uniform validity for any mixture of the source distributions, without requiring knowledge of the test point's source.

2. Efficiency Optimization via Dual Formulation

While the max-p aggregation ensures validity, naively applying it to standard single-source scores often results in overly conservative (large) prediction sets. To address this, the paper derives the population-level optimal prediction set that minimizes size subject to uniform coverage constraints.

Optimal Score: The authors show that the optimal prediction set corresponds to the superlevel set of a "shared" conformity score $h^*(x, y) = \sum_{k=1}^K \lambda_k^*(x) f_k(y|x)$ , where $f_k$ is the conditional density of source $k$ and $\lambda_k^*(x)$ are non-negative dual weights.
Learning the Weights: The weights $\lambda_k^*(x)$ are determined by solving a dual optimization problem that balances the coverage constraints across sources.
Empirical Algorithm: The paper proposes an end-to-end procedure to learn these weights:
1. Fit conditional density models (or classifiers) for each source and a pooled model.
2. Parameterize the weight functions $\lambda_k(x)$ (e.g., using splines or neural networks).
3. Optimize an empirical dual objective (derived from the Lagrangian of the population problem) to learn the parameters $\hat{\Theta}$ .
4. Construct the final MDCP set using the learned scores and the max-p aggregation.

Key Contributions

General Max-p Aggregation Scheme: Introduction of a simple, distribution-free mechanism that achieves finite-sample uniform coverage across multiple heterogeneous sources by aggregating single-source conformal p-values.
Characterization of Optimality: Theoretical derivation of the population-optimal prediction sets under uniform coverage constraints. The authors prove that the max-p aggregation is asymptotically optimal and tight when paired with the correctly learned conformity scores (specifically, the optimal shared score $h^*$ ).
End-to-End Algorithm: A practical pipeline for learning the optimal conformity scores for both classification and regression tasks, utilizing empirical risk minimization on a dual objective.
Empirical Validation: Extensive experiments demonstrating that MDCP achieves tight worst-case coverage while significantly reducing prediction set sizes compared to naive aggregation baselines.

Results

The paper evaluates MDCP on synthetic data and three real-world datasets:

Synthetic Experiments: Across linear, nonlinear, and varying temperature (heterogeneity) settings for both classification and regression, MDCP consistently maintains worst-case coverage close to the nominal level ( $1-\alpha$ ). In contrast, single-source baselines suffer from under-coverage on shifted data, and naive aggregation (BASELINE-AGG) is overly conservative. MDCP reduces prediction set sizes by approximately 34% in classification and 22% in regression compared to naive aggregation, often approaching the efficiency of an oracle with access to true densities.
Real-World Applications:
- FMoW (Satellite Imagery): Demonstrates robustness against geographic shifts across six regions. MDCP provides valid coverage across all regions with set sizes smaller than some single-source sets.
- PovertyMap: Addresses urban-rural shifts. MDCP maintains tight coverage where single-source methods fail (coverage as low as 0.6) and naive aggregation is too wide.
- MEPS (Medical Expenditure): Handles heterogeneity across sensitive groups (race) without requiring group labels at test time. MDCP achieves equalized coverage and shorter intervals compared to baselines.

Significance and Claims

The paper claims that MDCP offers a robust solution for uncertainty quantification in heterogeneous environments, bridging the gap between distributional robustness and fairness.

Fairness without Protected Attributes: The method enables group-conditional coverage guarantees (ensuring reliability for all subgroups) without needing to observe the sensitive group attribute at inference time, addressing a critical limitation in fair machine learning.
Subpopulation Shift Protection: It provides guarantees against arbitrary subpopulation shifts, ensuring reliability even when the test distribution is a mixture of sources different from the training mixture.
Efficiency-Validity Trade-off: The work demonstrates that it is possible to achieve uniform validity without the severe efficiency penalties typically associated with worst-case robustness, provided one learns an appropriate shared conformity score.

The authors position MDCP as a practical alternative to standard conformal prediction in multi-source settings, offering exact finite-sample guarantees and asymptotic optimality under mild consistency conditions.