SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Corollary 2.1 (Constraint on Data-Dependent Posterior). Fix  $\hat{\alpha} \leq \alpha$ , and a prior distribution  $P(\theta)$ . Given a calibration dataset  $D_{\mathrm{cal}}$  of  $N$  i.i.d. samples from  $\mathcal{D}$ , we have that with probability greater than  $1 - \delta$ , then simultaneously for all distributions  $Q(\theta)$  which satisfy
$$
\operatorname {K L} (Q \| P) \leq B (\alpha , \hat {\alpha}, \delta , N) := (N - 1) \mathrm {k l} \left(\frac {\left\lfloor (N + 1) \hat {\alpha} \right\rfloor - 1}{N - 1} \| \alpha\right) - \log \left(\frac {B (N)}{\delta}\right), \tag {10}
$$
it holds that  $\mathcal{L}_{\mathrm{cov}}(Q)\leq \alpha$
The proof for this corollary is provided in App. A. Importantly, this bound holds for all  $Q(\theta)$  that satisfy the KL constraint, including  $Q$  that depend on the calibration data  $D_N$ . Thus, we are free to use any optimization algorithm to search for a feasible  $Q$  that minimizes any auxiliary loss; with  $B(\alpha, \hat{\alpha}, \delta, N)$  serving as a "budget" for optimization, limiting the degree to which  $Q$  can deviate from the data-independent prior  $P$ . Fig. 2 visualizes this budget as a function of  $N$ , and  $\hat{\alpha}$  for a fixed  $\alpha$  and  $\delta$ . As the plots show, by choosing the threshold  $\tau$  to attain a more conservative  $\hat{\alpha}$ , we can afford more freedom to vary  $Q$ . The curve defined by  $B(\alpha, \hat{\alpha}, \delta, N) = 0$  indicates the maximum  $\hat{\alpha}$  for which our theory provides a guarantee on  $1 - \alpha$  test time coverage with probability greater than  $1 - \delta$  for a randomized data-independent score function ( $Q = P$ ). As visualized in the figure, this boundary aligns with maximum  $\hat{\alpha}$  implied by the PAC guarantee from Vovk [2012, Prop 2a], where differences are likely due to differences in bound derivation. However, there remains a gap between our analysis and the results from Vovk [2012, Prop 2b], suggesting room for future tighter bounds.
5.2 Optimizing efficiency via constrained stochastic gradient optimization
In order to select a  $Q$  from this set, we propose directly optimizing the efficiency objective $^{3}$  (7). However, direct optimization is difficult due to (i) the expectation over  $\pmb{x}$  and  $\theta$ , and (ii) the non-differentiability of the efficiency objective  $\mathcal{L}_{\mathrm{eff}}(C)$ . To address (i), we use minibatches of inputs sampled from  $D_{\mathrm{cal}}$  and a finite set of  $K$  samples of  $\theta$  from  $Q(\theta)$  to construct a Monte-Carlo approximation of the expectation, providing a stochastic estimate of the objective that can be used in a stochastic gradient descent algorithm. To ensure differentiability, we follow the strategy proposed in [Stutz et al., 2021] and replace any non-differentiable operations in the computation of  $\mathcal{L}_{\mathrm{eff}}$  with their "soft" differentiable counterparts. For example, the quantile operation needed to compute the optimal threshold  $\tau(\theta, D_{\mathrm{cal}}, \hat{\alpha})$  can be replaced with a soft quantile implementation leveraging a soft sorting
Alternatively, one could optimize the generalization bound (9) directly, which effectively introduces a penalty on the KL divergence. In our case, we already have a constraint on the KL which mitigates overfitting, so we opt not to add complexity to the optimization.
Algorithm 1 Optimal Conformal Prediction with Generalization Guarantees
Require: Parameterized score function  $s(\pmb {x},\pmb {y},\theta)$  , Calibration dataset  $D_{\mathrm{cal}}$  , Target coverage rate   
1-  $\alpha$  , Probability of correctness  $\delta$  , Differentiable efficiency objective  $\ell_{\mathrm{eff}}(C)$  , Prior distribution   
 $P(\theta) = \mathcal{N}(\theta ;\mu_0,\Sigma_0)$  , Posterior distribution family  $Q(\theta) = \mathcal{N}(\theta ;\mu ,\Sigma)$  , Empirical coverage  $1 - \hat{\alpha}$  Initialize posterior distribution parameters  $\mu ,\Sigma \gets \mu_0,\Sigma_0$    
