Learn then Test: Calibrating Predictive Algorithms to Achieve Risk Control

Anastasios N. Angelopoulos; Stephen Bates; Emmanuel J. Candès; Michael I. Jordan; Lihua Lei

Source-linked AI summary

Learn then Test: Calibrating Predictive Algorithms to Achieve Risk Control

Anastasios N. Angelopoulos, Stephen Bates, Emmanuel J. Candès, Michael I. Jordan, Lihua Lei

arXiv:2110.01052v5Oct 3, 2021cs.LGcs.AIcs.CVstat.MEstat.ML

TL;DR

The paper addresses the lack of finite-sample reliability guarantees for complex machine learning predictions. It calibrates pretrained models through Learn then Test, reframing risk control as multiple hypothesis testing, and reports tight risk control across several prediction tasks. The guarantees are marginal, and the paper leaves optimal testing procedures and balanced control across data strata open.

Problem
Complex machine learning models need statistical reliability guarantees, but standard tools do not apply broadly to structured predictions and multiple decisions.
Method
Learn then Test post-processes any pretrained predictive model with calibration data and uses multiple hypothesis testing to select risk-controlling parameters.
Results
The framework provides finite-sample risk guarantees across tasks including multi-label classification, selective prediction, OOD detection, and instance segmentation, with tight control from testing methods.
Takeaways & Limitations
Reframing risk control as hypothesis testing supports calibration of non-monotone risks without restricting predictions to confidence sets.
Takeaways & Limitations
The guarantees are marginal and do not ensure balanced errors across meaningful strata such as object class, race, sex, or illumination.

Abstract

from arXiv · show

We introduce a framework for calibrating machine learning models so that their predictions satisfy explicit, finite-sample statistical guarantees. Our calibration algorithms work with any underlying model and (unknown) data-generating distribution and do not require model refitting. The framework addresses, among other examples, false discovery rate control in multi-label classification, intersection-over-union control in instance segmentation, and the simultaneous control of the type-1 error of outlier detection and confidence set coverage in classification or regression. Our main insight is to reframe the risk-control problem as multiple hypothesis testing, enabling techniques and mathematical arguments different from those in the previous literature. We use the framework to provide new calibration methods for several core machine learning tasks, with detailed worked examples in computer vision and tabular medical data.

1 Introduction

The paper introduces Learn then Test (LTT), a modular calibration framework that gives complex predictive models finite-sample statistical guarantees without model refitting or assumptions about model internals or the data distribution. It reframes risk control as multiple hypothesis testing and applies the approach across structured prediction tasks.

Motivation: Complex machine learning models lack standard statistical tools for reliability guarantees, despite their widespread use in real-world systems.The motivation includes controlling false discovery rates across multiple decisions and understanding when models fail.
Framework: LTT provides finite-sample guarantees for any predictive model without assumptions on its inner workings or the underlying data distribution.Its modular design is intended to accommodate changing neural-network architectures, fitting schemes, and data types.
Framework: LTT post-processes a pretrained model with calibration data and a low-dimensional parameter to control a user-chosen statistical error rate.The model is not refit; calibration searches parameter values using the calibration data.
Applications: In multi-label classification, fixed sequence testing can certify a threshold with false discovery rate no more than 10%.The threshold determines which predicted classes are included, and the framework is also applied to selective classification, selective regression, OOD detection, and instance segmentation.
Novelty: The framework extends risk-controlled prediction beyond confidence sets by removing earlier restrictions to one-dimensional parameters and monotonic risk functions.This permits control of potentially non-monotone statistical risks.

2 Risk Control in Prediction

The Learn then Test procedure searches candidate post-processings, converts risk-control hypotheses into valid p-values, and uses family-wise error control to identify admissible settings. It supports several testing strategies, while uniform concentration can perform poorly in practice.

