Class Adaptive Conformal Training

2026-01-14

AI Summary

This work addresses the issue of unreliable probability estimates in deep neural networks, which often exhibit overconfidence and lack effective mechanisms for class-adaptive uncertainty quantification. To this end, the authors formulate conformal prediction as an augmented Lagrangian optimization problem, enabling, for the first time, end-to-end learning of class-conditional prediction sets without any distributional assumptions. The proposed method dynamically optimizes the shape and size of prediction sets for each class, significantly enhancing their informativeness and practical utility. Experimental results across multiple image recognition and text classification benchmarks demonstrate that the resulting prediction sets are substantially smaller and more discriminative while rigorously maintaining the desired coverage guarantee.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Stochastic Optimization
• Computer Vision: Learning & Optimization for CV

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems

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.

Problem

• Conformal Prediction
• Uncertainty Quantification
• Class-Conditional Prediction Sets
• Overconfidence
• Deep Neural Networks

Innovation

• Conformal Prediction
• Class-Adaptive Optimization
• Uncertainty Quantification
• Augmented Lagrangian
• Prediction Sets