Evaluating Predictive Distributions: Does Bayesian Deep Learning Work?
Year: 2021

Abstract:
Posterior predictive distributions quantify uncertainties ignored by point estimates. This paper introduces The Neural Testbed, which provides tools for the systematic evaluation of agents that generate such predictions. Crucially, these tools assess not only the quality of marginal predictions per input, but also joint predictions given many inputs. Joint distributions are often critical for useful uncertainty quantification, but they have been largely overlooked by the Bayesian deep learning community. We benchmark several approaches to uncertainty estimation using a neural-network-based data generating process. Our results reveal the importance of evaluation beyond marginal predictions. Further, they reconcile sources of confusion in the field, such as why Bayesian deep learning approaches that generate accurate marginal predictions perform poorly in sequential decision tasks, how incorporating priors can be helpful, and what roles epistemic versus aleatoric uncertainty play when evaluating performance. We also present experiments on real-world challenge datasets, which show a high correlation with testbed results, and that the importance of evaluating joint predictive distributions carries over to real data. As part of this effort, we opensource The Neural Testbed, including all implementations from this paper.

heading: 1 Introduction
text: Deep learning has emerged as the state-of-the-art approach across a number of application domains in which agents learn from large amounts of data (LeCun et al., 2015). Neural networks are increasingly used not only to predict outcomes but also to inform decisions. Common approaches in deep learning produce point estimates but not uncertainty estimates, which are often required for effective decision-making. Bayesian deep learning extends the methodology to produce such uncertainty estimates (MacKay, 1992; Neal, 2012). We consider agents that are trained on data pairs ((Xt, Yt+1) : t = 0, 1, . . . , T − 1) and subsequently generate predictions given new inputs. When presented with an input XT , a Bayesian neural network can generate a predictive distribution of the outcome YT+1 that is yet to be observed. This distribution characterizes the agent’s uncertainty about YT+1. We refer to such a prediction as marginal to distinguish it from a joint predictive distribution over a list (YT+1, . . . , YT+τ ) of prospective outcomes corresponding to inputs (XT , . . . , XT+τ−1). The importance of uncertainty estimation has motivated a great deal of research over recent years (Kendall & Gal, 2017). This research has produced a variety of agents that learn to generate predictive distributions. With this proliferation of alternatives, it is increasingly important to analyze and compare their performance (Filos et al., 2019; Nado et al., 2021). In this paper, we introduce new tools for systematic evaluation of such agents. Our tools overcome several limitations faced by previous methods of evaluation. First, by focusing purely on predictive distributions, we allow for a unified treatment of approaches developed within the Bayesian neural network community and beyond. This sidesteps the
Open source code available at https://anonymous.4open.science/r/neural-testbed-B839.
question of whether any approach ‘is really Bayesian’ (Wilson & Izmailov, 2020). Second, our tools evaluate the quality of higher-order joint predictions (τ > 1). Until now, the Bayesian deep learning literature has focused almost exclusively on evaluating marginal predictions (Wang et al., 2021). Finally, we develop a neural-network-based data generating process for Bayesian deep learning that can be used to drive insight and algorithm development. Where research has focused on a small set of challenge datasets, this might introduce bias through overfitting via multiple iterations of algorithm development. We use these tools to compare hundreds of agent variants. Further, we show that performance on our synthetic data generating process data is highly correlated with performance on real-world challenge datasets. We opensource all code used in this paper as The Neural Testbed. Our results reconcile several sources of confusion in the field. One concerns why particular approaches developed by the Bayesian deep learning community, such as Bayes-by-backprop, dropout, and deep ensembles, perform poorly in sequential decision tasks despite faring well based on evaluation metrics of that community (Osband et al., 2018). Our results demonstrate that, while such methods produce accurate marginal predictions, they are no longer competitive when it comes to high-order joint predictions. Joint predictions play a critical role in sequential decision-making (Lu et al., 2021). Another puzzling issue is that state-of-the-art methods do not employ domain-specific priors. Whether Bayesian deep learning approaches should at all is a subject of controversy (Wenzel et al., 2020). We show that the benefits of domain-specific priors can be pronounced when evaluating high-order joint predictions, even where they are negligible for marginals. We also help to resolve a point of philosophical debate within the deep learning community: the importance of epistemic versus aleatoric uncertainty1. The strangeness of this distinction has even made its way into wider popular culture, as satirized in the XKCD comic of Figure 1 (Munroe, 2021). For a given parametric model, we can clearly distinguish parameter uncertainty from noise, or reducible from irreducible uncertainty. However, from the perspective of a learning agent, the choice of model is subjective; different models can lead to the same marginal predictions. Our formulation provides a clear and objective way to assess the quality of predictive distributions, without reliance on this subjective distinction between knowledge and chance. Crucially, we show that this can be judged via the quality of joint predictions, but that marginals are not sufficient.
It is worth mentioning another notable contribution of this work. The quality of a predictive distribution is commonly assessed in terms of cross-entropy loss. While this measure is welldefined for both marginal and joint predictions, to the best of our knowledge, the literature has only addressed computation in the former case. For high-order joint predictions, the straightforward approach would require computing sums over exponentially many values. To render this computationally tractable, we developed a novel approximation algorithm that leverages a random partitioning operation and Monte Carlo simulation. While this approach is motivated by concepts from high-dimensional geometry (Kaski, 1998; Donoho, 2006), we leave its analysis as a topic for future theoretical research.
1Epistemic uncertainty relates to knowledge (ancient Greek episteme↔knowledge), as opposed to aleatoric uncertainty relating to chance (Latin alea↔dice) (Der Kiureghian & Ditlevsen, 2009).

heading: 2 Evaluating predictive distributions
text: In this section, we introduce notation for the standard supervised learning framework we will consider (classification) as well as our evaluation metric (the KL-loss). We also explain how we estimate the KL-loss for high-order joint predictions where exact computation is infeasible, using random partitions and Monte Carlo simulation.

