1 Introduction↩︎

Over the last 25 years, there have been several strong arguments favouring a Bayesian approach to deep learning [1]–[5]. Bayesian inference for neural networks promises improved predictions, reliable uncertainty estimates, and principled model comparison, naturally supporting active learning, continual learning, and decision-making under uncertainty. The Bayesian deep learning community has designed multiple successful practical methods inspired by the Bayesian approach [3], [6]–[11] with applications ranging from astrophysics [12] to automatic diagnosis of Diabetic Retinopathy [13], click-through rate prediction in advertising [14] and modeling of fluid dynamics [15].

However, inference with modern neural networks is distinctly challenging. We wish to compute a Bayesian model average corresponding to an integral over a multi-million dimensional multi-modal posterior, with unusual topological properties like mode-connectivity [16], [17], under severe computational constraints.

There are therefore many unresolved questions about Bayesian deep learning practice. Variational procedures typically provide unimodal Gaussian approximations to the multimodal posterior. Practically successful methods such as deep ensembles [18], [19] have a natural Bayesian interpretation [5], but only represent modes of the posterior. While Stochastic MCMC [7], [20], [21] is computationally convenient, it could be providing heavily biased estimates of posterior expectations. Moreover, [22] question the quality of standard Bayes posteriors, citing results where “cold posteriors”, raised to a power \(1/T\) with \(T<1\), improve performance.

Additionally, Bayesian deep learning methods are typically evaluated on their ability to generate useful, well-calibrated predictions on held-out or out-of-distribution data. However, strong performance on benchmark problems does not imply that the algorithm accurately approximates the true Bayesian model average (BMA).

In this paper, we investigate fundamental open questions in Bayesian deep learning, using multi-chain full-batch Hamiltonian Monte Carlo [23]. HMC is a highly-efficient and well-studied Markov Chain Monte Carlo (MCMC) method that is guaranteed to asymptotically produce samples from the true posterior. However it is enormously challenging to apply HMC to modern neural networks due to its extreme computational requirements: HMC can take tens of thousands of training epochs to produce a single sample from the posterior. To address this computational challenge, we parallelize the computation over hundreds of Tensor processing unit (TPU) devices.

We argue that full-batch HMC provides the most precise tool for studying the BNN posterior to date. We are not proposing HMC as a computationally efficient method for practical applications. Rather, using our implementation of HMC we are able to explore fundamental questions about posterior geometry, the performance of BNNs, approximate inference, effect of priors and posterior temperature.

In particular, we show: (1) BNNs can achieve significant performance gains over standard training and deep ensembles; (2) a single long HMC chain can provide a comparable representation of the posterior to multiple shorter chains; (3) in contrast to recent studies, we find posterior tempering is not needed for near-optimal performance, with little evidence for a “cold posterior” effect, which we show is largely an artifact of data augmentation; (4) BMA performance is robust to the choice of prior scale, and relatively similar for diagonal Gaussian, mixture of Gaussian, and logistic priors over weights. This result highlights the importance of architecture relative to parameter priors in specifying the prior over functions. (5) While Bayesian neural networks have good performance for OOD detection, they show surprisingly poor generalization under domain shift; (6) while cheaper alternatives such as deep ensembles and SGMCMC can provide good generalization, they provide distinct predictive distributions from HMC. Notably, deep ensemble predictive distributions are similarly close to HMC as standard SGLD, and closer than standard variational inference.

We additionally show how to effectively deploy full batch HMC on modern neural networks, including insights about how to tune crucial hyperparameters for good performance, and parallelize sampling over hundreds of TPUs. Our HMC samples and implementation will be a public resource. We hope this resource will both serve as a reference in evaluating and calibrating more practical alternatives to HMC, and aid in exploring questions for a better understanding of approximate inference in Bayesian deep learning.

2 Background↩︎

Bayesian neural networks The goal of classical learning is to find a single best setting of the parameters for the model, typically through maximum-likelihood optimization. In the Bayesian framework, the learner instead infers a posterior distribution \(p(w \vert \mathcal{D})\) over the parameters \(w\) of the model after observing the data \(\mathcal{D}\). The posterior distribution is given by Bayes’ rule: \(p(w \vert \mathcal{D}) \propto p(\mathcal{D} \vert w) p(w)\), where \(p(\mathcal{D} \vert w)\) is the likelihood of \(\mathcal{D}\) given by the model with parameters \(w\), and \(p(w)\) is the prior distribution over the parameters. The predictions of the model on a new test example \(x\) are then given by the Bayesian model average (BMA) \[\label{eq:bma} \textstyle p(y \vert x, \mathcal{D}) = \int_w p(y \vert x, w) p(w \vert \mathcal{D}) dw,\tag{1}\] where \(p(y \vert x, w)\) is the predictive distribution for a given value of the parameters \(w\). This BMA is particularly compelling in Bayesian deep learning, because the posterior over parameters for a modern neural network can represent many complementary solutions to a given problem, corresponding to different settings of parameters [5]. Unfortunately, the integral in Eq. 1 cannot be evaluated in closed form for neural networks, so one must resort to approximate inference. Moreover, approximating Eq. 1 is challenging due to a high dimensional and sophisticated posterior \(p(w \vert \mathcal{D})\). For a detailed discussion of Bayesian deep learning, see e.g. [5].

Markov Chain Monte Carlo. The integral in Eq. 1 can be approximated by sampling: \(p(y \vert x, \mathcal{D}) \approx \frac{1}{M} \sum_{i=1}^M p(y \vert x, w_i)\), where \(w_i \sim p(w \vert \mathcal{D})\) are samples drawn from the posterior. MCMC methods construct a Markov chain that, if simulated for long enough, generates approximate samples from the posterior. In this work, we focus on Hamiltonian Monte Carlo [23], a method that produces asymptotically exact samples assuming access to the unnormalized posterior density \(p(\mathcal{D} \vert w) p(w)\) and its gradient.

3 Related work↩︎

The bulk of work on Bayesian deep learning has focused on scalable approximate inference methods. These methods include stochastic variational inference [3], [24]–[33], dropout [6], [34]–[36], the Laplace approximation [8], [11], [37]–[39], expectation propagation [40], and leveraging the stochastic gradient descent (SGD) trajectory, either for a deterministic approximation, or sampling as in SGLD [5], [9], [41], [42]. [43] and [44] additionally consider the role of expressive posterior approximations in variational inference.

While these (and many other) methods often provide improved predictions or uncertainty estimates, to the best of our knowledge none of these methods have been directly evaluated on their ability to match the true posterior distribution using practical architectures and datasets. Moreover, many of these methods are often designed with train-time constraints in mind, to take roughly the same amount of compute as regular SGD training. To evaluate approximate inference procedures, and explore fundamental questions in Bayesian deep learning, we attempt to construct a posterior approximation of the highest possible quality, ignoring the practicality of the method.

The Monte Carlo literature for Bayesian neural networks has mainly focused on stochastic gradient-based methods [7], [20], [21], [45]–[49] for computational efficiency reasons. These methods are fundamentally biased: (1) they omit the Metropolis-Hastings correction step, and (2) the noise from subsampling the data perturbs their stationary distribution. In particular, [50] argues that HMC is incompatible with data subsampling. Notably, [51] recently proposed a second-order stochastic gradient MCMC method that is asymptotically exact.