Overview: LTT searches a collection of candidate functions and returns parameter values whose risks are guaranteed to be controlled.The selected value may be chosen data-dependently from the returned set.
Overview: For each candidate λ_j, LTT tests the null hypothesis H_j: R(λ_j) > α using a finite-sample valid p-value.Rejecting H_j corresponds to selecting a candidate whose risk is controlled.
Overview: Theorem 1 shows that any family-wise-error-rate-controlling algorithm applied to valid p-values yields a set of candidates that are risk-controlling predictions except with probability δ.Any selected λ from this set is an (α, δ)-risk-controlling prediction.
Valid p-values: For expected-loss risks, empirical calibration losses can be converted into valid p-values using concentration inequalities such as the hybridized Hoeffding-Bentkus bound.The bound applies when the loss is bounded in [0, 1]; asymptotically valid central-limit-theorem p-values are available for unbounded losses.
Multiple hypothesis testing: Bonferroni controls FWER but can lose performance for large hypothesis spaces, motivating methods that exploit dependence and structure among nearby candidates.The framework uses fixed sequence testing and sequential graphical testing as structured alternatives.
Multiple hypothesis testing: Fixed sequence testing and sequential graphical testing control FWER, with SGT redistributing rejected hypotheses’ error budgets to adjacent hypotheses.Fixed sequence testing is a special case of SGT, and under uniform null p-values its FWER is asymptotically δ.
Alternative approach: Uniform concentration has poor practical performance compared with hypothesis-testing approaches for the risk-control problem.The paper attributes this difference to uniform concentration solving a more difficult problem than the discretized testing formulation requires.

3 Examples

The paper applies Learn then Test to several vision and medical prediction tasks, calibrating pretrained models to control task-specific risks with finite-sample guarantees. Examples include false discovery rate, selective error or MSE, and simultaneous OOD type-1 error and coverage control.

Multi-label classification: In multi-label classification, thresholding estimated class probabilities produces prediction sets with FDR controlled at level (α, δ).On MS COCO, the experiment uses 4000 data points, 1000 validation points, α = 0.2, and δ = 0.1.
Selective classification: Selective classification on Imagenet abstains below a confidence threshold and evaluates error conditionally on predictions.The experiment uses α = 0.15, δ = 0.1, 5K calibration points, and 45K validation points, clipping thresholds with fewer than 25 predictions.
Selective regression: Selective regression abstains when estimated residual magnitude is large, targeting conditional MSE control on medical expenditure predictions.The MEPS experiment uses α = δ = 0.1 and gradient boosting models for the conditional mean and residual magnitude.

4 Example: Instance Segmentation with mIOU, Coverage, and Recall Guarantees

The instance-segmentation example calibrates three detector parameters to control recall, IOU, and classification coverage simultaneously. Experiments on MS COCO show controlled risks, with coverage tight and recall and IOU conservative under the parameter-selection order.

Problem and parameterization: Instance segmentation combines object localization, mask precision, and class assignment, evaluated through recall, IOU, and coverage.The method treats these as three risk-controlled subtasks.
Problem and parameterization: The three tuning parameters primarily control distinct subtasks: λ1 selects objects, λ2 adjusts bounding regions, and λ3 sets classification certainty.The parameters are λ1 for recall, λ2 for IOU, and λ3 for coverage.
Prediction construction: The detector constructs filtered objects, thresholded segmentation masks, and class prediction sets from estimated object, pixel, and class probabilities.Low-confidence detections can be discarded, while insufficient class confidence can produce null outputs.
Risk definitions: IOU aligns predicted and ground-truth objects through a maximizing permutation before contributing to the segmentation risk.The risk definitions also penalize spurious detections through the size of the retained prediction set.
Experiments: On MS COCO, fixed-testing experiments use α1 = 0.5, α2 = 0.5, α3 = 0.25, and δ = 0.1 to evaluate coverage, IOU, and recall.Figure 9 reports mean coverage, mean IOU, and recall over 1000 random splits.
Experiments: All risks remain controlled, but recall and IOU are conservative because λ3 is chosen before λ1 and λ2, making coverage tight.This ordering prioritizes small prediction sets and then optimizes the other risks subject to coverage control.