heading: 2.1 Kullback–Leibler loss
text: Consider a sequence of pairs ((Xt, Yt+1) : t = 0, 1, 2, . . .), where each Xt is a feature vector and each Yt+1 is its target label. This sequence is i.i.d. conditioned on the environment E , which produces the data, and which we view as a latent random variable. We consider an agent that is uncertain about the environment and predicts class labels YT+1:T+τ ≡ (YT+1, . . . , YT+τ ) given training data pairs DT ≡ ((Xt, Yt+1) : t = 0, 1, 2, . . . , T − 1) and unlabelled feature vectors XT :T+τ−1 ≡ (XT , . . . , XT+τ−1). From the agent’s perspective, each feature vector Xt is generated i.i.d from a fixed distribution P(Xt ∈ ·), and each class label Yt+1 is then drawn from P(Yt+1 ∈ ·|E , Xt). We describe the agent’s predictions in terms of a generative model, parameterized by a vector θT that the agent learns from the training data DT . For any inputs XT :T+τ−1, θT determines a predictive distribution, which could be used to sample imagined outcomes ŶT+1:T+τ . We define the τ th-order expected KL-loss by
dτKL =E [ dKL ( P (YT+1:T+τ ∈ ·|E , XT :T+τ−1)︸ ︷︷ ︸
environment likelihood
∥∥P(ŶT+1:T+τ ∈ ·|θT , XT :T+τ−1)︸ ︷︷ ︸ agent likelihood )] (1)
=−E [ log ( P ( ŶT+1:T+τ = YT+1:T+τ ∣∣∣θT , XT :T+τ−1, YT+1:T+τ))]︸ ︷︷ ︸ cross-entropy loss ≡ negative log-likelihood + C,
where C = E [log (P (YT+1:T+τ |E , XT :T+τ−1))] is independent of θT . The expectation is taken over all random variables, including the environment E , the parameters θT , XT :T+τ−1, and YT+1:T+τ . Note that dτKL is equivalent to the widely used notion of cross-entropy loss, though offset by a quantity that is independent of θT (Kullback & Leibler, 1951). For τ > 1, dτKL assesses joint rather than the marginal predictions.

heading: 2.2 Marginal Versus Joint Predictions
text: Evaluating an agent’s ability to estimate uncertainty on joint instead of marginal predictions can result in very different answers. We provide a simple example that illustrates the point. Suppose the data ((Xt, Yt+1) : t = 0, 1, 2, . . .) is generated by repeated tosses of a possibly biased coin with unknown probability p of heads.2 Let Xt = 0, to indicate that there is no input, and let each outcome Yt+1 be 0 or 1 to indicate tails or heads, respectively. Consider two agents that, without any training, predict outcomes. Agent 1 assumes p = 2/3 and models the outcome of each flip as pure chance. Agent 2 assumes that the coin is fully biased, meaning that p ∈ {0, 1}, but assigns probabilities 1/3 and 2/3 to 0 and 1.
Let Ŷ 1t+1 and Ŷ 2t+1 denote the outcomes imagined by the two agents. Despite their differing assumptions, the two agents generate identical marginal predictive distributions: P(Ŷ 1t+1 = 0) = P(Ŷ 2t+1 = 0) = 1/3. On the other hand, joint predictions greatly differ for large τ :
P(Ŷ 11 = 0, .., Ŷ 1τ = 0) = 1/3τ 1/3 = P(Ŷ 21 = 0, . . . , Ŷ 2τ = 0).
We can say that agent 1 attributes all uncertainty to aleatoric sources and agent 2, epistemic sources (although as Figure 1 alludes, there are many ways an agent can attribute sources of uncertainty). Evaluating marginal predictions cannot distinguish between the two possibilities, though for a specific prior distribution over p, one agent could be right and the other wrong. One must evaluate joint predictions to make this distinction.
2We consider this coin as an illustrative model of more complex binary outcomes, such as whether a user will click on an ad, or whether a given mortgage will default on payments.
When it comes to decision-making, this distinction can be critical (Lu et al., 2021). In a casino, under the first agent’s assumption, there is large upside and little risk on repeatedly betting on heads in the long run. However, if there is a 1/3 chance the coin will always land tails, as is the case in the second agent’s prediction, there is a ruinous risk to repeatedly betting heads. Evaluating joint predictions beyond marginals distinguishes these cases.

heading: 2.3 Computation of Kullback–Leibler loss
text: In contexts we will consider, it is not possible to compute dτKL exactly. As such, we will approximate dτKL via Monte Carlo simulation. This section provides a high level overview of our approach, we push the full details to Appendix A. Algorithm 1 outlines a basic approach to estimating dτKL with respect to a synthetic data generating process. The algorithm samples a set of environments and a training dataset for each environment. For each of these pairs, the agent is re-initialized, trained, and then tested on N independent test data τ -samples. Note that each test data τ -sample includes τ data pairs. For each test data τ -sample, the likelihood of the environment is computed exactly, but that of the agent’s belief distribution is approximated. The estimate of dτKL is taken to be the sample mean of the log-likelihood-ratios (Algorithm 2).
Algorithm 1 KL-Loss Computation 1: for j = 1, 2, . . . , J do 2: sample environment and training dataset, and train agent 3: for n = 1, 2, . . . , N do 4: sample a test data τ -sample with τ feature-label pairs 5: compute pj,n . likelihood of environment 6: compute p̂j,n . estimated likelihood of agent’s belief distribution 7: return 1JN ∑J j=1 ∑N n=1 log (pj,n/p̂j,n) . estimated log-likelihood-ratio
While the likelihood of an environment can be efficiently computed, that of an agent’s belief distribution poses a computational challenge. One approach is to estimate this likelihood via Monte Carlo simulation (Algorithm 3). This produces unbiased estimates, which can be accurate when τ is small. However, maintaining accuracy requires the number of samples to grow exponentially with τ , as discussed in Appendix A.1. To overcome this challenge, we propose a novel approach that estimates the likelihood of the agent’s beliefs via a combination of randomized partitioning and Monte Carlo simulation (Algorithm 4) (Kaski, 1998). We conjecture that, under suitable regularity conditions, this novel approach produces accurate estimates even when τ is large, but leave a formal analysis to future work. Even though Algorithm 1 is developed for a synthetic data generating process, it is straightforward to extend it to evaluate agents on real data. We outline our approach to real data in Section 5.1, with full details in Appendix A.2.

heading: 3 Benchmark agents
text: In this section we outline the baseline agents that we use to benchmark canonical approaches to uncertainty estimation in deep learning. Table 1 links to papers that introduce these agents, as well as the hyperparamters that we tuned to optimize their performance via gridsearch. In each case, we attempt to match ‘canonical’ implementations, which we open source at https://anonymous.4open.science/r/neural-testbed-B839. In addition to these agent implementations, our opensource project contains all the evaluation code to reproduce the results of this paper. Our code is written in Python and makes use of Jax internally (Bradbury et al., 2018). However, our evaluation procedure is framework agnostic, and can equally be used with any Python package including Tensorflow, Pytorch or even SKlearn. Over the course of this paper, we have made extensive use of parallel computation to facilitate large hyperparameter sweeps over many problems. Nevertheless, the overall computational cost is relatively low by modern deep learning standards and relies only on standard CPU. For reference, evaluating the mlp agent across all the problems in
our testbed and real data requires less than 3 CPU-hours. We view our opensource effort as one of the major contributions of this work. We provide clear and strong baselines, together with an objective and accessible method for assessing uncertainty estimates beyond marginal distributions.

