PAC Prediction Sets Under Covariate Shift

1 INTRODUCTION . A key challenge in machine learning is quantifying prediction uncertainty . A promising approach is via prediction sets , where the model predicts a set of labels instead of a single label . For example , prediction sets can be used by a robot to navigate to avoid regions where the prediction set includes an obstacle , or in healthcare to notify a doctor if the prediction set includes a problematic diagnosis . Various methods based on these approaches can provide probabilistic correctness guarantees ( i.e. , that the predicted set contains the true label with high probability ) when the training and test distributions are equal—formally , assuming the observations are exchangeable ( Vovk et al. , 2005 ; Papadopoulos et al. , 2002 ; Lei et al. , 2015 ) or i.i.d . ( Vovk , 2013 ; Park et al. , 2020a ; Bates et al. , 2021 ) . However , this assumption often fails to hold in practice due to covariate shift—i.e. , where the input distribution changes but the conditional label distribution remains the same . These shifts can be both natural ( e.g. , changes in color and lighting ) ( Hendrycks & Dietterich , 2019 ) or adversarial ( e.g. , ` ∞ attacks ) ( Szegedy et al. , 2014 ) . We consider unsupervised domain adaptation ( Ben-David et al. , 2007 ) , where we are given labeled examples from the source domain , but only unlabeled examples from the target ( covariate shifted ) domain . We propose an algorithm that constructs probably approximately correct ( PAC ) prediction sets under bounded covariate shifts ( Valiant , 1984 ) —i.e. , with high probability over the training data ( “ probably ” ) , the prediction set contains the true label for test instances ( “ approximately correct ” ) . Our algorithm uses importance weights to capture the likelihood of a source example under the target domain . When the importance weights are known , it uses rejection sampling ( von Neumann , 1951 ) to construct the prediction sets . Often , the importance weights must be estimated , in which case we need to account for estimation error . When only given confidence intervals around the importance weights , our algorithm constructs prediction sets that are robust to this uncertainty . We evaluate our approach in two settings . First , we consider rate shift , where the model is trained on a broad domain , but deployed on a narrow domain—e.g. , an autonomous car may be trained on both daytime and nighttime images but operate at night . Even if the model continues to perform well , this covariate shift may invalidate the prediction sets . We show that our approach constructs valid prediction sets under rate shifts constructed from DomainNet ( Peng et al. , 2019 ) , whereas several existing approaches do not . Second , we consider support shift , where the model is trained on one domain , but evaluated on another domain with shifted support . We train ResNet101 ( He et al. , 2016 ) using unsupervised domain adaptation ( Ganin et al. , 2016 ) on both ImageNet-C synthetic perturbations ( Hendrycks & Dietterich , 2019 ) and PGD adversarial perturbations ( Madry et al. , 2017 ) to ImageNet ( Russakovsky et al. , 2015 ) . Our PS-W algorithm satisfies the PAC constraint , and gives prediction sets with the smallest average normalized size among approaches that always satisfy the PAC constraint . See Figure 1 for a summary of our approach and results . Related work . Our work is related to results on conformal prediction ( Vovk et al. , 2005 ; Balasubramanian et al. , 2014 ) without a shift , where the goal is to construct models that predict sets of labels designed to include the true label with high probability ( instead of predicting individual labels ) under the assumption that the source and target distributions are the same . In particular , our setting is related to inductive ( or split ) conformal prediction ( Papadopoulos et al. , 2002 ; Lei et al. , 2015 ) , where the training set is split into a proper training set used to train a traditional predictive model , and a calibration set used to construct the prediction sets . The general approach is to train a model f : X ×Y → R that outputs conformity scores , and then to choose a threshold τ ∈ R that satisfies a correctness guarantee , where the corresponding prediction sets are C ( x ) = { y ∈ Y | f ( x , y ) ≥ τ } ( Park et al. , 2020a ; Gupta et al. , 2021 ) ; this notation is defined in Section 2 . Several kinds of correctness guarantees under no shift have been considered . One possibility is conditional ( Barber et al. , 2020 ) or object conditional ( Vovk , 2013 ) validity , which ensures correctness for all future covariates x , with high probability only over the conditional label distribution p ( y | x ) . This guarantee is very strong , and hard to ensure in practice . A weaker notion is approximate ( Barber et al. , 2020 ) or group ( Romano et al. , 2019 ) conditional validity , which ensures correctness with high probability over p ( y | x ) as well as some distribution p ( x ) over a subgroup . Finally , unconditional validity ensures only correctness over the joint distribution p ( x , y ) . We focus on unconditional validity , though our approach extends straightforwardly to group conditional validity . A separate issue , arising in the no shift setting , is how to condition on the calibration set Z . A conventional approach is to include the randomness in Z with the above , e.g. , for unconditional validity , the high-probability guarantee is now over p ( x , y , Z ) . We refer to this strategy as fully unconditional validity . However , the guarantee we consider uses a separate confidence level for the training data , which is called a training conditional guarantee ( Vovk , 2013 ) ; this correctness notion is equivalent to a PAC correctness guarantee ( Park et al. , 2020a ) , and is also related to the notion of tolerance regions from statistics ( Wilks , 1941 ) . We build on Park et al . ( 2020a ) , which formulates the problem of choosing τ as learning a binary classifier where the input and parameter spaces are both one-dimensional ; thus , the correctness guarantee corresponds to a PAC generalization bound . This approach is widely applicable since it can work with a variety of objectives ( Bates et al. , 2021 ) . Recent work has extended inductive conformal prediction to a setting with covariate shift ( Tibshirani et al. , 2019 ; Lei & Candès , 2020 ) ; similarly , Podkopaev & Ramdas ( 2021 ) considers conformal prediction under label shift ( Lipton et al. , 2018 ) , i.e. , assuming the conditional probabilities p ( x | y ) do not change . These approaches assume that the true importance weights ( IWs ) are known ( Tibshirani et al. , 2019 ; Podkopaev & Ramdas , 2021 ) , or assume some properties of the estimated IWs ( Lei & Candès , 2020 ) , whereas our approach considers uncertainty in the estimated IWs . Furthermore , they are focused on fully unconditional validity , whereas we obtain PAC prediction sets . In addition , Cauchois et al . ( 2020 ) designs confidence sets that are robust to all distribution shifts with bounded f -divergence ; in contrast , we consider the unsupervised learning setting where we have unlabeled examples from the target distribution . We provide additional related work in Appendix A . 2 BACKGROUND ON PAC PREDICTION SETS . We give background on PAC prediction sets ( Park et al. , 2020a ) ; we also re-interpret this approach using Clopper-Pearson confidence intervals ( Clopper & Pearson , 1934 ) , which aids our analysis . 2.1 PAC PREDICTION SETS ALGORITHM . Let x ∈ X be covariates and y ∈ Y be labels . We consider a source distribution P over X ×Y with probability density function ( PDF ) p ( x , y ) .1 A prediction set2 is a set-valued function C : X → 2Y . Inputs . We are given a held-out calibration set Sm of i.i.d . samples ( xi , yi ) ∼ P for i ∈ [ m ] : = { 1 , . . . , m } , written as Sm ∼ Pm , and a score function f : X × Y → R≥0 . For example , f ( x , y ) can be a prediction for the probability that y is the label for x . The score function can be arbitrary , but score functions assigning higher scores to likely labels yield smaller prediction sets . PAC prediction set . A PAC prediction set is a set-valued function C : X → 2Y satisfying the following two conditions . First , C is approximately correct in the sense that its expected error ( failing to contain the true label ) is bounded by a given ε ∈ ( 0 , 1 ) , i.e. , LP ( C ) : = E ( x , y ) ∼P [ 1 ( y /∈ C ( x ) ) ] = P ( x , y ) ∼P [ y /∈ C ( x ) ] ≤ ε . Second , noting that CSm is constructed based on a calibration set Sm ∼ Pm , the condition that CSm is approximately correct must be satisfied with high probability—i.e. , for given δ ∈ ( 0 , 1 ) , we have PSm∼Pm [ LP ( CSm ) ≤ ε ] ≥ 1− δ . Our goal is to devise an algorithm for constructing a PAC prediction set C. Letting C ( x ) = Y for all x ∈ X always satisfies both conditions above , but this extreme case would be uninformative if used as a prediction set . Therefore , we additionally want to minimize the expected sizeE [ S ( C ( x ) ) ] of the prediction sets C ( x ) , where S : 2Y → R≥0 is a size measure , which is application specific ( e.g. , the cardinality of a set in classification ) ; however , we only rely on the monotonicity of the size measure with respect to the prediction set parameterization in construction . Algorithm . To construct C , we first define the search space of possible prediction sets along with the size measure S. We parameterize C by a scalar τ ∈ T : = R≥0 as Cτ ( x ) = { y ∈ Y | f ( x , y ) ≥ τ } , i.e. , τ is the threshold on f ( x , y ) above which we include y in C ( x ) . Importantly , τ controls the tradeoff between size and expected error . The reason is that if τ1 ≤ τ2 , then Cτ2 ( x ) ⊆ Cτ1 ( x ) for all x ∈ X . Thus , size is monotonically decreasing in τ—i.e. , S ( Cτ2 ( x ) ) ≤ S ( Cτ1 ( x ) ) for all x ∈ X , and error is monotonically increasing in τ—i.e. , LP ( Cτ1 ) ≤ LP ( Cτ2 ) . See Figure 1b for an illustration , and Park et al . ( 2020a ) and Gupta et al . ( 2021 ) for details . As a consequence , a typical goal is to construct Cτ that provably contains the true label with high probability , while empirically minimizing size ( Vovk et al. , 2005 ; Gupta et al. , 2021 ) . In our setting , we want to maximize τ ( equivalently , minimize expected size ) subject to the constraint that 1All quantities that we define are measurable with respect to a fixed σ-algebra on X × Y ; for instance , p is the density induced by a fixed σ-finite measure . To be precise , we consider a probability measure P defined with respect to the base measureM on X ×Y ; then , p = dP/dM is the Radon-Nykodym derivative of P with respect to M . For classification , M is the product of a Lebesgue measure on X and a counting measure on Y . 2We use the term “ prediction set ” to denote both the set-valued function and a set output by this function . Cτ is PAC . Let L̄Sm ( Cτ ) : = ∑ ( x , y ) ∈Sm 1 ( y /∈ Cτ ( x ) ) be the empirical error count on Sm , and F ( k ; m , ε ) = ∑k i=0 ( m k ) εi ( 1−ε ) m−i be the cumulative distribution function ( CDF ) of the binomial distribution Binom ( m , ε ) with m trials and success probability ε . In prior work , Park et al . ( 2020a ) constructs Cτ̂ by solving the following problem : τ̂ = max τ∈T τ subj . to L̄Sm ( Cτ ) ≤ k ( m , ε , δ ) , ( 1 ) where k ( m , ε , δ ) = max k∈N∪ { 0 } k subj . to F ( k ; m , ε ) ≤ δ . Intuitively , the PAC guarantee via this construction is related to the Binomial distribution ; L̄Sm ( C ) has distribution Binom ( m , LP ( C ) ) ( given a fixed C ) , since 1 ( y 6∈ C ( x ) ) has a Bernoulli ( LP ( C ) ) distribution when ( x , y ) ∼ P . Thus , k ( m , ε , δ ) defines a “ confidence interval ” such that if L̄Sm ( C ) ≤ k ( m , ε , δ ) , then LP ( C ) ≤ ε with probability at least 1 − δ . Below , we formalize this intuition by drawing a connection to the Clopper-Pearson confidence interval .

Summary: Let $X$ be an instance space, $Y$ a set of labels, $D$ some underlying (hidden) distribution over $X \times Y$. This work studies a new method for converting the output of a probabilistic predictor (e.g. a deep net) to a good prediction set: that is a mapping from $C: X \to 2^Y$, such that for most samples $(x,y) \sim D$, $C(x)$ contains $y$. Formally, the authors study this problem in the PAC setting with covariate shift. The learner is given access to the score function, a labeled sample from the source distribution $P$ over $X \times Y$, and unlabeled samples from a target distribution $Q$ over $X \times Y$ whose marginal over $X$ may be shifted from the source. The goal is to output a prediction set which is as small as possible while still retaining PAC guarantees with respect to the shifted target distribution. This is a well-motivated model in practice. Probabilistic outputs of modern neural nets can be difficult to interpret, and small prediction sets may be useful in settings where one wishes to avoid or check more carefully a few marked outputs (e.g. a problematic medical diagnosis).  Assuming known, bounded importance weights, the authors provide an algorithm satisfying the PAC guarantee by maximizing a cutoff value for the score function that performs well over an empirical sample. They ensure their guarantee holds over the target rather than source distribution by rejection sampling to simulate the target distribution. The authors also extend this to settings where the importance weights are unknown but can be estimated from unlabeled samples. Finally, the authors provide experimental evidence over a couple common settings of covariate shift. They show that their algorithm outperforms baseline methods in the literature in the sense that it maintains PAC guarantees while outputting a smaller prediction set  in expectation.