SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Only the final term depends on the prediction  $\hat{P}$ . Its conditional expectation  $-\mathbb{E}[\ln \hat{P}(y_{1:\tau})|\mathcal{D}]$  is the log loss. Hence, minimizing expected KL divergence is equivalent to minimizing log loss. This result follows almost immediately from Pinsker's and Jensen's inequalities. A proof can be found in (Wen et al., 2022).
Theorem 2. [formal] For all data distributions, reward functions  $r$  with range  $[0,1]$ , and actions  $\hat{a} \in \arg \max_{a} \sum_{y_{1:\tau}} \hat{P}_{1:\tau}(y_{1:\tau})r(a,y_{1:\tau})$ ,
$$
\mathbb {E} [ r (\hat {a}, y _ {1: \tau}) | \mathcal {D}, x _ {1: \tau} ] \geq \max _ {a \in \mathcal {A}} \mathbb {E} [ r (a, y _ {1: \tau}) | \mathcal {D}, x _ {1: \tau} ] - \sqrt {2 \mathbb {E} [ \mathbf {d} _ {\mathrm {K L}} (P _ {1 : \tau} \| \hat {P} _ {1 : \tau}) | \mathcal {D} , x _ {1 : \tau} ]}.
$$
The KL-divergence is the difference between the log loss attained by  $\hat{P}_{1:\tau}$  and that attained by an optimal prediction  $P_{1:\tau}$ . The shortfall of action  $\hat{a}$  is hence bounded by a measure of joint prediction error. If this error is small,  $\hat{P}_{1:\tau}$  leads to good decisions regardless of the reward function.
C ENNs versus BNNs
In this appendix, we provide a proof of Theorem 3.
Proof. Consider a BNN  $(f,p)$ . Without loss of generality, take the parameter space of  $f$  to be  $\mathbb{R}^d$ . Let  $P_Z$  be an absolutely continuous reference distribution over  $\mathbb{R}^d$ . Via Knothe-Rosenblatt rearrangement (Knothe, 1957; Rosenblatt, 1952), for each  $\nu \in \mathbb{R}^d$ , there exists a transport map from  $P_Z$  to  $p_\nu$ . For each epistemic index  $z \in \mathbb{R}^d$  and BNN parameter vector  $\nu$ , let  $\hat{\theta}_{\nu,z}$  be the corresponding base network parameters generated by this transport map. Let  $f_{\nu}'(x,z) = f_{\hat{\theta}_{\nu,z}}(x)$ . It is easy to see that the ENN  $(f',P_Z)$  is defined with respect to  $f$  and expresses the BNN  $(f,p)$ .
For the second part of the theorem, consider a linear base network  $f_{\theta}(x) = \theta^{\top}x$ , a standard Gaussian epistemic index reference distribution, and an ENN
$$
f _ {\theta^ {\prime}} ^ {\prime} (x, z) = f _ {\theta} (x) + \theta^ {\prime \prime} z,
$$
where  $\theta' = (\theta, \theta'')$ . Clearly, there are ENN parameters  $\theta'$  and an epistemic index  $z$  such that no  $\hat{\theta}$  satisfies  $f_{\theta'}'(x, z) = f_{\hat{\theta}}(x)$ . If follows that no BNN defined with respect to  $f$  can express the ENN.
D Bayesian linear regression
In this appendix, we provide a proof of Theorem 4. First, we state without proof a standard result on Bayesian linear regression (Minka, 2000).
Lemma 1. Let  $\mathcal{D} = \{(x_i, y_i, i)\}_{i=1}^N$  be generated by a linear-Gaussian model. Conditioned on  $\mathcal{D}$ ,  $\nu_*$  is Gaussian and
$$
\mathbb {E} [ \nu_ {*} | \mathcal {D} ] = \left(\frac {1}{\sigma^ {2}} \sum_ {i = 1} ^ {N} x _ {i} x _ {i} ^ {\top} + \frac {1}{\sigma_ {0} ^ {2}} I\right) ^ {- 1} \left(\frac {1}{\sigma^ {2}} \sum_ {i = 1} ^ {N} x _ {i} y _ {i}\right), \quad \operatorname {C o v} [ \nu_ {*} | \mathcal {D} ] = \left(\frac {1}{\sigma^ {2}} \sum_ {i = 1} ^ {N} x _ {i} x _ {i} ^ {\top} + \frac {1}{\sigma_ {0} ^ {2}} I\right) ^ {- 1}. \tag {10}
$$
Next, we prove a result on the near-orthogonality of vectors sampled uniformly from a unit sphere. Note that a.s. abbreviates almost surely.
Lemma 2. Let  $b_n$  and  $c_n$  be independent vectors sampled uniformly from the  $n$ -dimensional unit sphere. Then,
$$
\lim  _ {n \to \infty} b _ {n} ^ {\top} c _ {n} \stackrel {{a. s.}} {=} 0.
$$
Proof. Note that unit random vectors generated by sampling normal vectors from  $N(0, I)$  and normalizing are uniformly distributed over the  $\mathbb{R}^d$  unit sphere.
Let  $\alpha_{n}$  and  $\beta_{n}$  be independent samples from a chi-squared distribution with  $n$  degrees of freedom. Let  $u_{n} = \alpha_{n}^{1 / 2}b_{n}$  and  $v_{n} = \beta_{n}^{1 / 2}c_{n}$ . Clearly the distribution of both  $u_{n}$  and  $v_{n}$  is
isotropic, and since  $\alpha_{n}$  and  $\beta_{n}$  are distributed chi-squared with  $n$  degrees of freedom (which is that of the length of  $n$ -dimensional standard normal),  $u_{n}$  and  $v_{n}$  are independent standard normal random vectors. Further,  $b_{n} = u_{n} / \| u_{n}\|_{2}$  and  $c_{n} = v_{n} / \| v_{n}\|_{2}$ .
For any  $n$
$$
\begin{array}{l} b _ {n} ^ {\top} c _ {n} = \frac {\sum_ {i = 1} ^ {n} u _ {i} v _ {i}}{\sqrt {\sum_ {i = 1} ^ {n} u _ {i} ^ {2} \sum_ {i = 1} ^ {n} v _ {i} ^ {2}}} \\ = \frac {\frac {\sum_ {i = 1} ^ {n} u _ {i} v _ {i}}{n}}{\sqrt {\frac {\sum_ {i = 1} ^ {n} u _ {i} ^ {2}}{n} \frac {\sum_ {i = 1} ^ {n} v _ {i} ^ {2}}{n}}} \\ \end{array}
$$
Since  $\{u_i\}_{i=1}^n$  and  $\{v_i\}_{i=1}^n$  are i.i.d  $N(0,1)$ , by the strong law of large numbers,
$$
\lim  _ {n \to \infty} \frac {\sum_ {i = 1} ^ {n} u _ {i} v _ {i}}{n} \stackrel {{\mathrm {a . s .}}} {{=}} E [ u _ {i} v _ {i} ] = 0,
$$
$$
\lim  _ {n \to \infty} \frac {\sum_ {i = 1} ^ {n} u _ {i} ^ {2}}{n} \stackrel {{\mathrm {a . s .}}} {{=}} E [ u _ {i} ^ {2} ] = 1,
$$
$$
\lim_{n\to \infty}\frac{\sum_{i = 1}^{n}v_{i}^{2}}{n}\stackrel {\text{a.s.}}{=}E[v_{i}^{2}] = 1.
$$
Hence, by the continuous mapping theorem,
$$
\lim  _ {n \to \infty} b _ {n} ^ {\top} c _ {n} \stackrel {{\text {a . s .}}} {=} 0.
$$
Theorem 4. Let data  $\mathcal{D} = \{(x_i,y_i,i)\}_{i = 1}^N$  be generated by a linear-Gaussian model and  $f$  be a linear-Gaussian epinet. Let  $\hat{\theta}\in \arg \min_{\theta}\sum_{i = 1}^{N}\int_{z}P_{Z}(dz)\ell_{\sigma ,\lambda}^{\mathrm{LSG}}(\theta ,z,x_{i},y_{i},i)$  with parameter  $\lambda = \sigma^2 /(N\sigma_0^2)$ . Then, conditioned on  $(\mathcal{D},c_{1:N},P_0)$ ,  $f_{\hat{\theta}}(\cdot ,z)$  converges in distribution to  $g_{\nu_*}$  as  $D_Z$  grows, almost surely.
Proof. The statement that, conditioned on  $(\mathcal{D}, c_{1:N}, P_0)$ ,  $f_{\hat{\theta}}(\cdot, z)$  converges in distribution to  $g_{\nu_*}$  as  $D_Z$  grows can be restated as
$$
\mathbb {P} \left(g _ {\nu_ {*}} \in F | \mathcal {D}\right) \stackrel {a. s.} {=} \lim  _ {D z \rightarrow \infty} \mathbb {P} \left(f _ {\hat {\theta}} (\cdot , z) \in F | \mathcal {D}, c _ {1: N}, P _ {0}\right), \tag {11}
$$
for all measurable sets  $F$  of functions mapping  $\mathbb{R}^D$  to  $\Re$ .
$$
\hat {\theta} = (\hat {\zeta}, \hat {\eta}) \in \underset {\zeta , \eta} {\arg \min } \sum_ {i = 1} ^ {N} \int_ {z} P _ {Z} (d z) \left(f _ {\zeta , \eta} \left(x _ {i}, z\right) - y - \sigma c _ {i} ^ {T} z\right) ^ {2} + \frac {\sigma^ {2}}{\sigma_ {0} ^ {2}} \left(\| \zeta \| _ {2} ^ {2} + \| \eta \| _ {2} ^ {2}\right).
$$
Since the optimization problem is strictly convex in  $\theta$ ,  $\hat{\theta}$  is the unique minimizer of this expression. Some simple algebra establishes that
$$
\hat {\zeta} = \hat {\Sigma} \left(\frac {1}{\sigma^ {2}} \sum_ {i = 1} ^ {N} x _ {i} y _ {i}\right) \text {a n d} \hat {\eta} = \left(\frac {1}{\sigma} \sum_ {i = 1} ^ {N} c _ {i} x _ {i} ^ {\top} + \frac {1}{\sigma_ {0}} P _ {0}\right) \hat {\Sigma} - P _ {0}, \tag {12}
$$
$$
\mathrm {w h e r e} \hat {\Sigma} = \left(\frac {1}{\sigma^ {2}} \sum_ {i = 1} ^ {N} x _ {i} x _ {i} ^ {\top} + \frac {1}{\sigma_ {0} ^ {2}} I\right) ^ {- 1}.
$$
For any  $x \in \mathbb{R}^D$  and  $z \in \mathbb{R}^{Dz}$ ,  $f_{\hat{\zeta},\hat{\eta}}(x,z) = \left(\hat{\zeta} + z^\top (\hat{\eta} + P_0)\right)^\top x$  and  $g_{\nu_*}(x) = \nu_*^\top x$ . Hence, (11) holds if and only if  $\hat{\zeta} + z^\top (\eta + P_0)$  converges in distribution to  $\nu_*$ , conditioned on  $(\mathcal{D},\{c_i\}_{i=1}^N,P_0)$ , almost surely. Since  $\nu_*$  and  $\hat{\zeta} + z^\top (\eta + P_0)$  are Gaussian, it is sufficient to show that
$$
\lim _ {D _ {Z} \to \infty} \hat {\zeta} = \mathbb {E} [ \nu_ {*} | \mathcal {D} ] \quad \mathrm {a n d} \quad \lim _ {D _ {Z} \to \infty} (\eta + P _ {0}) ^ {\top} (\eta + P _ {0}) = \operatorname {C o v} (\nu_ {*} | \mathcal {D}).
$$
Based on (10) and (12),  $\hat{\eta} = \mathbb{E}[\nu_{*}|\mathcal{D}]$  for any  $D_Z$ . Hence, in order to prove the theorem, it is sufficient to show
$$
\lim  _ {D _ {Z} \rightarrow \infty} (\eta + P _ {0}) ^ {\top} (\eta + P _ {0}) = \operatorname {C o v} \left(\nu_ {*} | \mathcal {D}\right). \tag {13}
$$
For any  $D_Z$
$$
\begin{array}{l} \left(\eta + P _ {0}\right) ^ {\top} \left(\eta + P _ {0}\right) = \hat {\Sigma} \left(\frac {1}{\sigma} \sum_ {i = 1} ^ {N} c _ {i} x _ {i} ^ {\top} + \frac {1}{\sigma_ {0}} P _ {0}\right) ^ {\top} \left(\frac {1}{\sigma} \sum_ {i = 1} ^ {N} c _ {i} x _ {i} ^ {\top} + \frac {1}{\sigma_ {0}} P _ {0}\right) \hat {\Sigma} \\ = \hat {\Sigma} \left(\frac {1}{\sigma^ {2}} \sum_ {i = 1} ^ {N} x _ {i} x _ {i} ^ {\top} + \frac {1}{\sigma_ {0} ^ {2}} I\right) \hat {\Sigma} + \hat {\Sigma} \left(\frac {1}{\sigma_ {0} ^ {2}} \left(P _ {0} ^ {\top} P _ {0} - I\right)\right) \hat {\Sigma} \\ + \hat {\Sigma} \left(\frac {1}{\sigma^ {2}} \sum_ {i = 1, j = 1, i \neq j} ^ {N} c _ {i} ^ {\top} c _ {j} x _ {i} x _ {j} ^ {\top} + \frac {1}{\sigma \sigma_ {0}} \sum_ {i = 1} ^ {N} \left(x _ {i} c _ {i} ^ {\top} P _ {0} + P _ {0} ^ {\top} c _ {i} x _ {i} ^ {\top}\right)\right) \hat {\Sigma} \\ \end{array}
$$
Recall that  $\{c_i\}_{i=1}^N$  and columns of  $P_0$  are all sampled i.i.d and uniformly from unit sphere in  $\mathbb{R}^{Dz}$ , by Lemma 2,
$$
\lim_{D_{Z}\to \infty}c_{i}^{\top}c_{j}\stackrel {\text{a.s.}}{=}0\forall i\neq j,
$$
$$
\lim_{D_{Z}\to \infty}P_{0}^{\top}c_{i}\stackrel {\text{a.s.}}{=}0\forall i
$$
$$
\lim_{D_{Z}\to \infty}P_{0}^{\top}P_{0}\stackrel {\text{a.s.}}{=}I.
$$
$$
\lim  _ {D _ {Z} \to \infty} \left(\eta + P _ {0}\right) ^ {\top} (\eta + P _ {0}) \hat {\Sigma} \stackrel {{\text {a . s .}}} {{=}} \hat {\Sigma} \left(\frac {1}{\sigma^ {2}} \sum_ {i = 1} ^ {N} x _ {i} x _ {i} ^ {\top} + \frac {1}{\sigma_ {0} ^ {2}} I\right) = \hat {\Sigma}.
$$
E Didactic examples
To offer some intuition for how epinets work and what they accomplish, we present a simple example specialized to a linear base model. A linear base model produces an output  $\mu_{\zeta}(x) = \zeta^T x$  given an input  $x$  and model parameters  $\zeta$ . It is natural to add to this a linear epinet  $\sigma_{\eta}(\phi_{\zeta}(x), z) = z^T \eta x$ . The combined architecture is equivalent to a linear hypermodel (Dwaracherla et al., 2020). To see this, note that,
$$
f _ {\theta} (x, z) = \mu_ {\zeta} (x) + \sigma_ {\eta} \left(\phi_ {\zeta} (x), z\right) = \zeta^ {T} x + z ^ {T} \eta x = (\zeta + \eta^ {T} z) ^ {T} x = \mu_ {\zeta + \eta^ {T} z} (x). \tag {14}
$$
As such, properties of linear hypermodels, such as their ability to implement exact Bayesian linear regression, carry over to such epinets.
Figures 7 illustrates predictive uncertainty estimates produced by linear epinets. In this figure, posterior credible intervals of an epinet are compared against exact Bayesian inference. Data is generated by a one-dimensional linear regression model with a Gaussian prior distribution and Gaussian noise. The loss function for this epinet follows prior work by includes Gaussian bootstrapping in regression (Lu and Van Roy, 2017; Dwaracherla et al., 2020),
$$
\ell (\theta , x _ {i}, y _ {i}, z) = \left(f _ {\theta} \left(x _ {i}, z\right) - y _ {i} + \sigma c _ {i} ^ {T} z\right) ^ {2}.
$$
Here  $c_{i}$  is a random signature drawn from the unit sphere in  $D_Z$  generated independently for each training example  $(x_{i},y_{i})$  and  $\sigma$  is the scale of the additive bootstrap noise. This figure indicates that the epinet outputs well-calibrated marginal predictive distributions.
We next consider classification with a two-dimensional input and two classes. Data is generated by a standard logistic regression model with parameters drawn from a Gaussian prior. Figure 8 presents standard deviations of marginal predictive distributions across the input space. We supplement a standard logistic regression model with a linear epinet. The
plots compare results against those generated via SGMCMC, which we expect in this case to closely approximate exact Bayesian inference. While these figures bear qualitative similarities, significant differences arise because the linear epinet architecture imposes symmetries that are not respected by exact posterior distributions. In particular, this epinet can be thought of as representing parameter uncertainty as Gaussian. While our data generating process assumes a Gaussian prior distribution, the posterior distributions, which are conditioned on binary outcomes, are not Gaussian. More complex, nonlinear, epinets should be able to more accurately represent the posterior distribution over classifiers.
Figure 7: Epinet predictions in Gaussian linear regression.
Figure 8: Epinet predictions in Logistic regression.
F Dyadic sampling
To evaluate the quality of joint predictions, we sample batches of inputs  $(x_{1},\ldots ,x_{\tau})$  and assess log-loss with respect to corresponding labels  $(y_{1},\dots ,y_{\tau})$ . With a high-dimensional input space, labels of inputs sampled uniformly at random are typically nearly independent. Hence, a large batch size  $\tau$  is required to distinguish joint predictions that effectively reflect interdependencies among labels. However, this becomes impractical because computational requirements for testing grow exponentially in  $\tau$ . Dyadic sampling serves as a practical heuristic that samples inputs more strategically so that effective agents are distinguished with a manageable batch size  $\tau$  (Osband et al., 2022b) even when inputs are high-dimensional.
F.1 Basic version
The basic version of dyadic sampling, for each batch, first samples two independent random anchor points,  $\tilde{x}_1$  and  $\tilde{x}_2$ , from the input distribution. Then, to form a batch of size  $\tau$ , sample  $\tau$  points independently and with equal probability from these two anchor points  $\{\tilde{x}_1,\tilde{x}_2\}$ . To assess an agent, its joint prediction of labels is evaluated for this batch of size  $\tau$ . Even with a moderate value of  $\tau = 10$ , a batch produced by dyadic sampling gives rise to labels that are likely to correlate – in particular, labels assigned to the same anchor point. Osband et al. (2022b) demonstrate that, across many problems, this sampling heuristic is effective in distinguishing the quality of joint predictions with modest computation.
On the Neural Testbed, the input distribution is standard normal. Thus, for each test batch, we sample anchor points  $x_{1}, x_{2} \sim N(0, I)$  and then re-sample  $\tau = 10$  points from these two anchor points to form a batch.
On ImageNet, for faster evaluation, rather than sampling anchor points from the evaluation set, we split the evaluation set into batches of size 2. We then iterate over these batches of size 2, re-sample  $\tau = 10$  points from each input pair, and evaluate the log-loss of an agent's joint predictions on these batches of size  $\tau = 10$ . Finally, we take the average of all the joint log-losses.
Figure 9(a) shows a few examples of these dyadic input batches of size  $\tau = 10$ . It may seem unsatisfactory that images are repeated exactly within a batch. Even though we view this metric as a unit test that an intelligent agent pass, it is conceivable to design a 'cheating' agent that takes advantage of this repeating structure. In the following section, we will consider a more robust version of dyadic sampling, where instead of repeating images exactly,
Figure 9: Examples of input batches used for evaluating joint predictions. Figure 9(a) is generated through dyadic sampling and Figure 9(b) includes additional perturbations.
we perturb images using standard data augmentation techniques to introduce more diversity within a dyadic batch.
F.2 Augmented dyadic sampling
In vanilla dyadic sampling, each dyadic batch has two unique elements (the anchor points), which are repeated multiple times within the batch. To make the metric more robust at distinguishing agents, we independently perturb each input in a dyadic batch using standard data augmentation techniques, so each input within a batch differs from others. For ImageNet, the perturbation takes the form of random cropping and flipping. We take the label of each perturbed image to be the label of its original image. Effective joint predictions indicate that images perturbed from the same anchor image are likely to be the same. Joint log-loss penalizes agents that do not recognize this. We call this sampling scheme augmented dyadic sampling. Figure 9(b) presents a few examples of augmented dyadic batches for Imagenet.
We evaluate our trained ResNet, ensemble, and epinet agents in Section 6 using augmented dyadic sampling, and we compare the joint log-loss with that obtained from basic dyadic sampling. In Figure 10, we see that the results from using these two sampling schemes are qualitatively similar. While the overall joint log-loss is higher for augmented dyadic sampling due to added perturbations, all agents benefit from increasing the model size. The joint log-loss of the ensemble agent improves more so than the ResNet agent with increasing model size. More importantly, the epinet agent outperforms both baselines by a huge margin under both sampling schemes. These results give us further confidence in the quality of joint predictions produced by the epinet agent.
Figure 10: An agent's joint log-loss under dyadic (left) and augmented dyadic sampling (right).
G Testbed experiments
This section provides details about the Neural Testbed experiments in Section 5. We begin with a review of the neural testbed as a benchmark problem, and the associated generative models. We then give an overview of the baseline agents we compare against in our evaluation. Next, we provide supplementary details for the hyperparameters and implementation details for the epinet agent as outlined in Section 5. Finally, we investigate the sensitivity our hyperparameter choices when evaluated across the testbed.
G.1 Neural testbed
The Neural Testbed (Osband et al., 2022a) is a collection of neural-network-based, synthetic classification problems that evaluate the quality of an agent's predictive distributions. We make use of the open-source code at https://github.com/deepmind/neural_testbed. The Testbed problems use random 2-layer MLPs with width 50 to generate training and testing data. The specific version we test our agents on entails binary classification, input dimension  $D \in \{2,10,100\}$ , number of training samples  $T = \lambda D$  for  $\lambda \in \{1,10,100,1000\}$ , temperature  $\rho \in \{0.01,0.1,0.5\}$  for controlling the signal-to-noise ratio, and 5 random seeds for generating different problems in each setting. The performance metrics are averaged across problems to give the final performance scores.
G.2 Benchmark agents
We follow Osband et al. (2022a) and consider the benchmark agents as in Table 1. We use the open-source implementation and hyperparameter sweeps at https://anonymous.4open.science/r/neural_testbed-8961/agents/factories. According to Osband et al. (2022a), the benchmark agents are carefully tuned on the Testbed problems, so we do not further tune these agents.
G.3 Epinet
We take the reference distribution of the epistemic index to be a standard Gaussian with dimension  $D_Z = 8$ . The base network  $\mu_{\zeta}$  has the same architecture as the baseline mlp agent, which is a 2-layer MLP with ReLU activation and 50 units in each hidden layer. The learnable part of the epinet  $\sigma_{\eta}^{L}$  takes  $\phi_{\zeta}(x)$  and index  $z$  as inputs, where  $\phi_{\zeta}(x)$  is the concatenation of  $x$  and the last-layer features of the base network. The learnable network has the form  $\sigma_{\eta}^{L}(\phi_{\zeta}(x), z) = g_{\eta}([\phi_{\zeta}(x), z])^{T} z$  where  $[\phi_{\zeta}(x), z]$  is the concatenation of  $\phi_{\zeta}(x)$  and  $z$ , and  $g_{\eta}(\cdot)$  is a 2-layer MLP with hidden width 15, ReLU activation, and outputs in  $\mathbb{R}^{D_Z \times C}$  for number of classes  $C = 2$ .
For the fixed prior  $\sigma^P$ , we consider an ensemble of  $D_Z$  networks. Each member of the ensemble is a small MLP with 2 hidden layers and 5 units in each hidden layer. Each MLP takes  $x$  as input and returns logits for the two classes. Let  $p^i(x) \in \mathbb{R}^C$  denote the output of the  $i^{\text{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$  and multiplying the sum by a tunable scaling factor  $\alpha$ . Thus, we can write the prior function  $\sigma^P(\phi_\zeta(x), z) = \alpha \sum_{i=1}^{D_Z} p^i(x) z_i$ .
We combine the base network and the epinet by adding their outputs and applying a stop-gradient operation on  $\phi_{\zeta}(x)$ ,
$$
f _ {\theta} (x, z) = \mu_ {\zeta} (x) + \sigma_ {\eta} ^ {L} (\operatorname {s g} [ \phi_ {\zeta} (x) ], z) + \sigma^ {P} (\operatorname {s g} [ \phi_ {\zeta} (x) ], z),
$$
where  $\theta = (\zeta, \eta)$  denotes all the parameters of the base net and epinet. We train the parameters  $\zeta$  and  $\eta$  jointly. The training loss takes the form as specified in (7), where we use log loss for the data loss and ridge regularization. We update  $\theta$  using Algorithm 1, with a batch size of 100 and number of epistemic index samples equal to the index dimension. We use Adam optimizer with learning rate 1e-3. The  $L2$  weight decay and prior scaling factor  $\alpha$  are roughly adjusted for different problem settings, taking in account the number of training samples and SNR.
Our implementation of the epinet agent can be found under the path /agents/factories/epinet.py in the anonymized neural testbed github.
G.4 Ablation studies
We run ablation experiments on the epinet agent that is trained simultaneously along with the base network. We sweep over various values for the index dimension, the number of hidden layers in the trainable epinet, the width of hidden layers of epinet, the ensemble prior's prior scale, the width of the hidden layers in the prior network, and  $L2$  weight decay. We keep all other hyperparameters fixed to the open-sourced default configuration while sweeping over one hyperparameter.
Our results are summarized in Figure 11. We see that a larger index dimension improves both joint and marginal kl estimates. The epinet performance is not sensitive to the number of hidden layers. We suspect that this is due to similar number of parameters across epinets with different number of hidden layers. We observe that epinet performance is not sensitive to width of the epinet hidden layers once the width is large enough. Performms of epinet seems sensitive to the prior scales. Smaller prior scale leads to better marginal kl, too small or too large prior scale degrades the joint kl. Increasing the width of the models in the ensemble prior network improves the epinet performance. However, the improvement seems marginal after a point. The epinet seems to be sensitive to  $L2$  weight decay. A very small or large weight decay degrades the performance of epinet.
Figure 11: Ablation studies of epinet with 2-layer mlp base model on the neural testbed.
H Image classification
This section gives an overview of our experiments on image classification problems outlined in Section 6. We begin with a review of the hyperparameter choices and design details for the agents as implemented in our experiments. Then, we include a comparison of these agents to benchmark implementations in the field as embodied by 'uncertainty baselines' (Nado et al., 2021). Next, we present results for an evaluation on both CIFAR-10 and CIFAR-100, and find that the overall results match that on ImageNet. We complement these results with an analysis of the computational cost in terms of FLOPs as well as memory on modern TPU architectures. Finally, we perform a suite of ablations to investigate the sensitivity of our results across ImageNet.
H.1 Epinet details