Epistemic Neural Networks
Authors: Ian Osband
Zheng Wen
Seyed Mohammad Asghari
Vikranth Dwaracherla
Morteza Ibrahimi
Xiuyuan Lu
Benjamin Van Roy
Venue: NeurIPS 2023 spotlight (2023)
Keywords: Uncertainty
Deep Learning
Neural Networks
TL;DR: Neural networks that know what they don't know, but are not necessarily Bayesian.

Abstract: Intelligence relies on an agent's knowledge of what it does not know.
This capability can be assessed based on the quality of joint predictions of labels across multiple inputs.
In principle, ensemble-based approaches can produce effective joint predictions, but the computational costs of large ensembles become prohibitive.
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 very large ensembles, consisting of hundreds or more particles, with orders of magnitude less computation.
The epinet does not fit the traditional framework of Bayesian neural networks.
To accommodate development of approaches beyond BNNs, such as the epinet, we introduce the epistemic neural network (ENN) as a general interface for models that produce joint predictions.

Epistemic Neural Networks
Ian Osband; Zheng Wen, Scyed Mohammad Asghari, Vikranth Dwaracherla, Morteza Ibrahimi, Xiuyuan Lu, and Benjamin Van Roy Google DeepMind, Efficient Agent Team, Mountain View
Alam.osband, m. ibrrairiii hogrmail com
sma.ggaaaenssssgaarr vikranthd 1xlu, benvanroyySgoogle.com
Abstract
Intelligent agents need to know what they don't know, and this capability be evaluated through the quality 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 supplement any conventional neural network, including large pretrained models, and can be trained with modest incremental computation to estimate uncertainty. With 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 newral network (ENN) as a general interface for models that generate joint predictions.
Introduction
Consider conventional neural network trained predict whether a random person would classify a drawing as a 'rabbit' or 'duck'. As illustrated in Figure given a single drawing, the network outputs a marginal prediction that assigns probabilities the probabilities are each 0.3, it remains unclear whether this is because labels sampled from random people are equally likely, whether the neural network would learn single class 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 joinl predictions (Wen et al. 2022).
the two classes.
The two tables to the right of Figure 1 represent. possible joint predictions that each consistent with the network's uniform marginal prediction. These joint predictions are over pairs of labels for the same image, (21) E { (J, For any such joint prediction, Bayes' rule defines a conditional prediction for yo given y1- 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 uncertaint;; conditioned on the first label, the prediction for the second label assigns all probability to the same outcome the first.
Figure presents the toy problem of predictions across two identical images 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). )
simple
Bayesian neural networks (BNNs) offer a statisticalll-prrinepleld way to make effective joint predictions, by maintaining an approximate posterior over the weights of a base neural network. Assymptotically these recover the exact posterior, but the computational costs
*Contact ian. oswanddgmail.ccm
37th Conference on Neural Information Processing Systems (NeurIPS 2023).

---

