Coverage-Guaranteed Prediction Sets for Out-of-Distribution Data

Recent years have witnessed the remarkable success of modern machine learning techniques in many applications. A fundamental assumption of most machine learning algorithms is that the training and test data are drawn from the same underlying distribution. However, this assumption is consistently violated in many practical applications. In reality, the test environment is influenced by a range of factors, such as the distributional shifts across photos caused by the use of different cameras in image classification tasks, the voices of different persons in voice recognition tasks, and the variations between scenes in self-driving tasks (Nagarajan, Andreassen, and Neyshabur 2021). As a result, there is now a rapidly growing body of research with a focus on generalizing to unseen target domains with the help of the source domains, namely OOD generalization (Shen et al. 2021).

Existing OOD generalization methods focus on improving worst-case performance on the target domains, i.e., improving the average test accuracy of the model on the worst target domain. However, in some systems that require high security (such as medical diagnosis), even a single mistake may have disastrous consequences. In these cases, it is important to quantify the uncertainty of the predictions. One way to perform uncertainty estimation (Amodei et al. 2016; Jiang et al. 2012, 2018; Angelopoulos et al. 2021) is to create confident prediction sets that provably contain the correct answer with high probability. Let $X_{n+1}\in\mathcal{X}$ be a new test example for which we would like to predict the corresponding label $Y_{n+1}\in\mathcal{Y}$, where $\mathcal{X}$ is the input space and $\mathcal{Y}$ is the label space. For any given $\alpha\in(0,1)$, the aim of confidence set prediction is to construct a set-valued output, $\mathcal{C}(X_{n+1})$, which contains the label $Y_{n+1}$ with distribution-free marginal coverage at a significance level $\alpha$, i.e., $\mathbb{P}\left(Y_{n+1}\in\mathcal{C}(X_{n+1})\right)\geq 1-\alpha$. A confidence set predictor $\mathcal{C}^{\alpha}$ is said to be valid if $\mathbb{P}\left(Y_{n+1}\in\mathcal{C}^{\alpha}(X_{n+1})\right)\geq 1-\alpha$ for any $\alpha\in(0,1)$, where $\alpha$ is a hyper-parameter of the predictor. To simplify the notation, we omit the superscript $\alpha$ in the remainder of this paper.

Conformal prediction (Vovk, Gammerman, and Shafer 2005; Shafer and Vovk 2008, CP) is a model-agnostic, non-parametric and distribution-free (the coverage guarantee holds for any distribution) framework for creating confident prediction sets. Split conformal prediction (Vovk 2013; Vovk, Gammerman, and Shafer 2005, SCP), a special type of CP, has been shown to be computationally efficient. SCP reserves a set of data as the calibration set, and then uses the relative value of scores of the calibration set and that of a new test example to construct the prediction set. The validity of SCP relies on the assumption that the examples are exchangeable. However, in the OOD setting, the distributional shift between the training and test distributions leads to the violation of the exchangeability assumption. We empirically evaluate the performance of SCP in the OOD setting in Section 4. Unfortunately, we find that trivially applying SCP results in a failure to maintain marginal coverage in the OOD setting.

To address this issue, we construct a set predictor based on the $f$-divergence (Alfréd 1961) between the test distribution (target domain) and the convex hull of the training distributions (source domains). We theoretically show that our set predictor is guaranteed to maintain the marginal coverage (Corollary 9). We then conduct simulation experiments to verify our theory.

The remainder of this article is structured as follows: §2 introduces some related works; §3 presents the notation definitions and preliminaries; §4 conducts experiments that show the failure of SCP in the OOD generalization setting; §5 creates corrected confidence set predictor in the OOD generalization setting; §6 provides our experimental results. §7 make discussions with the most related work. Finally, the conclusions are presented in §8. All of our proofs are attached in Appendix A.

