📊 statistics

Conformal Bayes under Label Shift: Post-Hoc Calibration vs. In-Training Adaptation

This paper presents a unified framework for handling label shift in Conformal Bayes by comparing post-hoc calibration, which adjusts the conformal threshold while keeping the parameter posterior fixed, against in-training adaptation, which reweights the parameter posterior itself, demonstrating that both achieve valid coverage but the latter offers superior efficiency in biased training regimes.

Original authors: Seungjin Choi

Published 2026-06-11

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Seungjin Choi

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 are a weather forecaster. You have built a model based on years of data from a specific region (let's call it "Source City"). Your model is great at predicting rain in Source City.

Now, you are asked to forecast for a new location, "Target City." However, there's a catch: Target City has a different climate profile. It rains much more often than Source City, but when it does rain, the weather patterns (clouds, wind, humidity) look exactly the same. This is what the paper calls Label Shift: the frequency of the outcome (rain vs. no rain) has changed, but the relationship between the weather signs and the rain hasn't.

If you just use your old Source City model on Target City without adjusting, your predictions will be wrong. You might say, "It's only 10% likely to rain," when in reality, it's 50%.

This paper introduces a method called Conformal Bayes to fix this. Instead of giving a single number (like "50% chance"), it gives a prediction set (e.g., "It will rain between 40% and 60% chance"). The goal is to make sure this range is accurate 90% of the time, even with the new climate.

The authors propose two different ways to fix the model, which they call Strategy A and Strategy B.

The Two Strategies

Strategy A: The "Post-Hoc" Tilt (The Quick Fix)

Think of this as taking your finished weather report and simply re-labeling the pages before handing it to the customer.

How it works: You keep your original model exactly as it is. You don't touch the internal math. Instead, when you calculate the final prediction range, you apply a "weight" to the data. If the new city rains more often, you mathematically "tilt" your final answer to account for that extra rain.
The Analogy: Imagine you have a map of Source City. You don't redraw the map. Instead, you put a transparent sheet over it with a filter that says, "Remember, it rains more here." You then draw your prediction circle based on this filtered view.
Pros: It's safe and works even if you don't know how the original model was built (a "black box").
Cons: If the original model was already biased (e.g., trained only on days when it rained heavily), this method just adjusts the final number but doesn't fix the underlying bias in the model's thinking.

Strategy B: The "In-Training" Tilt (The Deep Fix)

Think of this as rewriting the weather forecaster's training manual before they even start making predictions.

How it works: Before the model makes its final prediction, you go back into its "brain" (the parameters) and adjust its beliefs. You tell the model, "Hey, the data you learned from was skewed. Let's adjust your internal understanding of how often rain happens."
The Analogy: Instead of just putting a filter over the map, you actually redraw the map to reflect the new reality. You teach the forecaster that the "average day" in the new city is different.
Pros: If the original training data was "enriched" (e.g., the model only saw data from stormy days), this method actively de-biases the model. It fixes the root cause, leading to tighter, more precise prediction ranges.
Cons: You need access to the model's internal training process. If the model is a black box, you can't do this.

The Experiments: When to Use Which?

The authors ran two computer experiments to see which strategy wins.

Experiment 1: The "Fair" Training

Scenario: The model was trained on a perfectly balanced, fair dataset from Source City.
Result: Both strategies worked equally well. They both achieved the correct 90% accuracy. Strategy A was slightly simpler, but Strategy B didn't hurt anything.
Takeaway: If your training data is good and unbiased, you can use the simple "Post-Hoc" fix (Strategy A).

Experiment 2: The "Biased" Training (Lead Optimization)

Scenario: This simulates a real-world problem like drug discovery. Imagine you are training a model to find good drugs, but your training data only contains "high-activity" compounds (the best drugs). The model is biased because it thinks everything is a high-activity drug.
Result:
- Strategy A (the quick fix) still worked to get the right coverage, but the prediction ranges were a bit wide.
- Strategy B (the deep fix) shined here. By correcting the model's internal bias, it produced narrower, more precise prediction ranges while still keeping the 90% accuracy. It essentially "debiasing" the model, making it smarter about the new reality.
Takeaway: If your training data is skewed or biased (common in specialized fields like drug discovery), Strategy B is superior because it cleans up the model's internal logic, not just the final output.

Summary in Plain English

The Problem: Your AI model was trained on data that doesn't match the real world anymore (the "Label Shift").
The Solution: Use Conformal Bayes to create a safety net (a prediction range) that is guaranteed to be correct.
The Choice:
- If your model was trained on fair data, use Strategy A (adjust the final output). It's simple and effective.
- If your model was trained on biased or specialized data (like only the "best" examples), use Strategy B (fix the model's internal brain). It makes the predictions more precise and removes the bias.

The paper concludes that while both methods guarantee accuracy, Strategy B is the "de-biasing operator" that makes the model more efficient when the training data was imperfect.

Technical Summary: Conformal Bayes under Label Shift

1. Problem Statement

The paper addresses the challenge of constructing statistically valid prediction sets under label shift, a distributional shift where the marginal label distribution changes between source ( $p_s(y)$ ) and target ( $p_t(y)$ ) domains, while the conditional feature distribution remains invariant ( $p_t(x|y) = p_s(x|y)$ ).

In standard Conformal Bayes, prediction sets are constructed using the Bayesian posterior predictive $p(y|x, D_{tr})$ and a nonconformity score (typically negative log-predictive density, NLPD). Under label shift, calibration data drawn from the source domain and test data from the target domain are no longer exchangeable under the same measure. Consequently, prediction sets calibrated on source data fail to achieve the nominal coverage level on target test points. This issue is particularly relevant in domains like lead-optimization, where training data are enriched for high-activity compounds, but predictions are required over a full chemical library with a different label distribution.

2. Methodology

The authors propose a unified framework analyzing two complementary strategies to restore nominal target-domain coverage. Both strategies rely on importance-weighted (IW) conformal calibration to ensure validity but differ fundamentally in how they adjust the predictive geometry.

Core Assumptions

Label Shift (Conditional Invariance): $p_t(x|y) = p_s(x|y)$ , but $p_t(y) \neq p_s(y)$ .
Density Ratio: The shift is modeled via exponential tilting: $w(y) = p_t(y)/p_s(y) \propto \exp(\beta y)$ , where $\beta$ is a known parameter (or estimated).
Model: A frozen backbone $\phi(x)$ with a Gaussian linear head: $y|x, \theta \sim \mathcal{N}(\theta^\top \phi(x), \sigma^2)$ .

Strategy A: Post-Hoc Calibration (Predictive Tilting)

Mechanism: The parameter posterior $\pi_s(\theta|D_{tr})$ remains unchanged. The correction is applied at the predictive level.
Process:
1. The target posterior predictive is derived by tilting the source predictive: $p_A(y|x, D_{tr}) \propto p_s(y|x, D_{tr}) w(y)$ .
2. The nonconformity score is the negative log-density of this tilted predictive: $s_A(x, y) = -\log p_A(y|x, D_{tr})$ .
3. The conformal threshold is computed using an importance-weighted quantile of the calibration scores.
Geometry: The resulting prediction set center includes a shift proportional to the epistemic uncertainty ( $\phi(x)^\top \Sigma_\theta \phi(x)$ ). Inputs far from the training distribution receive wider intervals.
Requirements: Does not require access to the training pipeline; treats the source model as a black box.

Strategy B: In-Training Adaptation (Posterior Tilting)

Mechanism: The correction is applied at the parameter level before calibration.
Process:
1. A "target-tilted" posterior is constructed by adjusting the source posterior using per-observation likelihood ratios: $\pi_t(\theta|D_{tr}) \propto \pi_s(\theta|D_{tr}) \prod_{i} Z_\theta(x_i)^{-1}$ .
2. For Gaussian heads, this yields a closed-form Gaussian posterior with a shifted mean: $\hat{\theta}_t = \hat{\theta}_s - \beta n \Sigma_\theta \bar{\phi}$ , where $\bar{\phi}$ is the mean feature of the training set.
3. The nonconformity score is the NLPD of the predictive distribution derived from this corrected posterior.
4. The conformal threshold uses the same importance-weighted quantile as Strategy A.
Geometry: The shift in the prediction set center depends on the similarity between the test point and the training distribution ( $\phi(x)^\top \bar{\phi}$ ), rather than the epistemic uncertainty.
Requirements: Requires access to the training pipeline and the training features to compute the mean correction.

Algorithmic Pipeline

Both strategies follow a unified pipeline (Algorithm 1):

Compute source Bayesian linear regression parameters.
Compute the score function (either $s_A$ or $s_B$ ).
Score the calibration set.
Compute the importance-weighted quantile threshold.
Construct the prediction set $C(x) = \{y : s(x, y) \leq \hat{q}\}$ .

3. Key Contributions

Unified Framework: A separation of validity (governed by the importance-weighted quantile) from efficiency (governed by the choice of score geometry: predictive vs. posterior tilting).
Closed-Form Solutions: Derivation of a closed-form target-tilted posterior for Gaussian heads, enabling exact computation of the corrected predictive distribution.
Comparative Analysis: A head-to-head comparison identifying the regimes where each strategy is preferable:
- Unbiased Training: Both strategies achieve valid coverage, but Strategy A (Predictive Tilting) yields narrower intervals. Strategy B introduces unnecessary bias if the source data is already representative.
- Biased/Enriched Training (Lead-Optimization): Strategy B acts as a debiasing operator. By correcting the parameter posterior, it reduces parameter bias, predictive bias, and interval width simultaneously while maintaining valid coverage.

4. Experimental Results

The authors evaluate the methods on synthetic Gaussian linear regression data under two regimes:

Experiment 1 (Unbiased Regime): Training data is a uniform draw from the source distribution.
- Coverage: All methods using the IW quantile (CB-IW, Strategy A, Strategy B) maintain nominal 90% coverage regardless of label shift intensity ( $\beta$ ). The unweighted baseline fails, losing coverage as $\beta$ increases.
- Efficiency: Strategy A produces the narrowest intervals. Strategy B is wider because the MAP correction moves the posterior away from the truth in an unbiased setting.
Experiment 2 (Lead-Optimization Regime): Training data is enriched for high-response compounds (biased source).
- Coverage: All IW-based methods maintain valid 90% coverage. The unweighted baseline over-covers (91–95%) due to inflated calibration scores from the biased model.
- Debiasing: Strategy B reduces parameter bias by approximately 40% and predictive bias by 34% compared to the biased source posterior.
- Efficiency: Strategy B produces strictly narrower intervals than Strategy A (up to 16% reduction in width at $\beta=0.6$ ) while maintaining the same coverage. This demonstrates that correcting the parameter posterior is more efficient than correcting the predictive density when the source training data is biased.

5. Significance and Claims

The paper claims to provide a principled, unified view of handling label shift in Conformal Bayes. Its primary significance lies in distinguishing between calibration for validity and adaptation for efficiency:

Validity is guaranteed solely by the importance-weighted quantile rule, regardless of whether the score is corrected post-hoc or in-training.
Efficiency is context-dependent. In unbiased settings, post-hoc predictive tilting is superior as it avoids altering the parameter posterior. However, in scenarios with enriched/biased training data (common in lead-optimization), in-training posterior adaptation acts as a necessary debiasing step, improving both the accuracy of the underlying model and the efficiency of the prediction sets.

The authors note that the framework currently relies on strong assumptions (pure label shift, known tilt parameter, Gaussian linear heads) and synthetic validation. They suggest that extending to deep models requires approximate inference (e.g., Laplace approximation) and that real-world validation is left for future work. They also outline a two-step procedure for estimating $\beta$ without labeled target data using pseudo-labels, referencing a companion paper for molecular property applications.