heading: 4 The Neural Testbed
text: In this section we introduce the Neural Testbed, a system for assessing and comparing agent performance. The Testbed implements synthetic data generating processes and streamlines the process of sampling data, training agents, and evaluating test performance by estimating KL-loss for marginal and high-order joint predictions. Since independent data can be generated for each execution, the Testbed can drive insight and multiple iterations of algorithm development without risk of overfitting to a fixed dataset. We begin by describing the simple generative model based around a random 2-layer MLP. We then apply this testbed to evaluate a comprehensive set of benchmark agents.

heading: 4.1 Synthetic data generating processes
text: By data generating process, we do not mean only the conditional distribution of data pairs (Xt, Yt+1)|E but also the distribution of the environment E . The Testbed considers 2- dimensional inputs and binary classification problems, although the generating processes can be easily extended to any input dimension and number of classes. The Testbed offers three data generating processes distinguished by a “temperature” setting, which signifies the signal-to-noise ratio (SNR) regime of the generated data. The agent can be tuned separately for each setting. This reflects prior knowledge of whether the agent is operating in a high SNR regime such as image recognition or a low SNR regime such as weather forecasting. To generate a model, the Testbed samples a 2-hidden-layer ReLU MLP with 2 output units, which are scaled by 1/temperature and passed through a softmax function to produce class probabilities. The MLP is sampled according to standard Xavier initialization (Glorot & Bengio, 2010), with the exception that biases in the first layer are drawn from N(0, 12 ). The inputs (Xt : t = 0, 1, . . .) are drawn i.i.d. from N(0, I). The agent is provided with the data generating process as prior knowledge. In Section 2.1, we described KL-loss as a metric for evaluating performance of an agent. The Neural Testbed estimates KL-loss, with τ ∈ {1, 100}, for three temperature settings and several training dataset sizes. For each value of τ , the KL-losses are averaged to produce an aggregate performance measure. Further details concerning data generation and agent evaluation are offered in Appendix A.

heading: 4.2 Performance in marginal predictions
text: We begin our evaluation of benchmark approaches to Bayesian deep learning in marginal predictions (τ = 1). This setting has been the main focus of the Bayesian deep learning literature. Despite this focus, it is surprising to see in Figure 2 that none of the benchmark methods significantly outperform a well-tuned MLP baseline according to d1KL. Of course, there are many other metrics one might consider, but in this fundamental metric of prediction quality, the mlp agent presents a baseline that is difficult to outperform.
One of the keys to this result is that all of the agents are able to tune their hyperparameters, such as L2 weight decay, to the SNR regime and number of training points. This matches the way deep learning systems are typically implemented in practice, with extensive hyperparameter tuning on validation data. This methodology has led many practitioners to doubt the usefulness of automatic tuning procedures such as bootstrap sampling (Nixon et al., 2020). In Figure 3, we compare the performance of an ensemble+ agent that uses bootstrapping with and without the ability to tune the hyperparameters per problem setting. We see that bootstrap sampling is beneficial when the agent is expected to work robustly over a wide range of problem settings. However, the benefits are no longer apparent when the agent is allowed to tune its hyperparameters to individual tasks.

heading: 4.3 Performance beyond marginals
text: One of the key contributions of this paper is to evaluate predictive distributions beyond marginals. In Figure 2, the red bars show the results of benchmark agents evaluated on joint predictive distributions with τ = 100. Unlike when evaluating on marginal predictions, where no method significantly outperforms a well-tuned MLP, the potential benefits afforded by Bayesian deep learning become clear when examining higher-order predictive distributions. Our results refute prior works’ claims that examining dτKL beyond marginals provides little new information (Wang et al., 2021). Figure 2 shows that sgmcmc is the top-performing agent overall. This should be reassuring to the Bayesian deep learning community and beyond. In the limit of large compute this agent should recover the ‘gold-standard’ of Bayesian inference, and it does indeed perform best (Welling & Teh, 2011). However, some of the most popular approaches in this field (ensemble, dropout) do not actually provide good approximations to the predictive distribution in τ = 100. In fact, we see that even though Bayesian purists may deride ensemble+ and hypermodels as ‘not really Bayesian’, these methods actually provide much better approximations to the Bayesian posterior than ‘fully Bayesian’ VI approaches like bbb. We
note too that while sgmcmc performs best, it also requires orders of magnitude more computation than competitive methods even in this toy setting (see Appendix C.2). As we scale to more complex environments, it may therefore be worthwhile to consider alternative approaches to approximate Bayesian inference. For insight into where our top agents are able to outperform, we compare ensemble and ensemble+ under the medium SNR regime in Figures 4 and 5. These methods are identical, except for the addition of a randomized prior function (Osband et al., 2018). Figure 4 shows that, although these methods perform similarly in the quality of their marginal predictions (τ = 1), the addition of a prior function greatly improves the quality of joint predictive distributions (τ = 100) in the low data regime. Figure 5 provides additional intuition into how the randomized prior functions are able to drive improved performance. Figure 5a shows a sampled generative model from our Testbed, with the training data shown in red and blue circles. Figure 5b shows the mean predictions and 4 randomly sampled ensemble members from each agent (top=ensemble, bottom=ensemble+). We see that, although the agents mostly agree in their mean predictions, ensemble+ produces more diverse sampled outcomes enabled by the addition of randomized prior functions. In contrast, ensemble produces similar samples, which may explain why its performance is close to baseline mlp.

heading: 5 Performance on real data
text: Section 4 provides a simple, sanitized testbed for clear insight to the efficacy of Bayesian deep learning techniques. However, most deep learning research is not driven by these sorts of synthetic generative models, but the ultimate goal of performing well on real datasets. In this section, we apply the same benchmark agents to a selection of small challenge datasets. We find that, on average, tuning agents for the synthetic problems leads to better performance on real data. We also find that, just as the synthetic testbed, agents that perform similarly in marginal predictions may be distinguished in the quality of their joint predictions.

heading: 5.1 Datasets
text: We focus on 10 benchmark datasets (3 feature-based, 7 image from pixels) drawn from the literature including Iris, MNIST, and CIFAR-10 (TFD). This collection is not intended to be comprehensive, or to include the most challenging large-scale problems, but instead to represent some canonical real-world data that might reasonably be addressed with the MLP models of Section 4.1. We apply a basic pre-processing step to each dataset, normalizing input features and flattening observations. We push full details to Appendix D.1. To assess performance in real datasets, we follow a similar procedure as Algorithm 1. The only difference is that since it is impossible to compute the likelihood of environment for real datasets, we compute the negative log-likelihood (NLL) rather than dτKL. Appendix A.2 provides further details. Note that NLL and dτKL are equivalent for agent comparison since they differ by a constant (see Equation 1). Furthermore, to allow for more direct comparison with the synthetic testbed, we also consider variants of each dataset where the number of training pairs is limited to less than the ‘full’ dataset size.

