Distribution-Free, Risk-Controlling Prediction Sets

Stephen Bates; Anastasios Angelopoulos; Lihua Lei; Jitendra Malik; Michael I. Jordan

Source-linked AI summary

Distribution-Free, Risk-Controlling Prediction Sets

Stephen Bates, Anastasios Angelopoulos, Lihua Lei, Jitendra Malik, Michael I. Jordan

arXiv:2101.02703v3Jan 7, 2021cs.LGcs.AIcs.CVstat.MEstat.ML

TL;DR

Reliable uncertainty quantification is needed because black-box predictors can produce silent failures in consequential settings. The paper introduces risk-controlling prediction sets, which calibrate set-valued outputs using held-out data to control general loss with finite-sample guarantees. The framework is demonstrated across multiple machine-learning tasks and supports losses beyond binary coverage.

Problem
Black-box prediction systems are widely deployed without reliable uncertainty quantification, leaving prediction failures silent in consequential settings.
Method
The method calibrates set-valued predictions from black-box predictors to control a general loss with probability at least 1 − δ, using user-selected risk level α.
Results
The framework provides finite-sample risk control and is applied to classification with unequal penalties, multi-label and hierarchical classification, image segmentation, and 3D-structure prediction.
Takeaways & Limitations
RCPS represents uncertainty without retraining existing models and can support risk-sensitive decision-making, audits, and model-improvement protocols.
Takeaways & Limitations
The optimal predictor may require estimating the conditional response distribution, and its resulting sets may be difficult to compute.

Abstract

from arXiv · show

While improving prediction accuracy has been the focus of machine learning in recent years, this alone does not suffice for reliable decision-making. Deploying learning systems in consequential settings also requires calibrating and communicating the uncertainty of predictions. To convey instance-wise uncertainty for prediction tasks, we show how to generate set-valued predictions from a black-box predictor that control the expected loss on future test points at a user-specified level. Our approach provides explicit finite-sample guarantees for any dataset by using a holdout set to calibrate the size of the prediction sets. This framework enables simple, distribution-free, rigorous error control for many tasks, and we demonstrate it in five large-scale machine learning problems: (1) classification problems where some mistakes are more costly than others; (2) multi-label classification, where each observation has multiple associated labels; (3) classification problems where the labels have a hierarchical structure; (4) image segmentation, where we wish to predict a set of pixels containing an object of interest; and (5) protein structure prediction. Lastly, we discuss extensions to uncertainty quantification for ranking, metric learning and distributionally robust learning.

1 Introduction

Risk-controlling prediction sets modify black-box predictors to output interpretable sets of plausible responses while controlling expected loss at a user-specified level. The framework uses calibration and finite-sample guarantees to support uncertainty-aware decision-making across diverse prediction settings.

Black-box predictors often lack reliable uncertainty quantification, allowing prediction failures to remain silent in high-consequence settings.
Risk-controlling prediction sets (RCPS) modify black-box predictors to return plausible response sets that limit costly-error frequency at a user-chosen level.The sets provide distribution-free, finite-sample control of a general loss.
For MRI diagnosis, output sets encode differing mistake severities, with lower classifier accuracy or stricter risk requirements producing larger sets.The example assigns losses of 100 to the most severe mistake and 0.1 to the least severe mistake.
RCPS calibrates set-valued predictors so expected loss on future data remains below a desired level with high probability.The formal guarantee is that risk is at most α with probability at least 1 − δ.
RCPS uses data splitting: training data fit a predictive model, while held-out data calibrate a set-valued predictor based on that model.
The framework extends tolerance-region ideas to broader losses for set-valued predictors and covers classification, multi-label, hierarchical, segmentation, 3D-structure, ranking, and metric-learning problems.Its calibration procedure applies to any set-valued predictor under a specified monotonicity requirement.

2 Upper Confidence Bound Calibration

UCB calibration selects a nested set-valued predictor whose upper confidence bound on risk falls below the target level, yielding high-probability risk control. The guarantee relies on monotone losses and nested predictors, and can remain valid when training and calibration data come from different distributions as long as calibration and test data match.

