Class Adaptive Conformal Training

TL;DR

CaCT addresses the challenge of unreliable probability estimates by enabling class-conditioned uncertainty quantification. It reframes conformal training as an augmented Lagrangian optimization that learns per-class penalties, producing adaptive, compact prediction sets while preserving the $(1-\alpha)$ coverage guarantee. Across long-tailed and standard datasets, CaCT consistently outperforms existing conformal-training methods in efficiency and class-conditional coverage, with the ALM variant demonstrating robust performance and reduced hyperparameter sensitivity. The approach scales to many classes and practical imbalanced regimes, offering a principled, end-to-end framework for calibrated prediction in real-world, multi-class tasks.

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.

Figures (9)

• Figure 1: Existing conformal training methods fail to account for datasets with a large number of classes or long-tailed distributions. We compare CaCT against conformal training methods, and report coverage gap and set size across relevant benchmarks. Lower values for both metrics indicate better conformalization and more informative prediction sets (i.e., better efficiency) across all classes. Among all methods, CaCT exhibits superior performance.
• Figure 2: Impact of per-class conformal set size and coverage with class frequency ($\alpha=0.1$) in CIFAR10-LT. As classes become under-represented (toward the right of the x-axis), existing methods deteriorate their conformal performance, i.e., increase set sizes and decrease class-conditional coverage. In contrast, CaCT is less sensitive to class imbalance, yielding smaller class set sizes (left) while maintaining higher class coverage (right) for these classes.
• Figure 3: Discriminative results of conformal training methods. Top-$k$ ($k={1,3}$) results across several challenging datasets.
• Figure 4: Sensitivity to the balancing term $\lambda$. Selecting an appropriate $\lambda$ for existing methods, e.g., ConfTr stutz2022conftr and CUT einbinder2022cut is challenging, especially on complex datasets like ImageNet-LT.
• Figure 5: Left: Evolution of average soft set size (approximated using Sigmoid function) of CaCT and ConfTr (top), and values of multipliers $\boldsymbol{\lambda}$ for CaCT after each training epoch (bottom). Right: Effect of penalty functions and target size: Coverage gap and average set size on test set are shown across three popular penalty functions and different target size values. Ablation study on CIFAR100-LT using ResNet and THR score.