Since the classic work of [2], there have been a few recent attempts at using full-batch HMC in BNNs [22], [52]. These studies tend to use relatively short trajectory lengths (generally not considering a number of leapfrog steps greater than 100), and tend to focus on relatively small datasets and networks. We on the other hand experiment with practical architectures and datasets and use up to \(10^5\) leapfrog steps per iteration to ensure good mixing.

Our work is aimed at understanding the properties of true Bayesian neural networks. In a similar direction, [22] have recently explored the effect of the posterior temperature in Bayesian neural networks. We discuss their results in detail in 7, and provide our own exploration of the posterior temperature with a different result: we find that BNNs achieve strong performance at temperature \(1\) and do not require posterior tempering. Moreover, the scope of our paper extends well beyond temperature scaling, revealing for instance that while BNNs can provide strong in-domain generalization, they surprisingly suffer on the covariate shift problems that approximate inference methods are often applied towards. We also show that while deep ensembles are often treated as a non-Bayesian alternative, they in fact are providing higher fidelity approximations of the Bayesian model average than standard approximate inference procedures, as argued in [5]. We also explore several other key questions, including prior selection and posterior geometry.

4 HMC for deep neural networks↩︎

We use full-batch Hamiltonian Monte Carlo (HMC) to sample from the parameter posteriors for Bayesian neural networks. In this section, we show how to make HMC effective for modern Bayesian neural networks, discussing important details such as hyper-parameter specification. In the next sections, we use the HMC samples to explore fundamental questions about approximate inference in modern deep learning. We summarize HMC in Appendix 11 and [alg:leapfrog]. Intuitively, HMC is simulating the dynamics of a particle sliding on the plot of the density function that we are trying to sample from.¹

Implementation. To scale HMC to modern neural network architectures and for datasets like CIFAR-10 and IMDB, we parallelize the computation over 512 TPUv3 devices² [53]. We execute HMC in a single-program multiple-data (SPMD) configuration, wherein a dataset is sharded evenly over each of the devices and an identical HMC implementation is run on each device. Each device maintains a synchronized copy of the Markov chain state, where the full-batch gradients needed for leapfrog integration are computed using cross-device collectives. We release our JAX [54] implementation.

Neural network architectures. In our evaluation, following [22], we primarily focus on two architectures: ResNet-20-FRN and CNN-LSTM. ResNet-20-FRN is a residual architecture [55] of depth \(20\) with batch normalization layers [56] replaced with filter response normalization [57]. Batch normalization makes the likelihood harder to interpret by creating dependencies between training examples, whereas the outputs of FRN layers are independent across inputs. We use Swish (SiLU) activations [58]–[60] instead of ReLUs to ensure smoothness of the posterior density surface, which we found improves acceptance rates of HMC proposals without hurting the overall performance. The CNN-LSTM is a long short-term memory network [61] adapted from [22] without modifications.

Datasets and Data Augmentation. In our main evaluations we use the CIFAR image classification datasets [62] and the IMDB dataset [63] for sentiment analysis. We do not use any data augmentation, both because the random augmentations introduce stochasticity into the evaluation of the posterior log-density and its gradient, and because the expected randomly perturbed log-likelihood does not have a clean interpretation as a valid likelihood function [22]. In this section we perform ablations using ResNet-20-FRN on CIFAR-10 and CNN-LSTM on IMDB.

4.1 Trajectory length \(\tau\)↩︎

The trajectory length parameter \(\tau\) determines the length of the dynamics simulation on each HMC iteration. Effectively, it determines the correlation of subsequent samples produced by HMC. To suppress random-walk behavior and speed up mixing, we want the length of the trajectory to be relatively high. But increasing the length of the trajectory also leads to an increased computational cost: the number of evaluations of the gradient of the target density (evaluations of the gradient of the loss on the full dataset) is equal to the ratio \(\tau / \Delta\) of the trajectory length to the step size.

We suggest the following value of the trajectory length \(\tau\): \[\label{eqn:tau} \textstyle \hat{\tau} = \frac{\pi \alpha_{\text{prior}}}{2},\tag{2}\] where \(\alpha_{\text{prior}}\) is the standard deviation of the prior distribution over the parameters. If applied to a spherical Gaussian distribution, HMC with a small step size and this trajectory length will generate exact samples³. While we are interested in sampling from the posterior rather than from the spherical Gaussian prior, we argue that in large BNNs the prior tends to determine the scale of the posterior. We provide more detail and confirm this intuition empirically in the 13.

In order to test the validity of our recommended trajectory length, we perform an ablation and report the results in [fig:tau95ablation](a). As expected, longer trajectory lengths provide better performance in terms of accuracy and log-likelihood. Expected calibration error is generally low across the board. The trajectory length \(\hat{\tau}\) provides good performance in all three metrics. This result confirms that, despite the expense, when applying HMC to BNNs it is actually helpful to use tens of thousands of gradient evaluations per iteration.

4.2 Step size \(\Delta\)↩︎

The step size parameter \(\Delta\) determines the discretization step size of the Hamiltonian dynamics and consequently the number of leapfrog integrator steps. Lower step sizes lead to a better approximation of the dynamics and higher rates of proposal acceptance at the Metropolis-Hastings correction step. However, lower step sizes require more gradient evaluations per iteration to hold the trajectory length \(\tau\) constant.

Using ResNet-20-FRN on CIFAR-10, we run HMC for 50 iterations with step sizes of \(1 \cdot 10^{-5}\), \(5 \cdot 10^{-5}\), \(1 \cdot 10^{-4}\), and \(5 \cdot 10^{-4}\) respectively, ignoring the Metropolis-Hastings correction. We find the chains achieve average accept probabilities of \(72.2\%\), \(46.3\%\), \(22.2\%\), and \(12.5\%\), reflecting large drops in accept probability as step size is increased. We also observe BMA log-likelihoods of \(-0.331\), \(-0.3406\), \(-0.3407\), and \(-0.895\), indicating that higher accept rates result in higher likelihoods.

4.3 Number of HMC chains↩︎

We can improve the coverage of the posterior distribution by running multiple independent chains of HMC. Effectively, each chain is an independent run of the procedure using a different random initialization. Then, we combine the samples from the different chains. The computational requirements of running multiple chains are hence proportional to the number of chains.

We report the Bayesian model average performance as a function of the number of chains in [fig:tau95ablation](b). Holding compute budget fixed, using two or three chains is only slightly better than using one chain. This result notably shows that HMC is relatively unobstructed by energy barriers in the posterior surface that would otherwise require multiple chains to overcome. We explore this result further in 5.

5 How well does HMC mix?↩︎

The primary goal of our paper is to construct accurate samples from the posterior, and use them to understand the properties of Bayesian neural networks better. In this section we consider several diagnostics to evaluate whether our HMC sampler has converged, and discuss their implications to the posterior geometry. We consider mixing in both weight space and function space. A distribution over weights \(w\) combined with a neural network architecture \(f(x,w)\) induces a distribution over functions \(f(x)\). Ultimately, since we are using functions to make predictions, we care mostly about mixing in function space.

