Multi-Distribution Robust Conformal Prediction

AI summary

Overview

Research area: Conformal prediction, distributionally robust uncertainty quantification, multi-source (federated) learning, subpopulation shift, and group fairness in machine learning.

Technical level: Advanced. The paper is heavily theoretical, built on measure-theoretic population optimization, Lagrangian duality, complementary slackness, and sieve-based consistency arguments, with practical algorithms and experiments on top.

One-sentence scope: The paper introduces Multi-Distribution Conformal Prediction (MDCP), a framework that produces a single prediction set with finite-sample coverage guaranteed simultaneously across K heterogeneous source distributions, without observing which source a test point comes from.

What This Paper Is About

Standard conformal prediction calibrated on data pooled from several environments only guarantees coverage for the mixture distribution implied by the training sample, not for each environment individually. MDCP instead asks for a prediction set whose coverage exceeds 1 − α for every source distribution (and therefore for any mixture of them), which is the guarantee needed when test data may align with a single unknown member of the collection. The paper solves this by aggregating per-source conformal p-values with a max operation, then studies how to design conformity scores so the resulting set is as small as possible.

Key Contributions

1. A general max-p aggregation scheme. For each source, a conformal p-value is computed from any conformity score; taking the maximum p-value and inverting it yields a prediction set that is exactly the union of the per-source conformal sets and delivers finite-sample uniform coverage.

2. Population-level optimality and tightness characterization. The authors solve a pointwise size-minimization program under uniform conditional coverage, derive its dual, and prove that max-p aggregation, when paired with scores converging to the optimal dual-based score, is asymptotically equivalent to the oracle-optimal set and achieves (nearly) exact 1 − α coverage for at least one source.

3. An end-to-end learning pipeline. A dual-inspired empirical objective is used to learn the source weights λ(·), with concrete instantiation for classification (softplus-parameterized basis coefficients) and regression (Gaussian conditional density with gradient-boosted mean and scale functions).

4. Empirical validation. Synthetic and real-data experiments (described as satellite imagery and medical service datasets) show valid worst-case coverage while greatly reducing set size relative to naively applying max-p aggregation to single-source conformity scores.

Main Findings

• Max-p aggregation is exactly the union of single-source sets. By construction, the aggregated set equals ⋃{k=1}^{K} Ĉ^(k)(X{n+1}), and the paper proves in Theorem 1 that it satisfies P(Y_{n+1} ∈ Ĉ(X_{n+1})) ≥ 1 − α for any mixture P = Σ_k π_k P^(k) with weights summing to 1, which implies uniform coverage as a special case.

• The optimal set is a superlevel set of a shared score. Theorem 2 shows that for each x, the smallest set meeting uniform conditional coverage is C*(x) = {y : h_{λ*}(x, y) > 1} ∪ S(x), where h_{λ*}(x, y) = Σ_{k=1}^{K} λ_k*(x) f_k(y | x). The dual weights λ_k*(x) identify which sources are locally hardest to cover.

• Complementary slackness identifies which sources are binding. If λ_k*(x) > 0, the k-th source's conditional coverage equals 1 − α exactly; if λ_k*(x) = 0, its coverage is at least 1 − α. At least one λ_{k*}*(x) is strictly positive, which is the precise sense in which the construction is tight.

• Max-p is asymptotically oracle-optimal. Under the boundedness condition sup_{x,y} f_k(y|x) ≤ B_k, consistent density estimators, and uniform consistency of λ̂, Theorem 3 gives limsup | |Ĉ^(n)| − |C*| | ≤ ρ(T), where T = {(x, y) : h*(x, y) = 1}. When ρ(T) = 0, the symmetric difference between the estimated and optimal sets converges to zero in probability.

• The integrated dual objective recovers the right weights. Theorem 5 (informal) states that for any ν̃ covering the support of P^(k)(X), the maximizer of the integrated dual Φ(λ) coincides with the pointwise dual solution λ*(x) from Theorem 2, justifying the empirical objective used in practice.

• The ERM fitting procedure is consistent. Theorem 6 (informal) states that under smoothness conditions, a sieve approximation, and consistent density and pooled-density estimators, the empirical risk minimizer satisfies ||λ̂ − λ*||_∞ = o_P(1).

• Misspecification only costs efficiency, not validity. The estimation error of the density and dual-weight estimators affects the gap between MDCP set size and the oracle set size, but finite-sample uniform coverage is preserved regardless because it comes from the max-p aggregation itself.

• Empirical performance. In the synthetic and real-data experiments, MDCP attains valid worst-case coverage while greatly reducing set size compared to naively applying max-p aggregation to single-source conformity scores, and can be comparable in size to single-source prediction sets. (Specific numeric results, dataset names, and sample sizes are not reported in the available paper content.)

Methodology in Plain English

The paper starts from a simple idea. Split the labeled data from each of the K sources into a training fold and a calibration fold. On each training fold, fit whatever conformity score is convenient. On each calibration fold, compute a conformal p-value for every candidate label y at the test point — this is the usual rank of the test score among the calibration scores. A single prediction set valid for source k is obtained by keeping all y whose p-value clears α. To get a set valid across all sources, take, for each y, the largest p-value across the K sources, and keep y if that maximum clears α. Because a maximum clears the threshold only if some source's p-value does, the resulting set is exactly the union of the K single-source sets — and a union is never less valid than any of its parts, which is what gives the finite-sample uniform guarantee.