for each optimization iteration do   
 $\hat{\mathcal{L}}_{\mathrm{eff}}(\mu ,\Sigma)\gets \mathrm{DIFFERENTIABLELOSS}(\mu ,\Sigma ,\alpha)$ $c(\mu ,\Sigma)\gets \mathrm{KL}(\mathcal{N}(\mu ,\Sigma)\| \mathcal{N}(\mu_0,\Sigma_0)) - B(\alpha ,\hat{\alpha},\delta ,N)$  Update  $\mu ,\Sigma$  to minimize  $\mathcal{L}_{\mathrm{eff}}(\mu ,\Sigma)$  subject to  $c(\mu ,\Sigma) < 0$    
end for   
function DIFFERENTIABLELOSS  $(\mu ,\Sigma ,\alpha)$    
Sample mini-batch  $S = \{(x_j,y_j)\}_{j = 1}^J$  from  $D_{\mathrm{cal}}$    
Sample  $\zeta_1,\dots ,\zeta_K\sim \mathcal{N}(\mathbf{0}\in \mathbb{R}^d,I_d)$ $\theta_{k}\gets \mu +\Sigma \zeta_{k}$  for all  $k\in \{1,\ldots ,K\}$ $s_{j,k}\gets s(\pmb {x}_j,\pmb {y}_j;\theta_k)$  for all  $j\in \{1,\dots ,J\}$ $k\in \{1,\dots ,K\}$ $\tau_{k}\gets \mathrm{SOFTQUANTILE}(s_{[:,k]},[\left(J + 1\right)(1 - \hat{\alpha})] / J)$  for all  $k\in \{1,\dots ,K\}$    
Compute prediction sets  $C_{j,k} = C(x_j;\tau_k)$  for all  $j\in \{1,\dots ,J\}$ $k\in \{1,\dots ,K\}$    
return  $\hat{\mathcal{L}}_{\mathrm{eff}}(\mu ,\Sigma) = \frac{1}{JK}\sum_{k = 1}^{K}\sum_{j = 1}^{J}\ell_{\mathrm{eff}}(C_{j,k})$    
end function
algorithm [Cuturi et al., 2019, Grover et al., 2019]. As optimizing over the space of all probability distributions over  $\Theta$  is intractable, we restrict  $Q$  to be within a parametric class of distributions for which obtaining samples and evaluating the KL divergence with respect to the prior is tractable. In this work, we follow prior work in the PAC-Bayes and Bayesian Neural Network literature and fix both  $P$  and  $Q$  to be a Gaussian with a diagonal covariance matrix,  $\mathcal{N}(\boldsymbol{\mu}, \mathrm{diag}(\sigma^2))$ , where  $\boldsymbol{\mu} \in \mathbb{R}^d$  and  $\sigma^2 \in \mathbb{R}_+^d$ . This allows analytic evaluation of the KL divergence, and differentiable sampling via the reparametrization trick, as shown in Alg. 1. In this way, both the optimization objective and the constraint can be evaluated in a differentiable manner, and fed into a gradient based constrained optimization algorithm. We choose the  $Q$  with the lowest loss that still satisfies Corollary 2.1. More details on the optimization procedure can be found in Appendix B.
5.3 Practical considerations
Choice of prior. A prior  $P(\theta)$  may give high likelihood to values of  $\theta$  that yield good prediction sets, but also to  $\theta$  that perform poorly. To optimize efficiency,  $Q$  must shift probability mass away from poor  $\theta$  towards good  $\theta$ . However, due to the asymmetry of the KL divergence,  $\mathrm{KL}(Q||P)$  is larger if  $Q$  assigns probability mass where  $P$  does not than vice-versa. Thus, in order to effectively optimize efficiency under a tight KL budget, we must choose  $P$  that not only (i) assigns minimal probability density to bad values of  $\theta$ , but more importantly, (ii) assigns significant probability density to good choices of  $\theta$ . How do we choose an effective prior? In some cases, it suffices to choose  $P$  as an isotropic Gaussian distribution centered around a random initialization of  $\theta$ . However, especially when  $\theta$  correspond to parameters of a neural network, using a data-informed prior can lead to improved performance [Perez-Ortiz et al., 2021, Dziugaite and Roy, 2018]. Care is needed to ensure that a data-informed prior does not break the generalization guarantee; we find that a simple data-splitting approach similar to [Perez-Ortiz et al., 2021] is effective: We split  $D_{\mathrm{cal}}$  into two disjoint sets:  $D_0$ , used to optimize the prior without any constraints, and  $D_N$ , used to optimize the posterior according to Alg. 1. In our experiments, we consider two methods of optimizing a Gaussian prior  $P = \mathcal{N}(\boldsymbol{\mu}, \mathrm{diag}(\sigma^2))$ : First, we can optimize only the mean  $\boldsymbol{\mu}$  to minimize  $\mathcal{L}_{\mathrm{eff}}(\boldsymbol{\mu})$ , keeping  $\sigma^2$  fixed. This aims to shift  $P$  to assign more density to good values of  $\theta$ , but the fixed  $\sigma^2$  may still assign significant mass to bad choices of  $\theta$ . Alternatively, we consider optimizing both  $\boldsymbol{\mu}$  and  $\sigma^2$  to minimize  $\mathcal{L}_{\mathrm{eff}}(P)$ . By also optimizing the variance, we can minimize probability mass assigned to bad choices of  $\theta$ , but this comes with a risk of overfitting to  $D_0$ .
Test-time evaluation. The guarantees on coverage and efficiency generalization hold in expectation over both data sampled from the data distribution  $\mathcal{D}$  as well as model and score function parameters sampled from the learned posterior  $Q$ . Thus, in order to use such a predictor in practice and attain the desired coverage and efficiency, one would need to sample parameters  $\theta$  from  $Q$ , and then compute the
Figure 3: Standard ICP with a fixed score function (left) can yield over-conservative prediction sets. By using calibration data to also learn an uncertainty scaling factor in the nonconformity score function, both the learned (middle) and PAC-Bayes (right) approaches can yield more efficient prediction sets. When the calibration set size is large (bottom), both the learned and PAC-Bayes approaches do well. However, when calibration data is limited (top), our PAC-Bayes approach yields prediction sets with better test-time efficiency.
resulting prediction set  $C(\pmb{x};\tau)$ . However, as the threshold for this prediction set  $\tau = \tau^{*}(D_{\mathrm{cal}},\hat{\alpha};\theta)$  depends on  $D_{\mathrm{cal}}$ , this approach would require holding on to the calibration dataset to find the optimal threshold of each sample of  $\theta$ . In practice, once we optimize  $Q$ , we pre-sample  $K$  values of  $\theta$  from  $Q$ , and pre-compute the optimal threshold for each. Then, at test-time, we randomly select one  $\theta$ ,  $\tau$  pair from the set to evaluate each test input.
We evaluate our approach on an illustrative regression problem as well as on an MNIST classification scenario with simulated distributional shift. For each domain, we evaluate our PAC-Bayes approach against two baselines: (i) the standard ICP approach, where we use a fixed, data-independent score function  $s_{\mathrm{default}}(\boldsymbol{x}, \boldsymbol{y})$  and use the entirety of the calibration data to compute the critical threshold  $\tau$ ; and (ii) a learned ICP approach [Stutz et al., 2021] which uses a portion of  $D_{\mathrm{cal}}$  to optimize parameters  $\theta$  of a parametric score function  $s_{\mathrm{optim}}(\boldsymbol{x}, \boldsymbol{y}; \theta)$  to minimize an efficiency loss  $\mathcal{L}_{\mathrm{eff}}$ , and the remaining portion to re-estimate the threshold  $\tau$  to guarantee test-time coverage. For the PAC-Bayes method, we consider the same parametric form as the learned baseline, but instead randomize the score function by modeling a distribution over  $\theta$ . Both the learned and the PAC-Bayes baseline require choosing a data split ratio, i.e. what fraction is used for optimization (or prior tuning for the PAC-Bayes approach), and what fraction is held out for recalibration (or constrained optimization). For all methods, we target the same test-time coverage guarantee of  $1 - \alpha$  coverage with probability greater  $1 - \delta$ . For the standard and learned methods, the amount of data in the held-out recalibration set determines the empirical coverage level  $\hat{\alpha}$  needed to attain this guarantee, using either Vovk Prop 2a (4), or Vovk Prop 2b (5). We compare against both bounds in our experiments. For the PAC-Bayes method, the choice of  $\hat{\alpha}$  is a hyperparameter: choosing a lower value yields a larger budget for optimization, but entails using a more extreme quantile for the threshold  $\tau$ . As all methods provide guarantees on test-time coverage, our evaluation focuses on the efficiency of each method's predictions on held-out test data that is exchangeable with the calibration data. Additional results (including coverage results), as well as specific details for all experiments, can be found in App. B.
6.1 Illustrative demonstration: 1-D regression
As an illustrative example, we consider a 1-D regression problem where inputs  $\pmb{x} \in \mathbb{R}$  and targets  $\pmb{y} \in \mathbb{R}$  are drawn from a heteroskedastic distribution where noise increases with  $\pmb{x}$ , as shown in Figure 3. We train a two hidden layer fully connected neural network to minimize the mean squared error on this data, obtaining a base predictor  $f(\pmb{x})$  that produces a single point estimate for  $\pmb{y}$ . We aim
<sup>4</sup>We implemented our approach using PyTorch [Paszke et al., 2019] and Hydra [Yadan, 2019]; code to run all experiments is available at https://github.com/NVlabs/pac-bayes-conformal-prediction
Figure 4: Left: A base model is trained on clean MNIST digits (top), but calibrated and tested on corrupted digits (bottom). Middle: Average prediction set size on test data versus calibration set size  $(N)$  for a data split ratio of 0.5. Right: Prediction set size relative to standard ICP (Vovk 2a) as a function of the data split ratio, averaged over random seeds and across all calibration set sizes.
to obtain a set-valued predictor using a held-out set of calibration data with  $\alpha = 0.1$  and  $\delta = 0.05$ . For the sake of this example, we assume that we no longer have access to the training data, and only have calibration data and a pre-trained base predictor which only outputs point estimates and no input-dependent estimate of variance. Therefore, for the standard application of ICP, we must use the typical nonconformity function for regression,  $s_{\mathrm{default}}(\boldsymbol{x}, \boldsymbol{y}) = \|f(\boldsymbol{x}) - \boldsymbol{y}\|$ . Note that this choice yields prediction sets with a fixed, input-independent size, which is suboptimal in this heteroskedastic setting. To address this, we consider using the calibration dataset not only to determine the threshold  $\tau^{*}$ , but also to learn a parametric score function that can capture this heteroskedasticity. Specifically, we learn an input-dependent uncertainty score  $u(\boldsymbol{x}; \theta)$  which is used to scale the nonconformity score  $s_{\mathrm{optim}}(\boldsymbol{x}, \boldsymbol{y}; \theta) = \|f(\boldsymbol{x}) - \boldsymbol{y}\| / u(\boldsymbol{x}; \theta)$  [Angelopoulos and Bates, 2021]. We model  $u(\boldsymbol{x}, \theta)$  as a separate neural network with the same architecture as  $f$ , and optimize either  $\theta$  or  $Q(\theta)$  to minimize the volume of the resulting prediction sets. To do so, we minimize a loss proportional to the log of the radius of this set, which in this case can be written as  $\ell_{\mathrm{eff}}(C(\boldsymbol{x}, \tau)) = \log(u(\boldsymbol{x}; \theta) \cdot \tau)$ .
As can be seen in Figure 3, when  $N = 5000$ , both the learned baseline and our PAC-Bayes approach are able to learn an appropriate input-dependent uncertainty function which yields prediction sets which capture the heteroskedasticity of the data, and thus yield tighter sets in expectation over  $x$ . However, if we reduce the amount of calibration data available to  $N = 500$ , the learned approach starts to overfit on the split used to optimize efficiency. In contrast, the KL-constraint of the PAC-Bayes approach mitigates this overfitting, and outperforms both baselines in terms of test set efficiency for smaller amount of calibration data.
6.2 Corrupted MNIST
Next, we consider a more challenging distribution shift scenario, where we aim to deploy a model trained on "clean" data to a domain with distribution shift. As the base predictor, we use a LeNet convolutional neural network trained on a softmax objective to classify noise-free MNIST digits [LeCun et al., 1998]. Then, we simulate distribution shift by adding Gaussian noise and applying random rotations to a held-out set of the dataset: examples of the clean and corrupted digits are shown in Figure 4. Our goal is to use a calibration dataset of this corrupted data to design a set-valued predictor that achieves coverage on held-out test data with the same corruption distribution.
As is conventional, when applying ICP to probabilistic classifiers, we use the negative log probability as the nonconformity score  $s_{\mathrm{default}}(\pmb{x}, \pmb{y}) = [\log \mathrm{softmax}f(\pmb{x})]_{[\pmb{y}]}$  [Stutz et al., 2021, Angelopoulos and Bates, 2021]. This score function leverages the implicit uncertainty information encoded in the base model's probabilistic predictions; but these probabilities are far from calibrated, especially in the presence of distribution shift [Ovadia et al., 2019, Sharma et al., 2021]. Thus, we consider using the calibration dataset, which is representative of the corrupted test distribution, to fine-tune the predicted probabilities, defining  $s_{\mathrm{optim}}(\pmb{x}, \pmb{y}, \theta) = [\log \mathrm{softmax}f(\pmb{x}; \theta)]_{[\pmb{y}]}$ , where  $\theta$  are the parameters of the fully connected layers of the base predictor. As before, we measure efficiency as the size of the prediction set,  $\ell_{\mathrm{eff}}(C) = |C|$ . In contrast to the regression setting, here  $C$  is a discrete set over the set of class labels, and thus the size of this set is non-differentiable. To smooth this objective for training,
we follow the approach of [Stutz et al., 2021], and first use a sigmoid to compute soft set-membership assignment for each possible label, and then take the sum to estimate the size of the set,
$$
\hat {\ell} _ {\text {e f f}} (C (\boldsymbol {x}, \tau)) = \sum_ {\boldsymbol {y} \in \{0, \dots , 9 \}} \sigma \left(T ^ {- 1} \left(\tau - s (\boldsymbol {x}, \boldsymbol {y}, \theta)\right)\right), \tag {11}
$$
where the temperature  $T$  is a hyperparameter which controls the smoothness of the approximation; in our experiments we use  $T = 0.1$ .
We evaluate each method for different calibration set sizes  $N$  and data split ratios, repeating each experiment using 3 random seeds. For all approaches, we desire a coverage guarantee of  $\alpha = 0.1$  with  $\delta = 0.05$ . For the PAC-Bayes approach, we run the algorithm with  $\delta = 0.01$ , performing a grid search over 5 values of  $\hat{\alpha}$  evenly spaced between 0.2 and 0.8, and choosing the approach with the best efficiency generalization certificate (9); using the union bound for probability ensures that the efficiency for the best of the five runs will hold with probability at least  $1 - 0.05$ , i.e., an effective  $\delta$  of 0.05. For the PAC-Bayes approach, we optimize both the prior mean  $\mu$  and variance  $\sigma^2$  on  $D_0$ , as we found that optimizing the mean alone yielded a poor prior; ablations are provided in App. B. The results are summarized in Figure 4. In the middle figure, we hold the data split fixed at 0.5, and plot the average test set size as a function of  $N$ . First, we note that even for the standard ICP approach, increasing  $N$  leads to a smaller prediction set size, as we are able to calibrate to a larger  $\hat{\alpha}$  while still maintaining the desired test-time guarantee. In general, both the learned and PAC-Bayes approaches improve upon the standard baseline by fine-tuning on the corrupted calibration data. However, when data is limited ( $N = 1000$ ), the learned baseline overfits during optimization on the first half of the calibration data, and subsequent calibration yields set sizes that are even larger than the standard method. Our method mitigates this by using all the data to simultaneously fine-tune and calibrate while ensuring test-time generalization, thanks to the KL constraint. Empirically, the PAC-Bayes method yields prediction sets with efficiency comparable or higher than the learned baseline calibrated using Vovk 2a. However, there still remains a gap between our results and the learned baseline calibrated with tighter bound in Vovk 2b, as is evident in our analysis of our bound in Figure 2. In the right figure, we explore the sensitivity of both learned methods on the choice of data split ratio. We observe that the PAC-Bayes method is less sensitive to the choice of data split ratio.
In this paper, we introduced theory for optimizing conformal predictors while retaining PAC generalization guarantees on coverage and efficiency without the need to hold-out data for calibration. We achieved this by combining PAC-Bayes theory with ICP to furnish generalization bounds with the same data used for training. We translated our theory into a practical algorithm and demonstrated its efficacy on both regression and classification problems.
Limitations. As shown in Figure 2, our theory remains overconservative relative to the tight Vovk 2b PAC bound, which future work could aim to address. Practically, a key limitation of our approach is its dependence on the availability of a good prior. Although, we were able to mitigate this issue by training a prior on a portion of the calibration dataset, this strategy can struggle when datasets are small or the number of parameters to optimize are very large. Furthermore, while diagonal Gaussian distributions are convenient to optimize, they may not represent the optimal posterior. These limitations may explain why we obtained only modest improvements in efficiency with respect to the Vovk 2a baselines in our experiments; we are optimistic that future advances in prior selection and representation may lead to improvements. Finally, while PAC-Bayes theory requires sampling from the posterior  $Q$ , our set-up also requires computing the threshold  $\tau^{*}$  for each parameter sample. Future work could explore applying disintegrated PAC-Bayes theory [Viallard et al., 2023] to yield (perhaps more conservative) generalization bounds that hold for a single sample from the posterior.
Broader Impact. This work is building towards safer and trustworthy ML systems that can reason about uncertainty in their predictions. In particular, this work is a step towards developing calibrated uncertainty set predictors that are efficient enough for practical applications.
Future work. This work opens up exciting new theoretical and practical questions to pursue. On the theoretical front, we are excited to explore the development of an online learning approach that can leverage this theory to provably adapt to distribution shifts on the fly. On the practical front, we look forward to using this framework for providing reasonable uncertainty estimates in robot autonomy stacks to facilitate decision making and safety certification.
Anirudha Majumdar was partially supported by the NSF CAREER Award [#2044149] and the Office of Naval Research [N00014-23-1-2148]. Asher Hancock was supported by the National Science Foundation Graduate Research Fellowship Program under Grant No. DGE-2146755. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the National Science Foundation.