Class Adaptive Conformal Training

TL;DR

This paper introduces Class Adaptive Conformal Training (CaCT), a novel method that adaptively shapes class-conditional prediction sets in neural networks, improving uncertainty estimates without prior distribution knowledge.

Contribution

CaCT formulates conformal training as an augmented Lagrangian optimization, enabling class-conditional set shaping without distributional assumptions, outperforming prior methods.

Findings

CaCT produces smaller, more informative prediction sets.

CaCT maintains coverage guarantees across datasets.

CaCT outperforms existing conformal training methods.

Abstract

Deep neural networks have achieved remarkable success across a variety of tasks, yet they often suffer from unreliable probability estimates. As a result, they can be overconfident in their predictions. Conformal Prediction (CP) offers a principled framework for uncertainty quantification, yielding prediction sets with rigorous coverage guarantees. Existing conformal training methods optimize for overall set size, but shaping the prediction sets in a class-conditional manner is not straightforward and typically requires prior knowledge of the data distribution. In this work, we introduce Class Adaptive Conformal Training (CaCT), which formulates conformal training as an augmented Lagrangian optimization problem that adaptively learns to shape prediction sets class-conditionally without making any distributional assumptions. Experiments on multiple benchmark datasets, including standard and long-tailed image recognition as well as text classification, demonstrate that CaCT consistently outperforms prior conformal training methods, producing significantly smaller and more informative prediction sets while maintaining the desired coverage guarantees.