Assumptions: A prediction-set loss must decrease as the predicted set grows, while the predictor family must be nested so larger λ produces larger sets.These monotonicity conditions support selecting λ through an upper confidence bound for risk.
Calibration procedure: The procedure calibrates any nested family of set-valued predictors to achieve risk at most a user-specified level α with high probability.It searches over predictors indexed by λ using calibration data not used for model training.
Calibration procedure: The calibrated value λ̂ is the smallest λ for which the entire confidence region to its right lies below α.This choice produces a set-valued predictor with high-probability risk control.
Guarantee: Pointwise upper confidence bounds suffice for data-driven λ selection because risk monotonicity converts them into an RCPS guarantee.Without monotonicity, a uniform convergence result would be needed for a similar guarantee.
Scope: UCB calibration remains valid when the fitted predictive model uses data from another distribution, provided calibration and test data share a distribution.Continuity of the risk function is assumed for simplicity and can be removed with minor modifications.

3 Concentration Inequalities for the Upper Confidence Bound

The section develops finite-sample upper confidence bounds for bounded losses and compares them with asymptotically valid alternatives. The WSR bound is recommended for bounded losses overall, while exact binomial methods are preferred for binary losses and WSR performs best for non-binary losses in the reported experiments.

Overview: The WSR confidence bound is recommended for all bounded-loss cases after comparison with several alternative bounds.The paper describes it as the clear winner in the numerical comparisons.
Hoeffding bound: Hoeffding calibration yields an RCPS for losses bounded by one, but its simple empirical-risk threshold is presented only for illustration.The threshold subtracts log(1/δ)/2n from the target level, whereas tighter bounds are recommended in practice.
Hoeffding–Bentkus bound: The Hoeffding–Bentkus bound yields an RCPS for losses bounded by one and is nearly tight for binary losses.For binary loss, the exact binomial confidence result gives the most precise upper confidence bound.
Variance-adaptive bounds: Variance-adaptive bounds address the looseness of Hoeffding–Bentkus for non-binary losses, with WSR improving substantially over empirical Bernstein in the cited comparison.WSR is based on online inference and martingale tools rather than tail-probability inversion.
Numerical experiments: The numerical experiments report coverage and median upper-confidence-bound gap across one million replicates, measuring validity and power respectively.They compare Bernoulli and Beta losses across sample sizes, confidence levels, means, and variability settings.
Numerical experiments: WSR outperforms the other bounds for all tested Beta distributions and performs similarly to Hoeffding–Bentkus for Bernoulli distributions.The paper recommends WSR for non-binary bounded losses and exact binomial quantiles for binary losses.

3.2 Unbounded losses

For unbounded losses, finite-sample risk control requires distributional restrictions, while a CLT-based alternative provides only asymptotic validity. Experiments show the Pinelis–Utev bound preserves coverage under heavy tails, whereas the CLT bound is tighter for light-tailed losses and larger samples.

Unbounded losses: Finite-sample upper confidence bounds for nonnegative unbounded losses are impossible without additional restrictions.The paper therefore restricts attention to distributions satisfying regularity conditions such as bounded coefficient of variation.
Finite-sample bound: The Pinelis–Utev inequality yields an (α, δ)-RCPS when the coefficient of variation is bounded by a known constant.When the coefficient of variation is unknown, it can be treated as a sensitivity parameter or estimated from sample moments.
Numerical comparisons: The unbounded-loss experiments compare a coefficient-of-variation-based PU bound with a CLT bound across Gamma, square-t, and log-normal distributions.The distributions and parameter settings are summarized in Table 1.
Numerical comparisons: The CLT bound nearly achieves target coverage for light-tailed distributions but drastically undercovers for heavy-tailed distributions, while the PU bound has valid coverage in these settings.The PU bound is the finite-sample-valid alternative evaluated in the experiments.
Numerical comparisons: The CLT bound is much tighter in all cases, but the gap between the bounds shrinks as sample size grows.The paper recommends the CLT bound for believed light-tailed losses with moderately large samples, and the PU bound otherwise.
Calibration-set size: About 1,000 to 10,000 calibration points generally suffice for prediction sets that are not too conservative, depending on the loss setting.The bounded-loss examples require roughly 1,000, a few thousand, and 10,000 points as α decreases from 0.1 to 0.001; Gamma requires about 10,000.

4 Generating the Set-Valued Predictors

The paper constructs nested set-valued predictors from black-box models, using greedy risk-density ordering and proving optimality under important loss structures.