5.1 \(\hat{R}\) diagnostics↩︎

We apply the classic [64] “\(\hat{R}\)” potential-scale-reduction diagnostic to our HMC runs. Given two or more chains, \(\hat{R}\) estimates the ratio between the between-chain variance (i.e., the variance estimated by pooling samples from all chains) and the average within-chain variance (i.e., the variances estimated from each chain independently). The intuition is that, if the chains are stuck in isolated regions, then combining samples from multiple chains will yield greater diversity than taking samples from a single chain. For the precise mathematical definition of \(\hat{R}\), please see the 12.

We compute \(\hat{R}\) using TensorFlow Probability’s implementation⁴ [65] for both the weights and the test-set softmax predictions on CIFAR-10 with ResNet-20-FRN and on IMDB with CNN-LSTM. We report the results in 2. We observe that on both IMDB and CIFAR, the bulk of the function-space \(\hat{R}\) values is concentrated near \(1\), meaning intuitively that a single chain can capture the diversity of predictions on most of the test data points nearly as well as multiple chains. The mixing is especially good on the IMDB dataset, where only \(2\%\) of inputs correspond to \(\hat{R}\) larger than \(1.1\). In 15 we apply HMC to a synthetic regression problem and show that HMC can indeed mix in the prediction space: different HMC chains provide very similar predictions.

In weight space, although most parameters show no evidence of poor mixing, some have very large \(\hat{R}\)s, indicating that there are directions in which the chains fail to mix.

5.1.0.1 Implications for the Posterior Geometry.

The fact that a single HMC chain is able to mix well in function space (aka prediction space) suggests that the posterior contains connected regions which correspond to high functional diversity. Indeed, a single HMC chain is extremely unlikely to jump between isolated modes, but appears able to produce samples with diverse predictions. Moreover, HMC is able to navigate these regions effeciently. Prior work on mode connectivity [16], [17] has shown that there exist paths of high density connecting different modes of the posterior. Our observations suggest a stronger version of mode connectivity: not only do mode-connecting paths exist between functionally diverse modes, but also at least some of these paths can be leveraged by Monte Carlo methods to efficiently explore the posterior.

5.2 Posterior density visualizations↩︎

To further investigate how HMC is able to explore the posterior over the weights, we visualize a cross-section of the posterior density in subspaces of the parameter space containing the samples. Following [16], we study two-dimensional subspaces of the parameter space of the form \[\label{eq:subspace} \mathcal{S} = \{w \vert w = w_1 \cdot a + w_2 \cdot b + w_3 \cdot (1 - a - b)\}.\tag{3}\] \(\mathcal{S}\) is the unique two-dimensional affine subspace (plane) of the parameter space that includes parameter vectors \(w_1\), \(w_2\) and \(w_3\).

In 3 (a) we visualize the posterior log-density, log-likelihood and log-prior density of a ResNet-20-FRN on CIFAR-10. For the visualization, we use the subspace \(\mathcal{S}\) defined by the parameter vectors \(w_{1}, w_{51}\) and \(w_{101}\), the samples produced by HMC at iterations \(1, 51\) and \(101\) after burn-in respectively. We observe that HMC is able to navigate complex geometry: the samples fall in three seemingly isolated modes in our two-dimensional cross-section of the posterior. In other words, HMC samples from a single chain are not restricted to any specific convex Gaussian-like mode, and instead explore a region of high posterior density of a complex shape in the parameter space. We note that popular approximate inference procedures, such as variational methods, and Laplace approximations, are typically constrained to unimodal Gaussian approximations to the posterior, which we indeed expect to miss a large space of compelling solutions in the posterior.

In 3 (b) we provide a visualization for the samples produced by \(3\) different HMC chains at iteration \(51\) after burn-in. Comparing the visualizations for samples from the same chain and samples from independent chains in 3, we see that the shapes of the posterior surfaces are different, with the latter appearing more regular and symmetric. The qualitative differences between (a) and (b) suggest that while each HMC chain is able to navigate the posterior geometry the chains do not mix perfectly in the weight space, confirming our results in 5.1.

In 14 we provide additional details on our visualization procedure and visualizations using the CNN-LSTM on IMDB.

5.3 Convegence of the HMC chains↩︎

As another diagnostic, we look at the convergence of the performance of HMC BMA estimates and individual samples as a function of the length of the burn-in period. For a converged chain, the performance of the BMA and individual samples should be stationary not show any visible trends after a sufficiently long burn-in. We use the samples from \(3\) HMC chains, and evaluate performance of the ensemble of the first \(100\) HMC samples in each chain after discarding the first \(n_{bi}\) samples, where \(n_{bi}\) is the length of the burn-in. Additionally, we evaluate the performance of the individual HMC samples after \(n_{bi}\) iterations in each of the chains.

We report the results for ResNet-20-FRN on CIFAR-10 and CNN-LSTM on IMDB in 4. On IMDB, there is no visible trend in performance, so a burn-in of just \(10\) HMC iterations should be sufficient. On CIFAR-10, we observe a slowly rising trend that saturates at about \(50\) iterations, indicating that a longer burn-in period is needed compared to IMDB. We therefore use a burn-in period of \(50\) HMC iterations on both CIFAR and IMDB for the remainder of the paper.

6 Evaluating Bayesian neural networks↩︎

Now that we have a procedure for effective HMC sampling, we are primed to explore exciting questions about the fundamental behaviour of Bayesian neural networks, such as the role of tempering, the prior over parameters, generalization performance, and robustness to covariate shift. In this section we evaluate Bayesian neural networks in various problems using our implementation of HMC. Throughout the experiments, we use posterior temperature \(T=1\).

We emphasize that the main goal of our paper and this section in particular is to understand the behaviour of true Bayesian neural network posteriors using HMC as a precise tool, and not to argue for HMC as a practical method for Bayesian deep learning.

6.1 Regression on UCI datasets↩︎

Bayesian deep learning methods are often evaluated on small-scale regression problems using fully connected networks [9], [10], [31]. Following these works, we evaluate Bayesian neural networks using HMC on five UCI regression datasets: Concrete, Yacht, Boston, Energy and Naval. For each of these datasets, we construct \(20\) random \(90\)-to-\(10\) train-test splits and report the mean and standard deviation of performance over the splits. We use a fully connected neural network with a single hidden layer of size \(50\) and \(2\) outputs representing the predictive mean and standard deviation. For HMC we used a single chain with \(10\) burn-in iterations and \(90\) iterations of sampling. For more details, please see 11.

We report the results in 5. HMC typically outperforms all the baselines, often by a significant margin, both in test RMSE and log-likelihood. On the Boston dataset, HMC achieves a slightly higher average RMSE compared to the subspace inference and SWAG [9], [10] but outperforms both these methods significantly in terms of log-likelihood.

6.2 Image Classification on CIFAR↩︎

