Active, anytime-valid risk controlling prediction sets

TL;DR

This paper extends risk controlling prediction sets to an anytime-valid, sequential setting, enabling adaptive data collection with guarantees, and introduces active labeling strategies that optimize label usage and utility.

Contribution

It introduces a framework for risk controlling prediction sets in sequential, adaptive data collection, with active labeling and utility optimization, providing theoretical guarantees and practical algorithms.

Findings

Guarantees hold at all time steps in sequential data collection.

Active labeling policies reduce label usage while maintaining utility.

Empirical results show improved label efficiency over baselines.

Abstract

Rigorously establishing the safety of black-box machine learning models concerning critical risk measures is important for providing guarantees about model behavior. Recently, Bates et. al. (JACM '24) introduced the notion of a risk controlling prediction set (RCPS) for producing prediction sets that are statistically guaranteed low risk from machine learning models. Our method extends this notion to the sequential setting, where we provide guarantees even when the data is collected adaptively, and ensures that the risk guarantee is anytime-valid, i.e., simultaneously holds at all time steps. Further, we propose a framework for constructing RCPSes for active labeling, i.e., allowing one to use a labeling policy that chooses whether to query the true label for each received data point and ensures that the expected proportion of data points whose labels are queried are below a predetermined label budget. We also describe how to use predictors (i.e., the machine learning model for which we provide risk control guarantees) to further improve the utility of our RCPSes by estimating the expected risk conditioned on the covariates. We characterize the optimal choices of label policy and predictor under a fixed label budget and show a regret result that relates the estimation error of the optimal labeling policy and predictor to the wealth process that underlies our RCPSes. Lastly, we present practical ways of formulating label policies and empirically show that our label policies use fewer labels to reach higher utility than naive baseline labeling strategies on both simulations and real data.