5 Discussion

The framework provides statistical guarantees for complex machine-learning systems while leaving several important questions open. Its guarantees are marginal, and future work includes subgroup-balanced errors and improved testing procedures.

Discussion: LTT and related procedures can wrap around complex machine-learning deployments without directly analyzing how models are fitted.This modularity may help provide guarantees as architectures, fitting schemes, and data types evolve.
Discussion: Hypothesis testing combined with SGT controls risks more tightly than uniform concentration in the authors’ experiments.The authors suggest this may reflect the easier discretized problem solved by testing.
Limitations and open questions: All guarantees are marginal, so errors need not be balanced across strata such as object class, race, sex, or illumination.The authors identify approximately balanced error control across strata as an important open direction.
Limitations and open questions: The paper leaves optimal testing procedures unexplored, while noting that tightly controlling non-monotonic risks could support auditing and certification of algorithmic predictions.The SGTs used are described as relatively simple, hand-designed heuristics, and learning the graph from another data split is suggested as a possible improvement.

A.1 Numerical comparisons

On a correlated synthetic AR process, fixed sequence testing is less conservative than Bonferroni, while uniform concentration performs poorly. The comparison evaluates how far right risk-controlling procedures can extend their rejection sets relative to an empirical-risk baseline.

A.1 Numerical comparisons: The AR simulation uses a “V”-shaped risk curve, with losses generated independently across samples but correlated across λ through an autoregressive process.The risk falls below α near the center of Λ, and the simulated correlation reflects dependence among tests based on shared calibration data.
A.1 Numerical comparisons: The simulation applies Hoeffding-Bentkus p-values followed by Bonferroni and fixed sequence tests over the λ grid.The procedures test whether risk is below α at each candidate λ using calibration data.
A.1 Numerical comparisons: Uniform concentration performs badly in practice, whereas fixed sequence testing is less conservative than Bonferroni.Fixed sequence testing is designed to make the rejection set as wide as possible on the right side.
A.1 Numerical comparisons: The empirical-risk baseline selects the largest λ with empirical risk below α, serving as an upper bound rather than a risk-controlling procedure.Higher rightmost endpoints indicate less conservative procedures.

A.2 Additional experimental results for FDR control

Additional experiments report the multi-label classification procedure on MS COCO at the higher target FDR level α = 0.5. The section also describes the ODIN-based out-of-distribution score and its manual rescaling for numerical resolution.

A.2 Additional experimental results for FDR control: The MS COCO multi-label classification experiment is repeated with desired FDR level α = 0.5.The result is presented as the same figure as Figure 3 in the main text.
A.2 Additional experimental results for FDR control: ODIN perturbs input images to distinguish in-distribution from out-of-distribution examples using a temperature-scaled DenseNet softmax output.The perturbation magnitude is ϵ = 0.0014, the default value from the ODIN repository.
A.2 Additional experimental results for FDR control: ODIN improves on the score 1−f̂(X)(1) by exploiting the larger perturbation-induced growth of the top softmax output for in-distribution images.The paper describes this as an empirical observation underlying the improvement.
A.2 Additional experimental results for FDR control: On CIFAR-10, ODIN values fall between 0.890 and 0.899, so the authors manually rescale the score before grid-based calibration.With 1000 equally spaced points from 0 to 1000, only 9 points would land in that interval.

B CLT p-value

The CLT p-value construction uses the empirical mean and standard deviation of calibration losses and is asymptotically valid under finite mean and variance. A finite-sample Berry–Esseen correction is possible but not provided.

B CLT p-value: The CLT p-value is computed from the empirical standard deviation of calibration losses around their empirical mean.The displayed construction uses the standardized gap between the target level α and empirical risk, scaled by √n.
B CLT p-value: Under finite mean and variance of L(Tλj(Xi), Yi), the resulting p-value is asymptotically valid.The proposition states validity through a lim sup probability bound over u ∈ [0,1].
B CLT p-value: A standard finite-sample Berry–Esseen correction can be made, but the paper does not state it.

