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: This paper is concerned with learning prediction sets (in a PAC sense) under the covariate shift assumption, given a model $f(x,y)$. The form of the sets is restricted to $C_\tau(x) := \{y : f(x,y) \ge \tau\},$ and the problem is set up as learning a $\tau$, using a  labelled source dataset $S_m \sim P^{\otimes n}$ and an unlabelled target dataset $T_n \sim Q^{\otimes n}$ such that with high probability over $S_m, T_n,$ $Q(Y \in C_\tau(X)) \ge 1-\varepsilon,$ where $\varepsilon$ is a given target coverage level. It is desired that $\tau$ is as large as possible to minimise the size of the prediction sets learned.  The main scheme is presented modularly. The paper first describes how finding the maximum $\tau$ whilst ensuring that the number of points in $S_m$ it captures is large enough (as specified by using a Binomial tail inverse) gives valid sets when $Q = P$. Next, it is argued that when $Q$ is absolutely continuous with respect to $P$, and the derivative $w(x) = \frac{\mathrm{d}Q}{\mathrm{d}P}(x)$ is upper bounded and known, then one can importance sample the set $S_m$ to generate a sample from $Q$, which can then plug into the previous procedure. This step fundamentally uses the covariate shift assumption, and, to my understanding, is folklore. The following, then, constitute the main technical contributions.  Next, the assumption of exact knowledge of $w$ is relaxed, and it is argued that if instead for each $x_i \in S_m,$ bounds $\underline{w}_i \le w(x_i) \le \overline{w}_i$ are available, then one can produce a worst-case estimate of the coverage of $C_\tau$ for any $\tau$ (by taking points that were missed to have high weights, and points covered to have low weights), and if this pessimistic coverage also satisfies the constraint. So, for each $\tau$, the worst $w_\tau$ can be produced, which can then feed into the importance sampling procedure above.  Then, it is pointed out that such a confidence bound on the $w_i$s can be learned using a probabilistic classifier $s(\cdot|x)$ to separate data from $S_m$ and $T_n$, which leads to the main proposed algorithm. While it is roughly justified in the appendix that an accurate estimate can be obtained under smoothness assumptions with an appropriately fine gridding of the space (and a consequently huge sample complexity), the concrete proposal is to replace this step with a heuristic method for fitting bounds on the weights, and then plugging these into the above strategy. Note that this is not implemented as a direct optimisation over $\tau$ - instead a set $\mathcal{T}$ of possible values is pre-selected, and the procedure is executed for each $\tau$ (the algorithm recommends an increasing order on the same).  Finally, the paper presents experiments in the DomainNet dataset, and in the ImageNet dataset, where the shift in the latter corresponds to adversarial perturbations. The principal baselines are the weighted split conformal inference (WSCI) method, which in my opinion is an appropriate choice, and the "PS-C" method, which simply uses a prediction set that has $1-\varepsilon/b$ coverage on the source data, where $b$ is an estimated upper bound on $w$. The proposed method is ablatively presented, and it is seen that on just the DomainNet data, the method PS-R (which simply takes $(1-s)/s$ as an estimate of $w$) is both reliable and performs well, while the PS-M method, which further integrates samples in bins tends to be slightly optimistic on this dataset. Conversely, in the adversarially perturbed dataset, PS-R produces trivial prediction sets (since presumably this shift goes entirely outside the source domain's support), while PS-M performs well. The proposed method, PS-W, which further uses upper and lower bounds on $w$s after integrating them over bins, tends to be pessimistic (error-rates of $0.06-0.07$ are observed when only $0.1$ is demanded), but is not too much worse than either of these methods, and is effective in both types of shifts.