heading: 5.2 Synthetic data is predictive of real data
text: Recall that Figure 2 compares performance across an array of agents, assessed using our synthetic data generating process. Each agent’s hyperparameters were tuned by first enumerating a list of plausibly near-optimal choices and selecting the one that optimizes performance. Each of our real-world datasets can be viewed as generated by an environment sampled from an alternative data generating process. A natural question is whether performance on the synthetic data correlates with performance on the real-world data. The table of Figure 6a displays results pertaining to each of our agents. For each agent, performance for each candidate hyperparameter setting was assessed on synthetic and real data, and the correlation across these pairs is reported. The left and right columns restrict attention to datasets with low and high volumes of training data, respectively. If a correlation were equal to 1, the hyperparameter setting that optimizes agent performance on real data would be identical to that on synthetic data. It is reassuring that the correlations are high, reflecting a strong degree of alignment, with the exception of bbb in low data regime, for which there appear to be pathological outcomes distorting performance for small training sets. The values in parentheses express 5th and 95th percentile confidence bounds, measured via the statistical bootstrap. Figure 6b plots performance on real versus synthetic data for the high data regime. Each data point represents one agent-hyperparameter combination. If the correlation were equal to 1, the combination that performs best on the synthetic data would also perform best on the real data. It is reassuring that the correlation is large, and the confidence interval between the 5th and 95th percentiles small. Agent-hyperparameter combinations that perform better on the testbed tend to perform better on real data as well.

heading: 5.3 Higher order predictions and informative priors
text: Our synthetic testbed can be helpful in driving innovations that carry over to real data. Section 5.2 indicated that performance on the Testbed is correlated with that on realworld data. We now repeat the observation from Figure 4 on real data; additive prior functions can significantly improve the accuracy of joint predictive distributions generated by ensembles. We show this by comparing the performance of ensemble+ with different forms of prior functions on benchmark datasets. We evaluate an ensemble with no prior function (none), a random MLP prior (MLP), and a random linear function of a 2-dimensional latent representation as the prior, trained via variational autoencoder (VAE) (Kingma & Welling, 2014). We provide full details in Appendix D.3. Figure 7 plots the improvement in NLL for the ensemble agent relative to a baseline MLP (lower is better), and breaks out the result for datasets=MNIST,Iris and τ = 1, 100. We
can see that the results for Iris mirror our synthetic data almost exactly. The results for MNIST share some qualitative insights, but also reveal some important differences. For Iris τ = 1 none of the methods outperform the MLP baseline, but for τ = 100 we see significant benefits to an additive MLP prior in the low data regime. For MNIST τ = 1 we actually see benefits to ensembles, even without prior functions and even in the high data regime. This reveals some aspects of this real data that are not captured by our synthetic model, where we did not see this behaviour. For τ = 100 the random MLP prior gives a slight advantage, but the effect is much less pronounced. We hypothesize this is because, unlike the testbed, the MLP prior is not well-matched to the input image data. However, the VAE prior is able to provide significant benefit in the low data regime.3 These benefits also carry over to Iris, even where random MLPs already provided signficant value. Designing architectures that offer useful priors for learning agents is an exciting area for future work.

heading: 6 Conclusion
text: This paper highlights the need to evaluate predictive distributions beyond marginals. In addition to this conceptual contribution, we develop a suite of practical computational tools that can evaluate diverse approaches to uncertainty estimation. Together with these tools, we provide a neural-network-based data generating process that facilitates research and iteration beyond a small set of challenge datasets. We package these together as The Neural Testbed, including a variety of baseline agent implementations. We believe that this represents an exciting and valuable new benchmark for Bayesian deep learning and beyond. We have already used this testbed to generate several new insights in this paper. We have shown many popular Bayesian deep learning approaches perform similarly in marginal predictions but quite differently in joint predictions. We reveal the importance of bootstrapping for parameter robustness, and also help reconcile the observed lack of improvement when tuned to specific datasets. We have shown that these insights from synthetic data can carry over to real datasets; that performance in these settings is correlated, that agents with similar marginal predictions can be distinguished by their joint predictions, and that suitable prior functions can play an important role in driving good performance. The results in this paper are in some sense preliminary. The grand challenge for Bayesian deep learning is to provide effective uncertainty estimates in large, rich datasets. While we have demonstrated benefits to predictive evaluation beyond marginals only in the ‘low data’ regime and small-scale problems, we believe that they will extend more broadly to situations where new test inputs appear novel relative to training data. As such, we believe our core insights should carry over to the related problems of nonstationarity and covariate shift that plague modern deep learning systems. As an agent takes on more and more complex tasks, it will continue to run into new and unfamiliar settings and uncertain outcomes; as such, effective predictive distributions will be more important than ever.
3We hypothesize that appropriately initialized convnet architectures may be able to leverage image structure as noted in prior work (Ulyanov et al., 2018).

heading: A Testbed Pseudocode
text: We present the testbed pseudocode in this section. Specifically, Algorithm 2 is the pseudocode for our neural testbed, and Algorithm 3 and Algorithm 4 are two different approaches to estimate the likelihood of a test data τ -sample conditioned on an agent’s belief. Algorithm 3 is based on the standard Monte-Carlo estimation, while Algorithm 4 adopts a random partitioning approach. The presented testbed pseudocode works for any prior P(E ∈ ·) over the environment and any input distribution PX , including the ones described in Section 4.1. We also release full code and implementations at https://anonymous.4open.science/r/neural-testbed-B839. In addition to presenting the testbed pseudocode, we also discuss some core technical issues in the neural testbed design. Specifically, Appendix A.1 discusses how to estimate the likelihood of an agent’s belief distribution; Appendix A.2 discusses how to extend the testbed to agent evaluation on real data; finally, Appendix A.3 explains our choices of experiment parameters.
Algorithm 2 Neural Testbed Require: the testbed requires the following inputs
1. prior distribution over the environment P(E ∈ ·), input distribution PX 2. agent fθ 3. number of training data T , test distribution order τ 4. number of sampled problems J , number of test data samples N 5. parameters for agent likelihood estimation, as is specified in Algorithm 3 and 4
for j = 1, 2, . . . , J do Step 1: sample environment and training data
1. sample environment E ∼ P(E ∈ ·) 2. sample T inputs X0, X1, . . . , XT−1 i.i.d. from PX 3. sample the training labels Y1, . . . , YT conditionally i.i.d. as
Yt+1 ∼ P (Y ∈ ·|E , X = Xt) ∀t = 0, 1, . . . , T − 1
4. choose the training dataset as DT = {(Xt, Yt+1) , t = 0, . . . , T − 1} Step 2: train agent
train agent fθT based on training dataset DT Step 3: compute likelihoods
for n = 1, 2, . . . , N do 1. sample X(n)T , . . . , X (n) T+τ−1 i.i.d. from PX
2. generate Y (n)T+1, . . . , Y (n) T+τ conditionally independently as
Y (n) t+1 ∼ P
( Y ∈ · ∣∣∣E , X = X(n)t ) ∀t = T, T + 1, . . . , T + τ − 1 3. compute the likelihood under the environment E as
pj,n = P ( Y (n) T+1:T+τ ∣∣∣E , X(n)T :T+τ−1) = ∏T+τ−1t=T Pr(Y (n)t+1∣∣∣E , X(n)t ) 4. estimate the likelihood conditioned on the agent’s belief
p̂j,n ≈ P ( ŶT+1:T+τ = Y (n)T+1:T+τ ∣∣∣θT , X(n)T :T+τ−1, Y (n)T+1:T+τ) , based on Algorithm 3 or 4 with test data τ -sample ( X
(n) T :T+τ−1, Y (n) T+1:T+τ
) .
return 1JN ∑J j=1 ∑N n=1 log (pj,n/p̂j,n)
Algorithm 3 Monte Carlo Estimation of Likelihood of Agent’s Belief