Next, we evaluate Bayesian neural networks using HMC on image classification problems. We use the ResNet-20-FRN architecture on CIFAR-10 and CIFAR-100. We picked a random subset of 40960 of the 50000 images for each of the datasets to be able to evenly shard the data across the TPU devices; we use the same subset for both HMC and the baselines. We run 3 HMC chains using step size \(10^{-5}\) and a prior variance of \(1/5\), resulting in 70,248 leapfrog steps per sample. In each chain we discard the first 50 samples as burn-in, and then draw 240 samples (720 in total for 3 chains)⁵. For SGLD, we use a single chain with \(1000\) burn-in epochs and \(9000\) epochs of sampling producing \(900\) samples; we also report the performance of an ensemble of \(5\) independent SGLD chains. Next, we report the performance of a mean field variational inference (MFVI) solution; we initialize the mean of MFVI with a solution pre-trained with SGD and use an ensemble of \(50\) samples from the VI posterior at evaluation. Finally, we report the performance of a single SGD solution and a deep ensemble of \(50\) models. For more details, see 11.

We report the results in 6. Bayesian neural networks outperform all baselines in terms of accuracy and log-likelihood on both datasets. In terms of ECE, SGD provides the worst results across the board, and the rest of the methods are competitive; MFVI is particularly well-calibrated on CIFAR-100.

Out-of-distribution detection. Bayesian deep learning methods are often evaluated on out-of-distribution detection. In the ¿tbl:tab:ood? we report the performance of HMC-based Bayesian neural network on out-of-distribution (OOD) detection. To detect OOD data, we use the level of predicted confidence (value of the softmax class probability for the predicted class) from the HMC ensemble, measuring the area under the receiving operator characteristic curve (AUC-ROC). We train the methods on CIFAR-10 and use CIFAR-100 and SVHN as OOD data sources. We find that BNNs perform competitively with the specialized ODIN method in the challenging near-OOD detection setting (i.e. when the OOD data distribution is similar to the training data) of CIFAR-100, while underperforming in the easier far-OOD setting on SVHN relative to the baselines [66], [67].

Robustness to distribution shift Bayesian methods are often specifically applied to covariate shift problems [5], [33], [68]. We evaluate the the performance of HMC and Deep Ensemble-based Bayesian neural networks on the CIFAR-10-C dataset [69], which applies a set of corruptions to CIFAR-10 with varying intensities. Mimicking the setup in [68], we use the same 16 corruptions, evaluating the performance at all intensities. We report the results in 7. Surprisingly, we find that Deep Ensembles and SGLD are consistently more robust to distribution shift than HMC-based BNNs. For high corruption intensities, even a single SGD model outperforms the HMC ensemble.

In 16 we provide further exploration of this effect, where we see HMC samples are significantly less robust to many types of noise compared to conventionally-trained SGD models. Interestingly, the performance of HMC-based BNNs under data corruption can be significantly improved by using posterior tempering.

6.3 Language Classification on IMDB↩︎

We use a CNN-LSTM architecture on the IMDB binary text classification dataset. In 6 we report the results for HMC and the same baselines as in 6.2. We use HMC with a step size of \(10^{-5}\) and a prior variance of \(1/40\), resulting in 24,836 leapfrog steps per sample. We run \(3\) chains, burning-in for 50 samples, and drawing 400 samples per chain (1,200 total). For more details on the hyperparameters, please see 11. Analogously to the image classification experiments, HMC outperforms the baselines on accuracy and log-likelihood and provides competitive performance on ECE.

7 Do we need cold posteriors?↩︎

Multiple works have considered tempering the posterior in Bayesian neural networks [5], [21], [22], [70], [71]. Specifically, we can consider a distribution \[\label{eq:posterior95temp} p_T(w \vert \mathcal{D}) \propto \big ( p(\mathcal{D} \vert w) \cdot p(w) \big )^{1 / T},\tag{4}\] where \(w\) are the parameters of the network, \(\mathcal{D}\) is the training dataset, \(p(\mathcal{D} \vert w)\) is the likelihood of \(\mathcal{D}\) for the network with parameters \(w\) and \(T\) is the temperature. Note that at temperature \(T = 1\), \(p_T\) corresponds to the standard Bayesian posterior over the parameters of the network. Temperatures \(T < 1\) correspond to cold posteriors, distributions that are sharper than the Bayesian posterior. Similarly, temperatures \(T > 1\) correspond to warm posteriors which are softer than the Bayesian posterior. See Appendix 13 (d) for a visualization of the log-likelihood density surface at different temperatures.

[22] argue that Bayesian neural networks require a cold posterior, and the performance at temperature \(T=1\) is inferior to even a single model trained with SGD. The authors refer to this phenomenon as the cold posteriors effect. However, our results are different:

7.1 Testing the cold posteriors effect↩︎

[22] demonstrate the cold posteriors with two main experiments: ResNet-20 on CIFAR-10 and CNN-LSTM on IMDB. In these experiments the authors show poor performance at temperature \(T=1\), with strong benefits from decreasing the temperature. However, for the CIFAR-10 experiment, it is apparent [22] that the results at \(T = 1\) are near-optimal for the ResNet on CIFAR-10 if data augmentation is turned off and batch normalization is replaced with filter response normalization, which is in fact necessary for a clear Bayesian interpretation of the inference procedure.

Furthermore, in 6, we show that Bayesian neural networks can achieve performance superior to SGD and even deep ensembles at temperature \(T=1\), in particular using the same ResNet-20-FRN model on CIFAR-10 and CNN-LSTM model on IMDB used by [22].

To further understand the effect of posterior temperature \(T\), we compare the performance of the CNN-LSTM model at different \(T\) using our Hamiltonian Monte Carlo sampler. In all runs we used a fixed prior variance \(\alpha^2 = \frac{1}{40}\). We report the results in 8. We find that the performance of the BNN at \(T=1\) is better than the SGD baseline as well as a deep ensemble of \(50\) independent models. Moreover, the performance at \(T=1\) is better compared to all other temperatures we tested in terms of both test accuracy and log-likelihood.

We also note that while posterior tempering does not seem necessary for good predictive performance with BNNs, it may be helpful under distribution shift. In 16 we show that decreasing the temperature can significantly improve the robustness of BNN predictions to noise in the test inputs. [5] additionally argue that tempering may be a reasonable procedure in general, and is not necessarily at odds with Bayesian principles.

Role of data augmentation. Our results are in contrast with [22], who argue that cold posteriors are needed for good performance with BNNs. In 17 we provide an additional study of what may have caused the poor performance of BNNs in [22], using the code for inference provided by [22]. We identify data augmentation as the key factor responsible for the cold posterior effect, and also show that batch normalizing does not significantly influence this effect: when the data augmentation is turned off, we do not observe the cold posteriors effect. Data augmentation cannot be naively incorporated in the Bayesian neural network model (see the discussion in appendix K of [22]), and arguably it may be reasonable to decrease the temperature when using data augmentation: intuitively, data augmentation increases the amount of data observed by the model, and should lead to higher posterior contraction. We leave incorporating data augmentation in our HMC evaluation framework as an exciting direction of future work.

