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

Introduction

This paper addresses the challenge of out-of-distribution (OOD) generalization in machine learning, where test data comes from a different distribution than the training data. A key problem is that existing methods for uncertainty quantification, like split conformal prediction (SCP), fail to maintain reliable coverage guarantees under OOD conditions due to violations of the exchangeability assumption.

To address this, the paper proposes a new set predictor based on the f-divergence between the test distribution and the convex hull of the training distributions. This provably maintains the desired marginal coverage guarantee, even in the OOD setting. The paper presents theoretical analysis and simulation experiments to validate the proposed method.

Related Works

The text discusses out-of-distribution (OOD) generalization, which aims to train a model on data from source domains so that it can generalize well to unseen target domains. Several algorithms have been developed to improve OOD generalization, including minimizing discrepancies between source domains, using meta-learning approaches, and adversarial training.

The text then introduces conformal prediction, a framework that provides valid confidence set predictors, even when the exchangeability of examples is violated to some extent. Previous works have considered various situations where exchangeability is violated, such as adversarial attacks on test examples, known density ratios between target and source domains, few-shot learning settings, and online learning.

The current work differs from these previous works by considering the OOD generalization setting where the f-divergence between the target domain and the convex hull of the source domains is constrained. This setting is most related to the work by Cauchois et al., which studied the worst-case coverage guarantee of an f-divergence ball centered at a single source domain.

Preliminaries

Here is a summary of the provided text:

The paper begins by introducing notations and concepts related to out-of-distribution (OOD) generalization. The input space is denoted as X and the label space as Y. The goal is to learn a predictor h that minimizes the worst-case population risk on unseen target domains T from the set T.

To achieve high-security, the paper proposes outputting a prediction set C(x) with a marginal coverage guarantee - that for any target domain T, the probability that the true label Y is contained in C(X) is at least 1-α, where α is the pre-defined significance level.

The paper then introduces split conformal prediction (SCP) as a method to construct such prediction sets. SCP assumes exchangeability of the training examples and uses the notion of nonconformity score s(x,y) to rank the examples. The prediction set C^n(x) is then defined based on the quantile of the empirical distribution of the nonconformity scores.

However, the paper shows that the validity guarantee of SCP does not hold in the OOD setting, as the exchangeability assumption is violated when there is distribution shift between the training and test data. The paper states that it will present a new set predictor in Section 5 that provides marginal coverage guarantees for the OOD setting.

SCP Fails in the OOD Setting

This section presents a toy example to show that for the out-of-distribution (OOD) confidence set prediction problem, split conformal prediction (SCP) is no longer valid, even under a slight distributional shift. The authors consider a regression problem with X in R^l and Y in R.

The source domain S is defined by a linear predictor L(x) = <w*, x> + b*, where w* is in R^l and b* is in R. The marginal distribution of X and the conditional distribution of Y given X are specified as Gaussian distributions with different parameters between the source domain S and the target domain T.

The authors sample training examples from S to train a linear predictor L-hat(x) = <w-hat, x> + b-hat, and define the nonconformity score as s(x, y) = |L-hat(x) - y|. They then sample examples from S to construct the prediction set C-hat_n(x) using Equation 2, and sample test examples from T.

After running the experiment 1000 times with different random seeds, the results for the coverage and length of the prediction set are presented in box plots. The coverage is the ratio between the number of test examples such that y_i is in C-hat_n(x_i) and the size of the test set. The red lines stand for the desired marginal coverages.

The boxes in the left plot are below the red coverage lines, indicating that SCP fails to provide a prediction set with the desired coverage when there exists a distributional shift between the source and target domains.

Corrected SCP for OOD Data

The section considers correcting score-based conformal prediction (SCP) for out-of-distribution (OOD) data. First, it considers the case where the population distributions of the scores from the source domains are known. It defines a set of distributions T containing distributions that are close to the convex hull of the source domain distributions, as measured by an f-divergence. The paper then provides a way to construct a valid confidence set predictor that achieves the desired marginal coverage guarantee.

