Class Adaptive Conformal Training
Abstract
Deep neural networks have achieved remarkable success across a variety of tasks, yet they often suffer from unreliable probability estimates, leading to overconfident predictions. Conformal Prediction (CP) offers a principled framework for uncertainty quantification, yielding prediction sets with rigorous coverage guarantees. Recent work has proposed integrating conformal objectives into training, optimizing for overall set size. Nevertheless, shaping the prediction sets in a class-conditional manner is not straightforward and typically requires prior knowledge of the data distribution. While class-conditional efficiency has been studied at the post-hoc calibration level (e.g., via class-conditional scoring or calibration partitioning), no prior conformal training method explicitly addresses it. To fill this gap, we introduce Class Adaptive Conformal Training (CaCT), which formulates conformal training as an augmented Lagrangian optimization problem that adaptively shapes prediction sets class-conditionally without relying on parametric distributional assumptions, beyond a user-specified global coverage constraint. 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, reducing prediction set size while maintaining the desired coverage guarantees.