8 What is the effect of priors in Bayesian neural networks?↩︎

Bayesian deep learning is often criticized for the lack of intuitive priors over the parameters. For example, [22] hypothesize that the popular Gaussian priors of the form \(\mathcal{N}(0, \alpha^2 I)\) are inadequate and lead to poor performance. [72] propose a new prior for Bayesian neural networks inspired by Gaussian processes [73] based on this hypothesis. In concurrent work, [74] also explore several alternatives to standard Gaussian priors inspired by the cold posteriors effect. [5] on the other hand, argue that vague Gaussian priors in the parameter space induce useful function-space priors.

In 6 we have shown that Bayesian neural networks can achieve strong performance with vague Gaussian priors. In this section, we explore the sensitivity of BNNs to the choice of the prior scale as well as several alternative prior families, as a step towards a better understanding of the role of the prior in BNNs.

8.1 Effect of Gaussian prior scale↩︎

We use priors of the form \(\mathcal{N}(0, \alpha^2 I)\) and vary the prior variance \(\alpha^2\). For all cases, we use a single HMC chain producing \(40\) samples. These are much shorter chains than the ones we used in 6, so the results are not as good; the purpose of this section is to explore the relative performance of BNNs under different priors.

We report the results for the CIFAR-10 and IMDB datasets in 9. When the prior variance is too small, the regularization is too strong, hindering the performance. Setting the prior variance too large does not seem to hurt the performance as much. On both problems, the performance is fairly robust: a wide window of prior variances lead to strong performance. In particular, for all considered prior scales, the results are better than those of SGD training.

Why are BNNs so robust to the prior scale?One possible explanation for the relatively flat curves in 9 is that large prior variances imply a strong prior belief that the “true” classifier (i.e., the model that would be learned given infinite data) should make high-confidence predictions. Since the model is powerful enough to achieve any desired training accuracy, the likelihood does not overrule this prior belief, and so the posterior assigns most of its mass to very confident classifiers. Past a certain point, increasing the prior variance on the weights may have no effect on the classifiers’ already saturated probabilities. Consequently, nearly every member of the BMA may be highly overconfident. But the ensemble does not have to be overconfident—a mixture of overconfident experts can still make well-calibrated predictions. Appendix 17 provides some qualitative evidence for this explanation; for some CIFAR-10 test set images, the HMC chain oscillates between assigning the true label probabilities near 1 and probabilities near 0.

8.2 Non-Gaussian priors↩︎

In ¿tbl:tab:alternate-priors?, we report BMA accuracy, ECE, and log-likelihood for two non-Gaussian priors on the IMDB dataset: logistic and mixture of Gaussians (MoG). For the MoG prior we use a mixture of two Gaussians centered at \(0\), one with variance \(\frac{1}{40}\) and the other with variance \(\frac{1}{160}\). We pick prior scale of the logistic prior to have a variance of \(\frac{1}{40}\). We additionally provide the results for a Gaussian prior with variance \(\frac{1}{40}\). We approximate the BMA using 80 samples from a single HMC chain for each of the priors. We find that the heavier-tailed logistic prior performs slightly better than the Gaussian and MoG.

8.3 Importance of Architecture in Prior Specification↩︎

We often think of the prior narrowly in terms of a distribution over parameters \(p(w)\). But the prior that matters is the prior over functions \(p(f(x))\) that is induced when a prior over parameters \(p(w)\) is combined with the functional form of a neural network \(f(x,w)\). All of the results in this section point to the relative importance of the architecture in defining the prior over functions, compared to the prior over parameters. A vague prior over parameters is not necessarily vague in function-space. Moreover, while the details of the prior distribution over parameters \(p(w)\) have only a minor effect on performance, the choice of architecture certainly has a major effect on performance.

9 Do scalable BDL methods and HMC make similar predictions?↩︎

While HMC shows strong performance in our evaluation in 6, in most realistic BNN settings it is an impractical method. However, HMC can be used as a reference to evaluate and calibrate more scalable and practical alternatives. In this section, we evaluate the fidelity of SGMCMC, variational methods, and deep ensembles in representing the predictive distribution (Bayesian model average) given by our HMC reference.

9.1 Comparing the predictive distributions↩︎

We consider two primary metrics: agreement and total variation. We define the agreement between the predictive distributions \(\hat{p}\) of HMC and \(p\) of another method as the fraction of the test data points for which the top-1 predictions of \(\hat{p}\) and \(p\) are the same: \[\frac{1}{n} \sum_{i=1}^n I[\arg\max_j \hat{p} (y = j \vert x_i) = \arg\max_j p(y = j \vert x_i)],\] where \(I[\cdot]\) is the indicator function and \(n\) is the number of test data points \(x_i\). We define the total variation metric between \(\hat{p}\) and \(p\) as the total variation distance between the predictive distributions averaged over the test data points: \[\frac{1}{n} \sum_{i=1}^n \frac{1}{2}\sum_j \bigg| \hat{p}(y = j \vert x_i) - p(y = j \vert x_i)\bigg|.\] The agreement (higher is better) captures how well a method is able to capture the top-1 predictions of HMC, while the total variation (lower is better) compares the predictive probabilities for each of the classes.

In ¿tbl:tab:pred95comparison? we report the agreement and total variation metrics as well as the predictive accuracy on the CIFAR-10 and CIFAR-10-C test sets for a deep ensemble of \(50\) models and several SGLD variations: standard SGLD [7], SGLD with momentum (SGHMC) [20], SGLD with momentum and a cyclical learning rate schedule (SGHMC-CLR) [21] and SGLD with momentum, cyclical learning rate schedule and a preconditioner (SGHMC-CLR-Prec) [22]. All methods were trained on CIFAR-10. For more details, please see 11.

Overall, the absolute value of agreement achieved by all methods is fairly low on CIFAR-10 and especially on CIFAR-10-C. More advanced SGHMC-CLR and SGHMC-CLR-Prec methods provide a better fit of the HMC predictive distribution while not necessarily improving the accuracy. Notably, these methods are also less robust to the data corruptions in CIFAR-10-C, again suggesting that higher fidelity representations of the predictive distribution can lead to decreased robustness to covariate shift, as we found in section 6.2.

Deep ensembles, while not typically considered to be a Bayesian method, provide a reasonable approximation to the HMC predictive distribution. In particular, deep ensembles outperform both SGLD and SGHMC in terms of total variation on CIFAR-10 and in terms of agreement on CIFAR-10-C. These results support the argument that deep ensembles, while not typically characterized as a Bayesian method, provide a higher fidelity approximation to a Bayesian model average than methods that are conventionally accepted as Bayesian inference procedures in modern deep learning [5].

In 16 we explore the performance of HMC, SGD, deep ensembles, SGLD and SGHMC-CLR-Prec under different corruptions individually. Interestingly, the behavior of SGLD and SGHMC-CLR-Prec appears more similar to that of deep ensembles than that of HMC. So, while both SGMCMC and deep ensembles are very compelling practically, they provide relatively distinct predictive distributions from HMC. Mean-field variational inference methods are particularly far from the HMC predictive distribution. Thus, we should be very careful when making judgements about true Bayesian neural networks based on the SGMCMC or MFVI performance.

