Conformal Prediction With Conditional Guarantees

TL;DR

This work tackles distribution-free prediction with finite-sample conditional guarantees by reframing exact conditional coverage as coverage over a class of covariate shifts. It develops a general, calibration-based framework that extends split conformal prediction to finite-dimensional and infinite-dimensional function classes, achieving exact group-conditional and covariate-shift guarantees for finite-dimensional F and providing finite-sample error bounds for infinite-dimensional F (e.g., RKHS or Lipschitz classes). The proposed method integrates with black-box models via augmented quantile regression, employs a dual formulation for efficient computation, and delivers practical guidance through real-data experiments (Communities and Crime, RxRx1) with comparisons to CQR and localized conformal prediction. The results offer a controllable, distribution-free approach to quantify and manage uncertainty in modern predictive pipelines, even under distributional shifts.

Abstract

We consider the problem of constructing distribution-free prediction sets with finite-sample conditional guarantees. Prior work has shown that it is impossible to provide exact conditional coverage universally in finite samples. Thus, most popular methods only guarantee marginal coverage over the covariates or are restricted to a limited set of conditional targets, e.g. coverage over a finite set of pre-specified subgroups. This paper bridges this gap by defining a spectrum of problems that interpolate between marginal and conditional validity. We motivate these problems by reformulating conditional coverage as coverage over a class of covariate shifts. When the target class of shifts is finite-dimensional, we show how to simultaneously obtain exact finite-sample coverage over all possible shifts. For example, given a collection of subgroups, our prediction sets guarantee coverage over each group. For more flexible, infinite-dimensional classes where exact coverage is impossible, we provide a procedure for quantifying the coverage errors of our algorithm. Moreover, by tuning interpretable hyperparameters, we allow the practitioner to control the size of these errors across shifts of interest. Our methods can be incorporated into existing split conformal inference pipelines, and thus can be used to quantify the uncertainty of modern black-box algorithms without distributional assumptions.

Figures (16)

• Figure 1: Predicting with finite-sample guarantees: our conditionally valid pipeline vs. split conformal prediction.
• Figure 2: Comparison of split conformal prediction (blue, left-most panel) and the randomized implementation of our method (orange, center panel) on a simulated dataset first considered by Romano2019. Black curves denote an estimate of the conditional mean, while the blue and orange shaded regions indicate the fitted prediction intervals. For this experiment, our method is implemented using the procedure outlined in \ref{['sec:finite_dim']} with $\mathcal{F} := \{\sum_{G \in \mathcal{G}} \beta_G \mathbbm{1}\{x \in G\}: \beta \in \mathbb{R}^{|\mathcal{G}|}\}$. The rightmost panel shows the miscoverage of the two methods marginally over the x-axis and conditionally on x falling in the two grey shaded bands; the red line indicates the target level of $\alpha = 0.1$.
• Figure 3: Marginal calibration-conditional miscoverage (left panel) and length (right panel) of quantile regression (green) and the randomized (red) and unrandomized (orange) implementations of our conditional-calibration method on a simulated dataset. All methods are implemented in their two-sided form with conformity score $S(x,y) = y$, i.e. we estimate the $\alpha/2$ and $1-\alpha/2$ quantiles of $Y \mid X$ separately and define the prediction set to be the values of $y$ that fall between the two bounds (see \ref{['sec:app_two-sided']} for details). Data for this simulation are generated i.i.d. from $Y_i = X_i^\top w + \epsilon_i$ where $X_i \sim \mathcal{N}(0,I_d)$, $\epsilon_i \sim \mathcal{N}(0,1)$, and $w \sim \text{Unif}(\mathcal{S}^{d-1})$. We implement both vanilla quantile regression Jung2023 and our conditional-calibration methods on the function class $\mathcal{F} := \{\beta_0 + \sum_{i=1}^d \beta_i \mathbbm{1}\{x_i > 0\} : \beta \in \mathbb{R}^d\}$. Boxplots show empirical estimates obtained by averaging over 1000 test points for each of 100 calibration datasets. The red line in the left panel indicates the target coverage level of $1-\alpha = 0.9$.
• Figure 4: Comparison of split conformal prediction (blue, left-most panel) and the randomized implementation of our method (orange, center panel) on a simulated dataset first considered by Romano2019. Black curves denote an estimate of the conditional mean, while the blue and orange shaded regions indicate the fitted prediction intervals. We consider coverage under three scenarios; marginally over the whole x-axis and locally under two Gaussian tilts (denoted by $f_1$ and $f_2$ and plotted as grey dotted lines). For this experiment, our method is implemented using the procedure outlined in \ref{['sec:finite_dim']} with $\mathcal{F} := \{\beta_0 + \sum_{i = 1}^5 \beta_i w_i(x) : \beta \in \mathbb{R}^6\}$, where the $w_i$ corresponds to the Gaussian tilts with parameters $(\mu,\sigma) \in \{(0.5,1),(1.5,0.2),(2.5,1),(3.5,0.2),(4.5,1)\}$. The rightmost panel indicates the miscoverage of both methods under all three settings with a red line denoting the target level of $\alpha = 0.1$.
• Figure 5: Demonstration of the unrandomized implementation of our shift-agnostic method on a simulated dataset first considered by Romano2019. The orange shaded region in the left panel depicts the prediction interval output by our method when $\mathcal{F}$ is chosen to be the Gaussian reproducing kernel Hilbert space given by kernel $K(x,y) = \exp(-12.5|x-y|^2)$ and the hyperparameter $\lambda$ is set equal to $0.005$ (see \ref{['sec:infinite_dim']} for details). Hatched and solid bars in the right panel show the estimated coverage returned by our method and the true realized empirical coverage, respectively. Finally, the red line indicates the target level of $\alpha = 0.1$.

Theorems & Definitions (48)

• Theorem 1: Romano2019, Theorem 1, see also VovkBook
• Theorem 2
• Corollary 1
• proof : Proof sketch of \ref{['thm:finite_dim_result']}
• Theorem 3
• Proposition 1
• Proposition 2
• Proposition 3
• Proposition 4
• Theorem 4