Greedy construction: The method constructs nested prediction sets Tλ(x) from an estimated conditional risk density, re-computing remaining-element risks as sets grow.The greedy procedure is computationally convenient and approximately optimal for a broad class of losses.
Greedy construction: When infinitesimal risk is independent of the current set, the output contains responses whose conditional risk density exceeds a threshold.This special case is illustrated as the thresholding rule used by the optimal prediction-set construction.
Greedy construction: The greedy algorithm considers only the next infinitesimal risk portion, while look-ahead variants may be more expensive computationally.The paper identifies this restricted search as a design trade-off rather than a validity failure.
Optimality properties: For fixed class-specific losses, the greedy sets have the smallest average size among predictors with no greater risk, when the exact conditional density is known.The result uses Lebesgue measure for continuous responses and counting measure for discrete responses.
Optimality properties: The generalized optimal predictor includes z whenever its conditional expected omission cost E[ℓ(Y; z) | X = x] is at least −λ.This extends the fixed-loss construction to losses represented through a nonnegative function ℓ and a finite measure µ.
Optimality properties: In practice, the conditional response distribution must be estimated, and even the exactly optimal sets may be difficult to compute.Thus, the theoretical optimality result characterizes an important target rather than guaranteeing exact implementation.

5 Examples

RCPS is applied to classification, multi-label prediction, hierarchical classification, image segmentation, and protein structure prediction with task-specific losses and set constructions.

Overview: Across five prediction problems, the authors define task-specific losses and set-valued predictors, then evaluate upper-confidence-bound calibration.The experiments cover classification, multi-label classification, hierarchical classification, segmentation, and protein structure prediction.
Classification with a class-varying loss: On Imagenet, RCPS controls risk at the correct level while producing prediction sets of reasonable size across 100 random data splits.The experiment uses 30,000 calibration points and 20,000 evaluation points with a pretrained ResNet-152.
Multi-label classification: For MS COCO multi-label classification, RCPS allows less brittle losses such as false-negative rate, whereas conformal prediction requires all true labels to be covered and returns larger sets.The comparison highlights different coverage definitions rather than a common loss criterion.
Hierarchical classification: Hierarchical prediction represents uncertainty with an ancestor node, and its loss measures the normalized distance between predicted nodes and ancestors of the true label.The hierarchical structure is encoded as a rooted tree with finite depth and one leaf per label.
Image segmentation: For image segmentation, the method targets capturing at least a 1 − α fraction of object pixels within each connected object, rather than only averaging across the image.The construction detects local score peaks and connected regions before re-normalizing them to support per-object validity.
Protein structure prediction: For protein structure prediction, the paper reports that the CLT nearly controls risk in a small-sample regime, suggesting applicability with small calibration sets.The protein example uses prediction intervals for residue-pair distances, as illustrated for protein T0995.

6 Other Risk Functions

The calibration framework extends beyond mean pointwise risk to ranking, metric learning, and robustness settings by combining nested predictors with suitable concentration bounds.

General framework: UCB calibration requires a nested predictor family with monotone error and a pointwise upper confidence bound for each parameter value.These two ingredients support finite-sample error control when selecting the data-driven parameter.
Ranking: For any ranking loss bounded by one, RCPS provides R(Tλ̂) ≤ α with probability at least 1 − δ.The guarantee is obtained using concentration for bounded U-statistics and UCB calibration.
Ranking: For ranking, set-valued outputs encode uncertainty by including both signs and thereby abstaining on ambiguous pairs.Calibration seeks to minimize abstention while keeping non-abstained ranking errors below the user-specified level.
Metric learning: For metric learning, uncertainty is represented by set-valued distances, with an example using an interval around the estimated distance and a class-separation threshold at one.The framework controls the risk of predicting same-class points above the chosen distance boundary.
Metric learning: For metric learning with loss bounded by one, RCPS likewise guarantees R(Tλ̂) ≤ α with probability at least 1 − δ.The theorem applies to the specified nested set-valued metric family.
Robustness: For adversarial robustness, the framework can calibrate risk over perturbation balls, although implementing it would require computational innovations.The paper establishes statistical validity for this extension without carrying out the proposed procedure.

7 Discussion

RCPS provide model-agnostic, instance-wise uncertainty estimates by controlling user-specified losses, extending beyond coverage-based conformal prediction when concentration results are available.

RCPS represent uncertainty with prediction sets that can be applied to existing models without retraining.They are designed to be straightforward to use across many situations.
RCPS accommodate many loss functions, including false negative rate, rather than controlling only binary coverage risk.The multilabel classification setting illustrates this flexibility.
RCPS require access to a concentration result, whereas conformal prediction relies on exchangeability.This supports applications such as uncertainty quantification for ranking.
Standard train/validation/test splits estimate global uncertainty through overall accuracy, while RCPS return valid instance-wise uncertainty estimates.