Next, the section considers the more practical case where only the empirical distributions of the scores are available. It provides an error bound on estimating the minimum of the source distribution CDFs, and then uses this to derive a marginal coverage guarantee for the prediction set when using the empirical distributions. The paper provides a way to correct the prediction set to achieve the desired 1-α marginal coverage.

The section includes examples calculating the functions needed for the guarantees for specific f-divergences like chi-squared, total variation, and KL divergence. For KL divergence, the functions cannot be obtained in closed form, but can be efficiently computed using binary search.

Overall, the section derives marginal coverage guarantees for SCP in the OOD setting, both when the population distributions are known and when only empirical distributions are available.

Experiments

This section uses simulated data to verify the theory and validity of the constructed confidence set predictor called OOD-SCP. Two cases are considered:

1. The validity of OOD-SCP is verified using the same settings as in Section 4. The results show that unlike standard SCP, OOD-SCP is empirically valid for all values of the coverage parameter α.

2. A multi-source OOD confidence set prediction task is constructed. The results show that standard SCP is invalid in this case, while OOD-SCP remains valid, validating the theoretical claims.

The paper explains that real-world datasets are not used because the value of the parameter ρ is unknown for existing OOD datasets. The key claim is that when the target domain satisfies certain theoretical conditions, the coverage of the OOD-SCP method is guaranteed. This issue of setting ρ does not overshadow the main contributions, similar to how challenges in other areas like adversarial robustness and distributionally robust optimization do not undermine the value of those theoretical frameworks.

Discussions

The authors extend the work of Cauchois et al. (2020) to the multi-domain case. They identify two key issues with the trivial method of treating the mixture of source domains as a single domain:

1. The exact mixture weights of the source domains are unknown, so the set of points that the method provides a coverage guarantee for is unclear.
2. The method may not be able to provide a coverage guarantee for data from one of the source domains, as shown using KL-divergence as an example.

The authors' generalization does not suffer from these issues, even when using a coverage guarantee of 0. This makes their extension necessary.

The authors' proof techniques are more complex than those in the original work. They need to handle multiple input functions and take the infimum over the simplex of mixture weights. This leads to a more complicated set construction and analysis with four different cases to consider.

The authors' main results, Theorem 8 and Corollary 9, are novel and quite different from the results in Cauchois et al. (2020). The proof techniques also differ, as the examples are not exchangeable in the out-of-distribution setting, unlike in the original work.

Conclusion

The text describes a study on the confidence set prediction (SCP) problem in the out-of-distribution (OOD) generalization setting. The authors first show empirically that SCP is not valid in the OOD generalization setting. They then develop a method for forming valid confident prediction sets in the OOD setting and provide a theoretical proof of the validity of their proposed method. Finally, the authors conduct experiments on simulated data to verify the correctness of their theory and the validity of their proposed method.

Acknowledgements

This research was funded by grants from the National Key R&D Program of China, the National Natural Science Foundation of China, and the Fundamental Research Funds for the Central Universities. The grants supported this work.

Appendix A Proofs

The provided text presents proofs for several lemmas and theorems related to conformal prediction sets with f-divergence constraints. The key points are:

• Lemma 2 shows that the prediction set defined in Equation 2 satisfies the marginal coverage guarantee of at least 1-α.

• Lemma 3 establishes that the worst-case quantile function 𝒬~(α;𝒫f,ρ) can be expressed in terms of the function gf,ρ.

• Lemma 5 proves that the set 𝒫f,ρ is contained within the set 𝒫f,ρ.

• Theorem 6 shows how to efficiently compute the worst-case quantile function 𝒬~(α;𝒫f,ρ) using the function gf,ρ.

• Proposition 7 provides a finite-sample guarantee on the difference between the true and empirical worst-case quantile functions.

• Theorem 8 gives a marginal coverage guarantee for the empirical prediction set 𝒞~(x).

• Corollary 9 provides a way to correct the prediction set to achieve exact (1-α) marginal coverage.

• Lemma 10 shows that the inverse of gf,ρ can be efficiently computed.

The proofs rely on properties of f-divergences, convexity arguments, and concentration inequalities.

Appendix B B Additional experimental results