Distribution-free, Risk-controlling Prediction Sets

• Published in: Journal of the ACM Vol. 68; no. 6; pp. 1 - 34
• Main Authors: Bates, Stephen, Angelopoulos, Anastasios, Lei, Lihua, Malik, Jitendra, Jordan, Michael
• Format: Journal Article
• Language: English
• Published: New York Association for Computing Machinery 01.12.2021
• Subjects: Decision making; Image classification; Image segmentation; Labels; Machine learning; Risk communication; Risk management; Structural hierarchy; Uncertainty
• ISSN: 0004-5411; 1557-735X
• DOI: 10.1145/3478535

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 controls 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. Last, we discuss extensions to uncertainty quantification for ranking, metric learning, and distributionally robust learning.