SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We train one epinet for each ResNet-  $L$  baseline for  $L\in \{50,101,152,200\}$ . The ResNet baselines are open sourced in the ENN library under the path /networks/resnet/, and the checkpoints are available under /checkpoints/imagenet.py. For each epinet agent, we take the pre-trained ResNet as the base network. We do not update the base network during epinet training. The epinet network architecture together with checkpoint weights can be found in the ENN library under the paths /networks/epinet/ and /checkpoints/imagenet.py.
As discussed in Section 6, we choose the reference distribution of the epistemic index to be a standard Gaussian with dimension  $D_Z = 30$ . The input to the learnable part of the epinet  $\sigma_{\eta}^{L}$  includes the last-layer features of the base ResNet and the epistemic index  $z$ . Let  $C$  denote the number of classes. For last-layer features  $\phi$  and epistemic index  $z$ , the learnable network takes the form  $\sigma_{\eta}^{L}(\phi, z) = g_{\eta}([\phi, z])^{T}z$ , where  $[\phi, z]$  is the concatenation of  $\phi$  and  $z$ , and  $g_{\eta}(\cdot)$  is a 1-layer MLP with 50 hidden units, ReLU activation, and output  $\in \mathbb{R}^{D_Z \times C}$ .
The fixed prior  $\sigma^P$  is made up of two components. The outputs of the two parts are summed together to produce the prior output. In general, we could have tunable scaling factors (which we refer to as 'prior scales' in the ablation studies) for the output of each component before we add them together. However, for ImageNet, we find that scaling factors of 1 already work well. The first component of the prior is a network with the same architecture and initialization as the learnable network  $\sigma_{\eta}^{L}$ . The second component is an ensemble of small convolutional networks that act directly on the input images. The number of networks in the ensemble is equal to the index dimension. Each convolutional network has the number of channels  $(4,8,8)$ , kernel shapes  $(10\times 10,10\times 10,3\times 3)$ , and strides  $(5,5,2)$ . The outputs are flattened and taken through a linear layer to give a vector of dimension  $C$ . For input image  $x$ , let  $p^i (x)\in \mathbb{R}^C$  denote the output of the  $i^{\mathrm{th}}$  member of the ensemble. We combine the outputs of the ensemble members by taking the weighted sum  $\sum_{i = 1}^{D_Z}p^i (x)z_i$ .
The ResNet baselines are trained using log loss and ridge regularization. We optimize using SGD with a learning rate 0.1, a cosine learning rate decay schedule, and Nesterov momentum. We apply  $L2$  weight decay of strength 1e-4. We also incorporate label smoothing into the loss, where instead of one-hot labels, the incorrect classes receive a small weight of  $0.1 / C$ . We train the ResNet agents for 90 epochs on  $4 \times 4$  TPUs with a per-device batch size of 128.
We train the epinet using loss of the form (7) with log loss with ridge regularization. Similar to ResNet training, we apply  $L2$  weight decay of strength 1e-4 and incorporate label smoothing into the loss. We draw 5 epistemic index samples for each gradient step. We optimize the loss using SGD with a learning rate 0.1, Nesterov momentum and decay 0.9. The epinet is trained on the same hardware with the same batch size for 9 epochs.
H.2 Uncertainty baselines
In this section we compare our results to the open source 'uncertainty baselines', which provides a reference implementation for much work on uncertainty estimation in Bayesian deep learning (Nado et al., 2021). As part of our development, we upstream an optimized method for calculating joint log-loss, and contribute this to the community. We benchmark a few popular approaches to uncertainty estimation in terms of both marginal and joint predictions, and compare their results to ours. At a high level, our results mirror our own
results of Figure 2. After tuning, most of the agents appear on a roughly similar tradeoff in terms of marginal quality. However, the approaches are widely separated in terms of their quality on joint prediction, and here epinet performs much better than the alternatives.
Figure 6 repeats Figure 2 but adding a few new agents from uncertainty baselines, which we
- mimo: Multi-input Mulit-output ensemble (Havasi et al., 2020). This method appears to perform better than a baseline resnet in terms of marginal statistics, but provide no additional benefit to modeling joints. Note that the independent product of better marginals automatically does improve joints.
- dropout: Dropout as posterior approximation (Srivastava et al., 2014; Gal and Ghahramani, 2016). This approach seems to provide a slight improvement in marginal log-loss and a noticeable improvement in joint log-loss. However although dropout has a low computational cost in terms of parameters, these results required 10 forward passes of the network, and so actually underperform relative to an ensemble of similar inference cost (see Appendix H.4).
- sngp: Spectral-normalized Neural Gaussian Process (Liu et al., 2020). For this agent we introduced an additional temperature parameter rescaling logit samples, which we found was able to significantly improve joint log-loss without degrading marginal prediction quality. However, even with this tuning the quality of joint predictions cannot match ensemble of size 10. The classification accuracy of SNGP also appears to be significantly worse than other approaches benchmarked here.
- het: Heteroscedastic loss (Collier et al., 2020, 2021). This approach performs well in terms of joint log-loss, and achieves performance close to the ensemble of size=100. Although the results are still significantly worse than those of our epinet, it is interesting to note that the functional form of the resultant heteroscedastic agent can actually be written as a particular form of epinet. In future work, we would like to understand better the commonalities between these two approaches, and see if/when the algorithms can borrow from each others' strengths.
H.3 CIFAR-10 and CIFAR-100
In this section we reproduce a similar analysis to Section 6 applied to the CIFAR-10 and CIFAR-100 datasets. We find that, at a high level, our results mirror those when applying ResNet to ImageNet. In particular, we are able to produce results similar to Figure 2 for both CIFAR-10 and CIFAR-100. Relative to large ensembles, epinets greatly improve joint predictions at orders of magnitude lower computational cost.
For the experiments in this section, we mirror the ImageNet experiments but using the smaller ResNet architectures. The ResNet baselines are open sourced in the ENN library under the path /networks/resnet/, and the checkpoints are available under /checkpoints/cifar10.py and /checkpoints/cifar100.py. In particular, we tune ResNet-  $L$  for  $L\in$ $\{18,32,44,56,110\}$  over learning rate and weight decay. We did not include temperature rescaling in these sweeps, although this could further improve performance for all agents. After tuning hyperparameters we independently initialize and train 100 ResNet-18 models to serve as ensemble particles. These models are then used to form ensembles of size 1, 3, 10, 30, and 100.
For the epinet agent we take the pretrained ResNet as the base network and fix its weights. We then follow the same methodology as for ImageNet, described in Appendix H.1, but with a slightly smaller network. We use index dimension  $D_Z = 20$  and alter the convolutional prior to have channels (4,8,4) each with kernel size  $5 \times 5$  and stride 2 on account of the smaller image sizes in CIFAR-10 and CIFAR-100 (Krizhevsky, 2009).
The epinet network architecture together with checkpoint weights can be found in the ENN library under the paths /networks/epinet/, /checkpoints/cifar10.py, and /checkpoints/cifar100.py.
Figures 12 and 13 reproduce our scaling results for ImageNet when applied to these other datasets. At a high level, the key observations remain unchanged across datasets. We see that across all statistical losses the larger models generally perform better. When looking
at classification and marginal log-loss, epinets do not offer any particular advantage over baseline ResNets. However, when we look at the joint log-loss we can see that epinets offer huge improvements in performance, even when measured against very large ensembles.
Figure 12: Quality of marginal and joint predictions across models on CIFAR-10.
Figure 13: Quality of marginal and joint predictions across models on CIFAR-100.
H.4 Computational cost
This paper highlights a key tension in neural network development: balancing statistical loss against computational cost. Our main results on ImageNet (Figure 2) use 'memory' as a proxy for computational cost in large deep neural networks, for which this is often a hardware bottleneck (Kaplan et al., 2020). However, in the case of these models, the results are similar for many alternative measures.
Figure 14 reproduces the results of Figure 2 but with computational cost measured in total floating point operations at inference. This plot includes the total costs of one thousand forwards of the epinet. Even with these additional FLOPs, the overall cost of each epinet is still less than  $50\%$  of the total network, and the outperformance of the epinet in terms of joint log-loss is still remarkable. Further, on large modern TPU architectures these epinet operations can often be performed in parallel. This means that, in some cases, these extra FLOPs may require no extra time to forward on the device.
The results of Figure 14 hide an extra hyperparameter in epinet development: the number of independent samples of the index  $z$ , which we will call  $M$ . All of the results presented in this paper focus on the case  $M = 1000$ , however an agent designer may choose to vary this depending on their tradeoff between statistical loss and computational cost. Figure 15 shows this empirical tradeoff over  $M \in \{10, 30, 100, 300, 1000, 3000\}$  for each of the resnet variants. Once again, we see that in terms of marginal statistics, there is really no benefit to using an epinet. However, once you look at joint log-loss even a small number of epinet samples improves over the ResNet. Interestingly, these results continue to improve for  $M > 1000$  but at a higher computational cost.
H.5 Epinet ablations
We run ablation experiments on the epinet agent that builds on the pre-trained ResNet-50 base network. We sweep over various values for the index dimension, the number of hidden
Figure 14: Quality of marginal and joint predictions across models on ImageNet. Reproduces the results of Figure 2 but using inference FLOPs as measure of computation.
Figure 15: Comparing base ResNet against epinet for differing numbers of sampled indices  $z$ .
layers in the trainable epinet, the prior scale for the matched epinet prior, the prior scale for the ensemble of convolutional networks,  $L2$  weight decay, the number of index samples drawn for each gradient step, label smoothing, and temperature re-scaling post-training. We keep all other hyperparameters fixed to the open-sourced default configuration while sweeping over one hyperparameter.
Our results are summarized in Figure 16. We see that a larger index dimension improves the joint loss-loss, but does not necessarily improve the marginal log-loss and classification error. Adding another hidden layer to the trainable epinet makes the marginal and joint loss-loss worse, but we suspect that the performance could be improved with more tuning. The epinet performance seems sensitive to the prior scales. In the third and fourth rows, we see that the performance degrades quickly when the prior scales become too large. The epinet seems relatively robust to different values of  $L2$  weight decay. A large weight decay improves the joint log-loss but slightly worsens the marginal log-loss. The number of index samples and degree of label smoothing do not seem to affect the performance of epinet. Interestingly, we find that using a cold temperature  $< 1$  to re-scale the ENN output logits post-training improves the agent's performance during evaluation (Wenzel et al., 2020). The improvement is most dramatic in the joint log-loss.
Figure 16: Ablation studies of epinet with ResNet-50 base model on ImageNet.