SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ian Osband,\* Zheng Wen, Seyed Mohammad Asghari, Vikranth Dwaracherla, Morteza Ibrahimi, Xiuyuan Lu, and Benjamin Van Roy Google DeepMind, Efficient Agent Team, Mountain View {ian.osband, m.ibrahimi}@gmail.com {zhengwen,smasghari,vikranthd,lxlu,benvanroy}@google.com
Intelligent agents need to know what they don't know, and this capability can be evaluated through the quality of joint predictions. In principle, ensemble methods can produce effective joint predictions, but the compute costs are prohibitive for large models. We introduce the epinet: an architecture that can supplement any conventional neural network, including large pretrained models, and can be trained with modest incremental computation to estimate uncertainty. With an epinet, conventional neural networks outperform large ensembles of hundreds or more particles, and use orders of magnitude less computation. The epinet does not fit the traditional framework of Bayesian neural networks, so we introduce the epistemic neural network (ENN) as a general interface for models that generate joint predictions.
Consider a conventional neural network trained to predict whether a random person would classify a drawing as a 'rabbit' or a 'duck'. As illustrated in Figure 1, given a single drawing, the network outputs a marginal prediction that assigns probabilities to the two classes. If the probabilities are each 0.5, it remains unclear whether this is because labels sampled from random people are equally likely, or whether the neural network would learn a single class if trained on more data. Conventional neural networks do not distinguish these cases, even though it can be critical for decision making systems to know what they do not know. This capability can be assessed through the quality of joint predictions (Wen et al., 2022).
The two tables to the right of Figure 1 represent possible joint predictions that are each consistent with the network's uniform marginal prediction. These joint predictions are over pairs of labels for the same image,  $(y_{1},y_{2})\in \{R,D\} \times \{R,D\}$ . For any such joint prediction, Bayes' rule defines a conditional prediction for  $y_{2}$  given  $y_{1}$ . The first table indicates inevitable uncertainty that would not be resolved through training on additional data; conditioning on the first label does not alter the prediction for the second. The second table indicates that additional training should resolve uncertainty; conditioned on the first label, the prediction for the second label assigns all probability to the same outcome as the first.
Figure 1 presents the toy problem of predictions across two identical images as a simple illustration of these types of uncertainty. The observation that joint distributions express whether uncertainty is resolvable extends more generally to practical cases, where the inputs differ, or where there are more than two simultaneous predictions (Osband et al., 2022a).
Bayesian neural networks (BNNs) offer a statistically-principled way to make effective joint predictions, by maintaining an approximate posterior over the weights of a base neural network. Assymptotically these can recover the exact posterior, but the computational costs
*Contact ian.osband@gmail.com
Figure 1: Conventional neural nets generate marginal predictions, which do not distinguish genuine ambiguity from insufficiency of data. Joint predictions can make this distinction.
are prohibitive for large models (Welling and Teh, 2011). Ensemble-based BNNs offer a more practical approach by approximating the posterior distribution with an ensemble of statistically plausible networks that we call particles (Osband and Van Roy, 2015; Lakshminarayanan et al., 2017). While the quality of joint predictions improves with more particles, practical implementations are often limited to at most tens of particles due to computational constraints.
In this paper, we introduce an approach that outperforms ensembles of hundreds of particles at a computational cost less than that of two particles. Our key innovation is the epinet: a network architecture that can be added to any conventional neural network to estimate uncertainty. Figure 2 offers a preview of results presented in Section 6, where we compare these approaches on ImageNet. The quality of the ResNet's marginal predictions - measured by classification error or marginal log-loss - does not change much if supplemented with an epinet. However the epinet-enhanced ResNet dramatically improves the quality of joint predictions, as measured by the joint log-loss, outperforming the ensemble of 100 particles, with total parameters less than 2 particles. Prior work has shown the importance of joint predictions in driving effective decisions for a broad class of problems, including combinatorial decision problems and sequential decision problems (Wen et al., 2022; Osband et al., 2022a).
Figure 2: Quality of marginal and joint predictions across models on ImageNet (Section 6).
The epinet does not fit into the traditional framework of BNNs. In particular, it does not represent a distribution over base neural network parameters. To accommodate development of the epinet and other approaches that do not fit the BNN framework, we introduce the concept of epistemic neural networks (ENNs). We establish that all BNNs are ENNs, but there are useful ENNs such as the epinet, that are not BNNs.
Our research builds on the literature in Bayesian deep learning (Hinton and Van Camp, 1993; Neal, 2012). BNNs represent epistemic uncertainty via approximating the posterior distribution over parameters of a base neural network (Der Kiureghian and Ditlevsen, 2009; Kendall and Gal, 2017). A challenge is the computational cost of posterior inference, which becomes intractable even for small networks (MacKay, 1992), and even approximate SGMCMC becomes prohibitive for large scale models (Welling and Teh, 2011).
Tractable methods for approximate inference has renewed interest in BNNs. Variational approaches such as Bayes by backprop (Blundell et al., 2015) use an evidence-based lower bound (ELBO) to approximate the posterior distribution, and related approaches use the same objective with more expressive weight distributions (Louizos and Welling, 2017). One
influential line of work claims that MC dropout can be viewed as one such approach (Gal and Ghahramani, 2016), although subsequent papers have noted that the quality of this approximation can be very poor (Osband, 2016; Hron et al., 2017). As of 2022, perhaps the most popular approach is ensemble-based, with an ensemble of models, each referred to as a particle, trained in parallel so that they together approximate a posterior distribution (Osband and Van Roy, 2015; Lakshminarayanan et al., 2017).
Ensemble-based BNNs train multiple particles independently. This incurs computational cost that scales with the number of particles. A thriving literature has emerged that seeks the benefits of large ensembles at lower computational cost. Some approaches only ensemble parts of the network, rather than the whole (Osband et al., 2019; Havasi et al., 2020). Others introduce new architectures to directly incorporate uncertainty estimates, often inspired by connections to Gaussian processes (Malinin and Gales, 2018; Charpentier et al., 2020; Liu et al., 2020; van Amersfoort et al., 2021). Others perform Bayesian inference more directly in the function space in order to sidestep issues relating to overparameterization (Sun et al., 2019).
In general, research in Bayesian deep learning has focused more on developing methodology than unified evaluation (Osband et al., 2022a). Perhaps because the potential benefits of BNNs are so far-reaching, different papers have emphasized improvements in classification accuracy (Wilson, 2020), expected calibration error (Ovadia et al., 2019), OOD performance (Hendrycks and Dietterich, 2019), active learning (Gal et al., 2017) and decision making (Osband et al., 2019). However, in each of these settings it generally is possible to obtain improvements via methods that do not aim to approximate posterior distributions. Perhaps for this reason, there has been a recent effort to refocus evaluation on how well methods actually approximate 'gold standard' Bayes posteriors (Izmailov et al., 2021).
Our work on ENNs is motivated by the importance of joint predictions in driving decision, exploration, and adaptation (Wang et al., 2021; Wen et al., 2022; Osband et al., 2022a). This line of research, which we build on in Section 3.1, establishes a sense in which joint predictions are both necessary and sufficient to drive decisions. Effectiveness of ENN designs can be assessed through the quality of joint predictions. This perspective allows us to consider approaches beyond those accommodated by the BNN framework. As we will demonstrate, this can lead to significant improvements in performance.
3 Epistemic neural networks
A conventional neural network is specified by a parameterized function class  $f$ , which produces a vector-valued output  $f_{\theta}(x)$  given parameters  $\theta$  and an input  $x$ . The output  $f_{\theta}(x)$  assigns a corresponding probability  $\hat{P}(y) = \exp((f_{\theta}(x))_y) / \sum_{y'} \exp((f_{\theta}(x))_{y'})$  to each class  $y$ . For shorthand, we write such class probabilities as  $\hat{P}(y) = \text{softmax}(f_{\theta}(x))_y$ . We refer to a predictive class distribution  $\hat{P}$  produced in this way as a marginal prediction, as it pertains to a single input  $x$ .
An ENN architecture, on the other hand, is specified by a pair: a parameterized function class  $f$  and a reference distribution  $P_Z$ . The vector-valued output  $f_{\theta}(x,z)$  of an ENN depends additionally on an epistemic index  $z$ , which takes values in the support of  $P_Z$ . Typical choices of the reference distribution  $P_Z$  include a uniform distribution over a finite set or a standard Gaussian over a vector space. The index  $z$  is used to express epistemic uncertainty. In particular, variation of the network output with  $z$  indicates uncertainty that might be resolved by future data. As we will see, the introduction of an epistemic index allows us to represent the kind of uncertainty required to generate useful joint predictions.
Given inputs  $x_{1},\ldots ,x_{\tau}$ , a joint prediction assigns a probability  $\hat{P}_{1:\tau}(y_{1:\tau})$  to each class combination  $y_{1},\dots ,y_{\tau}$ . While conventional neural networks are not designed to provide joint predictions, joint predictions can be produced by multiplying marginal predictions:
$$
\hat {P} _ {1: \tau} ^ {\mathrm {N N}} \left(y _ {1: \tau}\right) = \prod_ {t = 1} ^ {\tau} \operatorname {s o f t m a x} \left(f _ {\theta} \left(x _ {t}\right)\right) _ {y _ {t}}. \tag {1}
$$
However, this representation models each outcome  $y_{1:\tau}$  as independent and so fails to distinguish ambiguity from insufficiency of data. ENNs address this by enabling more
expressive joint predictions through integrating over epistemic indices:
$$
\hat {P} _ {1: \tau} ^ {\mathrm {E N N}} \left(y _ {1: \tau}\right) = \int_ {z} P _ {Z} (d z) \prod_ {t = 1} ^ {\tau} \operatorname {s o f t m a x} \left(f _ {\theta} \left(x _ {t}, z\right)\right) _ {y _ {t}}. \tag {2}
$$
This integration introduces dependencies so that joint predictions are not necessarily just the product of marginals. Figure 3 provides a simple example of how two different ENNs can use the epistemic index to distinguish the sorts of uncertainty described in Figure 1.
In Figure 3(a) the ENN makes marginal predictions that do not vary with  $z$ , and so the resultant joint predictions are simply the independent product of marginals. This corresponds to an 'aleatoric' or 'irreducible' form of uncertainty that cannot be resolved with data. On the other hand, Figure 3(b) shows an ENN that makes predictions depending on the sign of the epistemic index. This corresponds to 'epistemic' or 'reducible' uncertainty that can be resolved with data. In this case, integrating the 2x2 matrix over  $z$  produces a diagonal matrix with 0.5 in each diagonal entry. As such, Figure 3 shows how an ENN can use the epistemic index to distinguish the two joint distributions of Figure 1.
(a) An ENN indicating an ambiguous image.
(b) An ENN indicating insufficient data.
Figure 3: An ENN can incorporate the epistemic index  $z \sim P_Z$  into its joint predictions. This allows an ENN to differentiate inevitable ambiguity from data insufficiency.
3.1 Evaluating ENN performance
Marginal log loss (also known as cross-entropy loss) is perhaps the most widely used evaluation metric in machine learning. For a single input  $x$ , if a neural network generates a prediction  $\hat{P}$  and the label turns out to be  $y$ , then the sample log loss is the form  $-\ln \hat{P}(y)$ . We say that this is a marginal loss because it only looks at the quality of a prediction over a single (input, output) pair. As we will discuss, minimizing marginal log loss does not generally lead to performant downstream decisions. Good decisions often require good joint predictions.
To formulate a generic decision problem, consider a reward function  $r$  that maps an action  $a \in \mathcal{A}$  and  $\tau$  labels  $y_{1:\tau}$  to a reward  $r(a,y_{1:\tau}) \in [0,1]$ . Given an exact posterior predictive  $P_{1:\tau}$ , consider as an objective maximization of the expected reward  $\sum_{y_{1:\tau}} P_{1:\tau}(y_{1:\tau})r(a,y_{1:\tau})$ . The optimal decision can be approximated based on an ENN's prediction  $\hat{P}_{1:\tau}$  by choosing an action  $a$  that maximizes  $\sum_{y_{1:\tau}} \hat{P}_{1:\tau}(y_{1:\tau})r(a,y_{1:\tau})$ . The following theorem is formalized and proved in Appendix B.
Theorem 1. [informal] There exists a decision problem and an ENN that attains small expected marginal log loss such that actions generated using the ENN perform no better than random guessing.
Theorem 1 indicates that minimizing marginal log loss does not suffice to support effective decisions. The key to this insight is that marginal predictions do not distinguish ambiguity from insufficiency of data. However, this can be addressed by instead considering the joint log loss. Given  $\tau$  data pairs and a joint prediction  $\hat{P}_{1:\tau}$ , we can consider the joint log loss  $-\ln \hat{P}_{1:\tau}(y_{1:\tau})$  in exactly the same way that we looked at the marginal log loss. We formalize our next result in Appendix B.
Theorem 2. [informal] For any decision problem, any ENN that attains small expected joint log loss leads to actions that attain near optimal expected reward.
Theorems 1 and 2 highlight the importance of joint predictions in driving decisions. Since we want machine learning systems to drive effective decisions, we will assess the performance
of ENNs by comparing their joint log loss.2 It is important to note that we will consider this joint loss as a method for assessing quality of a trained ENN. We have not yet discussed how particular forms of ENNs are trained, which will generally be up to the algorithm designer. Section 4 provides further detail on the specific architecture and training loss for the epinet ENN we develop in this paper.
3.2 ENNs versus BNNs
A base neural network  $f$  defines a class of functions. Each element  $f_{\theta}$  of this class is identified by a vector  $\theta$  of parameters, which specify weights and biases. An ENN or BNN is designed with respect to a specific base network, and seek to express uncertainty while learning a function in this class. We will formally define what it means for an ENN or BNN to be defined with respect to a base network. We will then establish results which indicate that, with respect to any base network, all BNNs can be expressed as ENNs but not vice versa.
Consider a base neural network  $f$  which, given parameters  $\theta$  and input  $x$ , produces an output  $f_{\theta}(x)$ . A typical BNN is specified by a pair: a base network  $f$  and a parameterized sampling distribution  $p$ . Given parameters  $\nu$ , a sample  $\hat{\theta}$  can be drawn from the distribution  $p_{\nu}$  to generate a function  $f_{\hat{\theta}}$ . Approaches such as stochastic gradient MCMC, deep ensembles, and dropout can all be framed in this way. For example, with a deep ensemble,  $\hat{\theta}$  comprises parameters of an ensemble particle and  $p_{\nu}$  is the distribution represented by the ensemble, which assigns probability to a finite set of vectors, each associated with one ensemble particle. For any inputs  $x_{1}, \ldots, x_{\tau}$ , by sampling many functions  $(f_{\hat{\theta}^{k}}: k = 1, \ldots, K)$ , a BNN can be used to approximate the corresponding joint distribution over labels  $y_{1}, \ldots, y_{\tau}$ , according to  $\hat{P}(y_{1:\tau}) = \frac{1}{K} \sum_{k=1}^{K} \mathbf{1}((f_{\hat{\theta}^{k}}(x_{1}), \ldots, f_{\hat{\theta}^{k}}(x_{\tau})) = y_{1:\tau})$ .
We say the BNN  $(f,p)$  is defined with respect to its base network  $f$ . We say an ENN  $(f^{\prime},P_{Z})$  is defined with respect to a base network  $f$  if, for any base network parameters  $\theta$ , there exist ENN parameters  $\theta^\prime$  such that  $f_{\theta '}^{\prime}(\cdot ,z) = f_{\theta}$  almost surely with respect to  $P_Z$ . Intuitively, being defined with respect to a particular base network means that the BNN or ENN is designed to learn any function within the class characterized by the base network.
We say that a BNN  $(f,p)$  is expressed as an ENN  $(f^{\prime},P_{Z})$  if, for all  $\nu ,\tau$  , and inputs  $x_{1:\tau}$  there exists  $\theta^\prime$  such that for  $\hat{\theta}\sim p_{\nu},z\sim P_{z}$
$$
\left(f _ {\hat {\theta}} \left(x _ {1}\right), \dots , f _ {\hat {\theta}} \left(x _ {\tau}\right)\right) \stackrel {{d}} {{=}} \left(f _ {\theta^ {\prime}} ^ {\prime} \left(x _ {1}, z\right), \dots , f _ {\theta^ {\prime}} ^ {\prime} \left(x _ {\tau}, z\right)\right). \tag {3}
$$
This condition means that the ENN and BNN make the same joint predictive distributions at all inputs.
We say that an ENN  $(f', P_Z)$  is expressed as a BNN if, for all  $\theta'$ ,  $\tau$ , and inputs  $x_1, \ldots, x_\tau$ , there exists a posterior distribution  $\nu$  such that (3) holds. Intuitively, one architecture is expressed as the other if the latter can represent the same distributions over functions. The following result, established in Appendix C, asserts that any BNN can be expressed as an ENN but not every ENN can be expressed as a BNN.
Theorem 3. For all base networks  $f$ , any BNN defined with respect to  $f$  can be expressed as an ENN defined with respect to  $f$ . However, there exists a base network  $f$  and ENN defined with respect to  $f$  that can not be expressed as a BNN defined with respect to  $f$ .
In supervised learning, de Finetti's Theorem implies that if a sequence of data pairs is exchangeable then they are i.i.d. conditioned on a latent random object (de Finetti, 1929). BNNs use base network parameters  $\theta$  as the object, while ENNs focus on the function  $g_{*}$  itself, without concerning the underlying parameters. ENNs serve as computational mechanisms to approximate the posterior distribution of  $g_{*}$ , allowing functions beyond the base network class to represent uncertainty and allowing better trade-offs between computation and prediction quality. The epinet is an example of an ENN that cannot be expressed as a BNN with the same base network, and showcases the benefits of the ENN interface beyond BNNs.
In problems with high-dimensional inputs, the number of inputs  $x_{1}, \ldots, x_{\tau}$  sampled uniformly from the input distribution, may have to be very large to distinguish ENNs in ways marginal log loss does not (Osband et al., 2022b). To sidestep prohibitive computational costs we use dyadic sampling in our empirical evaluation and review these details in Appendix F.
4 The epinet
This section introduces the epinet, which can supplement any conventional neural network to make a new kind of ENN architecture. Our approach reuses standard deep learning components and training algorithms, as outlined in Algorithm 1. As we will see, it is straightforward to add an epinet to any existing model, even one that has been pretrained. The key to successful application of the epinet comes in the design of the network architecture and loss functions.
Algorithm 1 ENN training via SGD
Inputs:
dataset	training examples D = {(xi,yi,i})N i=1
ENN	network f, reference PZ, initialization θ0
loss	ℓ evaluates example (xi,yi,i) for index z
batch size	data samples NB, index samples NZ
optimizer	update rule and number of iterations T
Returns:
θT	parameter estimates for the ENN.
1: for t in 0,...,T-1 do
2: sample data I = i1,...,inB ~ Unif{1,...,N}.
3: sample indices Z = z1,...,znZ ~ PZ.
4: compute grad ← ∇θ|θ=θt ∑z∈Z ∑i∈Iℓ(θ,z,xi,yi,i).
5: update θt+1← optimizer(θt,grad)
Figure 4: Epinet network architecture.
Consider a conventional neural network as the base network. Given base parameters  $\zeta$  and an input  $x$ , the output is  $\mu_{\zeta}(x)$ . For a classification model, the class probabilities would be  $\mathrm{softmax}(\mu_{\zeta}(x))$ . An epinet is a neural network with privileged access to inputs and outputs of activation units in the base network. A subset of these inputs and outputs, which we call features  $\phi_{\zeta}(x)$ , are taken as input to the epinet along with an epistemic index  $z$ . For example, these features might be the last hidden layer in a ResNet. The epistemic index is sampled from a standard Gaussian distribution in dimension  $D_Z$ . For epinet parameters  $\eta$ , the epinet outputs  $\sigma_{\eta}(\phi_{\zeta}(x), z)$ . To produce an ENN, the output of the epinet is added to that of the base network, though with a "stop gradient":3
$$
\underbrace {f _ {\theta} (x , z)} _ {\text {E N N}} = \underbrace {\mu_ {\zeta} (x)} _ {\text {b a s e n e t}} + \underbrace {\sigma_ {\eta} (\operatorname {s g} [ \phi_ {\zeta} (x) ] , z)} _ {\text {e p i n e t}}. \tag {4}
$$
We find that with this stop gradient, training dynamics more reliably produce models that perform well out of sample. The ENN parameters  $\theta = (\zeta, \eta)$  include those of the base network  $\zeta$  and epinet  $\eta$ . Due to the additive structure of the epinet, multiple samples of the ENN can be obtained with only one forward pass of the base network. Where the epinet is much smaller than the base network this can lead to significant computational savings.
Before training, variation of the ENN output  $f_{\theta}(x,z)$  as a function of  $z$  reflects prior uncertainty in predictions. Since the base network does not depend on  $z$ , this variation must derive from the epinet. In our experiments, we induce this initial variation using prior networks (Osband et al., 2018). In particular, for  $\tilde{x} \coloneqq \mathrm{sg}[\phi_{\zeta}(x)]$ , our epinets take the form
$$
\underbrace {\sigma_ {\eta} (\tilde {x} , z)} _ {\text {e p i n e t}} = \underbrace {\sigma_ {\eta} ^ {L} (\tilde {x} , z)} _ {\text {l e a r n a b l e}} + \underbrace {\sigma^ {P} (\tilde {x} , z)} _ {\text {p r i o r n e t}}. \tag {5}
$$
The prior network  $\sigma^P$  represents prior uncertainty and has no trainable parameters. The learnable network  $\sigma_{\eta}^{L}$  is typically initialized to output values close to zero, but is then trained so that the resultant sum  $\sigma_{\eta}$  produces statistically plausible predictions for all probable values of  $z$ . Variations of a prediction  $\sigma_{\eta} = \sigma_{\eta}^{L} + \sigma^{P}$  at an input  $x$  as a function of  $z$  indicate predictive epistemic uncertainty, just like the example from Figure 3(b).
Epinet architectures can be designed to encode inductive biases that are appropriate for the application at hand. In this paper we focus on a particularly simple form of architecture
<sup>3</sup>The "stop gradient" notation  $\mathrm{sg}[\cdot]$  indicates the argument is treated as fixed when computing a gradient. For example,  $\nabla_{\theta}f_{\theta}(x,z) = [\nabla_{\zeta}\mu_{\zeta}(x),\nabla_{\eta}\sigma_{\eta}(\phi_{\zeta}(x),z)]$
for  $\sigma_{\eta}^{L}$  based around standard multi-layered perceptron (MLP) with Glorot initialization (Glorot and Bengio, 2010):