Source-linked AI summary
Distribution-free uncertainty quantification for classification under label shift
Aleksandr Podkopaev, Aaditya Ramdas
TL;DR
Distribution-free uncertainty quantification has been studied under covariate shift, but label shift remains a gap despite shifts potentially invalidating UQ assumptions. This paper adapts conformal prediction and calibration to label shift; it argues that correction is needed in less idealized settings and identifies an efficiency cost for label-conditional conformal prediction when there are many classes.
Problem
Distribution shifts can violate assumptions behind uncertainty quantification, and distribution-free UQ under label shift remained an unaddressed gap.
Source: §1 · pp. 2–3Method
The paper adapts conformal prediction and post-hoc calibration to label shift through reweighting, using unlabeled target data under reasonable assumptions.
Source: §1 · p. 3Results
Label shift hurts calibration, and although its effect on prediction-set coverage is less obvious, correction is necessary in less idealized settings.
Source: §4 · p. 11Takeaways & Limitations
Adapting prediction sets and calibration via binning to label shift closes a gap in distribution-free UQ beyond the classic i.i.d. setting.
Source: §4 · p. 11Label-conditional conformal prediction splits calibration data into K parts, which can cause substantial statistical-efficiency losses when K is large.
Source: §2.2 · p. 7
Abstract
from arXiv · showhide
Trustworthy deployment of ML models requires a proper measure of uncertainty, especially in safety-critical applications. We focus on uncertainty quantification (UQ) for classification problems via two avenues -- prediction sets using conformal prediction and calibration of probabilistic predictors by post-hoc binning -- since these possess distribution-free guarantees for i.i.d. data. Two common ways of generalizing beyond the i.i.d. setting include handling covariate and label shift. Within the context of distribution-free UQ, the former has already received attention, but not the latter. It is known that label shift hurts prediction, and we first argue that it also hurts UQ, by showing degradation in coverage and calibration. Piggybacking on recent progress in addressing label shift (for better prediction), we examine the right way to achieve UQ by reweighting the aforementioned conformal and calibration procedures whenever some unlabeled data from the target distribution is available. We examine these techniques theoretically in a distribution-free framework and demonstrate their excellent practical performance.
1 Introduction
Distribution-free uncertainty quantification for classification uses prediction sets and calibrated probabilities, but distribution shifts can invalidate its guarantees. This work targets the comparatively unaddressed label-shift setting, adapting both approaches and considering target-unlabeled data.
1 Introduction
- Post-hoc calibration and conformal prediction use held-out data with a trained model to quantify predictive uncertainty without distributional assumptions about the data-generating mechanism.The paper notes that this generality comes at a cost discussed later.Source: §1 · p. 2
- Calibration may be more actionable in binary classification, whereas prediction sets become attractive with more labels; neither approach alone fully answers classification UQ.The paper presents the two approaches as complementary and practically efficient, theoretically grounded options.Source: §1 · p. 2
- Source-target distribution shifts can harm generalization and violate UQ assumptions; covariate shift changes q(x) but preserves q(y | x), while label shift changes q(y) but preserves q(x | y).The introduction gives disease prevalence changes during a pandemic as a plausible label-shift example.Source: §1 · p. 2
- Prior distribution-free UQ work addressed covariate shift in regression and binary classification, while this paper fills the label-shift gap by adapting conformal prediction and calibration using unlabeled target data under reasonable assumptions.The paper notes that label shift requires reweighting by label likelihood ratios rather than covariate likelihood ratios.Source: §1 · p. 3
2 Conformal classification
This section develops conformal prediction sets and adapts their calibration to label shift. Experiments show that uncorrected sets lose target coverage, while shift-corrected sets perform better.
2 Conformal classification
- Conformal prediction forms sets of labels using non-conformity scores, offering distribution-free marginal rather than conditional guarantees.The section focuses on probabilistic predictors, though the framework can also use scoring classifiers.Source: §2 · p. 3
2.1 Exchangeable conformal
- Under exchangeability, conformal prediction attains its stated coverage guarantee; when non-conformity scores are almost surely distinct, the coverage probability is upper bounded by 1−α+1/(n+1).The randomized sequence may yield empty, non-actionable prediction sets.Source: §2.1 · p. 4
2.2 Label-shifted conformal
- In toy and wine experiments, shift-corrected sets improve coverage over uncorrected sets; oracle weights match nominal coverage, while estimated weights are slightly worse.The toy experiment reports that uncorrected sets fail to achieve correct marginal coverage; on wine quality, estimated-weight sets have slightly downgraded performance relative to oracle-weight sets.Source: Figure 1 · p. 5, §2.2 · pp. 5–6
- With known label likelihood ratios, reweighting the calibration threshold yields prediction sets with a coverage guarantee under label shift.The guarantee uses weighted exchangeability; the experiments also apply BBSE-estimated importance weights.Source: §2.2 · pp. 5–6
- Plug-in weighted thresholds are consistent when the importance-weight estimator is consistent and the relevant discrete distribution has no jump at level 1−α.At a jump, the coin-toss example shows that the plug-in quantile estimator need not be consistent; the weight estimators rely on assumptions including identifiability and boundedness.Source: §2.2 · p. 6
- Label-conditional conformal splits calibration data by label, giving validity robust to changed class proportions without estimating importance weights.Its finite-sample guarantee comes at the cost of larger sets when classes are not well separated and potential efficiency loss when the number of classes is large.Source: §2.2 · p. 7
3 Calibration
Section 3 defines calibration for multiclass predictors and develops binning-based calibration for i.i.d. data and label shift. It shows how label-shift correction supports calibrated target probabilities, with error controlled by source calibration and weight estimation.
3 Calibration
- Canonical calibration concerns the full probability vector, but binning becomes costly as class counts grow because each bin needs enough calibration data.Class-wise calibration is one stated relaxation for multiclass problems.Source: §3 · pp. 7–8
3.1 Calibration for i.i.d. data
- For i.i.d. data, empirical class frequencies yield approximately calibrated probabilities when the least-populated bin has enough calibration points.The guarantee applies simultaneously across outputs of the recalibrated predictor.Source: §3.1 · p. 9
3.2 Label-shifted calibration
- A predictor calibrated on the source can become miscalibrated under label shift, so target-domain calibration requires correction.The example shows the source Bayes-optimal predictor is no longer calibrated on the target.Source: §3.2 · p. 9
- For Fisher’s LDA, uncorrected binning fails under label shift, while estimated- and oracle-weight corrections both yield near-perfect target calibration.The reliability curves for estimated and oracle weights are almost identical.Source: Figure 2 · p. 10, §3.2 · p. 11
- Target calibration error decomposes into source calibration error and importance-weight estimation error, yielding finite-sample guarantees when both are controlled.The cited results identify BBSE and RLLS as procedures with quantifiable weight-estimation error.Source: §3.2 · p. 11
4 Discussion
The discussion weighs label-shift uncertainty quantification against its scope and identifies full canonical multiclass calibration as an open direction.
4 Discussion
- The paper frames prediction sets and calibration as complementary routes to distribution-free uncertainty quantification for safety-critical deployment.It emphasizes the need for non-trivial guarantees that do not depend on unrealistic or unverifiable assumptions.Source: §4 · p. 11
- Adapting conformal sets and binned calibration addresses label-shift UQ without requiring labeled target data, though prediction-set coverage effects are less clear than calibration effects.In less idealized settings, the paper illustrates that label-shift correction is necessary.Source: §4 · p. 11
- Full canonical calibration guarantees for multiclass problems remain an open direction; the available theory covers binary calibration and multiclass class-wise calibration.The authors also note a gap between MLLS practice and its theoretical foundations when canonical calibration is hard to guarantee.Source: §4 · p. 11
A Importance weights estimation under label shift
This section describes label-shift importance-weight estimation from labeled source and unlabeled target data, along with assumptions and candidate procedures.
A Importance weights estimation under label shift
- BBSE and RLLS estimate label-shift weights from a source confusion matrix and can provide consistent estimates with quantifiable error under relatively mild assumptions [17].The estimation uses held-out labeled source data and unlabeled target data.Source: A Importance weights estimation under… · p. 14
- The setup assumes target labels have support in the source and that the expected source confusion matrix is full-rank.The text characterizes these as absolute continuity and identifiability conditions.Source: A Importance weights estimation under… · p. 14
- Under label shift, target predicted-label frequencies relate to source confusion probabilities through the label importance weights, yielding the BBSE plug-in estimator.The relation is expressed as μ = C_P(f)w.Source: A Importance weights estimation under… · p. 14
- MLLS and variants are alternatives, and the simulations use BBSE-soft while the modular UQ approach permits other weight estimators.The text cites for theoretical analysis of MLLS.Source: A Importance weights estimation under… · p. 14
B Conformal classification
The appendix discusses oracle prediction-set tie handling, including a randomized option for preserving conditional coverage.
B Conformal classification
- The appendix also covers proof details and a real-dataset simulation for the conformal-classification section.It separately discusses oracle-set tie-breaking and randomization.Source: B Conformal classification · p. 14
B.1 Tie-breaking rules for the oracle prediction set
- Ties among equal positive class-posterior probabilities can make an oracle prediction set conservative for some inputs.The oracle set excludes labels whose posterior probability is zero.Source: §B.1 · p. 15
- In the oracle case, ties can be broken arbitrarily while preserving conditional coverage, including by randomizing the order of tied labels.The randomized rule produces a smaller oracle prediction set.Source: §B.1 · p. 15
B.2 Note on randomization and conditional coverage
The note contrasts randomized and non-randomized conformal sequences, arguing that larger randomized sets can still improve conditional coverage. Its example also highlights a practical risk: randomization can produce empty sets when the most-likely label is the only label included.
B.2 Note on randomization and conditional coverage
- Randomized scores can better handle response-distribution heterogeneity across the sample space, unlike randomization of a conformal p-value aimed at conservative marginal coverage.Source: §B.2 · p. 15
- Scheme 1 randomizes all labels, whereas Scheme 2 randomizes all but the most-likely label; randomization can also turn a singleton most-likely set into an empty set.The authors say to apply this safeguard at prediction time only.Source: §B.2 · p. 16
- Only randomized Scheme 1 recovers the oracle prediction sets when the learned threshold is τ⋆=1−α, while all three conformal sets meet marginal coverage.The example uses the Bayes-optimal classifier, isolating the effects of conformalization.Source: §B.2 · p. 16
- Randomization can yield larger sets, but Scheme 1 has superior conditional coverage in this example rather than being inferior on cardinality alone.The authors report that increasing the calibration-set size does not recover the oracle sets for the other methods.Source: §B.2 · p. 16
B.3 Proofs
The proofs use exchangeability arguments to establish conformal coverage, including under label shift. In that setting, the test point receives a label-likelihood-ratio weight relative to the calibration points.
B.3 Proofs
- For fixed estimated class probabilities, exchangeability of the non-conformity scores yields the coverage result conditional on the training data.Source: §B.3 · p. 16
- With randomization, the scores are uniform, so the quantile argument gives nearly tight marginal coverage.Source: §B.3 · p. 16
- Under label shift, the augmented calibration and test data are weighted exchangeable, with test-point weight q(y)/p(y) and calibration-point weights equal to 1.Source: §B.3 · p. 19
B.4 Simulation on real data
The real-data simulation uses white wine data to construct a three-class problem and simulate label shift through resampling. It specifies separate data splits, a neural-network predictor, and BBSE-soft for estimating importance weights.
B.4 Simulation on real data
- The simulation uses 4,898 white-wine instances with 11 features, retaining classes 5, 6, and 7 to form a three-class problem.Less than 10% of the data points were removed.Source: §B.4 · p. 20
- Label shift is simulated by resampling class proportions, followed by separate training, calibration, weight-estimation, and test splits.The source-side splits use approximately 1,000 training, 100 calibration, and 700 weight-estimation instances; target-side splits use about 1,000 weight-estimation and 1,600 test instances.Source: §B.4 · p. 20
- The underlying model is a three-hidden-layer feed-forward neural network, and BBSE-soft estimates the label-shift importance weights.The hidden layers have 128, 64, and 32 neurons, with ℓ2-regularization.Source: §B.4 · p. 20
B.5 Marginal (standard) conformal versus label-conditional conformal
The section contrasts marginal and label-conditional conformal prediction, showing how conditional validity and calibration-set size affect prediction-set size, while label-conditional conformal is robust to changing class proportions.
B.5 Marginal (standard) conformal versus label-conditional conformal
- In the example, stronger class-specific validity makes many prediction sets larger and includes the least-populated class 1.The comparison uses the same randomized non-conformity scores and a shallow MLP.Source: §B.5 · p. 20, Figure 5 · p. 21
- After reweighting with true label likelihood ratios, marginal and label-conditional conformal produce similar average set sizes on the target distribution in 1000 simulations.Without changing class proportions, label-conditional conformal produces larger sets because of its stronger coverage requirement.Source: §B.5 · p. 20
- With approximately 100 calibration points, reduced statistical power to test label inclusion may yield larger prediction sets.The study also considers approximately 350 calibration points.Source: §B.5 · p. 21
- Label-conditional conformal is inherently robust to class-proportion changes and gives exact finite-sample guarantees without importance weights, but may be conservative where classes overlap and requires further calibration-set splitting.The splitting cost may be especially consequential when the number of classes is large.Source: §B.5 · p. 21
C Calibration
- Appendix C.1 contains proofs for Sections 3.1 and 3.2, while C.2 gives details of the real-dataset simulation from Section 3.2.Source: C Calibration · p. 21
C.1 Proofs
The proofs establish calibration guarantees by conditioning on bin assignments and use label shift with Bayes’ rule to express target probabilities; other bounds combine inequalities and a union bound.
C.1 Proofs
- Conditioning on the bin assignments fixes each bin’s calibration count; the proof then uses a high-probability event to establish its proposition.The cited proof fragments describe the event probability and conclude the proposition from event inclusion.Source: §C.1 · pp. 21, 23
- For the calibration guarantee, the proof handles merged bins by conditioning on the bin mapping and applying the tower rule, linearity of expectation, and Jensen’s inequality.Source: §C.1 · pp. 21, 23
- The Proposition 1 proof uses Bayes’ rule and label shift to express target conditional probabilities through source probabilities and target label proportions.It also uses normalization across labels within each bin.Source: §C.1 · p. 23
- Theorem 4’s proof bounds terms using triangle and Hölder inequalities, then combines the bounds and applies a union bound across bins.The cited fragments state the intermediate inequalities and the conditional probability guarantee.Source: §C.1 · pp. 21, 24, 25
C.2 Simulation on real data
A wine-quality simulation evaluates calibration under simulated label shift and finds that reweighting yields approximate target calibration, unlike the uncorrected procedure.
C.2 Simulation on real data
- The simulation uses logistic regression on a binary wine-quality task, with label proportions resampled to (0.8, 0.2) and (0.5, 0.5).Wine ratings of at least 7 out of 10 are treated as good.Source: §C.2 · pp. 25–26
- Proper reweighting leads to approximate calibration on the target domain, while the uncorrected procedure fails to do so.Reliability curves are shown for four random data splits.Source: §C.2 · p. 26
- The experiment uses approximately 1350 training and calibration instances, estimates importance weights from approximately 700 source and 400 target instances, and tests on approximately 1100 instances.Calibration uses uniform-mass binning with 10 bins.Source: §C.2 · p. 26
D Auxiliary results
The auxiliary material supplies supporting probability results and implementation details, alongside reliability-curve evidence that weighted calibration is close to ideal in the wine-quality simulation.
D Auxiliary results
- For conformal classification, the authors replace point masses at infinity with point masses at 1 in the cited lemmas.Source: D Auxiliary results · p. 26
- The auxiliary results include the Bretagnolle-Huber-Carol inequality for multinomial counts.Source: D Auxiliary results · p. 26
- The auxiliary lemmas state probability bounds for exchangeable and weighted-exchangeable variables, including a bound when exchangeable values are almost surely distinct.Source: D Auxiliary results · p. 26
- The stated conformal quantile is the ⌈β(m + 1)⌉-th smallest calibration value when β ≤ m/(m + 1), and otherwise is 1.Source: D Auxiliary results · p. 26
- On the wine-quality reliability curves, oracle and estimated importance-weight calibration are both close to the ideal diagonal, while uncorrected calibration is poor.Source: Figure 7 · p. 27