C Proofs

The proofs establish error-control guarantees by combining classical multiple-testing results, fixed-sequence arguments, and union bounds. They also cover procedures with multiple null families and sequential rejection rules.

The remaining proofs invoke established results, including hybrid inequalities, the central limit theorem, and classical multiple-testing theory.
For a single null family, the first encountered null bounds the probability of any false discovery by δ.
With multiple null families, parallel procedures run at level δ/|J|, and the union bound keeps the probability of any false rejection at most δ.
Fixed sequence testing rejects a null only when the first null index has p-value at most δ.

D Split Fixed Sequence Testing

Split fixed sequence testing learns a promising testing path from one calibration-data split and evaluates it on fresh data. This preserves rigorous FWER control while addressing graph-design challenges in large-scale systems.

Graph design becomes difficult when many parameter coordinates affect multiple risks, motivating an automated procedure that learns the graph from data.
Split fixed sequence testing selects a sequence of nearly equal p-values on one split, then applies fixed sequence testing on a fresh split.
The algorithm partitions calibration data into graph-selection and multiple-testing sets, restricting path learning to the former.
The path is parameterized by discretizing points whose risk-specific p-values are nearly equal, with duplicate sequence values removed in practice.
Figure 12 contrasts a fallback graph, which chains hypotheses with distributed initial error budgets, against the alternative graph design.
After selecting a controlled parameter set, the method can optimize jointly over its coordinates to obtain meaningful prediction sets, many detections, and high IOU.

E SGT methods used in Section 3.4

The paper develops SGT graph designs and a uniform-concentration alternative for risk-control calibration. The SGT methods exploit structured error-budget propagation, while the concentration approach is generally more conservative in experiments.

SGT methods: The fallback SGT spreads the initial error budget across λ1 levels and propagates it toward the Pareto frontier.
SGT methods: The Hamming graph allocates the initial error budget to the bottom-right node and cascades it outward with balanced-inflow weights.
SGT methods: Balanced inflow gives nodes equally distant from the root the potential to receive the same error budget, supporting movement across structured parameter spaces.
SGT methods: The Hamming SGT can outperform fallback when many lower-right-triangle hypotheses are rejected, but its structure can also be harmful in other configurations.
Uniform concentration alternative: Uniform concentration provides an alternative route to risk control, but the paper reports it as typically more conservative and less powerful than multiple testing.
Uniform concentration alternative: The uniform-concentration framework adapts self-normalized empirical-process results to bounded, non-binary risk summands and improves constants using sampling-without-replacement and Rademacher tools.
Uniform concentration alternative: For rare events with F_hat_n(w)=O(log n/n), the Anthony–Shaw-Taylor bound is O(log n/n), versus O(1/√n) for DKWM.
FDR calibration: The resulting FDR calibration algorithm returns a parameter intended to control FDR at level α using a uniform upper-confidence-bound construction.

F.5 Technical proofs

The technical appendix proves uniform concentration results for empirical risk processes using conditioning, sampling-without-replacement arguments, and concentration inequalities. The proofs combine several intermediate bounds into the stated theorems.

Theorem F.1 is proved through four steps involving fresh samples, empirical-process deviations, permutation arguments, and final combination of bounds.
The proof uses Bernstein-type control for bounded risk summands whose variance is bounded by their expectation.
Permutation conditioning treats calibration subsets as samples drawn without replacement from a fixed unordered collection.
The appendix states the uniform-bound results through generic empirical processes indexed by λ and growth-function quantities.
The technical arguments conclude by combining intermediate inequalities and taking infima over auxiliary parameters such as γ and n′.
Theorem F.2 further conditions on paired observations and uses Rademacher representations with Bentkus-Dzindzalieta, Pinelis, and Hoeffding inequalities.