SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\mathop{\mathbb{E}}_{\theta \sim Q}\left[\mathcal{L}_{\mathrm{eff}}(\theta ,\tau^{*}(D_{N},\hat{\alpha},\theta)) - \hat{\mathcal{L}}_{\mathrm{eff}}(\theta ,\tau^{*}(D_{N},\hat{\alpha},\theta),D_{N})\right]\leq \frac{2\beta L_{\tau}}{\sqrt{N}} +\sqrt{\frac{\mathrm{KL}(Q\|P) + \log\left(\frac{2N}{\gamma}\right)}{2(N - 1)}}
$$
$$
\begin{array}{l} \underset {\theta \sim Q} {\mathbb {E}} \left[ \mathcal {L} _ {\mathrm {e f f}} (\theta , \tau^ {*} (D _ {N}, \hat {\alpha}, \theta)) \right] \leq \underset {\theta \sim Q} {\mathbb {E}} \left[ \hat {\mathcal {L}} _ {\mathrm {e f f}} (\theta , \tau^ {*} (D _ {N}, \hat {\alpha}, \theta), D _ {N}) \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}} \\ \end{array}
$$
Noting that by definition,  $\mathcal{L}_{\mathrm{eff}}(Q) = \mathbb{E}_{\theta \sim Q}[\mathcal{L}_{\mathrm{eff}}(\theta ,\tau^{*}(D_{N},\hat{\alpha},\theta))]$ , which yields the theorem as stated, concluding the proof.
All experiments were performed on a workstation with a Intel Core i9-10980XE CPU with a NVIDIA GeForce RTX 3090 GPU.
B.1 Detailed Algorithm Overview
In this section, we give a detailed description of our general experimental method. Subsequent sections give specifics to each of the experimental domains we consider.
Model training. For each experiment, we first train a base neural network model on a training dataset. Next, for each calibration approach, we pair the neural network model with a score function which is used for calibration.
Model calibration. In the calibration phase, we first split the calibration data  $D_{\mathrm{cal}}$  into two random splits, a tuning dataset  $D_0$  and true calibration dataset  $D_N$ , where the fraction of data used for tuning (the data split) is a hyperparameter. In the standard baseline, the no optimization is performed so  $D_0$  is the empty set, and  $D_N = D_{\mathrm{cal}}$  is used to calibrate the model. In the learned baseline, parameters are optimized on only  $D_0$ , and  $D_N$  is used to calibrate by computing the threshold  $\tau^*$ . In our PAC-Bayes approach, we tune prior parameters on  $D_0$ , and then optimize the posterior following Algorithm 1 on  $D_N$ .
We aim for the same PAC guarantee of  $1 - \alpha$  test coverage rate with probability greater than  $1 - \delta$ . For the standard and learned baselines,  $N$ , i.e. the size of  $D_N$ , determines the required empirical miscoverage rate  $\hat{\alpha}$ . We can either use  $\hat{\alpha}$  equal to the value given by (4) [Vovk, 2012, Prop 2a], or by finding the value (e.g. by grid search) satisfying (5) [Vovk, 2012, Prop 2b]. For our PAC-Bayes approach,  $\hat{\alpha}$  is a hyperparameter which trades off conservatism in choosing  $\tau^*$  for the freedom to optimize score function parameters on  $D_N$  while retaining the PAC generalization guarantees. To select this hyperparameter, we choose  $K$  values for  $\hat{\alpha}$ , and for each, run our algorithm with a tighter  $\delta^i = \delta / K$ . This allows us to use a union bound argument to achieve the same PAC guarantee as the other approaches uniformly over the  $K$  predictors we obtain. We select between these predictors by choosing the predictor with the best generalization bound on efficiency (9).
**Optimization on  $D_0$ .** In both the learned and PAC-Bayes approaches, we optimize parameters on the tuning dataset to minimize the efficiency via gradient descent. Note that the efficiency of a conformal predictor depends on the score function evaluated at a given input  $\pmb{x}$ , as well as the threshold  $\tau^*$  which corresponds to the empirical quantile of the score function evaluated on the calibration dataset. We follow the approach of Stutz et al. [2021], and for each minibatch, we first compute  $\tau^*$  by evaluating the  $\hat{\alpha}$  quantile on the minibatch, and then use this threshold to evaluate the efficiency loss on the data. We use the softsort approach proposed by [Grover et al., 2019] to differentiate through the quantile function. The prior tuning in our PAC-Bayes approach follows the same protocol, with the addition of first sampling  $\theta \sim \mathcal{N}(\mu_0, \Sigma_0)$  (differentially, via the reparameterization trick), before computing the efficiency loss for each sample and averaging.
**Optimization on  $D_N$ .** In our PAC-Bayes approach, we further optimize the score function on the second split of data, this time holding the prior fixed and instead optimizing the posterior  $Q = \mathcal{N}(\mu, \Sigma)$  to minimize the efficiency loss (7) subject to the KL constraint (10). The efficiency loss is evaluated in the same way as above, and the KL constraint can be evaluated analytically. To perform the constrained optimization, we implement an augmented Lagrangian method. Specifically, we construct an unconstrained optimization problem with the objective
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {a u g}} (\mu , \Sigma , s; \lambda , \rho) = \hat {\mathcal {L}} _ {\mathrm {e f f}} (\mu , \Sigma) + \lambda c (\mu , \Sigma , s) + \frac {\rho}{2} c (\mu , \Sigma , s) ^ {2} \\ c (\mu , \Sigma , s) = \operatorname {K L} \left(\mathcal {N} \left(\mu_ {0}, \Sigma_ {0}\right) | | \mathcal {N} (\mu , \Sigma)\right) - B (\alpha , \hat {\alpha}, \delta , N) + s \\ \end{array}
$$
where  $B(\alpha, \hat{\alpha}, \delta, N)$  is the KL budget defined in (10) and  $s \in \mathbb{R}_+$  is a positive slack variable translating the inequality constraint in (10) to an equality constraint  $c(\mu, \Sigma, s) = 0$ . In each outer iteration, we approximately solve this unconstrained optimization problem for  $\mu, \Sigma$  and  $s$  using projected gradient descent (clamping  $s$  to 0 after any gradient step making it negative). After
optimization, we evaluate the constraint  $c(\mu, \Sigma, s)$ , and update the penalty scale according to
$$
\lambda \leftarrow \lambda + \rho c (\mu , \Sigma , s)
$$
before rerunning the inner optimization. We run 10 outer iterations, and return the  $\mu, \Sigma$  which yielded the best efficiency loss while satisfying the constraint.
Calibration on  $D_N$ . For all models, we calibrate using  $D_N$  by choosing  $\tau^{*}(D_{N},\hat{\alpha}ha)$  to be the  $q = \lceil (n + 1)(1 - \hat{\alpha})\rceil /n$  quantile of the set of score function values evaluated on  $D_N$ . In the case of the PAC-Bayes approach, the score function (and hence the threshold) depends on the sampled value of  $\theta$ . Therefore, after optimization, we first pre-sample  $M$  values of  $\theta \sim Q$ , and for each, compute  $\tau^{*}$  as above. We store the sampled  $\theta$  values along with their associated  $\tau^{*}$  values for use at test time.
Evaluation. To evaluate we construct prediction sets according to the learned score functions on a held-out test set of  $N_{\mathrm{test}}$  points. For each point, the prediction set is constructed according to (2) using the  $\theta, \tau^{*}$  pair obtained from the calibration phase. In the case of the PAC-Bayes baseline, each test sample is randomly assigned to one of the pre-sampled parameter/threshold pairs. We measure the rate of coverage as well as the efficiency  $\ell_{\mathrm{eff}}$  of each prediction set. We repeat the entire process from calibration and evaluation on multiple random seeds.
B.2 Illustrative example
Heteroskedastic data was generated by sampling  $\pmb{x} \sim \mathcal{U}(-1,1)$ , and computing the target  $\pmb{y} = \cos(5\pmb{x}) + 0.3*\epsilon_1 + 1.8\sigma(5\pmb{x})\epsilon_2$ , where  $\epsilon_1, \epsilon_2 \sim \mathcal{U}(-0.5,0.5)$ , and  $\sigma(\cdot)$  is the sigmoid function.
The base predictor was a 2 hidden layer feedforward network with ReLU activations, with layer width of 64 neurons. We trained the network on 100 randomly sampled datapoints, and held out  $N$  new datapoints to use for calibration. In the results in the paper, for all methods we targeted  $\alpha = 0.1$  with  $\delta = 0.05$ . We held the base predictor fixed, and optimize the score function
$$
s _ {\text {o p t i m}} (\boldsymbol {x}, \boldsymbol {y}; \theta) = \frac {| f (\boldsymbol {x}) - \boldsymbol {y} |}{1 + \sigma \left(\theta_ {g}\right) u (\boldsymbol {x} , \boldsymbol {y} ; \theta_ {u})} \tag {21}
$$
where  $\sigma$  is the sigmoid, and  $\theta_g$  functions as a gating parameter, and  $u(\pmb{x};\theta)$  is a learned input-dependent uncertainty function defined as
$$
u (\boldsymbol {x}, \boldsymbol {y}; \theta_ {u}) = - 1 + \text {s o f t p l u s} (\operatorname {F F} (\boldsymbol {x}, \boldsymbol {y}, \theta_ {u}) + 0. 6), \tag {22}
$$
where FF is a two-hidden layer feedforward network with tanh activations and a layer width of 128, and the softplus together with the offsets smoothly maps the output of the neural network to (approximately)  $\mathbb{R}_+$ .
For the learned model, we optimize the efficiency loss for 2000 steps with a learning rate of 1e-3 and a batch size of 100 samples. For the PAC-Bayes approach, we optimize using an augmented Lagrangian method, using 2000 steps of gradient descent with a learning rate of 1e-3 to solve the unconstrained penalized problem, running 7 outer iterations in which the penalty terms are updated and the unconstrained problem is re-solved. Each round of inner optimization requires around 1 minute of compute. In the results shown in the paper, we optimized only the prior mean  $\mu$ , and held the prior variance fixed at  $\sigma^2 = 0.02 / \sqrt{F}$  where  $F$  is the fan_in for that parameter.
B.3 Corrupted MNIST
For the corrupted MNIST experiments, we use 7000 images from the MNIST train split to train a base predictor with the LeNet-5 architecture. For the calibration and test data, we use the held-out test split from MNIST and apply random rotation, uniform between  $-30^{\circ}$  and  $30^{\circ}$  and Gaussian noise with standard deviation 1.3. We save these transformed data, and then split this data into two disjoint subsets for the calibration and test sets respectively, using a different split for each choice of random seed. We evaluate each approach on 1000 test samples.
Here, the score function for all methods is fixed to be the negative log softmax output by the base predictor, but for the learned and PAC-Bayes methods, we consider optimizing the fully-connected portion of the network. The efficiency loss we use is a differentiable approximation of the output set
Figure 5: Average prediction set size on held out test data as a function of calibration set size  $N$  for different data split ratios.
Figure 6: Average coverage rate on held-out test data as a function of calibration set size  $N$  for different data split ratios. Box plot shows variation in results over different seeds.
size. Note while the actual set-size of the resulting prediction sets is not Lipschitz continuous w.r.t. the threshold  $\tau$ , the smoothed loss is, as it is the sum of  $K$  sigmoid functions which are themselves Lipschitz continuous, where  $K$  is the number of classes. Here, we optimize the efficiency using the same optimization parameters as in the illustrative example. For the PAC-Bayes method, we initialize the prior with mean equal to the base predictor's original trained weights, and a variance with scale  $\sigma^2 = 0.01 / \sqrt{F}$  where once again,  $F$  corresponds to the is the fan_in for the layer that the parameter is in. In the body, we report results obtained by optimizing both the prior mean and variance to minimize the efficiency objective on  $D_0$  before optimizing the posterior subject to the KL constraint on  $D_N$ . We found that this led to better results, as selecting the variance using the simple initialization strategy we used in the illustrative example includes too many  $\theta$  values which lead to inefficient sets, hindering the quality of the posterior that can be found within the KL constraint.
B.3.1 Additional Results
In Figure 5, we show the complete set of prediction set size vs calibration dataset size results, for all choices of dataset split. In general, the PAC-Bayes approach yields equal or better sets than baselines across the choice of data split. Figure 6 shows the achieved test-time coverage rates for all methods. All methods target a generalization coverage rate of at least  $1 - \alpha = 0.9$ , but due to a finite test set size  $N_{\mathrm{test}} = 1000$ , the empirical coverage rate may appear lower. We show  $95\%$  confidence intervals for the empirical mean of 1000 draws from a Bernoulli random variable with probability 0.9. For our results to be consistent with the desired PAC guarantee, failures would only happen in the case of (a) our PAC guarantee failing to hold and the true coverage rate being lower than  $1 - \alpha$ , which should happen with probability less than  $\delta$ , or (b) the coverage rate being higher than  $\alpha$  but the empirical mean on the test set being lower than the  $95\%$  CI for  $\alpha = 0.9$ , which should occur with probability
Figure 7: Coverage vs efficiency plots on the MNIST experiment. There is an inherent trade-off between coverage and efficiency (larger sets inherently are more likely to contain the ground truth). Our PAC-Bayes approach consistently lies on the Pareto Frontier in these plots (high coverage and small set size).
less than 0.05. Indeed we see that all but one trial for the Learned (Vovk 2b) baseline yields an empirical coverage less than the  $95\%$  CI. We performed 30 total trials per method, so this observation is consistent with our PAC guarantee and finite test-set size. We find that the PAC-Bayes method is not significantly over-conservative compared to baselines using Vovk 2a to select the empirical coverage targets. However, Vovk 2b leads to much tighter results, indicating room for improvement in our PAC-Bayes results. The difference between the optimized methods and the standard ICP coverage rate is more pronounced for the 0.75:0.25 data split because the optimized method are using only a quarter of the data to guarantee test time coverage, while the standard method uses the whole dataset and thus can calibrate to a lower empirical miscoverage rate  $\hat{\alpha}$  following (4).
We also visualize efficiency (prediction set size) as a function of coverage in Figure 7. Indeed, these two quantities are linked: achieving higher coverage for a fixed score function requires increasing set size, coming at the cost of efficiency. The PAC-Bayes results lie on the Pareto frontier on this trade-off, highlighting that our method is able to learn good score functions that yield efficient test-sets on the test set. Future work tightening the generalization bound has the potential to improve practical utility by reducing over-conservatism and allowing calibrating to lower empirical coverage levels.
In Figure 8, we show the impact of different choices of prior optimization strategy, comparing no optimization (using the prior as initialized, using the whole calibration dataset to perform constrained optimization of the posterior) against both mean-only  $(\mu)$ , and mean-var  $(\mu, \sigma^2)$ . As can be seen, a crude, naive strategy of arbitrarily defining a prior by choosing a fixed variance around the initialization of the weights does not work well for this problem setting, highlighting the value of using some data to tune the prior. Furthermore, we find that optimizing both  $\mu$  and  $\sigma^2$  using a fraction of the data outperforms optimizing only  $\mu$ , for all choices of data split.
Figure 8: Average test-time prediction set size as a function calibration set size  $N$  for different prior optimization strategies and data split between  $D_0$  and  $D_N$ .