OOD generalization. OOD generalization aims to train a model with data from the source domains so that it is capable of generalizing to an unseen target domain. A large number of algorithms have been developed to improve OOD generalization. One series of works focuses on minimizing the discrepancies between the source domains (Li et al. 2018b; Ganin et al. 2016; Li et al. 2018c; Sun and Saenko 2016). Meta-learning domain generalization (Li et al. 2018a, MLDG) leverages the meta-learning approach and simulates train/test distributional shift during training by synthesizing virtual testing domains within each mini-batch. Another line of works (Xin et al. 2022; Wang et al. 2022) conducts adversarial training (Madry et al. 2018) to improve the OOD generalization performance. (Zou and Liu 2023) considers improving the adversarial robustness of the unseen target domain. Notably, the above works all focus on improving the average performance on the target domain; in contrast, we focus on designing valid confidence set predictors for data from the unseen target domains, as this is a crucial element of making high-stakes decisions in systems that require high security.

Conformal prediction. As introduced in §1, conformal prediction is a model-agnostic, non-parametric, and distribution-free framework that provides valid confidence set predictors. Generally speaking, examples are assumed to be exchangeable in a CP context. Most pertinent to our work, (Gendler et al. 2022; Tibshirani et al. 2019; Fisch et al. 2021; Cauchois et al. 2020; Gibbs and Candès 2021; Oliveira et al. 2022) all consider various situations in which the exchangeability of the examples is violated to some extent. (Gendler et al. 2022) considers the case in which the test examples may be adversarially attacked (Szegedy et al. 2014; Goodfellow, Shlens, and Szegedy 2015; Madry et al. 2018); (Tibshirani et al. 2019) investigates the situation in which the density ratio between the target domain and the source domain is known; (Fisch et al. 2021) studies the few-shot learning setting and assumes that the source domains and the target domain are independent and identically distributed (i.i.d.) from some distribution on the domains; (Gibbs and Candès 2021) considers an online learning setting and (Oliveira et al. 2022) provides results when the examples are mixing (Achim 2013; Xiaohong, Lars Peter, and Marine 2010; Bin 1994). Different from all the works discussed above, we consider the OOD generalization setting in which the $f$-divergence between the target domain and the convex hull of the source domains is constrained. The most related work among them is (Cauchois et al. 2020), which studies the worst-case coverage guarantee of a $f$-divergence ball centered at the single source domain. For the discussions about similarities and differences with (Cauchois et al. 2020), please refer to Section 7.

We begin with the OOD setups and a review of conformal prediction.

Notations. We denote $\{1,2,\dots,n\}$ by $[n]$ for $n\in\mathbb{N}_{+}$. For a distribution $P$ on $\mathbb{R}$, we define the quantile function of $P$ as $\mathcal{Q}(\beta;P)\coloneqq\inf\{s\in\mathbb{R}|P(S\leq s)\geq\beta\}$. Similarly, for a cumulative distribution function (c.d.f.) $F$ on $\mathbb{R}$, we define $\mathcal{Q}(\beta;F)\coloneqq\inf\{s\in\mathbb{R}|F(s)\geq\beta\}$. For $n$ distributions $P_{1},\dots,P_{n}$, we define $\mathcal{CH}\left(P_{1},\dots,P_{n}\right)\coloneqq\left\{\sum_{i=1}^{n}\lambda_{i}P_{i}|\lambda_{1},\dots,\lambda_{n}\geq 0;\sum_{i=1}^{n}\lambda_{i}=1\right\}$ as the convex hull of the distributions $P_{1},\dots,P_{n}$. We further define $\mathcal{N}(\mu,\Sigma)$ as the multi-variable Gaussian distribution with mean vector $\mu$ and covariance matrix $\Sigma$. For a set $A$, we define the indicator function as $\mathbb{I}_{A}(\cdot)$, where $\mathbb{I}_{A}(x)=1$ if $x\in A$ and $\mathbb{I}_{A}(x)=0$ otherwise.

In this section, we construct a toy example to show that for the OOD confidence set prediction problem, SCP is no longer valid, even under a slight distributional shift.