A Proofs

The appendix establishes the validity of RCPS calibration through concentration and confidence-bound arguments, while also showing that finite-sample valid upper confidence bounds are impossible for unrestricted unbounded losses.

Theorem A.1 calibrates a parameter by selecting the smallest value whose upper bound controls the monotone risk at level α.Its validity follows from the risk function’s continuity and monotonicity together with pointwise coverage of the upper bound.
The Hoeffding and Bentkus tail bounds extend from independent sums to the relevant order-two U-statistic using an effective sample size of floor(n/2).The proof conditions on a random permutation, whose paired summands are independent.
For all nonnegative finite-mean distributions, any finite-sample upper confidence bound for the mean with uniform coverage must equal infinity with probability at least 1−δ.This establishes an impossibility result for unrestricted unbounded losses.

B An Exact Bound for Binary Loss

For binary losses, RCPS calibration uses an exact binomial upper bound, computationally inverts it to select the calibration parameter, and retains the finite-sample risk-control guarantee.

Binary losses reduce the fixed-parameter risk to the mean of Bernoulli variables, enabling a tight upper confidence bound from a binomial quantile.
The binomial upper bound matches the Bentkus expression with an improvement by a factor of e.
The calibration parameter is obtained by computationally inverting the exact upper bound.
The resulting predictor is an (α, δ)-risk-controlling prediction set for binary losses.
In the binary case, the framework yields a classical tolerance region.

C Conformal Calibration

Conformal calibration selects a threshold from calibration scores using an empirical quantile, yielding a risk-control guarantee under the stated setup and exchangeability argument.

For binary loss, conformal calibration is the well-known split conformal prediction case when set-valued predictions use membership loss.The section also records a mild generalization to settings where Y′ differs from 2^Y.
Calibration defines a score for each observation and assumes the family of prediction sets has a minimal element with probability one.The text notes this is always true in practice when the parameter space is finite.
The threshold is chosen as the (1−α) empirical quantile of the calibration scores.
The resulting calibration procedure satisfies the stated risk-control guarantee.
The validity proof uses the usual conformal prediction exchangeability argument.

D Further Comparisons of Upper Confidence Bounds

Additional experiments compare upper confidence bounds across bounded and unbounded distributions, significance levels, and mean values. For binary loss, HB is often tighter in low-mean or low-δ settings, while WSR remains the best all-around choice for bounded losses.

Additional numerical comparisons: Figures 16–19 compare upper confidence bounds across bounded and unbounded distributions at δ = 0.01 and δ = 0.001.The evaluations use one million independent samples and vary distribution type, sample size, and mean.
Binary loss comparison: HB is more likely than WSR to be tighter when the mean µ or level δ is small.This pattern is reported for the binary loss case.
Binary loss comparison: The simple Hoeffding and empirical Bernstein bounds never win in the reported Bernoulli comparison.The symmetry between the HB and WSR curves reflects this result.
Binary loss comparison: WSR is not uniformly tighter than HB, but remains the best all-around choice for bounded losses.The paper presents this as the overall interpretation of the binary-loss comparison.
Binary loss comparison: Figure 20 reports which bound wins more often for Bernoulli losses, excluding CLT because it undercovers.Rows correspond to levels and columns to mean values.

E Adaptive Score Renormalization for Polyp Segmentation

The polyp segmentation predictor adaptively renormalizes raw pixel scores using image-processing operations that emphasize local peaks and separate them into regions. The resulting renormalization function is applied before the set predictor’s calibration step.

Score renormalization: The segmentation construction uses classical image-processing techniques to detect and emphasize local peaks in raw prediction scores.These operations are used to construct a better set predictor from the raw predictor.
Score renormalization: Gaussian blurring convolves a score matrix with a normalized kernel controlled by standard deviation σ and odd kernel size k.The kernel normalization constant C is chosen so its entries sum to one.
Score renormalization: Reconstruction by dilation retains local score peaks while setting other regions to zero.The operation is represented as dil applied to a matrix of scores.
Region extraction: Binarization thresholds the local peaks, and connected components divide the resulting binary image into disjoint regions.The connected-components operation is based on pixel adjacency in the binary matrix.
Pipeline composition: The final renormalization r applies blurring, peak extraction, binarization, connected components, bounding-box construction, and score renormalization in sequence.The resulting function is defined as r(M) = renorm(M, bboxes(E(A(bint(g(M, σ, k)))))).