💻 computer science

Class Adaptive Conformal Training

This paper introduces Class Adaptive Conformal Training (CaCT), a novel method that formulates conformal training as an augmented Lagrangian optimization problem to adaptively learn class-conditional prediction sets without distributional assumptions, thereby achieving smaller and more informative sets while maintaining rigorous coverage guarantees across various tasks.

Original authors: Badr-Eddine Marani, Julio Silva-Rodriguez, Ismail Ben Ayed, Maria Vakalopoulou, Stergios Christodoulidis, Jose Dolz

Published 2026-01-15

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Badr-Eddine Marani, Julio Silva-Rodriguez, Ismail Ben Ayed, Maria Vakalopoulou, Stergios Christodoulidis, Jose Dolz

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The Overconfident AI

Imagine you have a very smart AI assistant that looks at pictures and tells you what they are. It's usually right, but sometimes it gets overconfident. It might say, "I am 100% sure this is a cat," when it's actually a dog.

In high-stakes situations (like medical diagnosis or self-driving cars), we don't just want the AI to guess; we want it to say, "I'm not sure, so here are the top 3 possibilities." This is called a Prediction Set.

The Goal: The AI should be right 90% of the time (a safety guarantee).
The Catch: If the AI is too unsure, it might list every possible animal (cat, dog, bird, fish, etc.). This is a valid prediction set (it definitely includes the right answer), but it's useless because it's too big. We want the set to be as small as possible while still being safe.

The Old Way: The "One-Size-Fits-All" Penalty

Previous methods tried to teach the AI to make smaller lists by adding a "penalty" during training. Think of this like a coach yelling at a player: "If your list of guesses is too long, you get a penalty!"

However, the old methods used one single penalty rule for everyone.

The Flaw: Imagine a classroom with 100 students. 90 of them are math geniuses, and 10 are struggling with basic addition.
If the coach uses the same penalty for everyone, the geniuses might get punished too hard (making them too scared to guess), while the struggling students get off too easy (they keep guessing huge lists).
In AI terms, this means the model works great for common things (like "dog" or "car") but fails miserably for rare things (like a specific type of beetle or a rare disease). It either makes the list too big for rare items or misses the safety guarantee.

The New Solution: CaCT (Class Adaptive Conformal Training)

The authors introduce CaCT, which is like hiring a personalized coach for every single student (or every single category).

Instead of one rule for the whole class, CaCT gives the AI a custom penalty weight for every specific category.

For the "Geniuses" (Common Classes): The AI learns it can be confident and keep its list small.
For the "Struggling Students" (Rare Classes): The AI learns it needs to be extra careful. It might need to list a few more options to stay safe, but it learns exactly how many are needed without blowing the whole list up.

How It Works: The "Smart Adjuster"

You might ask, "How does the AI know what penalty to give each student? Do we have to guess?"

No! The paper uses a mathematical trick called Augmented Lagrangian Multipliers (ALM).

The Analogy: Imagine a thermostat that doesn't just have an "On/Off" switch. Instead, it has a smart sensor that constantly checks the temperature of every room in the house.
If the kitchen is too hot, the thermostat automatically turns the AC down just for the kitchen.
If the bedroom is too cold, it turns the heat up just for the bedroom.
In the AI: The system automatically learns the perfect "penalty weight" for every class during training. It doesn't need a human to tell it, "Hey, give the rare beetles a penalty of 5." The AI figures it out on its own to balance safety and efficiency.

The Results: Smarter, Tighter Lists

The paper tested this on many different datasets, including:

Standard Image Recognition: (Like identifying cats, dogs, and cars).
Long-Tailed Datasets: (Where some things appear thousands of times and others only a few times—like a library with 1,000 copies of one book and only 1 copy of 999 other books).
Text Classification: (Sorting news articles into categories).

The Findings:

Better Safety: The AI still keeps its promise of being right 90% of the time.
Smaller Lists: The lists of guesses are significantly shorter and more useful than before.
Fairness: The AI stops ignoring the "rare" items. It gives them the specific attention they need, so the rare items don't get huge, useless lists, and the common items don't get overly cautious lists.

Summary

Think of CaCT as moving from a rigid, one-size-fits-all rulebook to a dynamic, personalized training system. It allows the AI to be "smart" about how much uncertainty to show for different types of problems, resulting in predictions that are both safer and much more helpful for real-world use.

Technical Summary: Class Adaptive Conformal Training (CaCT)

Problem Statement

Deep neural networks frequently produce unreliable probability estimates, often exhibiting overconfidence in their predictions. While Conformal Prediction (CP) offers a principled framework for uncertainty quantification with rigorous, distribution-free coverage guarantees, existing methods face significant limitations in efficiency and adaptivity.

Standard CP methods, particularly those applied post-hoc, often generate prediction sets that are valid but uninformative (e.g., containing too many classes). To address this, Conformal Training (CT) methods like ConfTr integrate CP objectives directly into the training process to minimize prediction set sizes. However, current CT approaches suffer from two critical flaws:

Class-Agnostic Regularization: They typically employ a single global hyperparameter to penalize prediction set sizes across all classes. This rigid formulation fails to account for the intrinsic difficulty differences between classes.
Inefficiency in Imbalanced Scenarios: In long-tailed distributions (common in real-world applications), a global penalty leads to suboptimal performance. It often results in overly conservative (large) prediction sets for underrepresented (tail) classes while failing to adequately optimize for head classes, thereby reducing overall efficiency and adaptivity.

Methodology

The authors propose Class Adaptive Conformal Training (CaCT), a framework that reformulates conformal training as an augmented Lagrangian optimization problem. This approach adaptively learns class-specific penalty weights without requiring prior knowledge of the data distribution.

1. Class-Conditional Constraints

Instead of enforcing a marginal constraint on the average set size across all data points, CaCT introduces $K$ constraints (one for each class $k \in \mathcal{Y}$ ). The objective is to minimize the classification loss $L_{cls}(\theta)$ subject to the condition that the expected set size for each class $k$ does not exceed a target $\eta$ :
$\min_{\theta} L_{cls}(\theta) \quad \text{s.t.} \quad |C_\theta(X)| \le \eta, \quad \text{a.s. } X \sim P_{X|Y=k}$

2. Augmented Lagrangian Multiplier (ALM) Strategy

To solve this constrained optimization problem without manually tuning $K$ separate penalty weights (which is infeasible for large-scale datasets like ImageNet), CaCT utilizes a modified Augmented Lagrangian Method.

Objective Reformulation: The constrained problem is converted into an unconstrained one using a penalty function $P$ (specifically the PHR function) parameterized by Lagrange multipliers $\lambda_k$ and penalty coefficients $\rho_k$ for each class.
Adaptive Weight Learning:
- Inner Loop (Primal): The model parameters $\theta$ are updated via stochastic gradient descent to minimize the classification loss plus the current penalty terms.
- Outer Loop (Dual): The Lagrange multipliers $\lambda_k$ and penalty parameters $\rho_k$ are updated based on the violation of class-specific constraints observed on a validation set.
- Mechanism: If a specific class violates its size constraint (i.e., the average set size is too large), its corresponding multiplier $\lambda_k$ increases, strengthening the penalty for that class in subsequent iterations. Conversely, if the constraint is satisfied, the penalty decreases. This allows the model to learn an optimal trade-off between accuracy and conformalization for each class individually.

3. Differentiability

To enable end-to-end training, the non-differentiable components of CP (the indicator function for set size and the empirical quantile) are smoothed using a sigmoid function and differentiable sorting techniques. While this smoothing relaxes the marginal coverage guarantee during training, the authors note that the formal guarantee is restored at test time by reverting to the original non-smooth implementation.

Key Contributions

Class Adaptive Framework: Introduction of CaCT, which replaces global penalty weights with class-wise multipliers, enabling the model to adapt to the specific difficulty and frequency of each class.
Augmented Lagrangian Optimization: Development of a modified ALM algorithm that automatically learns the optimal class-specific penalty weights during training, eliminating the need for manual hyperparameter tuning or prior distributional assumptions.
Robustness to Imbalance: Demonstration that CaCT effectively mitigates the issues of long-tailed distributions, producing smaller prediction sets for underrepresented classes while maintaining coverage, a capability lacking in prior methods like ConfTr, CUT, and DPSM.

Experimental Results

The authors evaluated CaCT on multiple benchmarks, including standard image classification (MNIST, CIFAR-10/100, ImageNet), long-tailed variants (MNIST-LT, CIFAR-LT, ImageNet-LT), and text classification (20 Newsgroups).

Efficiency: CaCT consistently produces significantly smaller average prediction set sizes compared to state-of-the-art conformal training methods (ConfTr, CUT, InfoCTr, DPSM) across all datasets.
Coverage Gap: In long-tailed scenarios, CaCT achieves the lowest coverage gap (deviation from the target coverage $1-\alpha$ per class), indicating superior class-conditional coverage. Existing methods often fail to maintain coverage for tail classes, leading to large set sizes.
Discriminative Performance: CaCT maintains or improves Top-1 and Top-3 accuracy compared to baselines, demonstrating that reducing set size does not come at the cost of predictive power.
Robustness: The method shows robustness across varying imbalance factors ( $\gamma$ ) and different mis-coverage levels ( $\alpha$ ) at test time. It also performs well with different non-conformity scores (THR, APS, RAPS).
Hyperparameter Sensitivity: Unlike existing methods that are highly sensitive to the choice of a global regularization weight $\lambda$ , CaCT's adaptive learning mechanism renders it agnostic to this specific tuning, simplifying deployment.

Significance

The paper claims that CaCT addresses a fundamental limitation in conformal training: the inability to handle class-conditional efficiency in complex, imbalanced datasets. By formulating the problem as an adaptive optimization task, CaCT enables the construction of prediction sets that are both statistically valid and practically informative, particularly in high-stakes scenarios involving long-tailed distributions where minority classes are often neglected. The work bridges the gap between theoretical coverage guarantees and practical efficiency, offering a scalable solution for real-world deployment of uncertainty-aware deep learning systems.