The rest of the paper is about making that union small. The authors write down, for each fixed covariate value x, the problem of selecting the smallest subset of label space that captures at least 1 − α conditional probability under every source. This is a linear program whose decision variable is the set indicator. Solving its dual produces nonnegative weights λ_k(x) on the sources and a combined density h(x, y) = Σ_k λ_k(x) f_k(y|x). The optimal set turns out to be the region where this combined density exceeds 1. In plain terms, the weights decide how much each source's density gets to "pull" the boundary, and sources that are locally harder to cover get pulled more.

Since the true densities and true weights are unknown, the authors estimate them. They parameterize λ_k(x) with basis functions (splines or neural network features) passed through a softplus to keep values nonnegative, plug in fitted conditional densities, and minimize an empirical version of the dual objective — a hinge-like term penalizing when the combined score falls below 1, divided by the pooled density, plus a (1 − α) weighted sum of the λ's. This is a fairly standard ERM problem, solvable with gradient methods. In classification the densities are just fitted class probabilities; in regression the paper models the response as a location-scale Gaussian and fits the mean and standard deviation with gradient boosting. Finally, the learned combined score is used as the single conformity score for every source, and the max-p recipe is applied on the calibration folds. For regression, thresholding that score does not guarantee an interval, but with Gaussian components the set is a union of finitely many intervals, so a grid search gives an efficient superset.

Why This Matters

Impact on research. The paper gives a clean theoretical answer to an open question in multi-source conformal inference: what is the smallest prediction set that can be simultaneously valid for every source when test membership is unknown? It connects conformal prediction to the dual structure of group-wise distributionally robust optimization, and shows the max-p operation is not merely a safe heuristic but asymptotically the optimal aggregation. Prior multi-source and federated conformal work studied communication-limited calibration, borrowing strength across sources, and aggregating source-specific sets; prior group-conditional and multi-valid conformal work requires knowing the subgroup identity at inference. MDCP removes that requirement.

Real-world applications (drawn from the paper's own motivating scenarios):

• Clinical risk prediction across hospitals. A model trained on records from several sites must remain reliable when a new patient's record comes from one of these sites, without knowing which one.
• Fairness without protected attributes. When each P^(k) is a sensitive group, uniform coverage gives group-conditional coverage while avoiding reliance on race, gender, or other protected attributes that may be unavailable or sensitive at test time.
• Subpopulation shift. Because any test distribution of the form Σ_k π'_k P^(k) is covered, the method protects against arbitrary reweightings of subpopulations between training and deployment — for example when hospital or demographic mixture proportions change.
• Multi-source scientific and engineering data. Datasets collected under different protocols or environments can be aggregated with guarantees for each individual source, not just for the pooled mixture.

The paper reports experiments on satellite imagery and medical service datasets as the real-data applications.

Industry relevance. Any deployed model served across heterogeneous populations — clinical decision support, credit and insurance, content moderation, autonomous systems operating across regions — faces the question of whether its uncertainty guarantees survive the shift from training mixture to deployment mixture. MDCP provides a deployment-safe procedure that does not require knowing the deployment environment at inference, and its implementation uses off-the-shelf classifiers, gradient boosting, and neural network features rather than bespoke machinery.

Future Directions

1. Behavior when the boundary set has positive measure. Theorem 3 bounds the size gap by ρ(T), the measure of the tie region {(x, y) : h*(x, y) = 1}. Remark 4 characterizes limiting sets along subsequences in that case, but a practical recipe for choosing S*(x) so that |Ĉ^(n)| converges to |C*| even when ρ(T) > 0 is left open. In classification, where |T(x)| is non-negligible under the count measure, this directly affects set size.

2. Relaxing the consistency conditions. The asymptotic optimality requires consistent conditional density estimation for every source and uniform consistency of the learned weights, under the bounded density assumption sup_{x,y} f_k(y|x) ≤ B_k. How graceful the degradation is when density estimation is misspecified, or when one source has a small calibration fold, is not established.

3. Computation at scale. The regression set is only guaranteed to be a union of finitely many intervals, and the paper computes a superset via grid search with details deferred to an appendix. Whether the procedure scales to large label spaces or high-dimensional outputs — and how tight the grid-search superset is — remains to be examined.

4. Extensions beyond the studied settings. The paper focuses on classification and regression with a location-scale Gaussian working model for regression. Broader applications such as structured outputs, survival or time-to-event labels, and settings where the number of sources K grows with the sample size are natural next questions.

Target Audience

This paper is aimed at researchers and graduate students in conformal prediction, distributionally robust learning, and trustworthy machine learning, as well as statisticians working on uncertainty quantification under distribution shift. Practitioners deploying models across heterogeneous populations — clinical informatics teams, fairness and responsible-AI groups, and engineers working with multi-site or multi-protocol datasets — will find the algorithm and its validity guarantee directly useful, but should be comfortable with conformal prediction notation and optimization duality to follow the theoretical development.

Authors’ 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.