9.2 Predictive entropy and calibration curves↩︎

To provide an additional comparison of the predictive distributions between HMC and other methods, in 10 we visualize the distribution of predictive entropies and the calibration curves for HMC, SGD, deep ensembles, MFVI, SGLD and SGHMC-CLR-Prec on CIFAR-10 using ResNet-20-FRN.

All methods except fot SGD make conservative predictions: their confidences tend to underestimate their accuracies (10, right); SGD on the other hand is very over-confident [75]. Deep ensembles and MFVI provide the most calibrated predictions, while SGLD and SGHMC-CLR-Prec match the HMC entropy distribution and calibration curve closer.

10 Discussion↩︎

Despite the rapidly increasing popularity of approximate Bayesian inference in modern deep learning, little is known about the behaviour of truly Bayesian neural networks. To the best of our knowledge, our work provides the first realistic evaluation of Bayesian neural networks with precise and exhaustive posterior sampling. We establish several properties of Bayesian neural networks, including good generalization performance, lack of a cold posterior effect, and a surprising lack of robustness to covariate shift. We hope that our observations and the tools that we develop will facilitate fundamental progress in understanding the behaviour of Bayesian neural networks.

Is HMC Converging? In general, it is not possible to ensure that an MCMC method has converged to sampling from the true posterior distribution: theoretically, there may always remain regions of the posterior that cannot be discovered by the method but that contain most of the posterior mass. To maximize the performance of HMC, we choose the hyper-parameters that are the most likely to provide convergence: long trajectory lengths, and multiple long chains. In 5, we study the convergence of HMC using the available convergence diagnostics. We find that while HMC does not mix perfectly in weight space, in the space of predictions we cannot find evidence of non-mixing.

Should we use Bayesian neural networks? Our results show that Bayesian neural networks can improve the performance over models trained with SGD in a variety of settings, both in terms of uncertainty calibration and predictive accuracy. On most of the problems we considered in this work, the best results were achieved by Bayesian neural networks. We believe that our results provide motivation to use Bayesian neural networks with accurate posterior approximation in practical applications.

Should we use HMC in practice? For most realistic scenarios in Bayesian Deep Learning, HMC is an impractical method. On the image classification benchmarks, HMC takes orders of magnitude more compute than any of the baselines that we considered. We hope that our work will inspire the community to produce new accurate and scalable approximate inference methods for Bayesian deep learning.

10.1 Challenging conventional wisdom↩︎

A conventional wisdom has emerged that deep ensembles are a non-Bayesian alternative to variational methods, that standard priors for neural networks are poor, and that cold posteriors are a problematic result for Bayesian deep learning. Our results highlight that one should take care in uncritically repeating such claims. In fact, deep ensembles appear to provide a higher fidelity representation of the Bayesian predictive distribution than widely accepted approaches to approximate Bayesian inference. If anything, the takeaway from the relatively good performance of deep ensembles is that we would benefit from approximate inference being closer to the Bayesian ideal! Moreover, the details over the priors in weight space can have a relatively minor effect on performance, and there is no strong evidence that standard Gaussian priors are particularly bad. In fact, there are many reasons to believe these priors have useful properties [5]. Similarly, on close inspection, we found no evidence for a general cold posterior effect, which we identify as largely an artifact of data augmentation. Although we see here that tempering does not in fact seem to be required, as argued in [5] tempering is also not necessarily unreasonable or even divergent from Bayesian principles.

Even the results we found that are less favourable to Bayesian deep learning are contrary to the current orthodoxy. Indeed, higher fidelity Bayesian inference surprisingly appears to suffer more greatly from covariate shift, despite the popularity of approximate Bayesian inference procedures in this setting.

Acknowledgements↩︎

The authors would like to thank many people at Google for many helpful discussions and much helpful feedback, especially Rodolphe Jenatton, Rif A. Saurous, Jasper Snoek, Pavel Sountsov, Florian Wenzel, and the entire TensorFlow Probability team. This research is supported by an Amazon Research Award, NSF I-DISRE 193471, NIH R01DA048764-01A1, NSF IIS-1910266, and NSF 1922658NRT-HDR: FUTURE Foundations, Translation, and Responsibility for Data Science.

Appendix Outline↩︎

This appendix is organized as follows. We present the Hamiltonian Monte Carlo algorithm that we implement in the paper in 11, [alg:leapfrog]. In 11 we provide the details on hyper-parameters used in our experiments. In 12 we provide a description of the \(\hat{R}\) statistic used in 5.1. In 13 we study the marginals of the posterior over parameters estimated by HMC. In 14 we provide additional posterior density surface visualizations. In 15 we compare the BMA predictions using two independent HMC chains on a synthetic regression problem. In 16 we show that BNNs are not robust to distribution shift and discuss the reasons for this behavior. In 17 we provide a further discussion of the effect of posterior temperature. Finally, in 18 we visualize the predictions of HMC samples from a single chain on several CIFAR-10 test inputs.

11 Hyper-Parameters and Details↩︎

CIFAR and IMDB. In ¿tbl:tab:hyperparameters? we report the hyperparameters used by each of the methods in our main evaluation on CIFAR and IMDB datasets in 6. HMC was run on a cluster of 512 TPUs and the other baselines were run on a cluster of \(8\) TPUs. On CIFAR datasets the methods used a subset of \(40960\) datapoints. All methods were ran at posterior temperature \(1\). We tuned the hyper-parameters for all methods via cross-validation maximizing the accuracy on a validation set. For the step-sizes we considered an exponential grid with a step of \(\sqrt{10}\) with \(5\)-\(7\) different values, where the boundaries were selected for each method so it would not diverge. We considered weight decays \(1, 5, 10, 20, 40, 100\) and the corresponding prior variances. For batch sizes we considered values \(80\), \(200\), \(400\) and \(1000\); for all methods lower batch sizes resulted in the best performance. For HMC we set the trajectory length according to the strategy described in 4.1. For SGLD, we experimented with using a cosine learning rate schedule decaying to a non-zero final step size, but it did not improve upon a constant schedule. For MFVI we experimented with the SGD and Adam optimizers; we initialize the mean of the MFVI distribution with a pre-trained SGD solution, and the per-parameter variance with a value \(\sigma^\text{VI}_\text{Init}\); we tested values \(10^{-2}, 10^{-1}, 10^{0}\) for \(\sigma^\text{VI}_\text{Init}\). For all HMC hyper-parameters, we provide ablations illustrating their effect in 4. Producing a single sample with HMC on CIFAR datasets takes roughly one hour on our hardware, and on IMDB it takes \(105\) seconds; we can run up to three chains in parallel.