heading: Require:
text: 1. trained agent fθT and number of Monte Carlo samples M 2. test data τ -sample (XT :T+τ−1, YT+1:T+τ )
Step 1: sample M models Ê1, . . . , ÊM conditionally i.i.d. from P ( Ê ∈ · ∣∣∣fθT ) Step 2: estimate p̂ as
p̂ = 1 M M∑ m=1 P ( ŶT+1:T+τ = YT+1:T+τ ∣∣∣Êm, XT :T+τ−1, YT+1:T+τ) return p̂
Algorithm 4 Estimation of Likelihood of Agent’s Belief via Random Partitioning

heading: Require:
text: 1. trained agent fθT 2. number of Monte Carlo samples M 3. number of hyperplanes d 4. test data τ -sample (XT :T+τ−1, YT+1:T+τ )
Step 1: sample M models Ê1, . . . , ÊM conditionally i.i.d. from P(Ê ∈ ·|fθT ); for each model m = 1, . . . ,M , class k, and t = T, . . . , T + τ − 1, define
pm,t,k = P(Ŷ (m)t+1 = k| Êm, Xt),
and `m,t,k = Φ−1 (pm,t,k), where Φ(·) is the CDF of the standard normal function. For each model m, define a vector
`m = [`m,T,1, `m,T,2, . . . , `m,T+τ−1,K ] ∈ <Kτ
Step 2: sample a d × (Kτ) matrix A and a d-dimensional vector b, with each element/component sampled i.i.d. from N(0, 1). For each m = 1, . . . ,M , compute
ψm = 1 [A`m + b ≥ 0] ∈ {0, 1}d.
Step 3: partition the sampled models, with each cell indexed by ψ ∈ {0, 1}d and defined by Mψ = {m : ψm = ψ} and assign a probability to each cell:
qψ = |Mψ| M
Step 4: ∀ψ ∈ {0, 1}d and ∀t = T, T + 1, . . . , T + τ − 1, estimate the probability of predicting Ŷt+1 = k conditioned on the cell:
pψ,t,k = { 1 |Mψ| ∑ m∈Mψ pm,t,k if |Mψ| > 0
1 if |Mψ| = 0
Step 5: estimate Pr(Ŷt+1:T+τ = Yt+1:T+τ |θT , Xt:T+τ−1, Yt+1:T+τ ) as
p̂ = ∑
ψ∈{0,1}d qψ T+τ−1∏ t=T pψ,t,Yt+1
return p̂

heading: A.1 Estimating Likelihood of Agent’s Belief Distribution
text: We have presented two algorithms to estimate the likelihood of a test data τ -sample conditioned on a trained agent: Algorithm 3 is based on the standard Monte Carlo estimation, while Algorithm 4 adopts an approach combining random partitioning and Monte Carlo estimation. In this subsection, we briefly discuss the pros and cons between these two algorithms, and provide some general guidelines on how to choose between them. Algorithm 3 produces unbiased estimates of the likelihoods, which is usually accurate when τ is small (e.g. for τ ≤ 10). However, maintaining accuracy might require the number of samples M to grow exponentially with τ . The following is an illustrative example. Example 1 (Uniform belief over deterministic models): Consider a scenario where the number of class labels is K = 2. We say a model Ê is deterministic if for any feature vector Xt,
P(Ŷt+1 = 1 | Ê , Xt) ∈ {0, 1}.
Obviously, for any test data τ -sample (XT :T+τ−1, YT+1:T+τ ) with YT+1:T+τ ∈ {0, 1}τ , under a deterministic model Ê , we have
P ( ŶT+1:T+τ = YT+1:T+τ ∣∣∣ Ê , XT :T+τ−1, YT+1:T+τ) ∈ {0, 1}. When restricted to the inputs XT :T+τ−1, there are 2τ distinguishable deterministic models. Assume the agent’s belief distribution is uniform over these 2τ distinguishable deterministic models, then for any YT+1:T+τ ∈ {0, 1}τ , the likelihood of the agent’s belief distribution is
P ( ŶT+1:T+τ = YT+1:T+τ ∣∣∣ θT , XT :T+τ−1, YT+1:T+τ) = 2−τ . Now let’s consider Algorithm 3. When a model Êm is sampled from the agent’s belief distribution, with probability 2−τ ,
P ( ŶT+1:T+τ = YT+1:T+τ ∣∣∣ Êm, XT :T+τ−1, YT+1:T+τ) = 1, and with probability 1− 2−τ ,
P ( ŶT+1:T+τ = YT+1:T+τ ∣∣∣ Êm, XT :T+τ−1, YT+1:T+τ) = 0. Consequently, in expectation, we need the number of Monte Carlo samples M = Ω(2τ ) to ensure that the estimate p̂ returned by Algorithm 3 is non-zero. To overcome this challenge, we also propose a novel approach to estimate the likelihood of agent’s belief via a combination of randomized partitioning and Monte Carlo simulation, as is presented in Algorithm 4. This approach proceeds as follows. First, M models are sampled from the agent’s belief distribution. For each sampled model, each test data input Xt, and each class label k, a predictive probability pm,t,k and its probit `m,t,k = Φ−1(pm,t,k) are computed, where Φ(·) is the CDF of the standard normal distribution. For each sampled model, we also stack its probits into a probit vector `m ∈ <Kτ . Then, d random hyperplanes are sampled and used to partition <Kτ into 2d cells. Stacked probit vectors place models in cells. Predictive distributions of models in each cell are averaged, and the likelihood is calculated based on these averages, with each cell weighted according to the number of models it contains. The Neural Testbed applies Algorithm 4 with 2d M . Hence, some cells are assigned many models. We conjecture that, under suitable regularity conditions, models assigned to the same cell tend to generate similar predictions. If this is the case, this algorithm produces accurate estimates even when τ is large. We leave a formal analysis to future work. Finally, we briefly discuss how to choose between Algorithm 3 and Algorithm 4. As a rule of thumb, we recommend to choose Algorithm 3 for τ < 10 and Algorithm 4 with the number of hyperplanes d between 5 and 10 for τ ≥ 10.

heading: A.2 Agent Evaluation on Real Data
text: Algorithm 2 (and its simplified version Algorithm 1) is developed for a synthetic data generating processes. We now discuss how to extend it to agent evaluation on real data. Consider a scenario with J real datasets, and each dataset is further partitioned into a training dataset and a test dataset. The main difference between this scenario and a synthetic data generating process is that we cannot compute the likelihood of environment for real data. Thus, we compute the cross-entropy loss instead (see Equation 1). The computational approach is similar to Algorithm 1: for each real dataset, we use its training dataset to train an agent. Then, we sample N test data τ -samples from the test dataset, and estimate the likelihoods of the agent’s belief distribution. The estimate of the cross-entropy loss is taken to be the sample mean of the negative log-likelihoods. Note that when ranking agents, the cross-entropy loss and dτKL will lead to the same order of agents, since these two losses differ by a constant independent of the agent (see Equation 1).

heading: A.3 Choices of Experiment Parameters
text: To apply Algorithm 2, we need to specify an input distribution PX and a prior distribution on the environment P(E ∈ ·). Recall from Section 4.1 that we consider binary classification problems with input dimension 2. We choose PX = N(0, I), and we consider three environment priors distinguished by a temperature parameter that controls the signal-to-noise ratio (SNR) regime. We sweep over temperatures in {0.01, 0.1, 0.5}. The prior distribution P(E ∈ ·) is induced by a distribution over MLPs with 2 hidden layers and ReLU activation. The MLP is distributed according to standard Xavier initialization, except that biases in the first layer are drawn from N(0, 12 ). The MLP outputs two units, which are divided by the temperature parameter and passed through the softmax function to produce class probabilities. The implementation of this generative model is in our open source code under the path /generative/factories.py. We now describe the other parameters we use in the Testbed. In Algorithm 2, we pick the order of predictive distributions τ ∈ {1, 100}, training dataset size T ∈ {1, 3, 10, 30, 100, 300, 1000}, number of sampled problems J = 10, and number of testing data τ -samples N = 1000. We apply Algorithm 3 for evaluation of d1KL and Algorithm 4 for evaluation of d100KL . In both Algorithms 3 and 4, we sample M = 1000 models from the agent. In Algorithm 4, we set the number of hyperplanes d = 7. The specification of the testbed parameters is in our open soucre code under the path /leaderboard/sweep.py. On real datasets, we apply the same τ ∈ {1, 100}, N = 1000, and M = 1000. We set the number of hyperplanes d = 10 in Algorithm 4.

heading: B Agents
text: In this section, we describe the benchmark agents in Section 3 and the choice of various hyperparameters used in the implementation of these agents. The list of agents include MLP, ensemble, dropout, Bayes by backprop, stochastic Langevin MCMC, ensemble+ and hypermodel. We will also include other agents such as KNN, random forest, and deep kernel, but the performance of these agents was worse than the other benchmark agents, so we chose not to include them in the comparison in Section 4. In each case, we attempt to match the “canonical” implementation. The complete implementation of these agents including the hyperparameter sweeps used for the Testbed are available at https://anonymous.4open.science/r/neural-testbed-B839. We make use of the Epistemic Neural Networks notation from (Osband et al., 2021) in our code. We set the default hyperparameters of each agent to be the ones that minimize the aggregated KL score daggKL = d1KL + d100KL/100.

heading: B.1 MLP
text: The mlp agent learns a 2-layer MLP with 50 hidden units in each layer by minimizing the cross-entropy loss with L2 weight regularization. The L2 weight decay scale is chosen either to be λ 1T or λ d √ β
T , where d is the input dimension, β is the temperature of the generative process and T is the size of the training dataset. We sweep over λ ∈ {10−4, 10−3, 10−2, 10−1, 1, 10, 100}. We implement the MLP agent as a special case of a deep ensemble (B.2). The implementation and hyperparameter sweeps for the mlp agent can be found in our open source code, as a special case of the ensemble agent, under the path /agents/factories/ensemble.py.

heading: B.2 Ensemble
text: We implement the basic “deep ensembles” approach for posterior approximation (Lakshminarayanan et al., 2017). The ensemble agent learns an ensemble of MLPs by minimizing the cross-entropy loss with L2 weight regularization. The only difference between the ensemble members is their independently initialized network weights. We chose the L2 weight scale to be either λ 1MT or λ d √ β
MT , where M is the ensemble size, d is the input dimension, β is the temperature of the generative process, and T is the size of the training dataset. We sweep over ensemble size M ∈ {1, 3, 10, 30, 100} and λ ∈ {10−4, 10−3, 10−2, 10−1, 1, 10, 100}. We find that larger ensembles work better, but this effect is within margin of error after 10 elements. The implementation and hyperparameter sweeps for the ensemble agent can be found in our open source code under the path /agents/factories/ensemble.py.

heading: B.3 Dropout
text: We follow Gal & Ghahramani (2016) to build a droput agent for posterior approximation. The agent applies dropout on each layer of a fully connected MLP with ReLU activation and optimizes the network using the cross-entropy loss combined with L2 weight decay. The L2 weight decay scale is chosen to be either l 2 2T (1− pdrop) or d √ βl
T where pdrop is the dropping probability, d is the input dimension, β is the temperature of the data generating process, and T is the size of the training dataset. We sweep over dropout rate pdrop ∈ {0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8}, length scale (used for L2 weight decay) l ∈ {0.01, 0.1, 0.3, 1, 3, 10}, number of neural network layers ∈ {2, 3}, and hidden layer size ∈ {50, 100}. The implementation and hyperparameter sweeps for the dropout agent can be found in our open source code under the path /agents/factories/dropout.py.

heading: B.4 Bayes-by-backprop
text: We follow Blundell et al. (2015) to build a bbb agent for posterior approximation. We consider a scale mixture of two zero-mean Gaussian densities as the prior. The Gaussian densities have standard deviations σ1 and σ2, and they are mixed with probabilities p and 1− p,
respectively. We sweep over σ1 ∈ {1, 2, 4}, σ2 ∈ {0.25, 0.5, 0.75}, p ∈ {0, 0.25, 0.5, 0.75, 1}, learning rate ∈ {10−3, 3× 10−3}, number of training steps ∈ {500, 1000, 10000}, number of neural network layers ∈ {2, 3}, hidden layer size ∈ {50, 100}, and the ratio of the complexity cost to the likelihood cost ∈ {1, d √ β}, where d is the input dimension and β is the temperature of the data generating process. The implementation and hyperparameter sweeps for the bbb agent can be found in our open source code under the path /agents/factories/bbb.py.
B.5 Stochastic gradient Langevin dynamics
We follow Welling & Teh (2011) to implement a sgmcmc agent using stochastic gradient Langevin dynamics (SGLD). We consider two versions of SGLD, one with momentum and other without the momentum. We consider independent Gaussian prior on the neural network parameters where the prior variance is set to be
σ2 = λ T dβ ,
where λ is a hyperparameter that is swept over {0.01, 0.1, 0.5, 1}, d is the input dimension, β is the temperature of the data generating process, and T is the size of the training dataset. We consider a constant learning rate that is swept over {10−5, 5× 10−5, 10−4, 5× 10−4, 10−3, 5 × 10−3, 10−2}. For SGLD with momentum, the momentum decay term is always set to be 0.9. The number of training batches is 5 × 105 with burn-in time of 105 training batches. We save a model every 1000 steps after the burn-in time and use these models as an ensemble during the evaluation. The implementation and hyperparameter sweeps for the sgmcmc agent can be found in our open source code under the path /agents/ factories/sgmcmc.py.

heading: B.6 Ensemble+
text: We implement the ensemble+ agent using deep ensembles with randomized prior functions (Osband et al., 2018) and bootstrap sampling (Osband & Van Roy, 2015). Similar to the vanilla ensemble agent in Section B.2, we consider L2 weight scale to be either λ 1MT or λ d √ β
MT . We sweep over ensemble size M ∈ {1, 3, 10, 30, 100} and λ ∈ {10−4, 10−3, 10−2, 10−1, 1, 10, 100}. The randomized prior functions are sampled exactly from the data generating process, and we sweep over prior scaling ∈ {0, √ β, 1}. In addition, we sweep over bootstrap type (none, exponential, bernoulli). We find that the addition of randomized prior functions is crucial for improvement in performance over vanilla deep ensembles in terms of the quality of joint predictions. We also find that bootstrap sampling improves agent robustness, although the advantage is less apparent when one is allowed to tune the L2 weight decay for each task (see Figure 3). The implementation and hyperparameter sweeps for the ensemble+ agent can be found in our open source code under the path /agents/factories/ensemble_plus.py.

heading: B.7 Hypermodel
text: We follow Dwaracherla et al. (2020) to build a hypermodel agent for posterior approximation. We consider a linear hypermodel over a 2-layer MLP base model. We sweep over index dimension ∈ {1, 3, 5, 7}. The L2 weight decay is chosen to be either λ 1T or λ d √ β T with λ ∈ {0.1, 0.3, 1, 3, 10}, where d is the input dimension, β is the temperature of the data generating process, and T is the size of the training dataset. We chose three different bootstrapping methods of none, exponential, bernoulli. We use an additive prior which is a linear hypermodel prior over an MLP base model, which is similar to the generating process, with number of hidden layers in {1, 2}, 10 hidden units in each layer, and prior scale from {0, √ β, 1}. The implementation and hyperparameter sweeps for the hypermodel agent can be found in our open source code under the path /agents/factories/hypermodel.py.

heading: B.8 Non-parametric classifiers
text: K-nearest neighbors (k-NN) (Cover & Hart, 1967) and random forest classifiers (Friedman, 2017) are simple and cheap off-the-shelf non-parametric baselines (Murphy, 2012; Pedregosa et al., 2011). The ‘uncertainty’ in these classifiers arises merely from the fact that they produce distributions over the labels and as such we do not expect them to perform well relative to more principled approaches. Moreover, these methods have no capacity to model dτKL for τ > 1. For the knn agent we swept over the number of neighbors k ∈ {1, 5, 10, 30, 50, 100} and the weighting of the contribution of each neighbor as either uniform or based on distance. For the random forest agent we swept over the number of trees in the forest {10, 100, 1000}, and the splitting criterion which was either the Gini impurity coefficient or the information gain. To prevent infinite values in the KL we truncate the probabilities produced by these classifiers to be in the interval [0.01, 0.99]. The implementation and hyperparameter sweeps for the knn and random forest agents can be found in our open source code under the paths /agents/factories/knn.py and /agents/factories/random_forest.py.
B.9 Gaussian process with learned kernel
A neural network takes input Xt ∈ X and produces output Zt+1 = Wφθ(Xt) + b ∈ RK , where W ∈ RK×m is a matrix, b ∈ RK is a bias vector, and φθ : X → Rm is the output of the penultimate layer of the neural network. In the case of classification the output Zt+1 corresponds to the logits over the class labels, i.e., Ŷt+1 ∝ exp(Zt+1). The neural network should learn a function that maps the input into a space where the classes are linearly distinguishable. In other words, the mapping that the neural network is learning can be considered a form of kernel (Schölkopf & Smola, 2018), where the kernel function k : X ×X → R is simply k(X,X ′) = φθ(X)>φθ(X ′). With this in mind, we can take a trained neural network and consider the learned mapping to be the kernel in a Gaussian process (GP) (Rasmussen, 2003), from which we can obtain approximate uncertainty estimates. Concretely, let Φ0:T−1 ∈ RT×m be the matrix corresponding to the φθ(Xt), t = 0, . . . , T −1, vectors stacked row-wise and let ΦT :T+τ−1 ∈ Rτ×m denote the same quantity for the test set. Fix index i ∈ {0, . . . ,K − 1} to be a particular class index. A GP models the joint distribution over the dataset to be a multi-variate Gaussian, i.e.,[
Z (i) 1:T
Z (i) T+1:T+τ
] ∼ N ([ µ (i) 1:T
µ (i) T+1:T+τ
] , [ σ2I + Φ0:T−1Φ>0:T−1 ΦT :T+τ−1Φ>0:T−1
Φ0:T−1Φ>T :T+τ−1 ΦT :T+τ−1Φ>T :T+τ−1 ]) where σ > 0 models the noise in the training data measurement and µ(i)1:T , µ (i) T+1:T+τ are the means under the GP. The conditional distribution is given by
P (Z(i)T+1:T+τ | Z (i) 1:T , X0:T+τ−1) = N
( µ
(i) T+1:T+τ |1:T ,ΣT+1:T+τ |1:T ) where ΣT+1:T+τ |1:T = ΦT :T+τ−1Φ>T :T+τ−1 − ΦT :T+τ−1Φ>0:T−1(σ2I + Φ0:T−1Φ>0:T−1)−1Φ0:T−1Φ>T :T+τ−1.
and rather than use the GP to compute µ(i)T+1:T+τ |0:T (which would not be possible since we do not oberve the true logits) we just take it to be the output of the neural network when evaluated on the test dataset. The matrix being inverted in the expression for ΣT+1:T+τ |0:T has dimension T × T , which may be quite large. We use the Sherman-Morrison-Woodbury identity to rewrite it as follows (Woodbury, 1950)
ΣT+1:T+τ |0:T = ΦT :T+τ−1(I − Φ>0:T−1(σ2I + Φ0:T−1Φ>0:T−1)−1Φ0:T−1)Φ>T :T+τ−1 = σ2ΦT :T+τ−1(σ2I + Φ>0:T−1Φ0:T−1)−1Φ>T :T+τ−1,
which instead involves the inverse of an m×m matrix, which may be much smaller. If we perform a Cholesky factorization of positive definite matrix (σ2I + Φ>0:T−1Φ0:T−1) = LL> then the samples for all logits simultaneously can be drawn by first sampling ζ ∈ Rm×K , with each entry drawn IID from N (0, 1), then forming
ŶT+1:T+τ ∝ exp(µT+1:T+τ |1:T + σΦT :T+τ−1L−>ζ). The implementation and hyperparameter sweeps for the deep kernel agent can be found in our open source code under the path /agents/factories/deep_kernel.py.

heading: B.10 Other agents
text: In our paper we have made a concerted effort to include representative and canonical agents across different families of Bayesian deep learning and adjacent research. In addition to these implementations, we performed extensive tuning to make sure that each agent was given a fair shot. However, with the proliferation of research in this area, it was not possible for us to evaluate all competiting approaches. We hope that, by opensourcing the Neural Testbed, we can allow researchers in the field to easily assess and compare their agents to these baselines. For example, we highlight a few recent pieces of research that might be interesting to evaluate in our setting. Of course, there are many more methods to compare and benchmark. We leave this open as an exciting area for future research.
• Neural Tangent Kernel Prior Functions (He et al., 2020). Proposes a specific type of prior function in ensemble+ inspired by connections to the neural tangent kernel.
• Functional Variational Bayesian Neural Networks (Sun et al., 2019). Applies variational inference directly to the function outputs, rather than weights like bbb.
• Variational normalizing flows (Rezende & Mohamed, 2015). Applies variational inference over a more expressive family than bbb.
• No U-Turn Sampler (Hoffman et al., 2014). Another approach to sgmcmc that attempts to compute the posterior directly, computational costs can grow large.

heading: C Testbed results
text: In this section, we provide the complete results of the performance of benchmark agents on the Testbed, broken down by the temperature setting, which controls the SNR, and the size of the training dataset. We select the best performing agent within each agent family and plot d1KL and d100KL with the performance of an MLP agent as a reference. We also provide a plot comparing the training time of different agents.

heading: C.1 Performance breakdown
text: Figures 8 and 9 show the KL estimates evaluated on τ = 1 and τ = 100, respectively. For each agent, for each SNR regime, for each number of training points we plot the average KL estimate from the Testbed. In each plot, we include the “baseline” mlp agent as a black dashed line to allow for easy comparison across agents. A detailed description of these benchmark agents can be found in Appendix B.

heading: C.2 Training time
text: Figure 10 shows a plot comparing the d100KL and training time of different agents normalized with that of an MLP. We can see that sgmcmc agent has the best performance, but at the cost of more training time (computation). Both ensemble+ and hypermodel agents have similar performance as sgmcmc with lower training time. We trained our agents on CPU only systems.

heading: D Real data
text: This section provides supplementary details regarding the experiments in Section 5. As before, we include full implementation and source code at https://anonymous.4open. science/r/neural-testbed-B839.

heading: D.1 Datasets
text: Table 2 outlines the datasets included in our experiments. Unlike to the synthetic testbed, which evaluates agents over a range of SNR regimes, these datasets are generally all high
SNR regime. We can see this since the top-performing agents in the literature are able to obtain high levels of classification accuracy on held out data; something that is impossible if the underlying system has high levels of noise.
Each of these datasets is provided with a canonical training/test set of specific sizes. In order to examine performance in different data regimes we augment the default settings of Table 2 by also examining the performance of agents on these datasets with reduced training data. In a way that mirrors the testbed sweep of Section 4.1, we also look at settings where the training data is restricted to T = 1, 10, 100, 1000, 10000 data points respectively.

heading: D.2 Correlation
text: Figure 6 breaks down the correlation in performance between testbeds and real data. For the purposes of Table 6a we say that T = 1, 10 is the ‘low data’ regime, and the maximum training dataset size is the ‘high data’ regime. Our results show that, for each agent, for each data regime, performance of hyperparameters is correlated across settings. One concern might be that while performance on real data overall is highly correlated, that this might not necessarily be the case for any individual dataset. Or, alternatively, that this correlation is driven by extremely strong relationships in one dataset that are not present in others. Figure 11 shows that this is not the case. In fact, for each of the datasets considered we have strong and positive correlation over agent-hyperparameter pairs. This gives us confidence that the results of Figure 6b are robust not only to choice of agent, but also to some reasonable choice of datasets.

heading: D.3 Prior functions
text: We consider two different forms of prior functions for ensemble+: a random MLP of the input data and a random linear function of a 2-dimensional latent trained via variational autoencoder (VAE) (Kingma & Welling, 2014). For the MLP prior, we tried both linear (MLP with no hidden layer) and MLP with hidden layers, and observed that the linear prior works better. To train the 2-dimensional latent, we considered a 2-layer (128, 64) MLP for the Gaussian encoder and a 2-layer (64, 128) MLP for the Bernoulli decoder. We trained the VAE using all unsupervised training data available for each dataset. After training the VAE for 10,000 steps, we used the output mean of the Gaussian encoder as the latent.

Review summaries:
The authors discuss whether it is sufficient to consider the marginal posterior predictive vs considering a joint posterior predictive when evaluating Bayesian deep learning approaches. Introducing a set of experiments (_The Neural Testbed_), they demonstrate that the performance of common approaches can differ greatly depending on which of these predictive distributions are evaluated. Additionally, extensive code is provided for efficient implementation and evaluation of new models. 

This work advocates evaluating predictive distributions via joint predictions (rather than the standard practice of evaluating marginal predictions) and introduces The Neural Testbed, an open-source software which includes the testing suite, along with implementations of a handful of methods in uncertainty quantification (UQ). The core evaluation metric used is KL divergence (cross-entropy loss) between the predictive distribution and the true likelihood of the data-generating process, and this work proposes an algorithm to compute this metric with joint distributions. The empirical evaluations compare numerous standard UQ methods with the Testbed, with both synthetic and real datasets.

The paper proposes a simulation based framework to evaluate different techniques proposed for uncertainty estimation of predictive models. The approach relies on simulated data to control effects such as environments where data is collected, data and model uncertainty. This control enables generating interesting insights about various techniques.    The paper also go beyond evaluating marginal posterior predictive distributions and extend their benchmarking work to joint distributions capturing sequential decisions that can be made with such models. Some of the highlights from the results: 1) Their results show that Bayesian deep learning is impactful for capturing the joint predictive distributions. 2) Priors used in ensemble+ help with diversity and therefore enable better predictive distributions. 3) Bootstrapping help with robustness of predictions if the models hyperparameters are not tuned.