For simplicity, we consider a single-domain case. Specifically, we consider the regression problem and set $\mathcal{X}=\mathbb{R}^{l}$, $\mathcal{Y}=\mathbb{R}$. We define the source domain $S$ as follows: given a linear predictor $L(x)=\left<w^{\star},x\right>+b^{\star}$ where $w^{\star}\in\mathbb{R}^{l}$ and $b^{\star}\in\mathbb{R}$. The marginal distribution of $X$ and the conditional distribution of $Y$ given $X$ are defined as:

$$X\sim\mathcal{N}(\mu_{s},\sigma_{s,x}^{2}I_{l}),\ \ \ \ Y|X=x\sim\mathcal{N}(L(x),\sigma_{s,y}^{2}),$$

where $\mu_{s}\in\mathbb{R}^{l}$ is the mean vector of $X$, $\sigma_{s,x},\sigma_{s,y}$ are positive scalars, and $I_{l}\in\mathbb{R}^{l\times l}$ is an identity matrix. Similarly, for the target domain $T$, we define:

$$X\sim\mathcal{N}(\mu_{t},\sigma_{t,x}^{2}I_{l}),\ \ \ \ Y|X=x\sim\mathcal{N}(L(x),\sigma_{t,y}^{2}),$$

where $\mu_{t}\in\mathbb{R}^{l}$ is the mean vector of $X$ and $\sigma_{t,x},\sigma_{t,y}$ are positive scalars. For simplicity, we set $\mu_{s}=\mu_{t},\sigma_{s,x}=\sigma_{t,x}$ and $\sigma_{s,y}\neq\sigma_{t,y}$. We sample $m_{\text{train}}$ training examples from $S$ to train a linear predictor $\hat{L}(x)=\left<\hat{w},x\right>+\hat{b}$, where $\hat{w}\in\mathbb{R}^{l}$ and $\hat{b}\in\mathbb{R}$. We then define the nonconformity score as $s(x,y)=|\hat{L}(x)-y|$. We sample $n$ examples from $S$ to construct the prediction set $\widehat{\mathcal{C}}_{n}(x)$ in Equation 2 and sample $m_{\text{test}}$ examples from $T$ to form the test data. We run $1000$ times with different random seeds. The results for the coverage (left) and length (right) of the prediction set are presented in box plot form in Figure 1. Here, the coverage is the ratio between the number of test examples such that $y_{i}\in\widehat{\mathcal{C}}_{n}(x_{i})$ and the size of the test set. The red lines stand for the desired marginal coverages. Since the boxes are below the red coverage lines, we conclude that SCP fails to provide a prediction set with desired coverage when there exists a distributional shift between the source domain and the target domain.

In this section, we consider correcting SCP for OOD data. We first consider the case in which we have access to the population distributions of the scores from the source domains. We then consider the case in which we only have access to the empirical distributions and correct the prediction set to obtain a marginal coverage guarantee in Equation 3.

In this section, we use simulated data to verify our theory and the validity of our constructed confidence set predictor (referred to as OOD-SCP in the remainder of this paper). We consider two cases: first, we verify the validity of OOD-SCP using the same settings as in Section 4; then, we construct a multi-source OOD confidence set prediction task and show that OOD-SCP is valid for this task.

According to Figure 2, unlike standard SCP, for all values of $\alpha$, the violin for OOD-SCP is above the desired coverage line, which shows that OOD-SCP is empirically valid.

We next consider a multi-source OOD confidence set prediction task. Similar to Section 4, we consider the regression problem and set $\mathcal{X}=\mathbb{R}^{l},\mathcal{Y}=\mathbb{R}$. Define the oracle linear predictor $L:\mathcal{X}\xrightarrow{}\mathcal{Y}$ as $L(x)=\left<w^{\star},x\right>+b^{\star}$, where $w^{\star}\in\mathbb{R}^{l}$ and $b^{\star}\in\mathbb{R}$. We define the marginal distribution of $X$ for the source domains $S_{1}$ and $S_{2}$ as