Temperature scaling on IMDB. For the experiments in 7 we run a single HMC chain producing \(40\) samples after \(10\) burn-in epochs for each temperature. We used step-sizes \(5\cdot 10^{-5}\), \(3\cdot 10^{-5}\), \(10^{-5}\), \(3\cdot 10^{-6}\), \(10^{-6}\) and \(3\cdot 10^{-7}\) for temperatures \(10\), \(3\), \(1\), \(0.3\), \(0.1\) and \(0.03\) respectively, ensuring that the accept rates were close to \(100\%\). We used a prior variance of \(1 / 50\) in all experiments; the lower prior variance compared to ¿tbl:tab:hyperparameters? was chosen to reduce the number of leapfrog iterations, as we chose the trajectory length according to the strategy described in 4.1. We ran the experiments on \(8\) NVIDIA Tesla V-100 GPUs, as we found that sampling at low temperatures requires float64 precision which is not supported on TPUs.

UCI Datasets. In ¿tbl:tab:uci95hypers? we report the hyperparameters used by each of the methods in our main evaluation on UCI datasets in 6. For each datasets we construct \(20\) random splits with \(90\%\) of the data in the train and \(10\%\) of the data in the test split. In the evaluation, we report the mean and standard deviation of the results across the splits. We use another random split for cross-validation to tune the hyper-parameters. For all datasets we use a fully-connected network with a single hidden layer with \(50\) neurons and \(2\) outputs representing the predictive mean and variance for the given input. We use a Gaussian likelihood to train each of the methods. For the SGD and SGLD baselines, we did not use mini-batches: the gradients were computed over the entire dataset. We run each experiment on a single NVIDIA Tesla V-100 GPU.

SGMCMC Methods. In ¿tbl:tab:sgmcmc95hypers? we report the hyper-parameters of the SGMCMC methods on the CIFAR-10 dataset used in the evaluation in 9. We considered momenta in the set of \(\{0.9, 0.95, 0.99\}\) and step sizes in \(\{10^{-4}, 3\cdot 10^{-5}, 10^{-5}, 3\cdot 10^{-6}, 10^{-6}, 3\cdot 10^{-7}, 10^{-7}\}\). We selected the hyper-parameters with the best accuracy on the validation set. SGLD does not allow a momentum.

12 Description of \(\hat{R}\) Statistics↩︎

\(\hat{R}\) [64] is a popular MCMC convergence diagnostic. It is defined in terms of some scalar function \(\psi(\theta)\) of the Markov chain iterates \(\{\theta_{mn} | m\in\{1,\ldots,M\}, n\in \{1,\ldots,N\}\}\), where \(\theta_{mn}\) denotes the state of the \(m\)th of \(M\) chains at iteration \(n\) of \(N\). Letting \(\psi_{mn}\triangleq \psi(\theta_{mn})\), \(\hat{R}\) is defined as follows: \[\begin{gather} \bar \psi_{m \cdot} \triangleq \frac{1}{N}\sum_n \psi_{mn}; \quad \bar \psi_{\cdot \cdot} \triangleq \frac{1}{MN}\sum_{m,n} \psi_{mn}; \\ \frac{B}{N}\triangleq \frac{1}{M-1}\sum_m (\bar\psi_{m\cdot} - \bar\psi_{\cdot\cdot})^2; \\ W\triangleq \frac{1}{M(N-1)}\sum_{m,n} (\psi_{mn}-\bar\psi_{m\cdot})^2; \\ \hat{\sigma}^2_+ \triangleq \frac{N-1}{N}W + \frac{B}{N}; \\ \hat{R}\triangleq \frac{M+1}{M}\frac{\hat{\sigma}^2_+}{W} - \frac{N-1}{MN}. \end{gather}\] If the chains were initialized from their stationary distribution, then \(\hat{\sigma}^2_+\) would be an unbiased estimate of the stationary distribution’s variance. \(W\) is an estimate of the average within-chain variance; if the chains are stuck in isolated regions, then \(W\) should be smaller than \(\hat{\sigma}^2_+\), and \(\hat{R}\) will be clearly larger than 1. The \(\frac{M+1}{M}\) and \(\frac{N-1}{MN}\) terms are there to account for sampling variability—they vanish as \(N\) gets large if \(W\) approaches \(\hat{\sigma}^2_+\).

Since \(\hat{R}\) is defined in terms of a function of interest \(\psi\), we can compute it for many such functions. In 5.1 we evaluated it for each weight and each predicted softmax probability in the test set.

13 Marginal distributions of the weights↩︎

In 4.1, we argued for using a trajectory length \(\tau=\frac{\pi\sigma_\mathrm{prior}}{2}\) based on the intuition that the posterior scale is determined primarily by the prior scale. In 12 we examine this intuition. For each parameter, we estimate the marginal standard deviation of that parameter under the distribution sampled by HMC. Most of these marginal scales are close to the prior scale, and only a few are significantly larger (note logarithmic scale on y-axis), confirming that the posterior’s scale is determined by the prior.

14 Additional Posterior Visualizations↩︎

In 5.2 we study two-dimensional cross-sections of posterior log-density, log-likelihood and log-prior surfaces. We provide additional visualizations on the IMDB dataset in 13.

On IMDB, the posterior log-density is dominated by the prior, and the corresponding panels are virtually indistinguishable in 13. For the CNN-LSTM on IMDB the number of parameters is much larger than the number of data points, and hence the scale of the prior density values is much larger than the scale of the likelihood. Note that the likelihood still affects the posterior typical set, and the HMC samples land in the modes of the likelihood in the visualization. In contrast, on ResNet-20, the number of parameters is smaller and the number of data points is larger, so the posterior is dominated by the likelihood in 3. The log-likelihood panels for both datasets show that HMC is able to navigate complex geometry: the samples fall in three isolated modes in our two-dimensional cross-sections. On IMDB, the visualizations for samples from a single chain and for samples from three independent chains are qualitatively quite similar, hinting at better parameter-space mixing compared to CIFAR-10 (see 5.1).

In 13 (c), we visualize the likelihood cross-sections using our runs with varying posterior temperature on IMDB. The visualizations show that, as expected, low temperature leads to a sharp likelihood, while the high-temperature likelihood appear soft. In particular, the scale of the lowest likelihood values at \(T=10\) is only \(10^3\) while the scale at \(T=0.1\) is \(10^6\).

How are the visualizations created? To create the visualizations we pick the points in the parameter space corresponding to three HMC samples: \(w_1, w_2, w_3\). We construct a basis in the 2-dimensional affine subspace passing through these three points: \(u = w_2 - w_1\) and \(v = w_3 - w_1\). We then orthogonalize the basis: \(\hat{u} = u / \|u\|\), \(\hat{v} = (v - \hat{u}^T v) / \|v - \hat{u}^T v\|\). We construct a 2-dimensional uniform grid in the basis \(\hat{u}, \hat{v}\). Each point in the grid corresponds to a vector of parameters of the network. We evaluate the log-likelihood, log-prior and posterior log-density for each of the points in the grid, converting them to the corresponding network parameters. Finally, we produce contour plots using the collected values. The procedure is analogous to that used by [16]⁶.

15 HMC Predictive Distributions in Synthetic Regression↩︎

