Uncertainty quantification can be a critical tool for building reliable systems from machine learning components. For example, a medical decision support system can convey uncertainty to a doctor, or a robot can act conservatively with respect to uncertainty. These approaches are particularly important when the data distribution shifts as the predictive system is deployed, since they enable the decision-maker to react to degraded performance.
Conformal prediction (Vovk et al., 2005;Angelopoulos & Bates, 2021) is a promising approach to uncertainty quantification, aiming to outputs sets of labels instead of a single label. Under standard assumptions (i.i.d. or exchangeable data), it guarantees that the prediction set contains the true label with high probability. We consider probably approximately correct (PAC) (or calibration-setconditional) guarantees (Vovk, 2012;Park et al., 2019), which ensure high probability coverage over calibration datasets used to construct the prediction sets.
In this paper, we propose a novel prediction set algorithm that provides PAC guarantees under the label shift setting, where the distribution of the labels may shift, but the distribution of covariates conditioned on the labels remains fixed. For instance, during a pandemic, a disease may spread to a much larger fraction of the population, but the manifestations of the disease may remain the same. As another example, real-world data may have imbalanced classes, unlike the balanced classes typical of curated training datasets. We consider the unsupervised domain adaptation setting (Ben-David et al., 2006), where we are given labeled examples from a source domain, but only unlabeled examples from the target domain, and care about performance in the target domain.
A standard way to adapt conformal inference to handle distribution shift is by using importance weighting to "convert" data from the source distribution into data from the target distribution (Tibshirani et al., 2019). In the label shift setting, one possible way to express the importance weights is w * = C -1 P q * , where C P is the confusion matrix and q * is the distribution of predicted labels (Lipton et al., 2018); see details below. However, the estimation error for the unknown C P and q * breaks the PAC guarantee. An example of our approach on the ChestX-ray dataset. In the unshifted setting, standard PAC prediction sets guarantee high-probability coverage, but this guarantee fails under label shift. Our approach addresses this challenge and continues to work in the shifted environment.
Instead, we construct confidence intervals around C P and q * , and then devise a novel algorithm to propagate these intervals through the Gaussian elimination algorithm used to compute w * . Finally, we leverage an existing strategy for constructing PAC prediction sets when given confidence intervals for the importance weights (Park et al., 2021).
We empirically evaluate our approach on four datasets across three application domains: CIFAR-10 ( Krizhevsky et al., 2009) in the computer vision domain, the CDC Heart Dataset (Centers for Disease Control and Prevention (CDC), 1984) and ChestX-ray (National Institutes of Health and others, 2022) in the medical domain, and AGNews (Zhang et al., 2015) in the language domain.
Contributions. We propose a novel algorithm for constructing PAC prediction sets in the presence of label shift, which computes provably valid intervals around the true importance weights. Our algorithm is based on a technique for propagating confidence intervals through the updates of Gaussian elimination, which to our knowledge is a novel approach to uncertainty propagation in a prediction set construction setting. Finally, we empirically demonstrate that our approach satisfies the PAC guarantee while constructing smaller prediction sets than several baselines.
Example. Figure 1 illustrates a use case of our technique on the ChestX-ray dataset. In medical settings, prediction sets (denoted PS) provide a rigorous way to quantify uncertainty for making downstream decisions. In particular, they can guarantee that the prediction set contains the true label (in this case, a diagnosis) with high probability. However, label shift happens commonly in medical settings, for instance, many illnesses have varying rates of incidence over time even when the patient population remains the same. Unfortunately, label shift breaks the PAC coverage guarantee. Our approach (denoted PSW) corrects for the label shift via importance weighting; it does so in a provably correct way. The resulting prediction sets satisfy the PAC guarantee.
There has been recent interest in conformal inference under distribution shift, much of it focusing on covariate shift (Tibshirani et al., 2019;Lei & Candès, 2021;Qiu et al., 2022). (Podkopaev & Ramdas, 2021) develop methods for marginal coverage under label shift, whereas we are interested in training-set conditional-or PAC-guarantees. Furthermore, they formally assume that the true importance weights are known exactly, which is rarely the case. In the label shift setting, the importance weights can be estimated (Lipton et al., 2018), but as we show in our experiments, uncertainty in these estimates must be handled for the PAC guarantee to hold. elimination to propagate uncertainty through the computation of the weights w * = C -1 P q * . We give a more comprehensive discussion of related work in Appendix A.
2.1 BACKGROUND ON LABEL SHIFT Consider the goal of training a classifier g : X → Y, where X ⊆ R d is the covariate space, and Y = [K] = {1, ..., K} is the set of labels. We consider the setting where we train on one distribution P over X × Y-called the source-with a probability density function (PDF) p : (x, y) → p(x, y), and evaluate on a potentially different test distribution Q-called the target-with PDF q : (x, y) → q(x, y). We focus on the unsupervised domain adaptation setting (Ben-David et al., 2007), where we are given an i.i.d. sample S m ∼ P m of m labeled datapoints, and an i.i.d. sample of n unlabeled datapoints T n X ∼ Q n X . The label shift setting (Lipton et al., 2018) assumes that only the label distribution Q Y may be change from P Y , and the conditional covariate distributions remain the same:
We denote p(y) = P Y (Y = y) for all y ∈ Y and analogously for Q. (Lipton et al., 2018) consider two additional mild assumptions: Assumption 2.2. For all y ∈ Y such that q(y) > 0, we have p(y) > 0.
Next, given the trained classifier g :
This last assumption requires that the per-class expected predictor outputs be linearly independent; for instance, it is satisfied when g is reasonably accurate across all labels. In addition, one may test whether this assumption holds (Lipton et al., 2018).
Denoting the importance weights w * := (q(y)/p(y)) y∈Y ∈ R K , and ŷ := g(x), we will write p(ŷ|y) = P (X,Y )∼P [g(X) = ŷ|Y = y], and define p(ŷ, y), p ŷ as well as the corresponding expressions for q analogously. Since ŷ depends only on x, we have q(ŷ | y) = p(ŷ | y). Thus, see e.g., Lipton et al. (2018),
or in a matrix form, q * = C P w * , where q * := (q ŷ ) ŷ∈Y ∈ R K . As we assume C P is invertible,
Our algorithm uses this equation to approximate w * , and then use its approximation to construct PAC prediction sets that remain valid under label shift.
We are interested in constructing a prediction set C : X → 2 Y , which outputs a set of labels C(x) ⊆ Y for each given input x ∈ X rather than a single label. The benefit of outputting a set of labels is that we can obtain correctness guarantees such as:
where ε ∈ (0, 1) is a user-provided confidence level. Then, downstream decisions can be made in a way that accounts for all labels y ∈ C(x) rather than for a single label. Thus, prediction sets quantify uncertainty. Intuitively, equation 2 can be achieved if we output C(x) = Y for all x ∈ X , but this is not informative. Instead, the typical goal is to output prediction sets that are as small as possible.
The typical strategy for constructing prediction sets is to leverage an fixed existing model. In particular, we assume given a scoring function f : X × Y → R; most deep learning algorithms provide such scores in the form of predicted probabilities, with the corresponding classifier being g(x) = arg max y∈Y f (x, y). The scores do not need to be reliable in any way; if they are unreliable, the PAC prediction set algorithm will output larger sets. Then, we consider prediction sets parameterized by a real-valued threshold τ ∈ R:
In other words, we include all labels with score at least τ . First, we focus on correctness for P , in which case we only need S m , usually referred to as the calibration set. Then, a prediction set algorithm constructs a threshold τ (S m ) ∈ R and returns C τ (Sm) .
Finally, we want τ to satisfy (2); one caveat is that it may fail to do so due to randomness in S m . Thus, we allow an additional probability δ ∈ R of failure, resulting in the following desired guarantee:
Vovk (2012); Park et al. (2019) propose an algorithm τ that satisfies (3), see Appendix B.
Finally, we are interested in constructing PAC prediction sets in the label shift setting, using both the labeled calibration dataset S m ∼ P m from the source domain, and the unlabeled calibration dataset T X n ∼ Q n from the target distribution. Our goal is to construct τ (S m , T X n ) based on both S m and T X n , which satisfies the coverage guarantee over Q instead of P :
Importantly, the inner probability is over the shifted distribution Q instead of P .
To construct prediction sets valid under label shift, we first notice that it is enough to find elementwise confidence intervals for the importance weights w * . Suppose that we can construct
Then, when adapted to our setting, the results of Park et al. (2021)-originally for the covariate shift problem-provide an algorithm that returns a threshold τ (S m , V, W, b), where V ∼ Uniform([0, 1]) K is a vector of random variables, such that
(5) This is similar to equation 4 but it accounts for the randomness used by our algorithm-via V -in the outer probability. We give the details in Appendix C.
The key challenge is to construct W = k∈Y [w k , w k ] such that w * ∈ W with high probability. The approach from Park et al. (2021) for the covariate shift problem relies on training a source-target discriminator, which is not possible in our case since we do not have class labels from the target domain. Furthermore, Park et al. (2021)'s approach is does not provide conditions under which one can provide a valid confidence interval for the importance weights in their setting.
Our algorithm uses a novel approach, where we propagate intervals through the computation of importance weights. The weights w * are determined by the system of linear equations C P w * = q * . Since C P and q * are unknown, we start by constructing element-wise confidence intervals
with probability at least 1 -δ over our calibration datasets S m and T X n . We then propagate these confidence intervals through each step of Gaussian elimination, such that at the end of the algorithm, we obtain confidence intervals for its output-i.e., w * ≤ w * ≤ w * with probability at least 1 -δ.
Finally, we can use (7) with the algorithm from (Park et al., 2021) to construct PAC prediction sets under label shift. We describe our approach below.
Recall that C P = (c ij ) ij∈Y and q * = (q k ) k∈Y . Note that
Similarly, q k is the mean of 1(g(X) = k) over the randomness in X ∼ Q X . Thus, we can use the Clopper-Pearson (CP) intervals (Clopper & Pearson, 1934) for a Binomial success parameter to construct intervals around c ij and q k . Given a confidence level δ ∈ (0, 1) and the sample mean ĉij =foot_0 m (x,y)∈Sm 1(g(x) = i, y = j)-distributed as a scaled Binomial random variable-this is
Similarly, for q k , we can construct CP intervals based on qk = 1
for confidence levels δ ij and δ k chosen later, we obtain for all i, j, k ∈ [K],
Then, the following result holds by a union bound: Given any
(9)
We also need to set up notation for Gaussian elimination, which requires us to briefly recall the algorithm. To solve C P w * = q * , Gaussian elimination (see e.g., Golub & Van Loan, 2013) proceeds in two phases. Starting with c 0 = C P and q 0 = q * , on iteration t ≥ 1, Gaussian elimination uses row k = t to eliminate the kth column of rows i ∈ {k + 1, ..., K} (we introduce a separate variable k for clarity). In particular, if c t kk ̸ = 0, we denote
If c t kk = 0, but there is an element j > k in the kth column such that c t jk ̸ = 0, the kth and the jth rows are swapped and the above steps are executed. If no such element exists, the algorithm proceeds to the next step. At the end of the first phase, the matrix c K-1 has all elements below the diagonal equal to zero-i.e., c K-1 ij = 0 if j < i. In the second phase, the Gaussian elimination algorithm solves for w * i backwards from i = K to i = 1, introducing the following notation. For each i, if
, where
In our setting, we do not know c 0 and q 0 ; instead, we assume given entrywise confidence intervals as in equation 6, which amount to c 0 ≤ c 0 ≤ c 0 and q 0 ≤ q 0 ≤ q 0 . We now work on the event Ω that these bounds hold, and prove that our algorithm works on this event; later, we combine this result with Equation 9to obtain a high-probability guarantee. Then, our goal is to compute c t , c t , q t , q t such that for all iterations t ∈ {0, 1, ..., K -1}, we have elementwise confidence intervals specified by c t , c t , q t and c t for the outputs c t , q t of the Gaussian elimination algorithm:
The base case t = 0 holds by the assumption. Next, to propagate the uncertainty through the Gaussian elimination updates for each iteration t ∈ [K -1], our algorithm sets
for the lower bound, and computes
for the upper bound. The first case handles the fact that Gaussian elimination is guaranteed to zero out entries below the diagonal, and thus these entries have no uncertainty remaining. The second rule constructs confidence intervals based on the previous intervals and the algebraic update formulas used in Gaussian elimination for the entries for which i, j > k. For instance, the above confidence intervals use that on the event Ω, and by induction on t, if c t ij ≥ 0 and c t ii > 0 for all i, j ∈ [K] and for all t, the Gaussian elimination update c t+1 ij = c t ij -c t it c t tj /c t tt can be upper bounded as
The assumptions that c t ij ≥ 0 and c t ii > 0 for all i, j ∈ [K] and for all t may appear a little stringent, but the former can be removed at the cost of slightly larger intervals propagated to the next step, see Section D. The latter condition is satisfied by any classifier that obtains sufficient accuracy on all labels. We further discuss these conditions in Section D. The third rule in equation 11 and equation 12 handles the remaining entries, which do not change; and thus the confidence intervals from the previous step can be used. The rules for q are similar, and have a similar justification:
For these rules, our algorithm assumes q t i ≥ 0 for all i ∈ [K] and all t, and raises an error if this fails. As with the first condition above, this one can be straightforwardly relaxed; see Appendix D.
In the second phase, we compute w * i starting from i = K and iterating to i = 1. On iteration i, we assume we have the confidence intervals w * j ≤ w * j ≤ w * j for j > i. Then, we compute confidence intervals for the sum s i , with a similar justification based on the Gaussian elimination updates:
and show that they satisfy s i ≤ s i ≤ s i on the event Ω. Finally, we compute confidence intervals for w * i , assuming c K-1 ii > 0:
for which we can show that they satisfy w * i ≤ w * i ≤ w * i based on the Gaussian elimination updates. Letting W = {w | w * ≤ w ≤ w * }, we have the following (see Appendix E for a proof). Lemma 3.1 (Elementwise Confidence Interval for Importance Weights). If (6) holds, and for all i, j, t
We mention here that the idea of algorithmic uncertainty propagation may be of independent interest.
In future work, it may further be developed to other methods for solving linear systems (e.g., the LU decomposition, Golub & Van Loan ( 2013)), and other linear algebraic and numerical computations.
Algorithm 1 summarizes our approach. As can be seen, the coverage levels for the individual Clopper-Pearson intervals are chosen to satisfy the overall 1 -δ coverage guarantee. In particular, the PAC guarantee equation 4 follows from equation 5, equation 9, Lemma 3.1, and a union bound; we provide a more detailed explanation in Appendix F. Theorem 3.2 (PAC Prediction Sets under Label Shift). For any given ε, δ ∈ (0, 1), under Assumptions 2.1, 2.2 and 2.3, if ∀i, j, t ∈ [K], we have c t ij ≥ 0, c t ii > 0, and q t i ≥ 0, then Algorithm 1 satisfies
As discussed in Appendix D, we can remove the requirement that c t ij ≥ 0 and q t i ≥ 0 for i ̸ = j.
Algorithm 1 PS-W: PAC prediction sets in the label shift setting.
11: return c, c, q, q 12: procedure INTERVALGAUSSIANELIM(c 0 , c 0 , q 0 , q 0 ) 13:
Compute c t ij , c t ij using ( 11) & ( 12), and q t i , q t i using ( 14) 16:
if c t ij < 0 for some i ̸ = j or c t ii ≤ 0 for some i or q t i ≤ 0 for some i, then return ∅ 17:
for i ∈ [K, ..., 1] do 18:
Compute s t i , s t i using (15), and w i , wi using ( 16) 19:
4 EXPERIMENTS
Predictive models. We analyze four datasets: the CDC Heart dataset, CIFAR-10, Chest X-ray, and the AG News dataset; details are provided in Section 4.2. We use a two-layer MLP for the CDC Heart data with an SGD optimizer having a learning rate of 0.03 and a momentum of 0.9, using a batch size of 64 for 30 epochs. For CIFAR-10, we finetune a pretrained ResNet50 He et al. (2016), with a learning rate of 0.01 for 56 epochs. For the ChestX-ray14 dataset, we use a pre-trained CheXNet (Rajpurkar et al., 2017) with a DenseNet121 (Huang et al., 2017) backbone with a learning rate of 0.0003 for two epochs. For AGNews, we fine-tune a pre-trained Electra sequence classifier for one epoch with an AdamW optimizer using a learning rate of 0.00001.
Hyperparameter choices. There are two user-specified hyperparameters that control the guarantees, namely δ and ε. In our experiments, we choose δ = 5 × 10 -4 to ensure that, over 100 independent datasets S m , there is a 95% probability that the error rate is not exceeded. Specifically, this ensures that
] ≥ 1 -ε holds for all 100 trials, with probability approximately 1 -0.95 1/100 ≈ 5 × 10 -4 . We select ε for each dataset in a way that the resulting average prediction set size is greater than one. This ensures that the coverage guarantee is non-trivial, as a single-valued classifier should not achieve the desired coverage rate (as δ is small).
Dataset construction. We follow the label shift simulation strategies from previous work (Lipton et al., 2018). First, we split the full dataset into training data, and "base" source and target datasets.
We use the training dataset to fit a score function. Given label distributions P Y and Q Y , we generate the source dataset S m , target dataset T X n , and a labeled, size o target test dataset (sampled from Q) by sampling with replacement from the corresponding base datasets. We consider two choices of P Y and Q Y : (i) a tweak-one shift, where change the probability of one label, and keep the relative proportions of the other labels equally likely, and (ii) a general shift, where we shift each probability as described later.
Baselines. We compare our approach (PS-W) with several baselines (see Appendix G):
• PS: PAC prediction sets that do not account for label shift (Vovk, 2012;Park et al., 2019). This does not come with PAC guarantees under label shift. • WCP: Weighted conformal prediction under label shift, which targets marginal coverage (Podkopaev & Ramdas, 2021). This does not come with PAC guarantees under label shift either. • PS-R: PAC prediction sets that account for label shift but ignore uncertainty in the importance weights; again without guarantees. • PS-C: This addresses label shift via a conservative upper bound on the empirical loss (see Appendix G for details). This is the only baseline with a PAC guarantee under label shift.
We compare to other baselines in Appendix I.5, and to an oracle with the true weights in Appendix H.
Results for other hyperparameters are in Appendix I.
Metrics. We measure performance via the prediction set error, i.e., the fraction of (x, y) ∼ Q such that y / ∈ C τ (x); and the average prediction set size, i.e., the mean of |C τ (x)|, evaluated on the held-out test set. We report the results over 100 independent repetitions, randomizing both dataset generation and our algorithm.
CDC Heart. We use the CDC Heart dataset, a binary classification problem (Centers for Disease Control and Prevention (CDC), 1984). The goal is to predict the risk of heart attack given features such as level of exercise or weight. We use ε = 0.1 and δ = 5 × 10 -4 . We consider both large and small shifts. For the large shift, the label distributions-denoted (pos%, neg%)-are (94%, 6%) for source, and (63.6%, 36.4%) for target; results are in Figure 2b. We also consider a small shift with label distributions (94%, 6%) for source, and (91.3%, 8.7%) for target; results are in Figure 2a. As can be seen, our PS-W algorithm satisfies the PAC guarantee while achieving smaller prediction set size than PS-C, the only baseline to satisfy the PAC guarantee. The PS and PS-R algorithms violate the PAC guarantee. CIFAR-10. Next, we consider CIFAR-10 ( Krizhevsky et al., 2009), which has 10 labels. We use ε = 0.1 and δ = 5 × 10 -4 . We consider a large and a small tweak-one shift. For the large shift, the label probability is 10% for all labels in the source, 40.0% for the tweaked label, and 6.7% for other labels in the target; results are in Figure 3a. For small shifts, we use 10% for all labels for the source, 11.76% for the tweaked label, and 9.8% for other labels for the target; results are in Figure 3b. Under large shifts, our PS-W algorithm satisfies the PAC guarantee while outperforming PS-C by a large margin. When the shift is very small, PS-W still satisfies the PAC guarantee, but generates more conservative prediction sets similar in size to those of PS-C (e.g., Figure 3b) given the limited data. Results for a non-uniform source distribution and general shifts are shown in Appendix I.8.
AGNews. AG News is a subset of AG's corpus of news articles (Zhang et al., 2015). It is a text classification dataset with four labels: World, Sports, Business, and Sci/Tech. It contains 31,900 unique examples for each class. We use ε = 0.05 and δ = 5 × 10 -4 . We use tweak-one label shifts. We consider a large shift and a medium-sized calibration dataset, with label distributions equalling (30.8%, 30.8%, 7.7%, 30.8%) for the source, and (12.5%, 12.5%, 62.5%, 12.5%) for the target; results are in Figure 4a. As before, our PS-W approach satisfies the PAC guarantee while achieving smaller set sizes than PS-C.
ChestX-ray. ChestX-ray14 (Wang et al., 2017) is a medical imaging dataset containing about 112K frontal-view X-ray images of 30K unique patients with fourteen disease labels. This dataset contains PS PS-C PS-R PS-W WCP 0.00 0.02 0.04 0.06 0.08 0.10 0.12 0.14 prediction set error = 0.05 PS PS-C PS-R PS-W WCP 0.00 0.25 0.50 0.75 1.00 1.25 1.50 1.75 2.00 size (a) Prediction set error and size on the AGNews Dataset. Parameters are ε = 0.05, δ = 5 × 10 -4 , m = 26000, n = 12800, and t = 12800. PS PS-C PS-R PS-W WCP 0.0 0.1 0.2 0.3 0.4 prediction set error = 0.30 PS PS-C PS-R PS-W WCP 0 1 2 3 4 5 size (b) Prediction set error and size on the ChestX-ray dataset. Parameters are ε = 0.3, δ = 5 × 10 -4 , m = 67200, n = 35200, and o = 3520. instances with multiple labels, which we omit. We also omit classes with few positively labeled datapoints, leaving six classes: Atelectasis, Effusion, Infiltration, Mass, Nodule, Pneumothorax. We consider a general label shift. We consider a large tweak-one shift, with label distributions of (19.1%, . . . , 19.1%, 4.5%, 19.1%) for the source, and (11.1%, . . . , 11.1%, 44.5%, 11.1%) for the target. Results for ε = 0.3 are in Figure 4b. As before, our PS-W approach satisfies the PAC guarantee while outperforming PS-C. The PS-R and WCP methods violate the constraint.
Discussion. In all our experiments, our approach satisfies the PAC guarantee; furthermore, it produces smaller prediction set sizes than PS-C-the only baseline to consistently satisfy the PAC guaranteeexcept when the label shift is small and the calibration dataset is limited. In contrast, the PS baseline does not account for label shift, and the PS-R baseline does not account for uncertainty in the importance weights, so they do not satisfy the PAC guarantee. The WCP baseline is designed to target a different guarantee, and it does not satisfy the PAC guarantee. Thus, these results demonstrate the efficacy of our approach.
Limitations. Our approach is focused on problem settings where the label shift is not too small and sufficient calibration data is available; and may produce conservative prediction sets otherwise. This reflects the intrinsic difficulty of the problem in these settings. Importantly, our PAC coverage guarantees still hold.
We have proposed a PAC prediction set algorithm for the label shift setting, and illustrated its effectiveness in experiments. Directions for future work include improving performance when the calibration dataset or the label shift is small. Reproducibility statement. Our code is available at https://github.com/averysi224/ pac-ps-label-shift for reproducing our experiments.
Conformal Prediction. Our work falls into the broad areas of distribution-free uncertainty quantification, conformal prediction and tolerance regions (e.g., Guttman, 1970;Vovk et al., 2005;Balasubramanian et al., 2014;Angelopoulos & Bates, 2021;Kaur et al., 2022;Li et al., 2022;Dobriban & Yu, 2023, etc), which aim to construct prediction sets with finite sample coverage guarantees. The prediction sets are often realized by a setting a threshold on the values of a traditional single-label predictor (i.e. conformity/non-conformity score function) and predict all labels with scores above this threshold. Our setting is related to inductive conformal prediction (Papadopoulos et al., 2002;Vovk, 2012;Lei et al., 2015), where the dataset is split into a training set for fitting the scoring function and a calibration set for constructing the prediction sets.
A standard coverage guarantee for conformal prediction methods is marginal coverage over the training, calibration and test data. For inductive conformal prediction, training-data conditional correctness (Vovk, 2012;Park et al., 2019) aims to achieve an (ε, δ)-coverage guarantee, where coverage over test points exceeds the desired error rate level ε with probability at most δ. Sets satisfying this guarantee have been termed PAC prediction sets (Park et al., 2019). This guarantee is also equivalent to the classical coverage properties of tolerance regions (Wilks, 1941;Fraser, 1956). PAC coverage may be desired when a slight excess of the desired error rate cab cause a significant safety concern.
Label shift Importance weights, as a function of x, y q(x)/p(x) q(y)/p(y) Shared conditional distribution, for all x, y p(y
Table 1: Comparison of covariate shift and label shift.
Label shift. Label shift (Zadrozny, 2004;Huang et al., 2006;Sugiyama et al., 2007;Gretton et al., 2009) supposes that the conditional covariate distribution is fixed but the label distribution may change from a source to a target distribution; see Table 1. In more detail, let p and q denote the probability density functions of the source and target domains with respect to a common dominating measure, respectively. Label shift assumes that for all labels y, q(y) may differ from p(y), but for all features x and labels y, we have p(x | y) = q(x | y). Label shift can arise when the representations of the classes changes, for instance in scenarios like medical diagnosis and object recognition (Storkey et al., 2009;Saerens et al., 2002;Lipton et al., 2018).
Early solutions required the estimation of q and p, which may scale poorly with the dimension (Chan & Ng, 2005;Storkey et al., 2009;Zhang et al., 2013). More recently, two approaches achieved scalability by assuming an approximate relationship between ground truth labels y and the outputs ŷ of a classifier (Lipton et al., 2018;Azizzadenesheli et al., 2019;Saerens et al., 2002): Black Box Shift Estimation (BBSE) (Lipton et al., 2018) and RLLS (Azizzadenesheli et al., 2019) provided consistency results and finite-sample guarantees assuming the confusion matrix is invertible. Subsequent work developed a unified framework that decomposes the error in computing importance weights into miscalibration error and estimation error, with BBSE as a special case (Garg et al., 2020); this approach was extended to open set label shift domain adaptation (Garg et al., 2022). We also mention that label shift is a special case of a more general class of distribution shifts that allow the source and target distributions to be related by sharing sequential conditional distributions (Qiu et al., 2023).
Conformal methods and distribution shift. Conformal prediction has been adapted to distribution shifts such as covariate shift (Tibshirani et al., 2019;Lei & Candès, 2021;Qiu et al., 2022). (Podkopaev & Ramdas, 2021) considers label shift when the true importance weights are exactly known. Importance weights are typically estimated with uncertainty. We address this by developing an algorithm to find confidence intervals for the importance weights. More broadly, prediction sets have been studied in the meta-learning setting (Dunn et al., 2022;Park et al., 2022), as well as the setting of robustness to all distribution shifts with bounded f -divergence (Cauchois et al., 2020).
Class-conditional prediction sets. Although not designed specifically for solving the label shift problem, methods for class-conditional coverage (Sadinle et al., 2019) and adaptive prediction sets (APS) (Romano et al., 2020) can improve robustness to label shifts. However, class-conditional coverage is a stronger guarantee that leads to prediction sets larger than for our algorithm, while APS does not satisfy a PAC guarantee; we compare empirically to these approaches in Appendix I.5.
where k(m, ε, δ) = max
Their approach is equivalent to the method from Vovk (2012), but formulated in the language of learning theory. By viewing the prediction set as a binary classifier, the PAC guarantee via this construction can be connected to the Binomial distribution. Indeed, for fixed C, LSm (C) has distribution Bionm(m, L P (C)), since 1(y / ∈ C(x)) has a Bernoulli(L P (C)) distribution when (x, y) ∼ P . Thus, k(m, ε, δ) defines a bound such that if L P (C) ≤ ε, then LSm (C) ≤ k(m, ε, δ) with probability at least 1 -δ.
Here we demonstrate how to obtain prediction sets given intervals w * ≤ w * ≤ w * around the true importance weights. This approach is based closely on the strategy in (Park et al., 2021) for constructing prediction sets under covariate shift, but adapts it to the label shift setting (indeed, our setting is simpler since there are finitely many importance weights). The key challenge is computing the importance weight intervals, which we describe in detail below.
Given the true importance weights w * , one strategy would be to use rejection sampling (Von Neumann, 1951;Shapiro, 2003;Rubinstein & Kroese, 2016) to subsample S m to obtain a dataset that effectively consists of N ≤ m i.i.d. samples from Q (here, N is a random variable, but this turns out not to be an issue). Essentially, for each (x i , y i ) ∈ S m , we sample a random variable V i ∼ Uniform([0, 1]), and then accept samples where V i ≤ w * yi /b, where b is an upper bound on w * y :
In our setting, we can take b = max y∈Y w * y . Then, we return τ (T N (S m , V, w * , b)). Since T N (S m , V, w * , b) consists of an i.i.d. sample from Q, we obtain the desired PAC guarantee (4).
In practice, we do not know the true importance weights w * . Instead, suppose we can obtain intervals W y = [w * y , w * y ] such that w * y ∈ W y with high probability. We let W = y∈Y W y , and assume w * ∈ W with probability at least 1 -δ. The algorithm proposed in (Park et al., 2021) adjusts the above algorithm to conservatively account for this uncertainty-i.e., it chooses τ so the PAC guarantee (4) holds for any importance weights w ∈ W :
We minimize over τ since choosing smaller τ leads to larger prediction sets, which is more conservative. (Park et al., 2021) show how to compute (18) efficiently. We have the following guarantee: Theorem C.1 (Theorem 4 in (Park et al., 2021)). Assume that w * ∈ W . Letting U = Uniform([0, 1]),
In other words, the PAC guarantee (4) holds, with the modification that the outer probability includes the randomness over the samples V ∼ U m used by our algorithm.
The diagonal elements c kk of the confusion matrix of an accurate classifier, are typically much larger than the other elements. Indeed, for an accurate classifier, the probabilities of correct predictions c kk = P (g(x) = k, y = k) are higher than those of incorrect predictions c ik := P (g(x) = i, y = k), i ̸ = k. On the other hand, the Clopper-Pearson interval is expected to be short (for instance, the related Wald interval has length or order 1/ √ m, where m is the sample size). Thus, we expect that
Without loss of generality, we consider equation 11 as an example. In the Gaussian elimination process, recall that the update at step t is
Combining with equation 19, the factor by which the k-th row is multiplied is small, i.e., c t ik /c t kk ≪ 1. Thus the resulting i-th diagonal elements Balanced classifier. For a balanced classifier, when c ik and c jk are close for all i, j such that i ̸ = k, j ̸ = k, since the factor c t ik /c t kk is small, the lower bound in equation 20 is expected to be positive. Imbalanced classifier. For the more general case of a possibly imbalanced classifier, c ij and c kj may not be close. This could cause non-positive bounds at certain steps, so the confidence interval may not be valid at the next steps; e.g., equation 13 may fail. However, note that since
t ik ], c t kj ∈ [c t kj , c t kj ], we have c t ik c t kj ≤ max{|c t ik |, |c t ik |} • max{|c t kj |, |c t kj |} and hence c t ik c t kj c t kk ≤ max{|c t ik |, |c t ik |} • max{|c t kj |, |c t kj |} c t kk . In fact, one can derive the even tighter bound max c t ik c t kj c t kk , c t ik c t kj c t kk , c t ik c t kj c t kk , c t ik c t kj c t kk , c t ik c t kj c t kk , c t ik c t kj c t kk , c t ik c t kj c t kk , c t ik c t kj c t kk .
This can be checked by carefully going through all possible cases of positive and negative values of the bounds. Similar changes can be made to computing the upper bounds.
It is possible for our final interval W to contain negative lower bounds due to loose element-wise intervals or other factors. Since importance weights are non-negative, negative importance weight bounds act the same way as zero lower bounds in rejection sampling, and preserve our guarantees.
Finally, the requirement of an accurate classifier is already imposed by methods such as BBSE to ensure the invertibility of the confusion matrix. Therefore, our Gaussian elimination approach does not impose significantly stronger assumptions.
E PROOF OF LEMMA 3.1
For the first phase, we prove by induction on t that (10) holds for all t. The base case t = 0 holds by assumption. For the induction step, we focus on c t+1 ij ; the remaining bounds c t+1 ij , q t+1 k , and q t+1 k follow similarly. There are three sub-cases i, each corresponding to one of the update rules in (11).
For the first update rule c t+1 ij = 0, equation 10 follows since the Gaussian elimination algorithm guarantees that c t+1 ij = 0. For the second and third update rules, equation 10 follows by direct algebra and the induction hypothesis. For instance, for i, j > t, equation 13 holds, and similarly c t+1 ij ≥ c t+1 ij . For the second phase, the fact that s ≤ s ≤ s and w * ≤ w * ≤ w * follows by a similar induction argument. Since Gaussian elimination guarantees that w * = C -1 P q * , and we have shown that w * ∈ W = K i=1 [w i , w i ], the claim follows.
F PROOF SKETCH OF THEOREM 3.2 First, Theorem 3.2 follows from equation 5, assuming the given confidence intervals W i = [w i , w i ] for each importance weight w * i is valid-i.e., w * i ∈ W i . Equation 5follows by Theorem 4 in Park et al. (2021). Roughly speaking, if w * i is known, then we can use a standard rejection sampling procedure based on w * i to convert S m into a set of i.i.d. samples from Q. Then, the PAC guarantee would follow by standard conformal prediction results, e.g., Park et al. (2020). When w * i is not known, then intuitively, Park et al. (2021) takes the worst case over all w * i ∈ W i . They use a reparameterization trick to do so in a computationally efficient way. Finally, by Lemma 3.1, we have w * i ∈ W i for all i ∈ [n] with high probability, so Theorem 3.2 follows by a union bound.
We describe PS-C, the conservative baseline summarized in Algorithm 2. In particular, given an upper bound b ≥ w * on the importance weight, we use the upper bound c, c, q, q ← CPINTERVAL(S m , T X n , x → arg max y∈Y f (x, y), δ w ) 3:
W ← INTERVALGAUSSIANELIM(c, c, q, q)
Hence, using the PS prediction set algorithm with parameters (ε/b, δ), the output is (ε, δ)-PAC.
Here, we show comparisons to an oracle that is given the ground truth importance weights (which are unknown and must be estimated in most practical applications). It uses rejection sampling according to these weights rather than conservatively accounting for uncertainty in the weights. In contrast to our baselines (which need to estimate w), this oracle represents a "gold (yet not achievable) standard" to compare with. It enables us to quantify the increase in average prediction set size due to uncertainty in our estimates of w, stemming from (i) finite sample error, accounted for by Clopper-Pearson intervals, and (ii) label shift, accounted for by our interval Gaussian elimination algorithm.
The algorithm requires further discussion if c K-1 ii = 0 (Golub & Van Loan,
2013); this does not commonly happen in our motivating application so we will not consider this case. See Appendix D for details.