R
D
ambiguous image
0.25
0.25
0.25
conventional neural network
R
0.5
0.5
R
insufficient data
0.5
Figure 1: Conventional neural nets generate marginal predictions, which do not distinguish genuine ambiguity from insufficiency of data. Joint predictions can make this distinction.
Ensemble-based BNNs offer more practical approach by approxinating the posterior distribution with an ensemble of statistically plausible networks that we call particles (Osband and Van Roy 2015 Lakshmi, narayanan et al. 2017). While the quality of joint predictions improves with more particles, practical implementations often limited to at most. tens of particles due to computational constraints.
prohibitive for large models (Welling and Tch 2011
In this paper, 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: network architecture that can be added to any conventional neural network to estimate uncertainty. Figure 2 offers a preview of results presented in Section G where we compare these approaches ImageNet. The quality of the ResNet's marginal predictions measured by classification error or marginal log-loss does not change much if supplemented with epinet. However the epinet-enhanced ResNet dramatically improves the quality of joinl predictions, as measured by the joint log-loss, out oulperforming the ensemble of 100 particles, with total parameters less than 2 particles. Prior work has shown the importance of joint predictions driving effective decisions for a broad class of problems, including
combinatorial decision problems and sequential decision problems (Wen et al. 2022 et al. 2022a).
Osband
marginal log los-loss
Joint log
0.23
0.22
loss
ensemble
0.21
epinet
3e7
1e8
3e8
169
3e9
168
3e8
367 1eB
3e8
1e9
3G9
model (number of parameters)
Figure 2: Quality marginal and joint predictions across models on ImageNet (Section
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, introduce the concept of epistemic newral networks (ENNs). We establish that all BNNs are ENNs, but there are useful ENNs such the epinet, that are not BNNs.
2
Related work
Our research builds on the literature Bayesian deep learning (Hinonn and Van Camp
1993
Neal 2012). BNNs represents epistemic uncertainty via approximating the posterior distribution over parameters 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 Tch 2011).
Tractable methods for approximate inference has renewed interest BNNs. Variational
approaches such as Bayes by backprop (Blundell el 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 (Luuizos 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. 2017p. As of 2022, perhaps the most popular approach ensemble-based, with an ensemble of models, each referred to as a particle, trained in parallel so that they together approximate posterior distribution (Osband and Van Roy
2015 Lakshminarayanan et al. 2017).
Ensemble-basee 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 p. Others perform Bayesian inference more directly
in the lunction space in order to sidestep issucs relating 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.j 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 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. 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 load to significant improvements in performance.
2021 Wen et al. 2022 Osband al. 2022a). This
3 Epistemic neural networks
A conventional neural network is specified by parameterized function class f, which produces vector-valued output fo()) given parameters 0 and an input x. The output fo( assigns a corresponding probability P(y) = exp ((fo(=)),) Ey exp ((f(())))) to each class y. For shorthand, we write such class probabilities as P(y) (oftmax(fo( We refer a predictive class distribution P produced in this way as a marginal prediction, as it pertains to single input x.
An ENN architecture, on the other hand, specified by a pair: a parameterizod function class I and reference distribution Pz. The vector- valued output fa(x, 2) of an ENN depends additionally on an epistemic index z, which takes values in the support of Pz. Typical choices of the reference distribution Pz include uniform distribution over finite set or standard Gaussian over a vector space. The index 2 is used to express epistemic uncertainty. In particular, variation of the network output with 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 Zi,...,Zy, joint prediction assigns a probability (1:((1:7)) to each class combination ,f. While conventional neural networks are not designed to provide joint predictions, joint predictions can be produced by multiplying marginal predictions:
ANN (V1:))
II softmax (z:)),, t=1
(1)
However, this representation models each outcome Vi.r as independent and so fails to distinguish ambiguity from insufficiency of data. ENNs address this by enabling more
3

---

expressive joint predictions through integrating over epistemic indices:
PENN (v1=)) P2 P2((2) l
II softmax(foxxx.3))
(2)
This integration introduces dependencies so that joint predictions are not necessarily just the product of marginals. Figure provides simple example of how two different ENNs can use the epistemic index to distinguish the sorts of uncertainty described in Figure [
In Figure 3(a) the ENN makes marginal predictions that do not vary with and so the resultant joint predictions are simply the independent product 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 & produces diagonal matrix with 0.5 in each diagonal entry. As such, Figure shows how ENN can use the epistemic index to distinguish the two joint distributions of Figure
R
R
R
5
1|
ENN
ENN
1(2000)
0.25
0.25
[2)000)
ROOK NO0.
Index 8'0.3
(a) An ENN indicating an ambiguous image.
(b) ENN indicating insufficient data.
Figure 3: An ENN can incorporate the epistemic index Pz into its joint predictions. This allows an ENN to differentiale inevitable ambiguity from data insufliciency.
3.1 Evaluating ENN performance
Marginal log loss (also known 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 P and the label turns out to be y, then the sample log loss is the form - In P(v). We say that this is marginal loss because it only looks at the quality of a prediction over 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 E A and 7 labels V1.r to reward r(a, V1.7) [0, 1]. Given exact posterior predictive PE:,, consider as an objective maximization of the expected reward Σp.. (1: V1.r).
The optimal decision can be approximated based on an ENNN' prediction Pi.r by choosing an action that maximizes Σ yn.r (yr:)ra) y1.)). The following thcorem formalized and proved in Appendix B
Theorem 1. [informal] There exisls a decision problem and ENN that allains small expected marginal log loss such that actions generaled using the ENN perform no beller than random guessing.
Theorem indicates that minimizing marginal log loss does not suffice to support effective decisions. The key to this insight that marginal predictions do not distinguish ambiguity from insufficiency of data. However, this can can be addressed by instead considering the joint log loss. Given 7 data pairs and joint prediction P1., we can consider the joint log loss - In (y::() in exactly the same way that we looked at the marginal log loss. We formalize our next result. 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 highlight the importance of joint predictions driving decisions. Since want machine learning systems to drive effective decisions, will assess the performance

---

of ENNN by comparing their joint log loss 3 It is important to note that we will consider this joint loss as method for assessing quality of a trained ENN. We have not yet discussed how particular forms of ENNNs are trained, which will generally be up to the algorithm designer. Section | provides further detail on the specific architecture and training loss for the epinet ENN we develop in this paper.
3.2 ENNs versus BNNs
base neural network f defines a class of functions. Each element fo of this class is identified by vector I 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 0 and input x, produces an output fo(x). A typical BNN is specified by pair: a base network f and a parameterized sampling distribution p. Given parameters v, a sample i can be drawn from the distribution Pv to generate function fa. Approaches such as stochastic gradient MCMC, deep ensembles, and dropout can all be framed in this way. For example, with deep ensemble, i comprises parameters of an ensemble particle and Pv is the distribution represented by the ensemble, which assigns probability to finite set of vectors, each associated with one ensemble particle. For any inputs I1,...,Zy; Zr, by sampling many functions (fa= k = L.... K), a BNN can be used to approximate the corresponding joint distribution over labels 91..... Dr. according to P(y1:)) = L Σ 1(for (z1)..... .....ffo (x+)) = y1:7).
We say the BNN (f.P) is defined with respect to its base network f. We say an ENN Pz) is defined with respect to a base network f if, for any base network parameters 0, there exist ENN parameters 0' such that fo (-,=)= = fo almost surely with respect to Pz. 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 BNN (f.p) is expressed as an ENN (f'PP) if, for all v, T, and inputs Fi:r. there exists 0' such that for i ~ Pv. ~ Pai
(f((z)......f(
- (fo (21, (),. f. (x.: z)).
(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) P2) is expressed as a BNN if, for all 8', T, and inputs Z1,..., Tr. there exists a posterior distribution 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 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 base network 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 latent random object (de Finetti, 1929). BNNs use base network paramoters 0 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 9.: 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 nigh-dimensional inputs, the number of inputs Z1,.,, sampled uniformly from the input distribution, may have to be very large to distinguish ENNs in ways marginal log loss does not (Osann et al. 2022b). To sidestep prohibitive computational costs we use dyadic sampling in empirical evaluation and review these details in Appendix F
5

---

4 The epinet
This section introduces the epinet, which can supplement any conventional neural network to make new kind of ENN architecture. Our approach reuses standard deep learning components and training algorithms, as outlined in Algorithm [ 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 ENN training via SGD
Bac
Impule:
datasot
training examples D ((t.. 2.....)).... di).".
x
Bau rutwoees
Ouper (ok.
ENN
network f. reference Pa. initislization emaluats example (24) for index data sumples wa, Index sumples
loss
batch size
Fature
Prior network of
Splast
optimizer Returus:
update rule and number of iteraaions
Br
paramette estimates for the ENN.
Lanshic
2: sample dala 2 =iiodon ~ Unil(((...N)))
1: In O....TT do
Referace Z-P,
Pacarctio
3: sample indices 2 Eay
4: compute grad <- Vajaus, Lx Liz L.2 s.....
5: update B.41 optinizor gradd
Figure 4: Epinet network architecture.
4.1 Architecture
Consider a conventional neural network as the base network. Given base parameters Ç and an input x, the output is pę(x). For a classification model, the class probabilities would be (fftmax(((((z))) 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 (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 ResNet. The epistemic index is sampled from a standard Gaussian distribution in dimension Dz- For epinct parameters n, the epinct outputs o, (,)))(z)) 2). To produce ENN, the output of the epinct is added to that of the base network, though with a "stop gradient" "
f0( = (c(x) + (x)))
(4)
ENN
net
epinet
We find that with this stop gradient, training dynamics more reliably produce models that perform well out of sample. The ENN parameters 0 = (C, 17) include those of the base network S and epinet 7y. 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 fo(x, 2) as a function of reffects prior uncertainty in predictions. Since the base network does not depend on 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 ă. == sglø (x)], our epinets take the form
Tr (5, = of (i, 2) + 22
epinet
learnable
prior not
The prior network oP represents prior uncertainty and has no trainable parameters. The learnable network of is typically initialized to output values close to zero, but is then trained so that the resultant sum o, produces statistically plausible predictions for all probable values of Variations of prediction o,, = of + of at an input x as a function of = 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 particularly simple form of architecture
'The "stop gradient" notation sg(:] indicates the argument is treated as fixed when computing gradient. For example, 2 (x))(x), (n)r( 2)).

---

for of based around standard lulti-layered perceptron (MLP) with Glorot initialization
GGlorot and Bengio 2010
o (2, == m,, (() ((, RR
(6)
where mlp,, is an MLP with outputs in RDxxC and [ is flattened concatenation of B and 2. Depending on the choice of hidden units, of can represent highly nonlincar functions in ã, and thus allow for expressive joint predictions in high level features. The design, initialization and scaling the prior network o" allows an algorithm designer to encode prior beliels, and is essential for good performance in learning tasks. Typical choices might include of sampled from the same architecture as of but with different parameters.
4.2 Training loss function
While the epinet's novelty primarily lics in its architecture, the choice of loss function for training can also play role in its performance. In many classification problems, the standard regularized log loss suffices:
x EXENT (0, 2, F. 2., == In (softmax(fo(r;, z)) 1|10112
(7)
Here, À is a regularization penalty hyperparameter, while other notation are as defined in Sections 3 and 4.1 This is the loss function we use for the experiments with the Neural Testbed (Section and ImageNet (Section 6
Image classification benchmarks often exhibit a very high signal-to-noise ratio (SNR); identical images are almost always assigned the same label. As demonstrated by Dwaracherla et al. 2022 when the SNR is not so high and varies significantly across inputs, it can be beneficial to randomly perturb the loss function via versions of the statistical bootstrap (Efron and Tibshirani 1994). For example, a Bernoulli bootstrap omits each data pair with probability p E [0, IJ, giving rise to perturbed loss function:
Xx EXENT (O,NT 20, if 1<'=>0>>10 otherwise,
XXNT PXNNT(0, z, Ii, Vi, Vi., P.d
(8)
where each & is an independent random vector sampled uniformly from the unit sphere, and (()) is the cumulative distribution function of the standard normal distribution (O,, 1).
Note that the loss functions defined in equation and g only explicitly state the loss for single input-label pair (z,, 31) and a single epistemic index z. To compute a stochastic gradient, one needs to sample a batch of input-label pairs and a batch of epistemic indices, average the losses defined above, and then compute the gradient.
4.3
How can this work?
The epinet designed to produce effective joint predictions. As such, one might expect the training loss to explicitly reflect this and be surprised that we standard marginal loss functions such as & EXENT, Recall that a prediction fo =) produced by an epinet is given by a trainable component (x(xx) + (ā, perturbed by the prior function (ā, z), which has no trainable parameters. Minimizing PXENT can therefore be viewed as optimizing the learnable component (x) + to ()) with a perturbed loss function. Previous work has established that learning with prior functions can induce effective joint predictions (Osband
et al. 2018 He et al. 2020 Dwaracherla et al. 2020 2022). we extend this to the epinet.
To show this marginal loss function can lead effective joint predictions, Thcorem proves that this epinet training procedure can mimic exact Bayesian linear regression. Although this paper focuses on classification, in this subsection we consider a regression problem because its analytical tractability facilitates understanding. To establish this result, we introduce regularized squared loss perturbed by Gaussian bootstrapping:
(LSG (0, z, Fi, ye, 1) = (fe,( 2) - oc 2)² + x (IK<122 + llvl
(9)
Here, each 6 is a context vector sampled uniformly from the unit sphere in Dz dimensions, in the same manner as XENT

---

We say that a dataset D is generated by linear- Gaussian model if 9v. (a) = V2 x, Ve ~
(0, of,, and each data pair (Iz, ys) E 'D satisfies y: = gv. (x;) + & where E1NN are ii.d. according to o²). We say an ENN is a Einear-Gaussian epinet with parameters 0 = ((.7) its trainable component comprised of a linearly parameterized functions _ę((r) = C²
and o 2) = 2T næ with epinct input = 2; the prior function takes the form o² (z. 2) = Poz, where each column of the matrix Po is independently sampled uniformly from the unit sphere, and the reference distribution is taken to be P. N(O,TDz!
Theorem 4. Let dala D D= ((r,,,)))) be generated by a linear- Gaussian model and f be linear- Garasiann epinet. Let 0 Erggming Ein (ll_2(dd))) (0, (0,2,0i,Bi.3) with parameter \=o2)(Noo). Then, conditioned on (P)) fa(:2) converges in distribution to gu. as Dz grows, almost surely.
This result, proved in Appendix D serves as 'sanity check' that an epinet trained with a standard loss function can approach optimal joint predictions as the epistemic index dimension grows. Although our analysis is limited to Auear-Gaassian models, the loss function applies more broadly. Indeed, we next demonstrate that epinets and standard loss functions scale effectively to large and complex models.
Table 1: Summary of benchmark agents, full details in Appendix G
agent
description
hyperparameters
nlp
vanilla MLP
decay
ensemble
deep ensembles Lakshminaravauan al. 2017
La decay, eusemble size
dropout
Dropout [Gal and Ghabramam 2010
Lg decay, network, dropout rate
bbb
Bayes by baceprop (Eimuree et a. 2015
prior mixture, network, early stopping
hypermodel
hypermodel [Dwarscherla et al. 2020
La decay, prior, bootstrup, index dimension
ensemble*
ensemble + prior functions (Oshand r al. 2018
bg decay, ensemble sixe, prior scale, bootstrap learning rate, prior, momentium La decay, network, prior, index dimension
agnene
stochastic gradient MCMC (Weiing and Ieh 2011
epinet
MLP MLP epinet (this paper)
5 The neural testbed
The Newral Testbed is open-source benchmark that evaluates the quality of joint predictions in classification problems using synthetic data produced by neural-network based generative models (Osband et al. 2022a). We use this unit test to sanity-check learning algorithms in a controlled environment and compare the epinet against benchmark approaches.
Table | lists the agents that we study well as hyperparameters that tune via grid search. For our epinet agent, we use base network pe (x) that matches the baseline mlp agent. We take the features ((z) to be concatenation of the input x and the last hidden layer of the base network. We initialize the learnable epinet of according to (6) with hidden layers of 15 hidden units. The prior network oº is initialized as an ensemble of Dz 8 networks each with 2 hidden layers of 5 hidden units in each layer, and combine the output by dot-prood dot-product with index We push the details, together with open source code, to Appendix G
Figure examines the trade-offs between statistical loss and computational cost for epinet against benchmark agents. The bars indicate standard error, estimated over the testbed's random scods. After tuning, all agents perform similarly in marginal prediction, but the epinet is able to provide better joint predictions at lower computational cost. We first compare the performance of the epinet (blue) against that of an ensemble (red) as we grow the number of particles. We show the epinet is able to perform much better than an ensemble with 100 particles, with a model less than twice the size of a single particle. These results are still compelling when we compare against ensemblet, which includes prior functions, and which is necessary for good performance in low data regimes included in testbed evaluation. Included in this plot are the other agents bbb, dropout and hypermodel agents. The dashed line indicates performance of sgmemc agent, which can asymptotically obtain the Bayes optimal solution, but at much higher computational cost.
6 ImageNet
Benefits of the epinet scale up to more complex datasets and, in fact, become more substantial. This section focuses on experiments involving the ImageNet dataset (Deng

---

classification error
marginal log less
joint log less
enn
0.24
0.16
ensemble
ensemble+
opinet
0.22
0.14
bab
dropouu
0.13
hypermodel
0.20
104
1o5
1e6
104
105
1e6
164
1e5
1o6
model size (number of parameters)
Figure 5: Quality of marginal and joint prodictions across models on the Neural Testbed.
el. al. 2009 ; qualitatively similar results for both CIFAR-10 and CIFAR-100 are presented
in Appendix H We compare our epinet agent against ensemble approaches as well as the
uncertainty baselines of Nado al. (2021). Even after tuning optimize joint log-loss, none of these agents match epinet performance. We assess joint log-loss via dyadic sampling
(Osband al., 2022b), as explained in Appendix F
For our experiments, we first train several bascline ResNet architectures on ImageNet. We train each of the ResNet-L architectures for L E (50, 101, 152, 200} in the Jaxline framework (Babuschkin et al. 2020). We lune the learning rate, weight decay and temperature rescaling Wenzel al. 2020) Nes-et-50 and apply those settings other ResNets. The ensemble agent only uses the ResNet-50 architecture. A After tuning hyperparameters, we independently initialize and train 100 ResNet-50 models serve ensemble particles. These models
then used to form ensembles of sizes 1, 3, 10, 30, and 100.
The epinet takes pretrained ResNet as the base network with frozen weights. We fix the index dimension Dz = 30 and let the features be the last hidden layer of the RosNet. We use 1-layer MLP with 50 hidden units for the learnable network (6) The fixed prior oº consists of a network with the same architecture and initialization as o together with an ensemble of small random convolutional networks that directly take the image inputs. We push details on hyperparameters and evaluation, with open-source code, to Appendix
classification error
marginal loss
Joint log loss-lss
enn
resnet
0.90
ensemble
loss
0.22
spinat
0.85
mimo
dropout
sngp
het
3e8
3d9
367
169
3eg
168
169
model (number of parameters)
Figure 6: Marginal and joint predictions on ImageNet.
Figure 2 presents our key result of this paper: relative to large ensembles, epinets greatly improve joint predictions orders of magnitude lower compute cost. The figure plots performance of three agents with respect to three notions of loss as a function of model size in the spirit of Kaplan et al. (2020). The first two plots assess marginal prediction performance in terms of classification error and log-loss. Performance of the ResNet scales similarly whether or not supplemented with the epinet. Performance of the ensemble does not scale as well as that of the ResNet. The third plot pertains to performance of joint predictions and exhibits dramatic bonefits afforded by the epinct. While joint log-loss incurred by the ensemble agent improves with model size more than the ResNet, the epinet agent outperforms both alternatives by an enormous margin.
Figure 6 adds evaluation for the best single-model agents from uncertainty baselines (Nado al.) 2021): epinet also outperforms all of these methods. We tuned each uncertainty baseline agent to minimize joint log-loss, subject to the constraint

---

that their marginal log-loss does not degrade relative to published numbers. We can see that the epinet offers substantial improvements in joint prediction compared to sngp (Liu
et al. 2020), dropout (Gal and Ghahramani 20166) mimo (Havasi et al. 2020) and het Coolier et al. 2020 As demonstrated in Appendix H these qualitative observations remain
unchanged under alternative measures of computational cost, such as FLOOPs
7
Conclusion
This paper introduces ENNs as a new interface for uncertainty modeling in deep learning. We do this to facilitate the design of new approaches and the evaluation of joint predictions. Unlike BNNs, which focus on inferring unknown network parameters, ENNs focus on uncertainty that matters in the predictionn. The epinet novel ENN architecture that cannot be expressed as a BNN and significantly improves the tradeoffs in prediction quality and computation. For large models, the epinet enables joint predictions that outperform ensembles consisting of hundreds of particles at computational cost only slightly more than one particle. Importantly, you can add an epinet to large pretrained models with modest incremental computation.
ENNs enable agents to know what they do not know, thereby unlocking intelligent decision making capabilities. Specifically, ENNs allow agents to employ sophisticated exploration schemes, such as information-directer sampling (Russo and Van Roy 2014). when tackling complex online learning and reinforcement learning problems. A recent paper (Osband et al. 2023) has demonstrated efficacy of the epinet in several benchmark bandit and reinforcement learning problems. While ENNs can also be integrated into large language models, we defer this to future work.
References
Babuschkin, I., Baumli, K., Bell, A., Bhupatiraju, S., Bruce, J., Buchlovsky, P., Budden, D., Cai, T., Clark, A., Danihelka, L,, Fantacci, C., Godwin, J., Jones, C., Hennigan, T., Hessel, M., Kapturowski, S., Keck, T., Kemaev, L., King, M., Martens, L., Mikulik, V., Norman, T., Quan, J., Papamakarios, G., Ring, R., Ruiz, F., Sanchez, A., Schneider, R., Sezener, E., Spencer, S., Srinivasan, S., Stokowiec, W., and Viola, F. (2020). The DeepMind JAX Ecosystem.
Blundell, C., Cornebise, J., Kavukcuoglu, K., and Wierstra, D. (2015). Weight uncertainty in neural networks. International Conference on Machine Learning, pages 1613 1622. PMLR.
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G.,
Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q. (2018). JAX: composable transformations of fython-Numil programs.
Charpentier, B., Zügner, D., and Günnnemann, S. (2020). Posterior network: Uncertainty estimation without ood samples via density- based pseudo-counts. Advances in Newral Information Processing Systems, 33:1356-1367,
Collier, M., Mustafa, B., Kokiopoulou, E., Jenatton, R., and Berent, J. (2020). simple probabilistic method for deep classification under input-dependent label noise. arXiv preprint ar xixi:2000 06778.
Collier, M., Mustafa, B., Kokiopoulou, E., Jenatton, R., and Berent, J. (2021). Correlated input- dependent label noise in large-scale image classification. In Proceedings of the IEEE/CV Conference Computer Vision and Pattern Recognition, pages 1551-1560.
de Finetti, B. (1929). Funzione caratteristica di un fenomeno aleatorio. Atti del Congresso Internazionale dei Mulematici: Bologna, (1):1979 190.
Deng, J., Dong, W., Socher, R., Li, L.-J., Li, K., and Fei-Fei, L. (2009). Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 248-255, 255. lece.
Der Kiureghian, A. and Ditlevsen, O. (2009). Aleatory or epistemic? does it malter? Structural safety, 125
Dwaracherla, V., Lu, X., Ibrahimi, M., Osband, 1., Wen, Z., and Van Roy, B. (2020). Hypermodels for exploration. In International Conference on Learning Representations.
10

---

Dwaracherla, V., Wen, Z., Osband, I,, Lu, X., Asghari, S. M., and Van Roy, B. (2022). Ensem- bles for uncertainly estimation: Benefits of prior functions and bootstrapping. Xiv preprinl 2ix: 03633.
Efron, B. and Tibshirani, R. J. (1994). An introduction to the bootstrap. CRC press.
Gal, Y. and Chahramani, Z. (2016). Dropout as Bayesian approximation: Representing model uncertainty in deep learning. In International Conference Machine Learning.
Gal, Y., Islam, R., and Ghahramani, 7. (2017). Deep Bayesian active learning with image data. In International Conference on Machine Learning, pages 1183 1192. PMLR.
Glorot, X. and Bengio, Y. (2010). Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the 13th international conference artificial intelligence and statistics, pages 219 256.
Havasi, M., Jenatton, R., Fort, S., Liu, Z., Snoek, J., Lakshminarayanan, B., Dai, A. M., and Tran, D. (2020). Training independent subnetworks for robust. prediction. CoRR, abs/2010.00611
He, B., Takshminarayanan, B., and Teh, Y. W. (2020). Bayesian deep ensembles via the neural tangent kernel. Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H., editors, Ad- vances in Neural Information Processing Systems, volume 33, pages Inc.
Curran Associates,
Hendrycks, D. and Dietterich, T. (2019). Benchmarking neural network robustness to common corruptions and perturbations.
Hinton, G. E. and Van Camp, D. (1993). Keeping the neural networks simple by minimizing the description length of the weights. In Proceedings of the sirth annowal conference on Comyulalional learning theory, pages 13.
Hron, J., Matthews, A. G. d. G., and Ghahramani, Z. (2017). Variational Gaussian dropout is not Bayesian. ariiv preprint S.S::1711.0298.
Izmailov, P., Vikram, S., Hoffman, M. D., and Wilson, A. G. (2021). What are Bayesian neural network posteriors really like? arx preprint Xic:2104.1442.1
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. (2020). Scaling laws for neural language models. arXiv preprint arXio:2001. 08361.
computer vision? In Advances in Neural Information Processing Systems, volume 30. Kendall, A. and Gal, Y. (2017). What uncertainties do we need in Bayesian deep learning for
Knothe, H. (1957). Contributions the theory of convex bodies. Michigan Mathematical Journal, 1(1):39 52.
Krizhevsky, A. (2009). Learning multiple layers of features from tiny images. Technical Report University of Toronto, Toronto, Ontario.
Lakshminarayanan, B., Pritzel, A., and Blundell, C. (2017). Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in Neural Information Processing Systems, pages 6405- 6416.
Liu, J., Lin, Z., Padhy, S., Tran, D., Bedrax Weiss, T., and Lakshminarayanan, B. (2020). Simple and principled uncertainty estimation with deterministic deep learning via distance awareness. Advances in Neurul Information Processing Systems, 33:7418 7512.
Louizos, C. and Welling, M. (2017). Multiplicative normalizing flows for variational Bayesian neural networks. In International Conference on Machine Learning, pages 2218 2227. PMLR.
Lu, X. and Van Roy, B. (2017). Ensemble sampling. In Advonces in Neural Information Processing Systems, pages 3260 3268.
MacKay, D. J. (1992). A practical Bayesian framework for backpropagation networks. Neural computation, 4(3):448 472.
Malinin, A. and Gales, M. (2018). Predictive uncertainty estimation via prior networks. Advonces in neurol information processing systems, 31.
11

---

Minka, T. (2000). Bayesian linear regression. Technical report, Cileseer.
Nado, Z., Band, N., Collier, M., Djolonga, J., Dusenberry, M., Farquhar, S., Filos, A., Havasi, M., Jenatton, R., Jerfel, G., Liu, J., Mariet, Z., Nixon, J., Padlyy, S., Ren, J., Rudner, T., Wen, Y., Wenzel, F., Murphy, K., Sculley, D., Takshminarayanan, B., Snoek, J., Gal, Y., and Tran, D. (2021). Uneertainty Baselines: Benchmarks for uncertainty & robustness deep learning. arviv preprint arXtx2106600011:
Neal, R. M. (2012). Bayesion learning for neural networks, volume 118. Springer Science & Business Media.
Osband, T. (2016). Risk versus uncertainty in deep learning: Bayes, bootstrap and the dangers of dropout. In NIPS Workshop Bayesian Deep Learning, volume 192.
Osband, I,, Aslanides, J., and Cassirer, A. (2018). Randomized prior functions for deep reinforcement learning. In Bengio, S., Wallach, H., Larochelle, H., Grauman, K., Cesa-Biancha N., and Garnett, R., editors, Advonces Neural Information Processing Systems 91, pages 8617-8629. Curran Associates, Inc.
Osband, I. and Van Roy, B. (2015). Bootstrapped Thompson sampling and deep exploration. arXiv preprind arxi::1507 1.00000.
Osband, L., Van Roy, B., Russo, D. J., and Wen, Z. (2019). Deep exploration via randomized value functions. Journal of Machine Learning Research, 20(124):1-60 62.
Osband, I., Wen, Z., Asghari, S. M., Dwaracherla, V., Hao, B., Ibrahimi, M., Lawson, D., Lu, X., O'Donoguue, B., and Van Roy, B. (2022a). The neural testbed: Evaluating joint predictions. In Advances in Neural Information Processing Systems, volume 35. Curran Associates, Inc.
Oshand, I,, Wen, 7., Asghari, S. M., Dwaracherla, V., Thrahimi, M., Lu, X., and Van Roy, B. (2023). Approximate Thompson sampling via epistemic neural networks. In Evans, R. J. and Shpitser, I. editors, Proceedings of the Thirty-Ninth Conference on Uncertainty in Artificial Intelligence, volume 216 of Proceedings of Machine Learning Research, pages 1586 1595. PMLR.
Osband, I., Wen, Z., Asghari, S. M., Dwaracherla, V., Lu, X., and Van Roy, B. (2022b). Evaluating high-order predictive distributions in deep learning. In The 38th Conference on Uncertointy in Artificial Intelligence.
Ovadia, Y., Fertig, E., Ren, J., Nado, Z., Sculley, D., Nowozin, S., Dillon, J. V., Lakshminarayanan, B., and Snoek, J. (2019). Can you trust your model's uncertainty? evaluating predictive uncertainty under dataset shift. arXie preprint 1556:990000538
Rosenblatt, M. (1952). Remarks on multivariate transformation. The Annals of Mathematical Statistics, 23(3):470-477
Russo, D. and Van Roy, B. (2014). Learning to optimize via informatlion-directedd sampling. Advances in Neural Information Processing Systems, 27.
Srivastava, N., Hinton, G., Krizhevsky, A., Sutskever, L., and Salakhutdinov, R. (2014). Dropout: A simple way to prevent neural networks from overfitting. The Journal of Machine Learning Research, 15(1) (((1):1929199))
Sun, S., Zhang, G., Shi, J., and Grosse, R. (2019). Functional variational Bayesian neural networks. Xie preprint K6o:1903.055
van Amersfoort, J., Smith, T.., Jesson, A., Key, O., and Gal, Y. (2021). On feature collapse and deep kernel learning for single forward pass uncertainty. arXiv preprint arXim:2102.1 11409.
Wang, C., Sun, S., and Grosse, R. (2021). Beyond marginal uncertainty: How accurately can Bayesian regression models estimate posterior predictive correlations? In International Conference on Artificial Intelligence and Statistics, pages 2476-2484. PMLR
Welling, M. and Teh, Y. W. (2011). Bayesian learning via stochastic gradient Langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-11), pages 681-688. Citeseer.
Wen, Z., Osband, 1., Qin, C., Lu, X., Ibrahimi, M., Dwaracherla, V., Asghari, M., and Van Roy, B. (2022). From predictions to decisions: The importance of joint predictive distributions.
12

---

Wenzel, F., Roth, K., Veeling, B. S., Swiatkowski, J., Tran, L., Mandt, S., Snoek, J., Salimans, T., Jenatton, R., and Nowozin, S. (2020). How good is the Bayes posterior in deep neural networks really? Xrxi preprint ur Xrv:20002 02405.
Wilson, A. G. (2020). The case for Bayesian deep learning. arXiv preprint arXiv:2001.1099:
13

---

A
Open source code
Two related github repositories complement this paper:
enn: htte:://aaanpaaa...pprr..aaaaaar
neural_teatbed: htttpp///aaooyymuus.ooooo.oo..c
testbed-8961
These librarics contain the code neccssary to reproduce the key results our paper, divided into repositories based focus. Together with each repository, include several 'tutorial colabs' Jupyter notebooks that can be run in a browser without requiring any local installation. Each of these libraries is written in Python, and relies heavily JAX for scientific computing (Bradbury et al. 2018). We view this open-source effort as a major contribution of our paper.
The first library, enn, focuses on the design of epistemic neural networks and their training, This includes all of our network definitions and loss functions. Our library is built around Haiku (Babuschkin et al. 2020). The library provides the basis for all of the computational work reported in this paper.
The second library, testbbd, was introduced as part of The Neural Testbed (Osband et al. 2022a). We add the epinet agent, suitable for comparison with the existing agent implementations in that library.
B From predictions to decisions
Suppose a neural network has been trained on a dataset D and, then, for a random input I with label y generates a distributional prediction P. The expected marginal log loss is -I PD)) Note that this expectation is over x, y, and P. Results we establish point out that, while minimizing this produces optimal marginal predictions, that does not generally lead to performant downstream decisions.
To formulate generic decision problem, consider a reward function r that maps an action a E A and t labels yırr to a reward (1:)) E [0, 1]. Conditioned on training data D and random inputs Firr, an action generates expected reward
E [r(a, (i:r = E Pir (Mrr(
where (1=r())=7) = P(y = 1D, ));)) is the posterior predictive distribution. An optimal action a, can be determined by maximizing this conditional expectation. Tf an approximation Per is used instcad of Pı. then the objective becomes Σn,. Pi.r (f1:>)r(a y1.7)
Let P. P(y = (D, Epree-i P1: (31:7) be the marginal posterior predictive. Let P E PI.r (y1-7) be the approximate marginal. A joint prediction P1., minimizes marginal cross-entropy when P₂ P. for 1, I,..., The following result establishes that such predictions that minimize marginal log loss can result in decisions no better than random guesses. Note that when is clear from context, we use notation for a set, like A, to express cardinality of the set. The exchangeability requirement that the distribution of any set of data pairs be exchangeable ensures that we are in standard supervised learning setting.
Theorem [formal] For all T, there erists an action set A with [Al 2T, an ex- changeable data distribution, and a reward function with range [0, 1] such that if Pi.r (11..) I. PDyllD, x1.7) and à unif(arg max E Pr. (r(a, v1.))) then and Elr(â, vir ID, S E(r(ū, w D, a1rl < 1,
max Elr(a, (D,
where, conditioned on D and Tirr, ā ~ unif(A).
Proof. Let r be even; extending our argument to odd 7 is straightforward. Without loss of generality, let the input space be a singleton. Hence, P( -ID, P( ID).
14

---

Let labels be be in {-1,1} so that V1 E 1-1,11 Let A Let r(a, y1:7) = 1- (h Let P =... = y-1D) and, for all t, -1y-=11 1/2. Clearly, this data distribution is exchangeable, with uniform marginals. Hence, for any E A,
Σ D(a = (U.(( 1 V1.r
y017
Σ (U1:T)
Σ 1=1
y1
Mivv
=1 72
Pa(v))ei
1/T.
follows that any action maximizes En... Per (M (r(a, yir)), and therefore, selecting randomly from this set results in expected reward E[[(a, v1 ) x1r] = 1/T. Hence,
Bir(â, v1 ID=))D, =B(r(a, v 12, 21:7l
E
TzZit
D. x 1ir
27
=1 E
=1 1
E [(aa]]
=1 1/T.
For any action a such that Er= at = 0,
(1x x1.1 =P( -iDD)
1
-
1
Σ
a4
7
an))
P(y i|D)
1-
Σ
t=
=1.
The result follows.
The first displayed equation this theorem asserts that the optimal expected reward is one, while the second asserts that an action selected based on Piar, which minimizes marginal log loss, does no better than action chosen uniformly at random, which earns expected reward less than one.
While optimizing marginal predictions does not necessarily lead to performant downstream decisions, minimizing joint log loss does. Our next result formalizes this by bounding performance shortfall by function of the KL-divergence:
dRL. (PI-1I Σ (1:((1r:)) In (1s((11)))
Σ Pir(()))) PI. Vior
Only the final term depends on the prediction P. Its conditional expectation - EfIn P( 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
15

---

Theorem 2. [formal] For all data distributions, reward functions with range [0, 1], and actions à E arg m E Piir (ya
v1.r x1rT] max B1 81. (1,) - 2E|dKL (P.:rl111 x1rT).
The KL-divergence is the difference between the log loss attained by Par and that attained by an optimal prediction Pu.. The shortfall action hence bounded by a mcasure of joint prediction error. this error is small, Pi: leads to good decisions regardless the reward function.
C ENNs versus BNNs
In this appendix, we provide a proof of Theorem t
Theorem 3. For all base networks f, any BNN defined with respect to f can be expressed as an ENN defined with respect f. However, there erists base network and ENN defined BNN defined with respect to f.
with respect to that can not be expressed
Proof. Consider a BNN (S.p). Without loss of generality, take the parameter space of to be IRd, Let Pz be an absolutely continuous reference distribution over R Via Knothe- Rosenblatt rearrangement (Knothe 1957 Rosenblatt 1952), for each v E IR there exists transport map from Pz to Pv- For each epistemic index = E Rd and BNN parameter vector v, let dz.= be the corresponding base network parameters generated by this transport map. Let 2x) = fō.. (z)). It easy see that the ENN (f', is defined with respect to and expresses the BNN (f.p)-
For the second part of the theorem, consider linear base network fo(Ø) 0TT,, standard Gaussian epistemic index reference distribution, and an ENN
(x, 2) fo(z) + 0"2,
where 8' = (0,8"). Clearly, there are ENN parameters 8' and epistemic index such that no i satisfies for (zx,)) fz(x). If follows that no BNN defined with respect to f can express the ENN.
D Bayesian linear regression
In this appendix, we provide proof of Theorem 4. First, we state without proof standard result on Bayesian linear regression [Minka| 2000). Lemma 1. Let D= ( be generated by linear- Gaussion model. Conditioned on D, v, is Gaussian and
N
N
1 0
Ti
Tcyi
Cov[u,|D]=
+
02
ob
1=1
1=1
(10)
Next, we prove result on the ncar-orthogonality of vectors sampled uniformly from a unit sphere. Note that a.s. abbreviates almost surely. Lemma 2. Let by and en be independent vectors sampled uniformly from the n-dimensional unit sphere. Then,
lim D Go O. n-+00
Proof. Note that unit random vectors generated by sampling normal vectors from (0, and normalizing are uniformly distributed over the Rd unit sphere.
Let an and Bn be independent samples from a chi-squared distribution with n degrees of freedom. Let un an
1/2 b, and en & Cn. Clearly the distribution of both un and U, is
16

---

isotropic, and since an and 3n are distributed chi-squared with degrees of freedom (which that of the length of n-dimensional standard normal), un and 2n are independent standard normal random vectors. Further, b /n/l|uu l2 and un fn/l/en/12 For any n,
Σ
0 Gn
i=1 Ei=1
101
Since {ui]i=1 and (wi)2_1 are i.id N(0, 1), by the strong law of large numbers,
Σ 4.04 a.
lim
E Eluco]] 0,
n 20000
n
-i=1 up
lim
u.s.
Elu 1,
12+00
11
Σ i== Elu =
lim
Hence, by the continuous mapping theorem,
lim T cn n
4.8. 0.
Theorem 4. Let data D = {(1; Vs, 1) be generated by a linear- Ganssian model and f be linear-Gausiaan
linear- Gaussion epinet. Let Ô E arg ming E J. Pz (2z)) (0, 2, Ir, v1, with parometer
o² /(Not). Then, conditioned on (D, Po), fo(-.2) converges in distribution to gu. as Dz grows, almost surely.
Proof. The statement that, conditioned on (D,cN, Po), fz(;;2 converges in distribution to Iv, as Dz grows can be restated as
P (gv. E FID) a.e. lim
P(fa(,)) FID,cI,N. CINN Po),
(11)
D2
for all measurable sets F of functions mapping R R.
Note that
x² (1<111 + Illll))
i = (G, n) E arg min
Σ (2( y o00 G.7
Since the optimization problem is strictly convex in 0, is the unique minimizer of this expression. Some simple algebra establishes that
1
1 -
Ç Σ
Tilli i=1
and i
Ca
+
Po
E
Po.
(12)
0o
i=1
where b=
+ o 1
1=1
For any x E R and z E R22, (z ( (c+2 =T
Po)) x and Jv. (x) = =V_ Hence, holds if and only if ? + 2 TT( + Po) converges in distribution to conditioned on (c Po), almost surely. Since u. and Č & z (n Po) are Gaussian, it is sufficient to
11
show that
lim Č= E(u,ID] and D2-000
lim (n + Po) (m + Po) = Cov(u,1D). D2-000
17

---

Based on 10 and 12
E(v, [D] for any Dz. Hence, in order to prove the theorem, it is
sufficient show
lim (1 + Po) (1) + Po) = Cov(u.ID).
(13)
Dx
For any Dz,
N
T
N
1
(n+
(1)
Po)= $
C
Po
C.
-
Po E
00
1=1
1=1
1
N
1
1
(Po Po 11 11)
=
Σ
Σ+ 1
0
T
i=1
1
Σ
[ Po + Po exx,) E
000
Recall that {c)) and columns of Po are all sampled iid and uniformly from unit sphere in R22 by Lemma
lim D20
c c a. Ovi
+ j.
lim
PO as. Ovi
D2-100
lim
Po Po Po - I.
D2-x
Hence,
1 og 0
lim (n+ Po) (n - &
Σ
x
=.
D2
E Didactic examples
To offer some intuition for how epinets work and what they accomplish, we present simple example specialized to linear base model. A linear base model produces output ue(x)= = <"xx x given an input x and model parameters Ç. It is natural to add to this a linear epinet (x)) =) I.... The combined architecture is equivalent to a linear hypermodel (wwraccerla et al. 2020 To see this, note that,
fo(x, 2) Me (z) (,(()(()) = + =/)
(14)
As such, properties of linear hypermodels, such as their ability to implement exact Bayesian linear regression, carry over to such epinets.
Figures 7 illustraaee predictive uncertainty estimates produced by linear epinets. In this figure, posterior credible intervals of an epinet are compared against exact Bayesian inference. Data 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 regression (Lu and Van Roy 2017 Dwaracherla et al. 2020
E(0, yi, 2) = (2,, 2) - ys + oc 2)3.
Here c; is random signature drawn from the unit sphere in Dz generated independently for each training example (I;, Wi) and a is the scale of the additive bootstrap noisc. 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 generated by standard logistic regression model with parameters drawn from Gaussian prior. Figure presents standard deviations of marginal predictive distributions across the input space. We supplement standard logistic regression model with linear epinet. The
18

---

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.
std
0.05.0.3
0.15w0.2
R
D-
02:00:2
-2-
label
2
11
1
Figure 7: Epinet predictions Gaussian linear regression.
Figure 8: Epinct predictions in Logistic regression.
F Dyadic sampling
To evaluate the quality of joint predictions, we sample batches of inputs (di,...,30 and assess log-loss with respect to corresponding labels (y1,..., UU). With high-dimensional input space, labels of inputs sampled uniformly at random are typically nearly independent. Hence, large batch size is required to distinguish joint predictions that effectively reflect interdependencies among labels. However, this becomes impractical because computational requirements for testing grow exponcntially in Dyadic sampling serves heuristic that samples inputs more strategically so that effective agents are distinguished with manageable batch size (Osband et al. 2022b) even when inputs are high-dimensional.
practical
F.1
Basic version
The basic version of dyadic sampling, for each batch, first samples two independent random anchor points, ū and ã2, from the input distribution. Then, to form a batch of size sample points independently and with equal probability from these two anchor points {21, To assess an agent, its joint prediction of labels is evaluated for this batch of size т. Even with a moderate value of 10, batch produced by dyadic sampling gives rise to labels that are likely to correlate in particular, labels assigned to the same anchor point. Osband al. (2022b) demonstrate that, across many problems, this sampling hcuristic is effective in distinguishing the quality of joint predictions with modest computation.
On the Neural Testbed, the imput distribution is standard normal. Thus, for each test batch, we sample anchor points 21,02 ~ (0, and then re-sample 7 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, split the evaluation set into batches of size 2. We then iterate over these batches of size 2, re-sample 7 10 points from each input pair, and evaluate the log- of an agent's 's joint predictions these batches of size 10. Finally, we take the average of all the joint log-losses.
Figure ((a) shows a few examples of these dyadic input batches of size T 10. It may seem insatislactory that images are repeated exactly within a batch. Even though we view this metric as unit test. that intelligent agent pass, it is conceivable to design a 'cheating" agent that takes advantage this repeating structure. the following section, we will consider more robust version of dyadic sampling, where instead of repeating images exactly,
19

---

batth
baich
batch
(a) Examples of input batches generated by basic dyadic sampling. Each row is batch of size 10, which the joint log-loss evaluated.
baton
baten
(b) Examples imput batches generated using augmented dyadic sampling. Compared to Figure (a) each image is randomly cropped and flipped, introducing more diversity to the batch.
Figure 9: Examples of input batches used for evaluating joint predictions. Figure ((a)
is
generated through dyadic sampling and Figure 9(b) includes additional perturbations.
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 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 prodictions 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 [(b] presents 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 [0] 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 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 incrcasing model size. More importantly, the epinet agent outperforms both basclines by huge margin under both sampling schemes. These results give further confidence in the quality of joint predictions produced by the epinet agent.
20

---

dyadic sampling
sampling
10-
joint log-loss
enn
opinct
*
367
1e8
3e8
1e9
3e9
3e7
1e8
3e8
1e9
3e9
model size (number of parameters)
Figure 10: An agent's joint log-loss under dyadic (left) and augmented dyadic sampling (right).
21

---

G Testbed experiments
This section provides details about the Neural Testbed experiments in Section 5 with review of the neural testbed as 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 []] Finally, we investigate the sensitivity our hyperparameter choices when evaluated across the testbed.
We begin
G.1 Neural testbed
The Neural Testbed (Osband et al. 2022a) is collection of neural-network--bane synthetic classification problems that evaluate the quality of an agent's 's predictive distributions. We make of the open-source code https::://///.. //gtt///tttbbb wwt.coom/ceeee sateeppmind/neuuaaa 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 entails binary classification, input dimension D E (2, 10, 100), number of training samples = XD for E {1, 10, 100, 1000), temperature D {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 metries are averaged across problems to give the final performance scores.
G.2 Benchmark agents
We follow Osband al. (202aa) and consider the benchmark agents as in Table 1 use the open-source implementation and hyperparameter sweeps al hhtpss: https:///aonymmuu (aardd/-996//gggtt///ttteee
We
According 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 standard Gaussian with dimension Dz = 8. The base network jrc has the same architecture as the baseline mlp agent, which is 2-layer MLP with ReLU activation and 50 units in each hidden layer. The learnable part of the epinet of takes (x(x) and index as inputs, where (x(x) is the concatenation of x and the last-layer features of the base network. The learnable network has the form ())()),2) = 9n (z), 2)) where [x) 2] is the concatenation of &c (z) and z, and (-) is 2-layer MLP with hidden width 15, ReLU activation, and outputs in R D2xC for number of classes C' 2.
For the fixed prior of, consider ensemble of Dz networks. Each member of the ensemble is a small MLP with 2 hidden layers and 5 units in each hidden layer. Each MLP takes imput and returns logits for the two classes. Let p' pxxx) RC denote the output of the ith member of the ensemble. We combine the outputs of the ensemble members by taking the weighted sum, Σ.r (1));; and multiplying the sum by tunable scaling factor a. Thus, we can write the prior function of (+c(x);; 2) D2 ()))2
We combine the base network and the epinet by adding their outputs and applying stop-gradient operation on (x),,,
Hc((z) of (sg() + of (sgioc(z)).2),
where 0 = (C, m) denotes all the parameters of the base net and epinct. We train the parameters Ç and 17 jointly. The training loss takes the form as specified 77), where we use log loss for the data loss and ridge regularization. We update A using Algorithm| [ with a batch size of 100 and number of epistemic index samples equal the index dimension. We use Adam optimizer with learning rate 1e-3. The L2 weight decay and prior scaling factor are roughly adjusted for different problem settings, taking in account the number of training samples and SNR.
22

---

Our implementation of the epinet agent can be found under the path /agents/factories/ epinet.pp py in the anonymized neural testbed github.
G.4 Ablation studics
We run ablation experiments on the epinct agent that trained simultancously 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 epinct, 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 hyperparameter.
Our results are summarized in Figure I1 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 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 epinct hidden layers once the width large enough. Performs of opinct seoms sensitive to the prior scales. Smaller prior scale leads to better marginal kl, too small too large prior scale degrades the joint kl. Increasing the width of the models 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.
23

---

0.29
4.0
12.0
4.0
6.
12.0
index dimension
pent iI
0.42
1005
0.34
6.17
number of hidden layers
margfral
0.4.
0.185
0.180
0.175
5.38
8.170
0.265
5.0
10.0
15.0
20.0
30.0
$5.0
5.0
10.9
15.0
20.0
30.0
56.0
width opinet hidden layers
margiral
jint
0.50
0.4
0.3
0.2
0.03
0.1
0.3
2.0
3.0
0.03
0.1
0.3
3.0
ensemble prior scale
margiral
jiit 31
0.42
0.21
5.35
2.0
3.5
5.
15.0
70.0
36.0
52.2
2.0
3.5
5.3
15.0
20.0
30.0
55.0
width of prior networks
margnalli
[SH 2I
0.0
0.05
0.2
1.0
2.0
0.0
3.22
0.05
0.1
0.5
2.5
Figure 11: Ablation studies of epinet with 2-layer mlp base model the neural testbed.
24

---

H Image classification
This section gives an overview of our experiments on image classification problems outlined in Section 0 We begin with a review of the hyperparameter choices and design details for the agents 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, present results for evaluation 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 well as memory modern TPU architectures. Finally, we perform a suile of ablations to investigate the sensitivity of our results across ImageNet.
H.1 Epinet details
We train one epinet for each ResNet-L baseline for L (50, 101, 152, 200). The ResNet basolines are open sourced in the ENN library under the path /networks/resnet// and the checkpoints are available under / /checkpoints/Imagee 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 shetworks/epineee
py. For each epinct agent, we take
and / /checkpoints/immee htw.tttttimtmeiiii.eeaeeee py-
As discussed in Section G choose the reference distribution of the epistemic index to be standard Gaussian with dimension Dz = 30. The imput to the learnable part of the epinet of includes the last-layer features of the base ResNet and the epistemic index denote the number of classos. For last-layer features & and epistemic index 2, the learnable network takes the form of 2) ( z))² 2, where [ø) is the concatenation of 6 and 2, and 9-(-) is a 1-layer MLP with 50 hidden units, ReLU activation, and output E RDxxC
Let C
The fixed prior of 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 'prior scales' in the ablation studies) for the output of each component before we add them logether. However, for ImageNet, we find that scaling factors of 1 already work well. The first. component of the prior is network with the same architecture and initialization as the learnable network of-- 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 10, 10 x 10, x 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 (z) E RC denote the output of the fth member of the ensemble. We combine the outputs of the ensemble members by taking the weighted sum ΣP1 ('))))
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 schecule, 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 x 4 TPUs with 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 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 epochs.
H.2 Uncertainty baselines
In this scction we compare our results to the open source 'uncortainty basclines', 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 high level, our results mirror our own
25

---

results of Figure 2 After tuning, most of the agents appear on a roughly similar tradeoff in
terms of marginal quality. However, the approachs are widely separated in terms of their quality on joint prediction, and here epinet performs much better than the alternatives.
Figure 0 repeats Figure 2 but adding few new agents from uncertainty baselines, which we
mimo:
Multi-input oulit-output ensemble (Havasi et al. 2020). This method appears perform better than baseline resnet in terms of marginal statisties, but provide no additional benefit to modeling joints. Note that the independent product of better marginals automatically does improve joints.
dropout: Dropout posterior approximation (Srivastava et al., 2014 Gal and Ghahramani 2016). This approach seems to provide a slight improvement in marginal log-loss and noticeable improvement in joint log-loss. However although dropout has a low computational cost. in terms of paromelers, these results required 10 forward passes of the network, and so actually underperfom 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 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 like to understand better the commonalities between these two approaches, and see if/when the algorithms can borrow from each others' strengths.
particular form of epinct. In future work, we would
H.3 CIFAR-10 and CIFAR-100
In this section we reproduce a similar analysis Section 6 applied to the CIFAR-10 10 and CIFAR-100 datascts. We find that, at a high level, our results mirror those when applying ResNet ImageNet. In particular, we are able to produce results similar to Figure 2 for both CIFAR-10 and CIFAR-100. Relative to large ensembles, epincts 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 /netxorks/reent// and the checkpoints are available under /checkpoints/ cifar10. py and piteckpointssciffarll
py. In particular, we tune ResNet-L for L e {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.]] but with slightly smaller network. We use index dimension Dz = 20 and alter the convolutional prior to have channels (4, 4) each with kernel size 5 x 5 and stride 2 on account of the smaller image sizes in CIFAR-10 and CIFAR-100 (Krizhevsky (209)).
The epinct network architecture together with checkpoint weights can be found in the ENN library under the paths /nptworss/epinet/ /checkpoints/ciffrl// .y, and /checkpointss/cfarr000
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
26

---

at classification and marginal log-loss, epinets do not offer any particular advantage over baseline ResNets. However, when we look the joint log-loss we can see that epinets offer huge improvements in performance, even when measured against very large ensembles.
classification eiror
marginal log loss
J nt log-loss
0.22
lo5s
0.06
erismbbee
net
0.05
*
0.5
365
166
366
167
3,55
106
306
167
365
166
306
167
model size (number parameters)
Figure 12: Quality of marginal and joint predictions across models on CIFAR-10.
Cassificatin error
marginal log ioss
joint logloss
0.325
10-
0.300
loss
0.275
fedinet
entsemble
0.250
4
305
365
model size (number of parameters)
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 noural 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, 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 1 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 vary this depending on their tradeoff between statistical loss and computational cost. Figure 15 shows this empirical tradeoff over M (10,30, 30, 100, 300, 1000,3000), 3000} for each of the resnet variants. Once again, 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 Res-et-50 base network. We sweep over various values for the index dimension, the number of hidden
27

---

clsssfication errer
marginal log 1o55
Jo nt log 1055
0.23
*
0.22
reshet
#
insemble
0.21
5-
0.80
4-
369
1610
3e10
1611
3611
369
1e10
3630
1el1
Je11
369
1630
3e10
lel1
Je11
computational cost (FLOPs)
Figure 14: Quality of marginal and joint predictions across models on ImageNet. Reproduces the results of Figure 2
but using inference FLOPs as mcasure of computation.
cassification
0.23
spinet
resnet
1685
resnet_ resnet_vvriant
variant
101
0.85
200
10
15
25
20
10 20
computational cost (gFLOPs)
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 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 los-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, 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, to re-scale the ENN output logits
we find that using a cold temperature improves the agent's performance during evaluation (Wezzel et al. 2020 The improvement is most dramatic the joint log-loss.
28

---

1.04
0.235
6.230
20
30
60
20
30
20
30
indexx ditwirsion
2.240
5.235
3.925
number hidden layers
clasiiactton 6722
2.4
margral log lom
Jann log
0.74
30-
0.33
1.2
0.22-
matched
0ie5 OFFY
TRE NOS
J40 (8) 055
34
12
4.50-
0.23-
4.2
0.72
335
34
01
ensemble ble of COTV nets prior Scale
3.075
0.225
3.033
0,200
1002
16:05
3r/66
1603
1e42
16-05 1-bt 3r:8 16
Leas 1e:tt 1690 3r:42
12 weight decay
mergral
19s1
0.240
0.35
0.225
0.02
0.92
5.220
number index
Cloaiiiaainn ent
merginal kay Ini
1911 log los
0.240
0.92
4.5
0.41
0.225
0.92
42
0.92
smoothing
Canefeett on InTe
marginal loge
Fine ea looe
0.235
1055
0230
Samperature rescaling
Figure 16: Ablation studies of epinet with ResNet-50 base model on ImageNet.


Reviews:
[{"review_id": "YN4CKdS9Jx", "reviewer": "Reviewer_UZs5", "summary": "This paper proposes epistemic neural networks (ENNs) as a novel framework for uncertainty estimation in neural network predictions. ENNs introduce an epistemic index that expresses uncertainty and correlations across multiple inputs via joint predictive distributions. The paper argues that joint predictions are critical for effectively evaluating uncertainty quality and enabling good decision making. It is argued that minimizing joint log-loss leads to near optimal actions while marginal log-loss does not. ENNs generalize Bayesian neural networks, as any BNN can be expressed as an ENN but the reverse does not hold. A novel ENN architecture called epinet is introduced, which adds a small auxiliary network to any existing neural network to produce uncertainty estimates. Experiments demonstrate that epinets can match the joint prediction performance of large ensembles while adding relatively minor computational overhead. ", "questions": "Theoretical Analysis\n\nCould you provide a more thorough theoretical analysis quantifying the advantages of joint modeling for decision making? Any regret bounds or performance guarantees?\n\nIs it possible to better characterize the properties of the proposed training objectives? Do they provably optimize calibration of uncertainties?\n\nExperimental Evaluation\n\nHave you considered evaluating on tasks more directly relevant to decision making such as experiment design, active learning, or contextual bandits?\n\n\nImpact\n\nCould you better highlight the broader potential impact enabled by the approach? What new directions are unlocked?", "limitations": "NA", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Originality:\n\nProposes a new conceptual framework of epistemic neural networks that expands uncertainty estimation options beyond Bayesian neural networks. \n\nIntroduces a practically effective and scalable approach via epinets that represents a novel architecture and training methodology.\n\nQuality:\n\nTechnically sound way of modeling uncertainty through joint predictive distributions.\n\nStrong empirical results outperforming baselines.\n\nClarity:\n\nMotivates the limitations of marginal predictions and need for joint modeling.\n\nExplains and visualizes key concepts effectively through examples. \n\nResults presented clearly through tables/plots.\n\nSignificance: \n\nUncertainty estimation is a fundamental challenge for deploying reliable neural networks.\n\nCan enable progress on critical applications like exploration, experiment design, and robust decision making.\n\n", "weaknesses": "Theoretical Analysis:\n\nThe connections between joint log-loss and decision performance could be bolstered with more rigorous analysis of regret bounds or performance guarantees. This would strengthen claims about utility for decision making.\n\nAdditional theoretical characterization of the proposed training objectives and epinet architecture properties could provide better insights into why the approach works.\n\nExperimental Evaluation:\n\nExperiments focus on image classification. Evaluating on decision-making benchmarks more directly relevant to the motivations could better highlight benefits.\n\nCould ablate design decisions like network architectures and priors more thoroughly to understand their impact. This provides guidance on best practices.\n\nNovelty and Impact:\n\nThe high-level ENN concept builds closely on established perspectives like Bayesian learning. Examining more novel implications would strengthen novelty.\n\nWhile solid incremental gains are shown, the broader advancements enabled by this approach could be better highlighted. Articulating impact on future work would increase significance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their time and effort in reviewing the paper. We address the comments below:\n\n**> Theoretical analysis regarding regret bound: Could you provide a more thorough theoretical analysis quantifying the advantages of joint modeling for decision making? Any regret bounds or performance guarantees?**\n\nThanks for raising this. We would like to point the author to Theorem 2, which has a formal version and a proof in Appendix B. The theorem establishes a rigorous regret bound in bandit settings. It is an adaptation of the theory presented by Wen et al. (https://arxiv.org/pdf/2107.09224.pdf), which we plan to cite in the main paper as well.\n\n**> Theory regarding the choice of training objective: Is it possible to better characterize the properties of the proposed training objectives? Do they provably optimize calibration of uncertainties?**\n\nWe would like to point the reviewer to Theorem 4 and the associated lemmas and proof in Appendix D. In Theorem 4, we show that the distribution approximated by an epinet converges to the posterior in linear regression setting, under appropriate technical conditions. We plan to extend this theory beyond linear regression in future work.\n\n**> Have you considered evaluating on tasks more directly relevant to decision making such as experiment design, active learning, or contextual bandits?**\n\nWe would like to point the reviewer to a relevant paper: (https://arxiv.org/pdf/2302.09205.pdf)  which presents empirical results that improved joint predictions via epinet lead to better performance in bandit and reinforcement learning tasks. \n\n**> Could you better highlight the broader potential impact enabled by the approach? What new directions are unlocked?**\n\nOur main motivation for this work is to enable principled and practical methods for scalable uncertainty estimation. Our work, especially epinet, will enable uncertainty estimation for large scale models (such as large language models). Furthermore, scalable uncertainty estimation can also unlock practical implementations of efficient exploration algorithms, such as Thompson sampling and information-directed sampling, which may dramatically improve data efficiency in sequential decision-making tasks. Finally, our ENN framework and open source library will allow researchers to iterate and develop better networks and approaches to uncertainty modeling. \n"}, {"review_id": "z78FXCpfxP", "reviewer": "Reviewer_TShA", "summary": "In this work, the authors introduce a new framework \"Epistemic Networks\" to better capture uncertainty by integrating over epistemic indices when computing the joint distribution over multiple inputs. This approach is flexible and can be added to any existing neural network architecture with some differences in the training algorithm. The authors show through experimental evaluations how their approach shows promising results when compared against other baseline methods for uncertainty quantification, including BNNs and deep ensembles on joint log-loss. ", "questions": "- Can you highlight the parallels between the expected reward under joint prediction to the work in the Bayesian decision theory area? In Bayesian decision theory too our goal is to get better decision utility for downstream decision-making task. [1-3]\n\n- Building upon the previous point, why should one focus on the decision making framework presented in this paper versus other Bayesian decision theoretic applications that look at marginals? [1-3]\n\n- Could you you describe how this approach captures aleatoric uncertainty more clearly? \n\nReferences\n\n[1] Vadera, Meet, Soumya Ghosh, Kenney Ng, and Benjamin M. Marlin. \"Post-hoc loss-calibration for Bayesian neural networks.\" In UAI (2021).\n\n[2] Cobb, Adam D., Stephen J. Roberts, and Yarin Gal. \"Loss-calibrated approximate inference in Bayesian neural networks.\" arXiv preprint arXiv:1805.03901 (2018).\n\n[3] Simon Lacoste-Julien, Ferenc Huszár, and Zoubin Ghahra\u0002mani. Approximate inference for the loss-calibrated\nbayesian. In AISTATS (2011)\n", "limitations": "Highlighted above in the review. The authors do not clearly specify the limitations in their submission. ", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- This paper presents a new approach that's flexible and doesn't add a lot of parameters to the existing model. \n\n- The paper evaluates the joint log-loss performance rather than just marginal log-loss and thus helps us understand how other baseline methods fare when looking at the joint prediction performance. ", "weaknesses": "- I think the authors need to do a better job to motivate the case for joint log-loss. Why should one care for joint log-loss performance as opposed to marginal? Can you describe any applications where this might be useful? \n\n- The presentation can improve a bit. The authors should describe section 4 to make sure it's easier to understand for folks. For e.g., what are all the different terms in sec 4.2 (eq 7, 8) and what do they signify? ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank Reviewer TShA for reviewing our paper. The following are our point-to-point responses:\n\n**I think the authors need to do a better job to motivate the case for joint log-loss. Why should one care for joint log-loss performance as opposed to marginal? Can you describe any applications where this might be useful?**\n\nWe would like to clarify that, as a recent paper [Wen et al. 2022] has shown, for a broad class of decision problems, such as combinatorial decision problems, sequential prediction problems, and multi-armed bandit problems, accurate joint predictions are required to deliver good performance, while accurate marginal predictions alone are insufficient to guarantee good performance. We will further explain and discuss this in the revision.\n\nWen et al. 2022, “From predictions to decisions: the importance of joint predictive distributions ”, https://arxiv.org/abs/2107.09224.\n\n**The presentation can improve a bit. The authors should describe section 4 to make sure it's easier to understand for folks. For e.g., what are all the different terms in sec 4.2 (eq 7, 8) and what do they signify?**\n\nThanks for raising this. We will further improve the writing of Section 4.2. The notations in equations 7 and 8 follow the definitions in Sections 3 and 4.1. In particular: \\theta is the parameters of the ENN, z is the ENN index, (x_i, y_i) is an input-label pair. f_\\theta(x, z_i) is the output of the ENN, which is the logit vector. The softmax transforms the logit vector into a probability vector, and the subscript y_i  in eq 7 denotes the (y_i)-th component of this probability vector. We will clarify that in eq 8, \\Phi is the CDF of a standard Gaussian random variable.\n\nWe recognize that Section 4.2 only explicitly states the loss for a single input-label pair and a single epistemic index. We will clarify in the paper that for each stochastic gradient step, the method samples a batch of input-label pairs and a batch of indices and averages over the losses.\n\n**Can you highlight the parallels between the expected reward under joint prediction to the work in the Bayesian decision theory area? In Bayesian decision theory too our goal is to get better decision utility for downstream decision-making task. [1-3]**\n\nBayesian decision problems are typically framed as maximizing expected utility.  Our approach offers tools for such problems.  Theorem 2 ensures that minimizing joint log-loss leads to effective decisions for problems framed in that manner.\n\nPlease also note that, as Theorem 2 indicates, for the special case when the reward only depends on one label (i.e. tau=1 in the formal version of Theorem 2 in Appendix B), joint loss reduces to marginal loss. For that special case, minimizing marginal log-loss is sufficient.\n\n**Building upon the previous point, why should one focus on the decision making framework presented in this paper versus other Bayesian decision theoretic applications that look at marginals? [1-3]**\n\nAs we have discussed above, our approach is consistent with Bayesian decision theory.  Joint log-loss serves as a unit test to ensure that our tool will serve the needs of Bayesian decision making.\n\n**Could you you describe how this approach captures aleatoric uncertainty more clearly?**\n\nVariation across epistemic indices models epistemic uncertainty.  For a fixed epistemic index, label probabilities produced as model outputs express aleatoric uncertainty.\n\n"}, {"review_id": "LIeqy2P67R", "reviewer": "Reviewer_jzvR", "summary": "The paper introduces epinets as part of Epistemic Neural Networks (ENNs), a novel approach to uncertainty estimation in deep learning models. Epinets extend neural networks to create ENNs, which can be used to estimate uncertainty. The paper then proceeds to present experiments on image classification tasks, demonstrating that epinets outperform both uncertainty baseline models and ensemble methods in terms of joint log-loss, while maintaining comparable performance in terms of marginal log-loss. At the same time, the computational cost of these networks is orders of magnitude less than larger ensembles.", "questions": "1. Has the epinet framework been tested on real-world datasets outside of the experiments mentioned in the paper?\n2. How does the interpretability of the uncertainty estimates by epinets compare to that of BNNs?\n3. In the ablation study in the appendix, it is shown that larger values for the index dimension improve both the joint and marginal kl estimates. Is there a point where this is no longer the case? What is the tradeoff?", "limitations": "To my knowledge, the limitations of epistemic neural networks are not addressed in the paper. The limitations and possible directions for future work could be addressed to improve the understanding and practicality of epinets.\n\nEdit: I thank the authors for addressing my questions and comments in their rebuttal.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper presents a clear distinction between epistemic neural networks and Bayesian neural networks (BNNs). This distinction helps readers grasp the unique characteristics and advantages of the proposed framework. \n- The results showcase the benefits of epinet compared to both ensemble approaches and uncertainty baselines. More comprehensive results in terms of the computational cost in FLOPs are provided in the appendix.\n- The explanations provided throughout the paper are balanced between intuitive and easy-to-follow explanations and rigorous theorems. The figures provided in the document also effectively illustrate both the core principles of the epinet framework and the key results obtained. These visuals enhance the understanding of the concepts discussed in the paper.", "weaknesses": "- Even though the framework is rather general, the focus of the experiments seems rather narrow and confined to image classification tasks. It would be useful to see how well the epinet generalizes to other domains and whether there are scenarios or domains where it is more or less effective.\n- Related to the previous point, the limitations of the work are not discussed in the conclusion. While epistemic neural networks are a general framework, addressing such limitations could provide insights on future directions to take.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for reviewing our paper. Our point-to-point responses are as follows:\n\n**Even though the framework is rather general, the focus of the experiments seems rather narrow and confined to image classification tasks. It would be useful to see how well the epinet generalizes to other domains and whether there are scenarios or domains where it is more or less effective.**\n\nWe thank the reviewer for this comment. We would like to point the reviewer to a recent paper [Osband et al. 2023], on approximate Thompson sampling via epistemic neural networks. [Osband et al. 2023] has provided more experiment results for epinet, in both the neural bandit problems and reinforcement learning problems. Epinet has been found to be effective in these two problems.\n\nOsband et al. “Approximate Thompson sampling via Epistemic Neural Networks” https://arxiv.org/pdf/2302.09205.pdf\n\n**Related to the previous point, the limitations of the work are not discussed in the conclusion. While epistemic neural networks are a general framework, addressing such limitations could provide insights on future directions to take.**\n\nWe thank the reviewer for this comment. One limitation of the current ENN framework is that it does not consider problems with sequential inputs and outputs (e.g. the language translation problems). We will clarify and discuss this limitation in the revision. A possible future direction is to build ENNs for these sequential problems.\n\n**Has the epinet framework been tested on real-world datasets outside of the experiments mentioned in the paper?**\n\nAs we have mentioned above, the epinet framework has been tested in other problems, including the neural bandit problems and reinforcement learning problems.\n\n**How does the interpretability of the uncertainty estimates by epinets compare to that of BNNs?**\n\nOur understanding is that interpreting uncertainty modeling in complex prediction/decision problems might be too challenging. It is better to compare them based on the effectiveness of resulting predictions/decisions. In this paper, we have used the joint log-loss to measure the uncertainty estimates from different agents, including epinets and BNNs. As Theorem 2 of this paper shows, minimizing joint log-loss leads to effective decisions.\n\n**In the ablation study in the appendix, it is shown that larger values for the index dimension improve both the joint and marginal kl estimates. Is there a point where this is no longer the case? What is the tradeoff?**\n\nThanks for the comment. Usually, increasing the ENN index dimension will improve both the joint and marginal kl estimates (performance). However, it will also require more computation. Thus, there exists a computation-performance tradeoff here.\n\n\n"}, {"review_id": "R4dT09hq7X", "reviewer": "Reviewer_zTnm", "summary": "In this paper, the authors present a new approach for uncertainty estimation utilizing joint predictions. They introduce the concept of an Epistemic Neural Network (ENN). Within this framework, they propose an innovative architecture called 'epinet,' which supplements any conventional neural network. This additional architecture helps conventional neural networks outperform large ensembles in terms of the joint loglikelihood at the cost of a minimal computational increase. This approach is shown to be effective and practical through numerical experiments, including large-scale datasets like ImageNet.", "questions": "Please address the weaknesses I mentioned above.\n\nEdit: I would like to thank the authors for their answers during rebuttal period.", "limitations": "Authors adequately addressed the limitations, there is no negative societal impact.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper addresses an important problem of uncertainty quantification and disentanglement (aleatoric and epistemic) and introduces a novel class of models, ENN, to tackle it.\n2. The paper is well-written and easy to follow.\n3. The code for the reproduction of experiments is provided.", "weaknesses": "I think that there is a missing opportunity to show the performance of the proposed approach on direct uncertainty quantification. The motivation example in the introduction section showed that it is important to realise where uncertainty came from -- is it due to the noise in data (aleatoric) or due to the lack of knowledge (epistemic). \nIt would be helpful to see the authors run standard experiments, like checking how well the model can tell the difference between one dataset and another. This would give us a clearer picture of how well the model understands what it doesn't know. One such experiment might be to train a model on one dataset (say CIFAR100) and compare some epistemic uncertainty measure on another dataset (say LSUN). And then compute ROCAUC.\nRight now, the results are just about how well the model can classify data and predict losses based on the same training dataset, which I believe doesn't tell us the whole story.\n\nMinor concerns:\n1) In the related work section, line 74, there are cited several papers, and mentioned their connection to Gaussian Processes. But two out of three cited papers - PriorNets and PostNets -- actually focus more on Dirichlet parametrization and don't have a clear link to Gaussian Processes.  I think SNGP[1] should be cited here.\n\n\n[1] Liu, Jeremiah, et al. \"Simple and principled uncertainty estimation with deterministic deep learning via distance awareness.\" Advances in Neural Information Processing Systems 33 (2020): 7498-7512.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and effort in reviewing this paper. Our responses are:\n\n> I think that there is a missing opportunity to show the performance of the proposed approach on direct uncertainty quantification... It would be helpful to see the authors run standard experiments, like checking how well the model can tell the difference between one dataset and another...\n\nThanks for the great suggestion!  We plan to include a visualization of the uncertainty estimates produced by our methods on some 2D OOD-like classification problems. This should give more intuition on how epinets are performing outside of the train dataset. It will be interesting to investigate more about how our methods will perform on more general OOD tasks  in the future.\n\n> In the related work section, line 74, there are cited several papers, and mentioned their connection to Gaussian Processes. But two out of three cited papers - PriorNets and PostNets -- actually focus more on Dirichlet parametrization and don't have a clear link to Gaussian Processes. I think SNGP[1] should be cited here.\n\nThanks for pointing us to the SNGP paper [1]. We will make appropriate modifications and include the paper."}, {"review_id": "fmM6Qv2gwI", "reviewer": "Reviewer_tvti", "summary": "This paper introduces *Epistemic Neural Networks (ENN)* as a general approach for uncertainty estimation for deep learning, and then the *epinet*, which is a particular ENN instantiation. The paper argues that *joint* predictions are essential for sequential decision making problem, and shows *epinets* offer similar marginal performance but better joint predictive performance than conventional BNNs/DNNs.", "questions": "1. L42-43 \"evaluating each model on out-of-sample data\", is this OOD data e.g., from ImageNet-C, or just the normal test set? Please specify?\n\n2. How would one actually use an ENN on sequential decision making problems or active learning problems? Update the epinet parameters only? Retrain the whole thing? Perform inference over Z? Given that this is the aim of the paper, this is an important topic to discuss in the main text.\n\n3. Is the Glorot initialisation used for the MLP important?", "limitations": "Looks good to me.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 1, "contribution": 3, "strengths": "1. Interesting ideas, particularly around setting aside getting good weight-space posteriors and instead focusing on getting good predictive performance (whether that be marginal or joint predictions). I think this is good, and is of interest to the community.\n\n2. Solid experimental results, especially around joint predictions. Nice job.", "weaknesses": "Although I think the paper is of interest to the community, I think the paper could be substantially improved. I have some concerns about the presentation and writing of the paper that limit how excited I am about this work being published at NeurIPS. On balance, I would like to see the suggestions I've made below incorporated into the writing of the text, and this would enable me to increase my score. \n\n# 1. Poor Technical Presentation, Writing, Technical Soundness\nI have a number of specific concerns here. Some of these might be seen as nit-picky, I admit, but I remain concerned here.\n1. L7-9. \"With an epinet, CNNs outperform large ensembles of hundred of more particles, and use orders of magnitude less compute.\" This claim should be qualified because it only applies robustly to joint-predictions. This applies elsewhere in the paper.\n\n2. L31-32 \"These approaches can almost match the posterior distribution.\" I think citing Welling and Teh for this doesn't make sense, and I don't think this is true. See, for example, the analysis in [1].\n\n3. L36-37. \"Practical large scale implementations are often limited to ten or fewer [particles] due to computational reasons\". The way you've defined \"ensemble based BNNs\", sampling based inference or even VI when sampling from the approximate posterior counts as an ensemble. And some approaches use more than 10 samples, e.g., sampling 30+ particles from an approximate posterior should be fine.\n\n4. In a number of places in the paper, intuition could be provided that would allow the reader to understand the points you are making more clearly. e.g., L50-51. \"All BNNs are ENNs, but there are useful ENNs such as the epinet that are not BNNs\". Explain why please. ENNs have the extra index, which mean that they have a wider class of functions (as far as I can tell), but I'm not sure this is right. L215-216 \"with this stop gradient, training dynamics more reliably produce models that perform out ot sample\". Please provide more intuition.\n\n5. The writing in Section 3 only addresses classification problems. It should be more general.\n\n6. The notation in 3.2 is confusing. Both $\\nu$ and $\\theta$ correspond to parameters.\n\n7. It is worth noting that recent work argues that BNNs do not actually need to maintain distributions over all of their parameters for good predictive performance [1]. This is important because it's mentioned in the paper that BNNs learn distributions over all parameters, but this is not completely true, some BNNs do not do this. This characterisation seems wrong.\n\n8. L163: \"A BNN is specified by a pair: a base network $f$ and a parameterised sampling distribution $p$.\" I don't think this is quite right. I think a BNN is specified by a likelihood and prior, which defines a posterior distribution, at least technically speaking. Just because we have a distribution over network parameters does not mean that the network is a BNN. That might be contentious, but the sampling distribution certainly does not need to be parametric! For example, HMC has no parametric sampling distribution!\n\n9. Eq (3): it seems to be the definition should integrate over $z$ and examine the distribution over outputs when doing that. The actual equation currently seems to be a stronger condition: correct joint predictions for every value of $z$. Am I missing something here?\n\n10. Theorem 3. \"any BNN defined with respect to $f$ can be expressed as an ENN defined with respect to $f$\". I find this is a bit confused. Are there no technical conditions here? If there are technical conditions, please improve them. For instance, if there is a BNN with 10 million parameters, how do we represent this with an ENN that has one dimensional $Z$. \n\n11. Similarly, notation in 4.1 is a bit confusing, with different symbols used for different parameters etc.\n\n12. I think it is incredibly confusing that the paper introduces ENNs and epinets, but that these are not the same thing! I would suggest changing the name of one of them, probably the epinet.\n\n# 2. Evaluations\n1. What does the reduction in joint-log loss actually correspond to here? Does it yield better predictions in sequential decision making problems? Can you test this? For example, on an active learning or bandit problem? I appreciate this might be a tough ask, but I think given that decision making problems are a central focus for this work, this would be a substantial improvement to the submission. I'd also like to see the performance of an ENN on a 1D regression problem to understand whether ENNs offer coherent updates when observing data, since BNNs (ideally) would offer this. This could be done by retraining, or by just updating the posterior over the epistemic index.\n\n2. I'd like to see results on ImageNet-C i.e., OOD robustness. BNNs and other approaches are often used here. I'm curious to see if the epinet helps here. Maybe this is what you are doing already, it wasn't clear to me reading the main paper alone, and it should be.\n\n3. Are there error bars on Figure 2? It seems not? Please add them, if possible! \n\n4. It is not clear what the reduction in joint log-loss corresponds to, the number is hard to interpret or understand. As a reader, I cannot understand how significant this is. Benchmarking on a sequential decision making problem and demonstrating improved decision making there would be a much more significant result. I appreciate there is theory here, but I'm more interested in the practical implications of the work.\n\n5. Ideally, there would be examples of networks having good marginal predictions but poor joint predictions that lead to problems in sequential decision making problems.\n\n[1] Sharma, Mrinank, et al. \"Do Bayesian Neural Networks Need To Be Fully Stochastic?.\" International Conference on Artificial Intelligence and Statistics. PMLR, 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for carefully reviewing our paper and writing detailed reviews. Our responses are as follows. Due to the length limit of the rebuttal, we will omit parts of the questions and provide concise answers.\n\n**Regarding your questions:**\n\n**1. L42-43 \"evaluating each model on out-of-sample data\", is this OOD data e.g., from ImageNet-C, or just the normal test set? Please specify?**\n\nThis is a typical test dataset. We will clarify.\n\n**2. How would one actually use an ENN on sequential decision making problems or active learning problems? Update the epinet parameters only? Retrain the whole thing? Perform inference over Z? Given that this is the aim of the paper, this is an important topic to discuss in the main text.**\n\nThanks. In case of sequential decision making, on obtaining new data, both base model and epinet needs to be trained (as the data provides information for both base model and epinet). This can be done either concurrently or by first training the base network and then training the epinet in an iterative manner. We believe that both of these should be equivalent. \n\n**3. Is the Glorot initialisation used for the MLP important?**\n\nWe think that Glorot initialisation would lead to outputs roughly in N(0,1) when inputs are from N(0, I). This helps us by not having to tune the scale for different depths and widths of the neural networks. But techniques presented in our paper should also work for other initialization schemes. \n\n**Regarding your comments on technical presentation, writing, and technical soundness:**\n\n**L7-9. \"With an epinet, CNNs ... This applies elsewhere in the paper.**\n\nWe will polish the wording in the revision.\n\n**L31-32 \"These approaches can almost ... See, for example, the analysis in [1].**\n\nThanks for pointing us to [1], we were not aware of this recent paper. We will make modifications accordingly.\n\n**L36-37. \"Practical large scale ... e.g., sampling 30+ particles from an approximate posterior should be fine.**\n\nWhen we refer to “ensemble”, we just mean an ensemble of particles, not more general sampling based inference. We will change “ten or fewer particles” to “at most tens of particles”.\n\n**In a number of places in the paper ... Please provide more intuition.**\n\nThanks for the comments! We will provide more intuitions in the revision to further improve the writing.\n\n**The writing in Section 3 only addresses classification problems. It should be more general.**\n\nWe fully agree that it can be more general, in particular, it can also include the problem formulations for regression problems. We only include the classification formulation since we would like to keep the paper short and simpler. \n\n**The notation in 3.2 is confusing. Both \\nu and \\theta correspond to parameters.**\n\nNote that \\theta corresponds to the parameters of the neural network, and \\nu corresponds to the parameters of the sampling distribution.\n\n**It is worth noting that recent work ... some BNNs do not do this. This characterisation seems wrong.**\n\nThanks for pointing us to this very recent work, we were not aware of it. We would like to clarify that we just would like to give a high-level overview of how classical BNNs model uncertainty. We will change the wording accordingly.\n\n**L163: \"A BNN is specified by a pair... For example, HMC has no parametric sampling distribution!**\n\nWe will clarify in the revision that the sampling distribution is supposed to be an approximation of the posterior distribution.\n\n**Eq (3): it seems to be the definition ... Am I missing something here?**\n\nNote that `z` is a random variable, so in Equation (3), we imply that the distribution of a BNN can be matched by the distribution generated by an ENN. \n\n**Theorem 3. \"any BNN defined with respect to `f` can be expressed ... how do we represent this with an ENN that has one dimensional `Z`.**\n\nTheorem 3 doesn’t require any additional technical conditions. This can be observed from the proof of Theorem 3 in Appendix C. \n\n**Similarly, notation in 4.1 is a bit confusing, with different symbols used for different parameters etc.**\n\t\nWe used different symbols to distinguish different parameters to avoid confusion.\n\n**I think it is incredibly confusing that the paper ..., probably the epinet.**\n\nENNs are a more general framework. Epinet is one specific kind of ENN. For example, ENNs are to NNs as Epinet is to an MLP.  \n\n**Regarding your comments on evaluations:**\n\n**What does the reduction in joint-log loss actually correspond to here? ...  or by just updating the posterior over the epistemic index.**\n\nThanks for the great suggestion. We will include these results with 1D regression in the paper.\n\n**I'd like to see results on ImageNet-C i.e., OOD robustness. ..., it wasn't clear to me reading the main paper alone, and it should be.**\n\nThanks for this suggestion! We will examine performance on Imagenet-C and other OOD datasets in our future work.\n\n**Are there error bars on Figure 2? It seems not? Please add them, if possible!**\n\nWe weren't able to add the error bars because of the compute requirements for some of the agents (like ensembles), but we will try to include them in the final version.\n\n**It is not clear what the reduction in joint log-loss corresponds to, ... but I'm more interested in the practical implications of the work.**\n\nWe would like to point the reviewer to a recent paper [1] that examines performance of epinet, in both the neural bandit problems and reinforcement learning problems.\n\n[1]  “Approximate Thompson sampling via Epistemic Neural Networks” https://openreview.net/pdf?id=xampQmrqD8U\n\n**Ideally, there would be examples of networks having good marginal predictions but poor joint predictions that lead to problems in sequential decision making problems.**\n\nBased on the paper [1] discussed above, it is indeed the case that the  agents which have good marginal predictions and poor joint predictions  perform poorly in sequential decision problems.\n"}]

Author rebuttal:

Meta-review: The reviewers praise the significance, originality, strong empirical results, technical soundness, and flexibility and scalability of the approach. The main criticisms are the lack of theoretical insights, limited scope of experiments, lack of motivation, lack of clarity in the presentation, and lack of OOD experiments. However, these concerns have been mostly addressed by the rebuttal.
Decision: Accept (spotlight)