We consider a one-dimensional synthetic regression problem. We follow the general setup of [10] and [5]. We generate the training inputs as a uniform grid with \(40\) points in each of the following intervals (\(120\) datapoints in total): \([-10, -6], [6, 10]\) and \([14, 18]\). We construct the ground truth target values using a neural network with 3 hidden layers, each of dimension \(100\), one output and two inputs: following [10], for each datapoint \(x\) we pass \(x\) and \(x^2\) as inputs to the network to enlarge the class of functions that the network can represent. We draw the parameters of the network from a Gaussian distribution with mean \(0\) and standard deviation \(0.1\). We show the sample function used to generate the target values as a black line in each of the panels in 14. We then add Gaussian noise with mean \(0\) and standard deviation \(0.02\) to each of the target values. The final dataset used in the experiment is shown with red circles in 14.

For inference, we use the same model architecture that was used to generate the data. We sample the initialization parameters of the network from a Gaussian distribution with mean \(0\) and standard deviation \(0.005\). We use a Gaussian distribution with mean zero and standard deviation \(0.1\) as the prior over the parameters, same as the distribution used to sample the parameters of the ground truth solution. We use a Gaussian likelihood with standard deviation \(0.02\), same as the noise distribution in the data. We run two HMC chains from different random initializations. Each chain uses a step-size of \(10^{-5}\) and the trajectory length is set according to the strategy described in 4.1, resulting in \(15708\) leapfrog steps per HMC iteration. We run each chain for \(100\) HMC iterations and collect the predictions corresponding to all the accepted samples, resulting in \(89\) and \(82\) samples for the first and second chain respectively. We discard the first samples and only use the last \(70\) samples from each chain. For each input point we compute the mean and standard deviation of the predictions.

We report the results in 14. In panels (a), (b) we show the predictive distributions for each of the chains, and in panel (c) we show them overlaid on top of each other. Both chains provide high uncertainty away from the data, and low uncertainty near the data as desired [76]. Moreover, the true data-generating function lies in the \(3\sigma\)-region of the predictive distribution for each chain. Finally, the predictive distributions for the two chains are almost identical. This result suggests that on the synthetic problem under consideration HMC is able to mix in the space of predictions, and provides similar results independent of initialization and random seed. We come to the same conclusion for more realistic problems in 5.

16 BNNs are not Robust to Domain Shift↩︎

In 6.2, 7 we have seen that surprisingly BNNs via HMC underperform significantly on corrupted data from CIFAR-10-C compared to SGLD, deep ensembles and even MFVI and SGD. We provide detailed results in 15. HMC shows surprisingly poor robustness in terms of accuracy and log-likelihood across the corruptions. The ECE results are mixed. In most cases, the HMC ensemble of \(720\) models loses to a single SGD solution!

The poor performance of HMC on OOD data is surprising. Bayesian methods average the predictions over multiple models for the data, and faithfully represent uncertainty. Hence, Bayesian deep learning methods are expected to be robust to noise in the data, and are often explicitly evaluated on CIFAR-10-C [5], [33]. Our results suggest that the improvements achieved by Bayesian methods on corrupted data may be a sign of poor posterior approximation.

To further understand the robustness results, we reproduce the same effect on a small fully-connected network with two hidden layers of width \(256\) on MNIST. We run HMC at temperatures \(T=1\) and \(T=10^{-3}\) and SGD and report the results for both the BMA ensembles and individual samples in 16. For all methods, we train the models on the original MNIST training set, and evaluate on the test set with random Gaussian noise \(\mathcal{N}(0, \sigma^2 I)\) of varying scale \(\sigma\). We report the test accuracy as a function of \(\sigma\). We find that while the performance on the original test set is very close for all methods, the accuracy of HMC at \(T=1\) drops much quicker compared to that of SGD as we increase the noise scale.

Notably, the individual sample performance of \(T=1\) HMC is especially poor compared to SGD. For example, at noise scale \(\sigma=3\) the SGD accuracy is near \(60\%\) while the HMC sample only achieves around \(20\%\) accuracy!

HMC can be though of as sampling points at a certain sub-optimal level of the training loss, significantly lower than that of SGD solutions. As a result, HMC samples are individually inferior to SGD solutions. On the original test data ensembling the HMC samples leads to strong performance significantly outperforming SGD (see 6). However, as we apply noise to the test data, ensembling can no longer close the gap to the SGD solutions. To provide evidence for this explanation, we run evaluate HMC at a very low temperature \(T=10^{-3}\), as low temperature posteriors concentrate on high-performance solutions similar to the ones found by SGD. We find that at this temperature, HMC performs comparably with SGD, closing the gap in robustness We have also experimented with varying the prior scale but were unable to close the gap in robustness at temperature \(T=1\).

We hypothesize that using a lower temperature with HMC would also significantly improve robustness on CIFAR-10-C. Verifying this hypothesis, and generally understanding the robustness of BNNs further is an exciting direction of future work.

17 Further discussion of cold posteriors↩︎

In 7 we have seen that the cold posteriors are not needed to achieve strong performance with BNNs. We have even shown that cold (as well as warm) posteriors may hurt the performance. On the other hand, in 16 we have shown that lowering the temperature can improve robustness under the distribution shift, at least for a small MLP on MNIST. Here, we discuss the potential reasons for why the cold posteriors effect was observed in [22].

17.1 What causes the difference with [22]?↩︎

There are several key differences between the experiments in our study and [22].

First of all, the predictive distributions of SGLD (a version of which was used in [22]) are highly dependent on the hyper-parameters such as the batch size and learning rate, and are inherently biased: SGLD with a non-vanishing step size samples from a perturbed version of the posterior, both because it omits a Metropolis-Hastings accept-reject step and because its updates include minibatch noise. Both of these perturbations should tend to make the entropy of SGLD’s stationary distribution increase with its step size; we might expect this to translate to approximations to the BMA that are overdispersed.

Furthermore, [22] show in Figure 6 that with a high batch size they achieve good performance at \(T=1\) for the CNN-LSTM. Using the code provided by the autors⁷ with default hyper-parameters we achieved strong performance at \(T=1\) for the CNN-LSTM (accuracy of \(0.855\) and cross-entropy of \(0.35\), compared to \(0.81\) and \(0.45\) reported in Figure 1 of [22]); we were, however, able to reproduce the cold posteriors effect on CIFAR-10 using the same code.

On CIFAR-10, the main difference between our setup and the configuration in [22] is the use of batch normalization and data augmentation. In the appendix K and Figure 28 of [22], the authors show that if both the data augmentation and batch normalization are turned off, we no longer observe the cold posteriors effect. In ¿tbl:tab:app95cold95posteriors? we confirm using the code provided by the authors that in fact it is sufficient to turn off just the data augmentation to remove the cold posteriors effect. It is thus likely that the results in [22] are at least partly affected by the use of data augmentation.

18 Visualizing prediction variations in HMC samples↩︎

In 17 we visualize the predicted class probability of the true class for \(100\) HMC samples on eight different input images. While on some images the predicted class probability is always close to \(1\), on other inputs it is close to \(1\) for some of the samples (confidently correct), close to \(0\) for some of the samples (confidently wrong) and in between \(0\) and \(1\) for the remaining samples (unconfident). So, some of the samples are individually over-confident, but the ensemble is well-calibrated as other samples assign high probabilities to the correct class.

References↩︎