📊 statistics

Audited Conformal Prediction for Classification under Unknown Distribution Shift

This paper introduces Audited Conformal Prediction (ACP), a novel framework that utilizes a small labeled target dataset to train an auxiliary audit model, thereby enhancing conditional coverage and providing rigorous theoretical guarantees for uncertainty quantification in pretrained classifiers under unknown distribution shifts.

Original authors: Yanfei Zhou, Rizal Fathony, Nam H. Nguyen, Matteo Sesia

Published 2026-06-16

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yanfei Zhou, Rizal Fathony, Nam H. Nguyen, Matteo Sesia

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

Imagine you have a very experienced, old-school detective (the "Legacy Model") who has solved thousands of cases based on historical data. This detective is usually great, but the world is changing. New types of crimes are appearing, or the way criminals act has shifted (this is what the paper calls "Distribution Shift").

The problem is that while the detective might still guess the right answer sometimes, they often don't realize when they are guessing blindly. They might say, "I'm 100% sure this is a cat," when it's actually a dog, simply because the lighting is different. In machine learning, we call this a lack of uncertainty quantification. We need a way to say, "I'm not sure," or "I'm only 60% sure," so we don't make dangerous mistakes.

The Old Way: "Just Trust the Average"

Standard methods (called Conformal Prediction) try to fix this by saying, "Okay, we need to be right 90% of the time on average." To do this, they look at the detective's past mistakes and draw a safety net.

The Flaw: This safety net is like a blanket that covers everyone equally. If the detective is great at identifying cats but terrible at identifying dogs in the new environment, the blanket might be too loose for the dogs (letting them slip through) and too tight for the cats (wasting time). It guarantees you are right overall, but it doesn't guarantee you are right for specific, tricky cases.

The New Solution: "Audited Conformal Prediction" (ACP)

The authors propose a new system called Audited Conformal Prediction (ACP). Instead of just trusting the old detective, they hire a new, specialized manager (the "Audit Model").

Here is how the analogy works:

The Legacy Detective: The original AI model that makes the predictions. It's fast and knows a lot, but it might be out of touch with the new world.
The Audit Manager: A smaller, simpler AI trained on a tiny amount of new, labeled data. Its only job is not to solve the crime. Its only job is to look at the Detective's work and say, "Hey, I think the Detective is going to mess this one up," or "This one looks safe."
- Think of the Audit Manager as a reliability radar. It doesn't need to know the answer; it just needs to know when the Detective is likely to be wrong.
The Process:
- The Detective makes a guess.
- The Audit Manager checks the guess. If the Manager says, "This looks risky," the system automatically widens the safety net (the prediction set) to include more possibilities, just to be safe.
- If the Manager says, "This looks safe," the system keeps the safety net tight and efficient.

Why is this better?

The paper claims this method achieves two things that other methods struggle with:

The "Fairness" Guarantee: It ensures that the system is right 90% of the time overall (Marginal Coverage).
The "Smart" Guarantee: It actually gets much better at being right for the hard cases (Conditional Coverage).

The Analogy of the Umbrella:
Imagine it's raining.

Standard Method: Gives everyone a giant, heavy umbrella. It keeps everyone dry on average, but it's clumsy and heavy for people who are just standing under a tree (easy cases), and it might still be too small for people in a storm (hard cases).
ACP Method: The Audit Manager looks at the sky. If it's a light drizzle, it gives you a small, light umbrella. If it's a thunderstorm, it immediately hands you a massive, reinforced umbrella. You stay dry in both situations, but you aren't weighed down when you don't need to be.

The "Secret Sauce": Two Strategies

The paper describes two ways to use this Audit Manager:

The "Smart Adjuster" (ACP-MC): This version tweaks the Detective's confidence scores. If the Audit Manager thinks the Detective is overconfident, it lowers the confidence. This keeps the safety net small but makes it much more accurate for tricky cases.
The "Group Protector" (ACP-ACC): This version groups cases by how "risky" the Audit Manager thinks they are. It guarantees that for the "High Risk" group, the system is right 90% of the time, and for the "Low Risk" group, it's also right 90% of the time. This is like having a specific safety protocol for the stormy areas and a different one for the sunny areas.

What the Experiments Showed

The authors tested this on:

Fake Data: Where they knew exactly where the "tricky" cases were.
Real Data:
- Medical Images: Detecting tumors in hospital scans from different hospitals (where the equipment and patients differ).
- Car Photos: Identifying cars in images that have been corrupted by rain, fog, or blur.

The Result:
In all these tests, the ACP method was able to:

Keep the overall error rate low (meeting the 90% guarantee).
Significantly improve the accuracy on the "hard" or "unreliable" samples compared to just retraining the model or using standard methods.
Avoid making the predictions too vague (the "umbrellas" didn't get unnecessarily huge).

The Bottom Line

The paper argues that when an AI model is deployed in a changing world, you don't always need to throw away the old model and retrain it from scratch (which is expensive and slow). Instead, you can add a small, smart "audit" layer that knows when the old model is likely to fail. This allows the system to be safer on the hard stuff without becoming clumsy on the easy stuff.

The authors explicitly state that this method works well even when you only have a small amount of new labeled data, which is a common real-world problem. They also note that the method's success depends on how good the "Audit Manager" is at spotting the risks.

Technical Summary: Audited Conformal Prediction for Classification under Unknown Distribution Shift

Problem Statement

The paper addresses the challenge of uncertainty quantification for pre-trained ("legacy") classification models deployed in environments subject to unknown distribution shifts. While standard conformal prediction guarantees marginal coverage (i.e., the true label is contained in the prediction set with probability $1-\alpha$ on average over the population), this guarantee often masks severe conditional failures. Specifically, when the training and target distributions differ, standard methods may produce overly conservative prediction sets for some inputs while systematically under-covering others (e.g., samples where the legacy model is likely to fail).

Existing approaches to mitigate this either require retraining the model (which is resource-intensive and often suboptimal with limited new labeled data) or rely on assumptions about the nature of the distribution shift. The authors note that while retraining is intuitive, it can be inefficient when new labeled data is scarce relative to historical datasets. Furthermore, standard conformal methods do not inherently distinguish between "easy" and "hard" samples under shift, leading to uninformative prediction sets precisely where reliability is most needed.

Methodology: Audited Conformal Prediction (ACP)

The proposed solution, Audited Conformal Prediction (ACP), introduces an auxiliary audit model to identify inputs where the legacy model is likely to fail, without requiring retraining of the legacy model or assumptions about the shift structure.

Core Components

The Audit Model ( $\hat{r}$ ): A binary classifier trained on a small subset of labeled target data. Its task is to predict the correctness of the legacy model's point prediction ( $f_0$ ) for a given input $X$ .
- Training: The audit model is trained to predict the binary label $R_i = \mathbb{1}[f_0(X_i) = Y_i]$ using features derived from the input $X_i$ and signals from the legacy model (e.g., softmax probabilities, entropy, embedding distances).
- Oracle Concept: Theoretically, the method relies on an oracle audit model $r^*(x) = P(f_0(X)=Y | X=x)$ , which represents the true conditional reliability of the legacy model.
Integration Strategies: The paper proposes two complementary strategies for integrating the audit model into the conformal framework:
- Strategy 1: ACP with Marginal Coverage (ACP-MC)
  - Mechanism: The audit model $\hat{r}$ is used to refine the legacy model's estimated conditional label distribution. For binary classification, if $f_0(x)=1$ , the estimated probability is $(\hat{r}(x), 1-\hat{r}(x))$ ; otherwise, it is $(1-\hat{r}(x), \hat{r}(x))$ .
  - Goal: This refined distribution is used to compute nonconformity scores. The method guarantees standard marginal coverage ( $1-\alpha$ ) while empirically improving conditional coverage by adapting the prediction set size based on the estimated reliability of the input.
- Strategy 2: ACP with Audited Conditional Coverage (ACP-ACC/AEC)
  - Mechanism: This approach explicitly targets group-conditional coverage. The feature space is partitioned based on the audit model's output (e.g., binning $\hat{r}(x)$ into high-confidence and low-confidence groups).
  - Goal: It enforces coverage guarantees conditional on these audit-defined groups. The paper establishes a finite-sample theorem showing that if the audit model consistently estimates the oracle reliability, the method achieves coverage guarantees for the true reliability groups, with the error bounded by the discrepancy between the estimated and true audit groups.
Adaptive ACP (AACP): A data-driven procedure that adaptively selects between ACP and a retrained baseline. It splits the calibration data to evaluate which method yields higher conditional coverage on "unreliable" groups (identified by the audit model) and selects the winner.

Key Contributions

Novel Framework: The introduction of ACP, which leverages a small labeled target dataset to train an audit model that predicts legacy model failure, rather than retraining the legacy model itself.
Theoretical Guarantees:
- Proof that ACP-MC maintains marginal coverage guarantees.
- Derivation of a finite-sample bound for ACP-ACC, demonstrating that conditional coverage approaches the oracle guarantee as the audit model's estimation error vanishes.
Practical Efficiency: The method avoids the computational cost of full model retraining and the data hunger associated with fine-tuning complex models, making it suitable for scenarios with limited new labeled data.
Empirical Validation: Extensive experiments on synthetic data (covering concept shift, covariate shift, and label noise) and real-world datasets (Camelyon17 and CIFAR-10/CIFAR-10-C) demonstrate that ACP variants achieve significantly higher conditional coverage on difficult samples compared to standard conformal prediction, AdaTS (Adaptive Temperature Scaling), and retraining baselines, without inflating the average prediction set size.

Results

Synthetic Data: In scenarios with intrinsic label uncertainty (concept shift) or increased proportions of hard samples (covariate shift), standard methods and retraining often fail to distinguish between easy and hard samples, leading to under-coverage on hard bins. ACP variants successfully assign larger prediction sets to hard samples and smaller sets to easy ones, achieving near-perfect conditional coverage on hard bins while maintaining compact average set sizes.
Real-World Data: On Camelyon17 (histopathology) and CIFAR-10-C (corrupted images), ACP consistently outperformed baselines in conditional coverage across reliability bins. While retraining often produced overly conservative sets that failed to differentiate reliability levels, ACP maintained informative set sizes.
Trade-offs: The experiments highlight a trade-off between the strictness of conditional guarantees and set size. ACP-ACC (enforcing coverage over multiple overlapping groups) is the most conservative but offers the strongest conditional guarantees, while ACP-MC offers a balance with smaller set sizes.

Significance and Claims

The paper claims that ACP provides a practical and statistically principled approach to uncertainty-aware classification under arbitrary distribution shifts. Its primary significance lies in its ability to improve conditional coverage—ensuring reliability where it matters most—without the prohibitive data requirements of retraining complex models.

The authors emphasize that their contribution is not the design of a new learning procedure for the audit model (which can be any standard classifier), but rather the integration of audit predictions into the conformal framework. This allows for a direct reliability measure that does not require access to the original training data or assumptions about the shift mechanism.

The paper concludes modestly, acknowledging that the method's effectiveness depends on the quality of the audit model, which can be challenging in very low-data regimes. Future work suggested by the authors includes more data-efficient extensions (e.g., via cross-validation conformal methods) and adaptations for online or streaming settings.