$$S_{1X}=\mathcal{N}(\mu_{1},\sigma_{s,x}^{2}I_{l}),\ \ \ \ S_{2X}=\mathcal{N}(\mu_{2},\sigma_{s,x}^{2}I_{l})$$

respectively, where $\mu_{1},\mu_{2}\in\mathbb{R}^{l}$ are the mean vectors, $\sigma_{s,x}>0$ is a scalar, and $I_{l}\in\mathbb{R}^{l\times l}$ is the identity matrix with dimension $l\times l$. We define $Y|X=x\sim\mathcal{N}(L(x),\sigma_{s,y}^{2})$ for both $S_{1}$ and $S_{2}$. For the target domain $T$, we define the marginal distribution of $X$ as $T_{X}=\frac{S_{1X}+S_{2X}}{2}$ and the conditional distribution of $Y$ given $X$ as $Y|X=x\sim\mathcal{N}(L(x),\sigma_{t,y}^{2})$. Here, $\sigma_{s,y},\sigma_{t,y}>0$ are the standard deviations and $\sigma_{s,y}\neq\sigma_{t,y}$.

Similar to Section 4, we sample $\frac{m_{\text{train}}}{2}$ examples from $S_{1}$ and $\frac{m_{\text{train}}}{2}$ examples from $S_{2}$ to train a linear predictor $\hat{L}(x)=\left<\hat{w},x\right>+\hat{b}$, where $\hat{w}\in\mathbb{R}^{l}$ and $\hat{b}\in\mathbb{R}$. We then define the nonconformity score as $s(x,y)=|\hat{L}(x)-y|$. We sample $\frac{n}{2}$ examples from $S_{1}$, and $\frac{n}{2}$ examples from $S_{2}$ to construct the prediction set $\widetilde{\mathcal{C}}(x)$ and sample $m_{\text{test}}$ examples from $T$ to form the test data.

Figure 3 shows the results for the multi-source OOD confidence set prediction task. From the figure, we can see that the violins for the standard SCP are under the desired coverage lines, which means that the standard SCP is invalid in this case. By contrast, the violins for OOD-SCP are above the desired coverage lines, indicating that OOD-SCP is valid, which validates Corollary 9.

The reason we do not do experiments on real datasets is that do not know how to set the value of $\rho$ for the existing OOD datasets. Our main claim is that when the target domain satisfies $T\in\mathcal{T}_{f,\rho}$, the coverage of our method is guaranteed. However, we claim it is acceptable. In many fields, we face the same problem.

In adversarial robustness (Szegedy et al. 2014), the theories (for example (Montasser, Hanneke, and Srebro 2019)) provide an upper bound of the test adversarial robustness $\underset{(x,y)\sim D}{\mathbb{P}}[\exists\|\delta\|\leq\epsilon:h(x+\delta)\neq y]$, where $h$ is a classifier. The results just tell us that we have the guarantee for the test accuracy if the test perturbation $\delta$ satisfies $\|\delta\|\leq\epsilon$. However, what if $\|\delta\|>\epsilon$? It is out of the scope of their theories.

For distributional robustness optimization (DRO), the theories (Lee and Raginsky 2018) prove that if the test distribution is in a Wasserstein ball with radius $r$, then the test risk can be upper bounded. Formally, $\underset{D\in W(r)}{\max}\underset{(x,y)\sim D}{\mathbb{P}}[h(x)\neq y]$ is upper bounded, where $W(r)$ is a Wasserstein ball with radius. They do not know how to set $r$ to make $W(r)$ contain the test distribution either, however, this does not overshadow their contribution to the DRO community. In other words, the issues of $\rho$ do not overshadow our contribution to the OOD community.

Our work is an extension of (Cauchois et al. 2020) to the multi-domain case. In this section, we discuss the differences between our work and (Cauchois et al. 2020).