SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Apoorva Sharma
apoorvas@nvidia.com
Sushant Veer
sveer@nvidia.com
Asher Hancock
ah4775@princeton.edu
Heng Yang
& NVIDIA Research
hengy@nvidia.com
Marco Pavone
mpavone@nvidia.com
Anirudha Majumdar
ani.majumdar@princeton.edu
Inductive Conformal Prediction (ICP) provides a practical and effective approach for equipping deep learning models with uncertainty estimates in the form of set-valued predictions which are guaranteed to contain the ground truth with high probability. Despite the appeal of this coverage guarantee, these sets may not be efficient: the size and contents of the prediction sets are not directly controlled, and instead depend on the underlying model and choice of score function. To remedy this, recent work has proposed learning model and score function parameters using data to directly optimize the efficiency of the ICP prediction sets. While appealing, the generalization theory for such an approach is lacking: direct optimization of empirical efficiency may yield prediction sets that are either no longer efficient on test data, or no longer obtain the required coverage on test data. In this work, we use PAC-Bayes theory to obtain generalization bounds on both the coverage and the efficiency of set-valued predictors which can be directly optimized to maximize efficiency while satisfying a desired test coverage. In contrast to prior work, our framework allows us to utilize the entire calibration dataset to learn the parameters of the model and score function, instead of requiring a separate hold-out set for obtaining test-time coverage guarantees. We leverage these theoretical results to provide a practical algorithm for using calibration data to simultaneously fine-tune the parameters of a model and score function while guaranteeing test-time coverage and efficiency of the resulting prediction sets. We evaluate the approach on regression and classification tasks, and outperform baselines calibrated using a Hoeffding bound-based PAC guarantee on ICP, especially in the low-data regime.
Machine learning (ML) models have rapidly proliferated across numerous applications, including safety-critical ones such as autonomous vehicles [Schwarting et al., 2018, Waymo, 2021], medical diagnosis [Ahsan et al., 2022], and drug discovery [Vamathevan et al., 2019]. Due to the severity of outcomes in these applications, decision making cannot solely hinge on "point" predictions from the ML model, but must also encapsulate a measure of the uncertainty in the predictions. As a result, accurate uncertainty estimation is a cornerstone of robust and trustworthy ML systems; however, overly-optimistic or overly-conservative estimates limit their usefulness. In this paper, we will present a method that builds upon inductive conformal prediction (ICP) [Vovk et al., 2005] and probably approximately correct (PAC) Bayes theory [McAllester, 1998] to furnish uncertainty estimates that provably control the uncertainty estimate's conservatism while achieving desired correctness rates.
Figure 1: Standard ICP uses a calibration dataset to guarantee coverage on exchangeable test data, but the efficiency of the resulting prediction sets is not directly optimized. Using the same data to select score function and model parameters as well as to perform ICP no longer provides guarantees on test-time coverage and efficiency. Therefore, prior work on optimal ICP has typically relied on holding out some calibration data for a post-optimization recalibration step to retain coverage guarantees. In this work, we leverage PAC-Bayes theory to optimize efficiency and provide generalization guarantees using the same dataset. Practically, we demonstrate that a hybrid approach using some data to tune a data-dependent prior, and use the rest to further optimize and obtain generalization guarantees yields the best performance.
ICP has emerged as a practical approach for equipping deep learned models with uncertainty estimates in the form of set-valued predictions with high-probability coverage guarantees, i.e., with high probability, at test time, the ground truth labels will lie within the predicted set. However, post-hoc application of ICP can often lead to overly-conservative set sizes; the tightness of the predicted set (referred to as efficiency) highly depends on the underlying model as well as the choice of the score function used in ICP. Various approaches have worked towards alleviating this challenge either by developing new application-specific score functions [Romano et al., 2020, Yang and Pavone, 2023, Lindemann et al., 2023] or by optimizing the model and / or score function directly [Yang and Kuchibhotla, 2021, Cleaveland et al., 2023, Bai et al., 2022]. The direct optimization approach is appealing as a general application-agnostic method for obtaining tight uncertainty estimates; however, it requires re-calibration on held-out data in order to retain generalization guarantees on coverage. By coalescing ICP with PAC-Bayes, we are able to train using all available data, while retaining generalization guarantees for coverage and efficiency; see Fig. 1 for more details.
Contributions: In this work, we make the following core contributions:
1. We use PAC-Bayes theory to obtain generalization bounds on both the coverage and efficiency of set-valued predictors. In contrast to prior work (see Fig. 1 and Sec. 3), our framework allows us to utilize the entire calibration dataset to learn the model parameters and score function (instead of requiring a held-out set for obtaining a guarantee on coverage).
2. We leverage these theoretical results to provide a practical algorithm (see Alg. 1) for utilizing calibration data to fine-tune the parameters of the model and score function while guaranteeing test-time coverage and efficiency. This algorithm allows the user to specify a desired coverage level, which in turn specifies the degree to which efficiency can be optimized (since there is a trade-off between the desired coverage level and efficiency).
3. We evaluate our approach on regression and image classification problems with neural network-based score functions. We show that our method can yield more efficient predictors than prior work on learned conformal prediction calibrated using a Hoeffding-bound based PAC guarantee on coverage [Vovk, 2012, Prop 2a].
2 Background: Inductive Conformal Prediction
We consider a supervised learning setup, wherein our goal is to predict labels  $\pmb{y} \in \mathcal{V}$  given inputs  $\pmb{x} \in \mathcal{X}$ . In inductive conformal prediction (ICP) [Vovk et al., 2005], our goal is to develop a set-valued predictor which maps inputs  $\pmb{x}$  to a subset of the label space  $C(\pmb{x}) \subseteq \mathcal{V}$ . Specifically, we want this set-valued function to satisfy a coverage guarantee, ensuring that the chance that the prediction set fails to contain the ground truth label is bounded to a user-specified level  $\alpha$ , i.e.
$$
\underset {\boldsymbol {x}, \boldsymbol {y} \sim \mathcal {D}} {\mathbb {P}} (\boldsymbol {y} \notin C (\boldsymbol {x})) \leq \alpha , \tag {1}
$$
where  $\mathcal{D}$  is the joint distribution over inputs and labels that test data are drawn from. In practice,  $\mathcal{D}$  is not known, but we assume we have access to a calibration dataset  $D_{\mathrm{cal}} = \{(x_i,y_i)\}_{i = 1}^N$  where each example  $(\pmb {x}_i,\pmb {y}_i)$  is exchangeable with the test data. Given this set-up, ICP constructs a set-valued predictor of the form
$$
C (\boldsymbol {x}; \tau) := \left\{\boldsymbol {y} \in \mathcal {Y} \mid s (\boldsymbol {x}, \boldsymbol {y}) \leq \tau \right\}, \tag {2}
$$
i.e., the  $\tau$  sub-level set of the nonconformity function  $s: \mathcal{X} \times \mathcal{Y} \to \mathbb{R}$  evaluated at the input  $\pmb{x}$ . The nonconformity function assigns a scalar value quantifying how poorly an example  $\pmb{x}, \pmb{y}$  conforms to the training dataset. It is common to define a nonconformity function that depends on the learned model  $f$ , choosing  $s(\pmb{x}, \pmb{y}) = \ell(f(\pmb{x}), \pmb{y})$ , measuring how poorly the prediction  $f(\pmb{x})$  aligns with the label  $\pmb{y}$  [Angelopoulos and Bates, 2021]. Armed with the calibration dataset as well as the nonconformity function, all that remains to construct the set-valued predictor is to choose the proper threshold  $\tau$ . ICP defines a calibration strategy to choose  $\tau$  that guarantees a desired miscoverage rate  $\alpha$ . Specifically, let  $S_{\mathrm{cal}} = \{s(\pmb{x}, \pmb{y}) \mid \pmb{x}, \pmb{y} \in D_{\mathrm{cal}}\}$  and choose  $\tau^{*}(D_{\mathrm{cal}}, \alpha)$  to be the  $q = \lceil (n + 1)(1 - \alpha) \rceil / n$  quantile of the set  $S$ . Then, so long as the score function  $s$  is independent from  $D_{\mathrm{cal}}$ , we have the following guarantee on the prediction sets defined by (2):
$$
\alpha - \frac {1}{N + 1} <   \mathbb {P} _ {D _ {\mathrm {c a l}} \sim \mathcal {D} ^ {N}, (\boldsymbol {x}, \boldsymbol {y}) \sim \mathcal {D}} (\boldsymbol {y} \notin C (\boldsymbol {x}; \tau^ {*} (D _ {\mathrm {c a l}}, \alpha)) \leq \alpha . \tag {3}
$$
Importantly, this probabilistic guarantee is marginal over the sampling of the test data point as well as the calibration dataset. In practice, we are given a fixed calibration dataset, and would like to bound the miscoverage rate conditioned on observing  $D_{\mathrm{cal}}$ . As shown in Vovk [2012, Prop 2a], we can obtain a probably-approximately-correct (PAC) style guarantee on the ICP predictor of the form
$$
\mathbb {P} _ {D _ {\mathrm {c a l}} \sim \mathcal {D}} \left(\mathbb {P} _ {\boldsymbol {x}, \boldsymbol {y} \sim \mathcal {D}} (\boldsymbol {y} \notin C (\boldsymbol {x}; \tau^ {*} (D _ {\mathrm {c a l}}, \alpha)) <   \alpha + \sqrt {\frac {- \log \delta}{2 N}}\right) \geq 1 - \delta . \tag {4}
$$
In other words, with probability at least  $1 - \delta$  over the sampling of a fixed calibration dataset, we can ensure the test-time miscoverage rate is bounded by  $\alpha$  by calibrating for a marginal coverage rate of  $\hat{\alpha} = \alpha -\sqrt{-\log\delta / 2N}$ . Alternatively, one can achieve the same PAC guarantee using the tighter bound in [Vovk, 2012, Prop 2b] by choosing the largest  $\hat{\alpha}\in (0,\alpha)$  for which
$$
\delta \geq I _ {1 - \alpha} (N - \lfloor \hat {\alpha} (N + 1) - 1 \rfloor , \lfloor \hat {\alpha} (N + 1) - 1 \rfloor + 1) \tag {5}
$$
where  $I_{x}(a,b)$  is the regularized incomplete beta distribution. This lacks an analytic solution for the optimal  $\hat{\alpha}$ , but is less conservative.
The strength of ICP lies in its ability to guarantee test-time coverage for any base predictor and nonconformity score function. However, ICP does not directly control the efficiency of the resulting prediction sets; the size and make-up of the prediction set for any given input  $\pmb{x}$  is highly dependent on the score function  $s$  (and thus also the base prediction model  $f$ ). For this reason, naive application of ICP can often yield prediction sets that are too "loose" to be useful downstream. Often, both the model and the score function may depend on various parameters  $\theta \in \mathbb{R}^d$ , e.g., the weights of a neural network base prediction model. The values of these parameters influence the ICP procedure, from the computation of the score function  $s(\pmb{x},\pmb{y};\theta)$ , the resulting threshold  $\tau^{*}(D_{\mathrm{cal}},\alpha ;\theta)$ , and finally, the prediction sets themselves  $C(\pmb{x};\tau^{*},\theta)$ . Our goal in this work is to investigate how we can use calibration data to fine-tune these parameters to optimize an efficiency objective while retaining guarantees on test-time coverage.
Handcrafted score functions. In classification  $(\mathcal{Y} = [K])$ , suppose  $f(\pmb{x})$  predicts the softmax class probabilities; typically the score function is chosen as either  $s(\pmb{x},y) = 1 - f(\pmb{x})_y$  or  $s(\pmb{x},y) = \sum_{k\in [K]}\{f(\pmb{x})_k\mid f(\pmb{x})_k\geq f(\pmb{x})_y\}$ , where  $f(\pmb{x})_y$  denotes the softmax probability of the groundtruth label. While the former score function produces prediction sets with the smallest average size [Sadinle et al., 2019], the latter produces prediction sets whose size can adapt to the difficulty of the problem [Romano et al., 2020]. In 1-D regression  $(\mathcal{Y} = \mathbb{R})$ , [Romano et al., 2019] recommend training  $f(\pmb{x}) = [f_{\alpha /2}(\pmb{x}),f_{1 - \alpha}(\pmb{x})]$  using quantile regression to output the (heuristic)  $\alpha /2$  and  $1 - \alpha /2$  quantiles, and then set the score function as  $s(x,y) = \max \{f_{\alpha /2}(x) - y,y - f_{1 - \alpha}(x)\}$
to compute the distance from the groundtruth  $y$  to the prediction interval  $f(\pmb{x})$ . In  $n$ -D regression  $(\mathcal{Y} = \mathbb{R}^n)$ , it is common to design  $s(f(\pmb{x}), \pmb{y}) = \frac{\|\mu(\pmb{x}) - \pmb{y}\|}{u(\pmb{x})}$ , where  $f(\pmb{x}) = (\mu(\pmb{x}), u(\pmb{x}))$  outputs both the mean of the label  $\mu(\pmb{x})$  and a heuristic notion of uncertainty  $u(\pmb{x})$  [Yang and Pavone, 2023, Lindemann et al., 2023]. While this leads to a ball-shaped prediction set, designing  $s(f(\pmb{x}), \pmb{y}) = \max_{i \in [n]} \{|\mu(\pmb{x})_i - y_i|\}$  can produce a box-shaped prediction set [Bai et al., 2022].
Learning score functions. [Yang and Kuchibhotla, 2021] propose selection algorithms to yield the smallest conformal prediction intervals given a family of learning algorithms. [Cleaveland et al., 2023] consider a time series prediction setup and leverage linear complementarity programming to optimize a score function parameterized over multiple time steps. [Stutz et al., 2021] develop conformal training to shrink the prediction set for classification, which learns the prediction function and the score function end-to-end by differentiating through and simulating the conformal prediction procedure during training. [Bai et al., 2022] considers optimizing the efficiency of a parameterized score function subject to coverage constraints on a calibration dataset, and derives generalization bounds depending on the complexity of the parametrization. [Einbinder et al., 2022] focus the classification-specific adaptive predictive set (APS) score function [Romano et al., 2020] score function, and propose an auxiliary loss term to train the model such that its predictions more closely meet the conditions under which the APS score function yields optimal efficiency. The drawback of [Cleaveland et al., 2023, Stutz et al., 2021, Bai et al., 2022, Einbinder et al., 2022] is that a separate held-out dataset is required to recalibrate the optimized score function to obtain test-time coverage guarantees. In this paper, we leverage PAC-Bayes theory to alleviate this drawback and allow learning the parameters of the model and score function using the entire calibration dataset, while offering efficiency and coverage guarantees.
Generalization theory and PAC-Bayes. Generalization theory seeks to quantify how well a given model will generalize to examples beyond the training set, and how to learn models that will generalize well. Early work includes VC theory [Vapnik and Chervonenkis, 1968], Rademacher complexity [Shalev-Shwartz and Ben-David, 2014], and the minimum description length principle [Rissanen, 1989, Blumer et al., 1987]. In this work, we make use of PAC-Bayes theory [McAllester, 1998]. In PAC-Bayes learning, one fixes a data-independent prior over models and then obtains a bound on the expected loss that holds for any (potentially data-dependent) choice of posterior distribution over models. One can then optimize the PAC-Bayes bound via the choice of posterior in order to obtain a certificate on generalization. In contrast to bounds based on VC theory and Rademacher complexity, PAC-Bayes provides numerically strong generalization bounds for deep neural networks for supervised learning [Dziugaite and Roy, 2017, Neyshabur et al., 2017a,b, Bartlett et al., 2017, Arora et al., 2018, Rivasplata et al., 2019, Perez-Ortiz et al., 2021, Jiang et al., 2020, Lotfi et al., 2022] and reinforcement learning [Fard et al., 2012, Majumdar et al., 2021, Veer and Majumdar, 2020, Ren et al., 2021]. We highlight that the standard framework of generalization theory (including PAC-Bayes) provides bounds for point predictors (i.e., models that output a single prediction for a given input). Here, we utilize PAC-Bayes in the context of conformal prediction to learn set-valued predictors with guarantees on coverage and efficiency.
4 PAC-Bayes Generalization Bounds for Inductive Conformal Prediction
Optimizing efficiency with ICP requires splitting the dataset to train on one part and calibrate on the other. The reduced data for calibration can result in weaker miscoverage guarantees. Drawing from PAC-Bayes generalization theory, we will develop a theory that facilitates simultaneous calibration and efficiency optimization for ICP using the entire calibration dataset.
Our generalization bounds are derived by randomizing the parameters of the score function  $\theta$ . Let  $Q(\theta)$  be a distribution over parameters  $\theta$ . Note that, for a fixed target miscoverage rate  $\hat{\alpha}$ , every sample from  $\theta$  induces a different prediction set  $C(\boldsymbol{x}; \tau^{*}(D_{\mathrm{cal}}, \hat{\alpha}; \theta), \theta)$ . Test-time coverage for such a randomized prediction set corresponds to the probability of miscoverage, marginalizing over the sampling of  $\theta$ :
$$
\mathcal {L} _ {\mathrm {c o v}} (Q) := \underset {\theta \sim Q} {\mathbb {P}} \left(\underset {\boldsymbol {x}, \boldsymbol {y} \sim \mathcal {D}} {\mathbb {P}} (\boldsymbol {y} \notin C (\boldsymbol {x}; \tau^ {*} (D _ {\mathrm {c a l}}, \hat {\alpha}; \theta), \theta))\right). \tag {6}
$$
The efficiency  $\ell_{\mathrm{eff}}$  of a predicted set  $C$  is a metric that can encode task-specific preferences (e.g., the volume of the predicted set) that we wish to minimize while satisfying the desired miscoverage rate.
Similar to (6), we can define the efficiency measure  $\mathcal{L}_{\mathrm{eff}}$  on randomized prediction sets by taking an expectation of a single-set efficiency measure:
$$
\mathcal {L} _ {\text {e f f}} (Q) := \underset {\theta \sim Q} {\mathbb {E}} \left[ \underset {\boldsymbol {x}, \boldsymbol {y} \sim \mathcal {D}} {\mathbb {E}} \left[ \ell_ {\text {e f f}} \left(C (\boldsymbol {x}; \tau^ {*} \left(D _ {\text {c a l}}, \hat {\alpha}; \theta\right), \theta\right), \boldsymbol {y}) \right] \right]. \tag {7}
$$
Our core theoretical contributions are to provide bounds on  $\mathcal{L}_{\mathrm{cov}}(Q)$  and  $\mathcal{L}_{\mathrm{eff}}(Q)$  that can be computed using  $D_{\mathrm{cal}}$  and, critically, hold uniformly for all choices of  $Q$ , even those that depend on  $D_{\mathrm{cal}}$ . This allows us to learn the distribution  $Q$  over score function parameters using  $D_{\mathrm{cal}}$ . First, we state our generalization bound for test-time coverage.
Theorem 1 (PAC-Bayes Bound on Coverage of ICP Prediction Sets). Let  $P(\theta)$  be a (data-independent) prior distribution over the parameters of the score function. Let  $D_N$  be a set of  $N$  samples from  $\mathcal{D}$ . Choose empirical coverage parameter  $\hat{\alpha} \in (0,1)$  such that  $N > 1 / \hat{\alpha} - 1$ . Then, with probability greater than  $1 - \delta$  over the sampling of  $D_N$ , the following holds simultaneously for all distributions  $Q(\theta)$ :
$$
\operatorname {k l} \left(\frac {\left\lfloor (N + 1) \hat {\alpha} \right\rfloor - 1}{N - 1} \| \mathcal {L} _ {\mathrm {c o v}} (Q)\right) \leq \frac {\mathrm {K L} (Q \| P) + \log \left(\frac {B (N)}{\delta}\right)}{N - 1}, \tag {8}
$$
where  $\mathrm{kl}(p||q)$  is the KL divergence between Bernoulli distributions with success probability  $p$  and  $q$  respectively, and  $B(N) := \operatorname{Beta}\left(\frac{k - 1}{N - 1}, k, N + 1 - k = O\left(\sqrt{N / (\hat{\alpha}(1 - \hat{\alpha}))}\right)\right)$ , where  $k = \lfloor (N + 1)\hat{\alpha} \rfloor$
The proof for this theorem is presented in App. A $^1$ . This theorem states that applying the calibration procedure will yield a test-time coverage rate  $\mathcal{L}_{\mathrm{cov}}(Q)$  close to  $\frac{\lfloor(N + 1)\hat{\alpha}\rfloor - 1}{N - 1}$ , and the amount that it can differ shrinks as the cardinality  $N$  of the calibration set increases, but grows as  $Q$  deviates from  $P$ . Finally, we note that an upper bound on  $\mathcal{L}_{\mathrm{cov}}(Q)$  can be computed by inverting the KL bound [Dziugaite and Roy, 2017] using convex optimization [Majumdar et al., 2021, Sec. 3.1.1].
Next, we will provide a generalization bound on efficiency.
Theorem 2 (PAC-Bayes Bound on Efficiency of ICP Prediction Sets). Let  $P(\theta)$  be a (data-independent) prior distribution over the parameters of the score function. Let  $D_N$  be a set of  $N$  samples from  $\mathcal{D}$  and  $\gamma > 0$ . Let the efficiency  $\ell_{\mathrm{eff}}$  of a predicted set  $C$  always lie within  $[0,1]$ . Assume the score function is bounded,  $s(\boldsymbol{x},\boldsymbol{y};\theta) < \beta$ , and the efficiency loss  $\tau \rightarrow \ell_{\mathrm{eff}}(C(\boldsymbol{x};\tau,\theta))$  is  $L_{\tau}$ -Lipschitz continuous in  $\tau$  for all values of  $\theta$ . Then, with probability greater than  $1 - \gamma$ , we have for all  $Q$ ,
$$
\mathcal {L} _ {\text {e f f}} (Q) \leq \underset {\theta \sim Q} {\mathbb {E}} \left[ \frac {1}{N} \sum_ {i = 1} ^ {N} \ell_ {\text {e f f}} \left(C \left(\boldsymbol {x} _ {i}; \tau^ {*} \left(D _ {N}, \hat {\alpha}; \theta\right), \theta\right)\right) \right] + \frac {2 \beta L _ {\tau}}{\sqrt {N}} + \frac {\sqrt {\frac {1}{2} \mathrm {K L} (Q | | P) + \frac {1}{2} \log \left(\frac {2 N}{\gamma}\right)}}{\sqrt {N - 1}} \tag {9}
$$
The proof for this theorem follows a similar approach to standard PAC-Bayes bounds, but is complicated by the dependence of  $\tau^{*}$  on  $D_N$ . To mitigate this, we instead consider the generalization error between empirical and test-time efficiency for the worst-case choice of  $\tau$ , which requires assuming  $\tau$  is bounded and the efficiency objective smooth w.r.t.  $\tau$ . The full proof is presented in App. A.
5 Practical Algorithmic Implementation
The generalization bounds above suggest a practical algorithm for using calibration data to simultaneously conformalize a base predictor while also fine tuning model and score function parameters. Specifically, we leverage the theory to formulate a constrained optimization problem, where the objective aims to minimize an efficiency loss, while the constraint ensures that test-time coverage can still be guaranteed. The overall algorithm is summarized in Alg. 1.
<sup>1</sup>Our bound takes a similar form to the Maurer-Langford-Seeger PAC-Bayesian bound [Seeger et al., 2001, Maurer, 2004] on generalization risk given empirical risk. However, while the risk on two samples is independent for a fixed hypothesis, the coverage for a particular sample depends the threshold  $\tau^{*}$ , which is a function of the entire calibration dataset  $D_N$ . This key difference precludes a direct application existing bounds and necessitates a novel theoretical approach.
2 Any uniformly bounded efficiency measure can be mapped to [0,1] by a scaling factor.
Figure 2: Visualization of the KL budget allowed by the PAC-Bayes generalization bound to achieve a target test-time coverage of  $\alpha = 0.1$ , plotted as a function of  $\hat{\alpha}$  and calibration set size  $N$  (left) as well as guarantee failure probability  $\delta$  (right). Note the log scale on the x axis on both plots.
5.1 Guaranteeing test-time coverage via a KL constraint
Using the bound in Theorem 1, we derive a constraint over data-dependent posteriors  $Q(\theta)$  such that choosing a threshold using the conformal prediction procedure  $\tau(\theta, D_N, \hat{\alpha})$  will guarantee a test-time coverage greater than  $1 - \alpha$ , where the coverage we enforce over  $D_N$ ,  $(1 - \hat{\alpha})$ , maybe different from the coverage we wish to achieve at test time,  $1 - \alpha$ .