Posterior Reﬁnement Improves Sample Efﬁciency in Bayesian Neural Networks
Abstract
Monte Carlo (MC) integration is the de facto method for approximating the predictive distribution of Bayesian neural networks (BNNs).
But, even with many MC samples, Gaussian-based BNNs could still yield bad predictive performance due to the posterior approximation’s error.
Meanwhile, alternatives to MC integration tend to be more expensive and biased.
In this work, we experimentally show that the key to good MC-approximated predictive distributions is the quality of the approximate posterior itself.
However, previous methods for obtaining accurate posterior approximations are expensive and non-trivial to implement.
We, therefore, propose to reﬁne Gaussian approximate posteriors with normalizing ﬂows.
When applied to last-layer BNNs, it yields a simple post hoc method for improving pre-existing parametric approximations.
We show that the resulting posterior approximation is competitive with even the gold-standard full-batch Hamiltonian Monte Carlo.
Introduction
Predictive uncertainty is crucial in safety-critical systems [1].
Yet, commonly-used neural network (NN) predictive systems are overconﬁdent [2, 3].
Approximate Bayesian inference of NNs, resulting in Bayesian neural networks (BNNs) is a principled way of mitigating this issue [4].
Indeed, even crude approximations of BNNs’ posteriors, such as the Laplace approximation [5] can lead to good predictive uncertainty quantiﬁcation (UQ) performance [6] provided the predictive distributions are correctly computed [7].
A prediction in a BNN amounts to an integration of the likelihood w.r.t.
the (approximate) posterior measure.
Due to the non-linearity of NNs, no analytic solution to the integral exists, even when the likelihood and the approximate posterior are both Gaussian.
A low-cost, unbiased, stochastic approximation can be obtained via Monte Carlo (MC) integration: obtain S samples from the approximate posterior and then compute the empirical expectation of the likelihood w.r.t.
these samples.
While MC integration is accurate for large S , because of the sheer size of modern (B)NNs, virtually all BNNs use small S (typically 10 to 30 , see [8–14, etc.]).
Due to its well-known error scaling of Θ(1 / √ S ) , intuitively, MC integration with a small S is inadequate for an accurate prediction—yet, this has not been studied in depth for BNNs.
Furthermore, while linearization of an
NN around a point estimate in the parameter space is an alternative to MC integration [7, 15–17], it is generally costly due to the computation of per-example Jacobian matrices.
In this work, we study the quality of MC integration for making predictions in BNNs.
We show that few-sample MC integration is inaccurate for even an “easy” integral such as when the domain of 36th Conference on Neural Information Processing Systems (NeurIPS 2022).
S = 30 ; Error = 0 .
S
The MC Standard Error
to
R σ ( x ) N ( x | 5 , 10 2 ) dx ∈ [0 , 1]
ECE ↓
NLL ↓0.
Metrics
Calibration on F-MNIST
LA( S = 10000 ) HMC( S = 10 )
HMC( S = 600 ) integration is the output space of a binary-classiﬁcation BNNs, cf.
Fig. 1 (left). Further, we show that its alternative, the network linearization, disagrees with large-sample MC integration, making it unsuitable as a general-purpose predictive approximation method, i.e. when the Gauss-Newton matrix is not employed. Meanwhile, as indicated by full-batch Markov Chain Monte Carlo methods [18, 19] or even (the Bayesian interpretation of) ensemble methods [20], few-sample MC integration might still be useful for making predictions, provided that the approximate posterior measure is “close enough” to the true posterior—cf.
Fig. 1 (right). This implies that one should focus on improving the accuracy of posterior approximations.
Nevertheless, prior methods for obtaining expressive posteriors [e.g., 10, 21] require either signiﬁcant modiﬁcation to the NN or storing many copies of the parameters.
Moreover, they require training from scratch and thus introduce a signiﬁcant overhead, which can be undesirable in practical applications.
Therefore, we propose a post hoc method for “reﬁning” a Gaussian approximate posterior by leveraging normalizing ﬂows [22].
Contrary to the existing normalizing ﬂow methods with a priori base distribution (e.g. N (0 , I ) ), the proposed reﬁnement method converges faster with shorter ﬂows, making it cheaper than the naïve application of normalizing ﬂows in BNNs.
When used in conjunction with last-layer BNNs, which have been shown to be competitive to their all-layer counterparts [6], the proposed method is simple, cheap, yet competitive to even the gold-standard Hamiltonian Monte
Carlo in terms of predictive performance.
To summarize, our contributions are as follows:
(i)
We highlight the deﬁciencies of both few- and many-sample MC integration for computing BNNs’ predictive distributions.
(ii)
We argue that one must be careful when applying analytic alternatives to MC integration, such as linearization and the probit approximation, since they can yield unintended effects.
(iii)
We propose a widely-applicable technique for reﬁning parametric posterior approximations by leveraging normalizing ﬂows, which yields a cost-efﬁcient post hoc method when used in conjunction with last-layer BNNs.
(iv) We validate the method via extensive experiments and show that reﬁned posteriors are competitive with the much more expensive full-batch Hamiltonian Monte Carlo.
Preliminaries 2.1 Bayesian neural networks
Let D := { ( x i , y i ) } mi =1 be an i.i.d.
dataset, sampled from some distributions p ( x ) on R n and p ( y | x ) on R c .
A Bayesian neural network (BNN) is a neural network (NN)
f θ : R n → R c with a random parameter θ ∼ p ( θ ) on R d and a likelihood p ( y | f θ ( x )) , along with the associated posterior p ( θ | D ) ∝ p ( θ ) (cid:81)
mi
=1 p ( y i | f θ ( x i )) .
However, the exact posterior p ( θ | D ) is intractable in general due to its normalization constant.
Approximate BNNs , which approximate p ( θ | D ) with only its samples or with a simpler parametric distributions, must thus be employed in most cases.
2.1.1 Posterior approximations
The Laplace approximation [LA, 5] is one of the simplest approximation methods for BNNs.
The main idea is to construct a local Gaussian approximation to p ( θ | D ) , centered at a maximum a posteriori (MAP) estimate θ MAP :
= arg max θ p ( θ | D ) , with the covariance matrix given by the negative inverse-Hessian ( −∇ 2 θ log p ( θ | D ) | θ MAP ) − 1 .
Since the MAP estimation is the standard training procedure for NNs, the LA can be efﬁciently applied on top of pre-trained NNs [6], especially due to recent efﬁcient second-order optimization libraries [23, 24].
Variational Bayes [VB, 25] assumes a family of simple approximate distributions over the parameter space, e.g. M :=
{ q ( θ ) :
=
N ( θ | µ, Σ ) : ( µ, Σ ) ∈ R d + d 2 } , and ﬁnds a q ∈ M that is the closest to p ( θ | D ) by minimizing the reverse Kullback-Leibler (KL) divergence D KL ( q ( θ ) , p ( θ | D )) over M .
Unlike the LA, VB is not constrained to capture just the local mode of the log-posterior landscape.
However, VB is not post hoc and generally not as straightforward to implement.
Unlike the previous two approximations, Markov Chain Monte Carlo (MCMC) methods do not obtain parametric approximations to the posterior.
Instead, they collect samples from asymptotically the true posterior and use them to approximate integrals under the posterior measure.
A popular MCMC method is the Hamiltonian Monte Carlo (HMC) method [18, 19], where sampling processes are casted as Hamiltonian dynamics.
Nevertheless, MCMC methods are generally expensive since they require full-batches of data and storage of many copies of the networks’ parameters.
2.1.2 Predictive approximations
Let q ( θ ) be an approximate posterior of a NN f θ under D .
Given a test point x ∗ , the prediction y ∗ is distributed as p ( y ∗ | x ∗ , D ) :
= (cid:82) R d p ( y ∗ | f θ ( x ∗ ))
q ( θ ) dθ .
However, even when both p ( y ∗ | f θ ( x )) and q ( θ ) are Gaussians, this integral does not have an analytic solution due to the nonlinearity of f θ .
Further approximations on top of the posterior approximation are thus necessary.
Monte Carlo integration The Monte Carlo (MC) integration is a simple technique to approximate integrals under probability measures.
Speciﬁcally, in our case,
[Equation]
where S is the number of samples from q ( θ ) .
It is easy to show that MC integration is unbiased with error given by the standard error around the mean that scales like Θ(1 / √ S ) , see e.g. Murphy [26,
Section 2.7.3].
That is, MC integration (slowly) becomes more accurate as the number of samples S increases.
In the realm of BNNs, however, S is often chosen to be a small number, e.g. S = 20 , due to the computational cost of evaluating f θ .
Analytic approximations Since a Gaussian N ( θ | µ, Σ ) is the de facto choice for q ( θ ) , an analytic approximation of the marginal output distribution 2 p ( f ( x ) | x, D ) can be obtained by linearizing f θ around µ .
Speciﬁcally, we perform a ﬁrst-order Taylor expansion f θ ( x ) ≈ f µ ( x )
+ J f ( x ) · ( θ − µ ) where J f ( x ) ∈ R c × d is the Jacobian matrix of the network output w.r.t.
θ at µ and · denotes matrix multiplication.
Under this scheme, denoting f ∗ =
f ( x ∗ ) , we thus have for each test point x ∗
[Equation]
One can then use this distribution to obtain the predictive distribution p ( y ∗ | x ∗ , D ) ≈ (cid:82) R c p ( y | f ∗ ) p ( f ∗ | x ∗ , D ) df ∗ , which again can be approximated via MC integration.
Nevertheless, there exist analytic approximations to this integral [27, 15, 28, 29].
Speciﬁcally for binary classiﬁcation where p ( y | f ∗ ) =
σ ( f ∗ ) for the logistic function σ , one can use the so-called probit approximation [27, 15]: p ( y = 1 | x ∗ , D )
≈ σ ( f µ ( x ∗ ) / (cid:112)
+ π/ 8 S ( x ∗ )) .
Moreover, for
Here, S ( x ∗ ) ∈ R > 0 since f ( x ∗ ) is a real-valued random variable.
[Figure] [Figure]
Probit 0.
the multi- case, we can use the
[Equation]
where the vector division is taken component-wise.
2.2 Normalizing ﬂows for Bayesian inference
Let F : R d → R d be a diffeomorphism 4 in the sense of C k ( R d ) for a ﬁxed k ≥ 1 ∈ N .
If p is a density on R d , 5 then the density (cid:101) p of the random variable (cid:101) θ = F ( θ ) is given by (cid:101) p ( (cid:101) θ )
=
p ( F − 1 ( (cid:101) θ ))
| det J F − 1 ( (cid:101) θ ) | , where J F − 1 =
J − 1 F is the Jacobian matrix of F − 1 .
A normalizing ﬂow (NF) is a method exploiting this relation [22, 30, 31].
Speciﬁcally, it is a way of constructing a sophisticated diffeomorphism F , by composing several simple parametric ones.
The resulting diffeomorphism thus transforms a simple density into a complicated one in a tractable manner.
Let F φ l : R d → R d be a diffeomorphism parametrized by φ l ∈ R k with a known inverse F − 1 φ l and a known d × d Jacobian matrix J F φl for l = 1 , . . .
, (cid:96) .
Writing F φ :
= F φ (cid:96) ◦ · · · ◦ F φ 1 and φ := ( φ l ) (cid:96)l =1 , then the change-of-density formula becomes
[Equation]
Given a target posterior density p ∗ , the goal of a normalizing ﬂow is then to estimate φ s.t. (cid:101) p φ is close to p ∗ .
In Bayesian inference, the reverse KL-divergence D KL ( (cid:101) p φ , p ∗ ) is often used.
Pitfalls of BNNs’ Approximate Predictive Distributions
In this section we study the failure modes of BNNs, especially last-layer ones, for predictive uncertainty calibration. We shall show that there are two parts contributing to the inaccuracy in the predictive distribution of a BNN: (i) the weight-space approximation and (ii) the integration method to do Bayesian model averaging.
We shall observe that while the accuracy of the integration method is impactful (i.e., the number of samples in MC integration), it is less important than the quality of the posterior approximation itself.
That is, higher-quality weight-space approximations—even in last-layer BNNs—yield samples that are more efﬁcient: Few-sample MC integration can already yield good predictive performance in this case.
3.1 Accurate MC integration requires many samples
We begin our analysis with the simplest yet practically relevant case.
Let x ∗ ∈ R n and f ∗ :
= f ( x ∗ ) , with p ( f ∗ | x ∗ , D ) =
N ( f ∗ | m, s 2 ) for some m ∈ R and s 2 ∈ R > 0 .
We are interested in computing the integral (cf.
Fig. 3)
[Equation]
Note that this integral is prevalent in practice, e.g. in linearized or last-layer classiﬁcation BNNs [4, 7, 17, 32]. I ( m,s ) σ ( f ∗ ) N ( f ∗ | m,s 2 )
We compare a MC integration with S = 100 against the probit approximation, using a trapezoid quadrature with 20000 evaluation points to represent a gold-standard baseline.
That is, we compute the discrepancy | (cid:101) I ( m, s ) − I ∗ ( m, s ) | where (cid:101) I ( m, s ) is I ( m, s ) computed either with MC integration or the probit approximation, and I ∗ ( m, s ) obtained via the trapezoid method.
The results are in Fig.
Even with S = 100 samples— larger than the usual S = 10 - 30 —MC integration is inaccurate.
Its error can be as high as 0 .
18 , which is substantial considering I ( m, s ) ∈ [0 , 1] .
As S increases beyond 10000 , MC integration improves and eventually overtakes the probit approximation, as to be expected given its theoretical guarantees.
This highlights the ﬂaw of few-sample MC integration—accurate MC integration generally requires large number of samples.
3.2 Many-sample MC integration is not sufﬁcient
[Table]
However, even with a large S , MC integration can still fail to yield good predictive performance in BNNs.
This can happen when the approximation q ( θ ) used in (1) is an inaccurate approximation of p ( θ | D ) —virtually the case for every parametric BNN.
Thankfully, there is some evidence that the error of MC integration might be relatively small in comparison to the error generated by crude posterior approximations.
As an extreme example, Deep Ensemble [20] and its variants, even though they perform MC integration with a small number of samples (usually S = 5 ), generally yield better approximations to the predictive distributions.
This is perhaps due to their multimodality, i.e. due to their ﬁner-grained posterior approximations.
To show this more concretely, consider the following experiment.
We take Fashion-MNIST (FMNIST) pre-trained LeNet and CIFAR-10 pre-trained WideResNet-16-4.
For each case, we perform a last-layer Laplace approximation with the exact Hessian and a full-batch NUTS-HMC [19], under the same prior and likelihood.
We ﬁnd in Table 1 that HMC, even with few samples, yields bettercalibrated predictive distribution than the LA with three orders of magnitude more samples.
This ﬁnding validates the widely-believed wisdoms [9, 10, 14, 32, etc.] that highly accurate posterior approximations are most important for BNNs.
Furthermore, this also shows that with a ﬁne-grained posterior approximation, the predictive performance of a BNN is less sensitive to the number of MC samples, leading to better test-time efﬁciency.
3.3 Analytic alternatives to MC integration are not the deﬁnitive answers
Of course, ﬁne-grained posterior approximations are expensive.
So, a natural question is whether one can keep a cheap but crude posterior approximation by replacing MC integration.
Network linearization seems to be a prime candidate to replace MC integration in the case of Gaussian approximations [7, 15].
But it poses several problems: First, it requires relatively expensive computation of the per-example Jacobian J f ( x ∗ ) , where for each test point x ∗ , one must store the associated c × d matrix, where d could easily be tens of millions (e.g. in ResNets) and c could be in the order of thousands (e.g. ImageNet).
Second, while Immer et al.
[7] argued that network linearization is the correct way to make predictions in Gauss-Newton-based Gaussian approximations, it is not generally applicable.
We show this in Fig.
4: Everything else being equal, MC integration ( S = 10000 ) and linearization yield different results especially in terms of predictive uncertainty. Since one should -    -   (a) Predictive distributions. (b) Densities of pred. means.
MC
Lin. (c) Densities of pred. std. dev. prefer MC integration in this many-sample regime, linearization is thus not accurate for the general cases.
Methods [Table] Moreover, we show that the multi-class probit approximation (MPA), which is often used on top of a linearized output distribution p ( f ( x ∗ ) | x ∗ , D ) to obtain the predictive distribution p ( y ∗ | x ∗ , D ) [6, 17, 32], can also obscure the predictive performance, albeit often for the better [6, 32].
To that end, we use the same experiment setting as Section 3.2 and Table 1, and compare the MPA against MC integration.
The results are in Table 2.
Indeed, we see that the MPA yields better-calibrated predictive distributions.
However, this results should be taken with a grain of salt: The MPA ignores the off-diagonal elements of the covariance of p ( f ( x ∗ ) | x ∗ , D ) , as shown in (2)—further detail is in Appendix A.
That is, the MPA “biases” the predictive distribution p ( y | x ∗ , D ) since it generally assumes that f ∗ has lower uncertainty than it actually has.
And since p ( f ∗ | x ∗ , D ) is usually induced by the LA or VB which are often underconﬁdent on large networks [7, 33, 34], the bias of MPA towards overconﬁdence thus counterbalances the underconﬁdence of p ( f ∗ | x ∗ , D ) .
Therefore, in this case, the MPA can yield better-calibrated predictive distributions than MC integration [32, 6].
Nevertheless, it can also fail even in simple cases such as in Fig.
5, 6 where the MPA yields underconﬁdent predictions even near the training data (notice the lighter shade in Fig.
5, bottom). Moreover, the structured nature of the error of the MPA (cf.
Fig.
2 for the binary case) might also contribute to the cases where the MPA differs from MC integration.
Both examples above thus highlight the need of careful consideration when analytic alternatives to MC integration is employed.
Reﬁning Gaussian Approximate Posteriors
The previous analysis indicates the importance of of accurate approximate posteriors q ( θ ) for BNNs’ predictive distributions.
However, the current de facto way of obtaining accurate posterior approximations, HMC [18, 19], are very expensive since they require a full-batch of data in their updates [35].
While mini-batch versions of HMC exist, they do not seem to yield as good of results as their full-batch counterparts. Indeed, Daxberger et al.
[6] showed that a well-tuned last-layer LA can outperform a state-of-the-art all-layer stochastic-gradient MCMC method [21].
Furthermore, these samplebased methods—both the full- and mini-batch versions, along with deep ensembles and their variants—are costly in terms of storage since one effectively must store S copies of the network’s parameters.
We, therefore, propose a simple post hoc technique for “reﬁning”
Gaussian approximations using normalizing ﬂows.
The resulting method is thus analytic yet can produce high-quality samples.
Let f θ be a NN equipped with a Gaussian approximate posterior q ( θ ) =
N ( θ | µ, Σ ) (e [Equation] g., via a
LA, VB, or SWAG [13]) on the parameter space under a dataset D .
Given a NF F φ of length (cid:96) with parameter φ , we obtain the reﬁned posterior by
[Equation]
Then, a reﬁnement of q ( θ ) amounts to minimizing the reverse KL-divergence to the true posterior, using the evidence lower bound (ELBO) as a proxy ( H below is the entropy functional):
[Equation]
Given a reﬁned posterior q φ ∗ ( (cid:101) θ ) and a test point x ∗ , we can obtain the predictive distribution via MC integration:
[Equation]
; where [Equation] Due to the expressiveness of NFs [36], we can expect based on the previous analysis that a large S is not necessary here to obtain good predictive performance.
We shall validate this in Section 6.
Last but not least, this reﬁnement technique is especially useful for last-layer BNNs since their parameter spaces typically have manageable dimensions, cf.
Section 6.4.[-
-]
Related Work
Normalizing ﬂows have previously been used for approximate Bayesian inference in BNNs.
An obvious way to do so, based on the ﬂexibility of NFs in approximating any density [36], would be to apply a NF on top of the standard normal distribution N ( θ | 0 , I ) to approximate p ( θ | D ) , see e.g.
Izmailov et al.
[38] and the default implementation of variational approximation with NF in Pyro [39].
We show in Section 6.3 that the subtle difference that we make—using an approximate posterior instead of an a priori distribution—is more cost-effective.
In a more sophisticated model, Louizos and Welling [10] combine VB with NF by assuming a compound distribution on each NN’s weight matrix and use the NF to obtain an expressive mixing distribution.
However, their method requires training both the BNN and the NF jointly from scratch.
In an adjacent ﬁeld, Maroñas et al.
[40] use
NFs to transform Gaussian process priors.
Posterior reﬁnement in approximate Bayesian inference has recently been studied.
Immer et al.
[7] proposes to reﬁne the LA using a Gaussian-based VB and Gaussian processes.
However, this implies that they still assume a Gaussian posterior.
To obtain a non-Gaussian posterior, Miller et al.
[41] form a mixture-of-Gaussians approximation by iteratively adding component distributions.
But, at every iteration, their methods require a full ELBO optimization, making it costly for BNNs.
A lower-cost
LA-based alternative to their work has also been proposed by Eschenhagen et al.
These methods have high storage costs since they must store many high-dimensional Gaussians.
By contrast, our method only does an ELBO optimization once and only requires storage of the base Gaussian and the parameters of a NF.
In a similar vein, Havasi et al.
[42] perform an iterative ELBO optimization for reﬁning a sample of a variational approximation.
But, this inner-loop optimization needs to be conducted each time a sample is drawn, e.g., during MC integration at test time.
Our method, on the other hand, only requires sampling from a Gaussian and evaluations of a NF, cf.
While our results also reafﬁrm the widely held belief that single-mode Gaussian approximations are inferior to more ﬁne-grained counterparts, our conclusion differs slightly from that of Wilson and
Izmailov [43]: They argue “[...]
Ultimately, the goal is to accurately compute the predictive distribution, rather than ﬁnd a generally accurate representation of the posterior.
[...]” and subsequently propose a multi-modal approximation to the true posterior.
Here, we show that ﬁnding an accurate (non-Gaussian) weight-space posterior approximation can still be a worthy goal, even when only utilizing a single mode of the true posterior. [Figure]  Experiments All implementations and experiments are available at: https://github.com/wiseodd/laplace_ refinement .
See Appendix B and Appendix C for additional details.
Datasets We validate our method using standard classiﬁcation datasets: Fashion-MNIST (FMNIST), CIFAR-10, and CIFAR-100.
For the out-of-distribution (OOD) detection task, we use three standard OOD test sets for each in-distribution dataset, see Appendix B.2 for the full list.
Finally, for the toy logistic regression experiment, a dataset of size 50 is generated by sampling from a bivariate, bimodal Gaussian.
Network architectures For the F-MNIST experiments, we use the LeNet-5 architecture [44].
Meanwhile, for the CIFAR experiments, we use the WideResNet architecture with a depth ofand widen factor of 4 [WRN-16-4, 37].
For the NF, we use the radial ﬂow [22] which is among the simplest and cheapest non-trivial NF architectures.
Baselines We focus on last-layer Bayesian methods to validate the reﬁnement technique.
9 We use the LA to obtain the base distribution for our reﬁnement method—following recent practice [6], we tune the prior precision via post hoc marginal likelihood maximization.
The No-U-TurnSampler Hamiltonian Monte Carlo [HMC, 18, 19] with 600 samples is used as a gold-standard baseline—all HMC results presented in this paper are well-converged in term of the Gelman-Rubin diagnostic [45], see Appendix B. Furthermore, we compare our method against recent, all-layer BNN baselines: variational Bayes with the Flipout estimator [VB, 46] and cyclical stochastic-gradient HMC [CSGHMC, 21].
For all methods, we use MC integration with 20 samples to obtain the predictive distribution, except for HMC and CSGHMC where we use S = 600 and S = 12 , respectively.
We use prior precisions of 510 and 40 for the last-layer F-MNIST and CIFAR experiments, respectively.
These prior precisions are obtained via grid search on the respective HMC baseline, maximizing validation log-likelihood.
Additionally, we use MAP and its temperature-scaled version (MAP-Temp) as non-Bayesian baselines.
More implementation details are in Appendix B.
Metrics We use standard metrics to measure both calibration and OOD-detection performance.
For the former, we use the negative log-likelihood (NLL) and the expected calibration error [ECE, 47].
For the latter, we employ the false-positive rate at 95% true-positive rate (FPR95).
Finally, to measure the closeness between an approximate posterior to the true posterior, we use the maximum-mean discrepancy [MMD, 48] distance between the said approximation’s samples to the HMC samples, i.e. we use HMC as a proxy to the true posterior.
6.2 Toy example
We visualize different approximations to the posterior of the toy logistic regression problem in Fig.
Note that the true posterior density is non-symmetric and non-Gaussian.
[Table] The LA matches the weight-space posterior mean, but inaccurate the further away from it.
It even assigns probability mass on what are supposed to be low-density regions.
While VB yields a more accurate result than the LA, it still assigns some probability mass on low-density regions due to the symmetry of the Gaussian approximation.
Furthermore, it is unable to match the posterior’s mode well.
The proposed reﬁnement method, on the other hand, is able to make the LA more accurate—it yields a skewed, non-Gaussian approximation, similar to HMC [Table]
We further quantify the previous observation using the MMD distance between each approximation’s samples and HMC’s samples.
The LA, as expected, obtain the worst weight-space MMD.
While VB is better than the LA with an MMD, the reﬁned LA achieve even better MMD.
This quantiﬁes the previous visual observation.
6.3 Image classiﬁcation
[Table]
We present the calibration results in Table 3 using the LA and HMC as baselines, which represents a two “extremes” in the continuum of posterior approximations.
As has previously shown by e.g. Guo et al.
[49], the vanilla MAP approximation yields uncalibrated, low-quality predictive distributions in terms of NLL and ECE.
While the LA is a cheap way to improve MAP predictive, its predictive performance is still lagging behind HMC.
By reﬁning it with a NF, the LA becomes even better and closer to the
HMC predictive.
We also note that one does not need a complicated nor long (thus expensive) NF to achieve these improvements.
Furthermore, we observe a positive correlation between posterior-approximation quality (measured via MMD) and the predictive quality.
Considering that S = 20 is used, this validates our hypothesis that one can “get away” with fewer number of MC samples when accurate weight-space posterior approximations are employed.
[Figure]
Moreover, we present out-of-distribution (OOD) data detection in Table 4.
We observe that while the last-layer LA baseline can already be better than all-layer baselines—as also observed by Daxberger et al.
[6]—reﬁning it can yield better results: Even with a small NF, e.g., (cid:96) = 5 , the OOD detection performance of the reﬁned LA is close to that of the gold-standard HMC.
Finally, we show that it is indeed desirable to do reﬁnement , i.e. using an approximate posterior as the base distribution of the NF, instead of starting from scratch, i.e. starting from a data-independent distribution such as N (0 , I ) .
As shown in Fig.
7, starting from an approximate posterior yields better predictive distributions faster than when N (0 , I ) is used as the base distribution of the NF.
This is particularly important since the computational cost of a NF depends on its length: We see a 53% increase in training time from (cid:96) =
5 to (cid:96) = 10 —the latter is required for the a priori NF approximation to yield similar predictive performance to the reﬁned posterior.
6.4 Costs
The proposed reﬁnement technique is post hoc and cheap when applied to last-layer BNNs.
Suppose one already has a last-layer Gaussian approximate posterior.
Using a standard consumer GPU (Nvidia RTX 2080Ti), each epoch a length- 5 NF’s optimization takes around 3 . 4 seconds.
From our experiments, we found that a low number of epochs (we use 20 ) is already sufﬁcient for improving a crude approximate posterior.
Thus, the entire reﬁnement process is quick, especially when compared to MAP estimation, ELBO optimization, or HMC sampling. Additional details in Appendix C.3.
Concluding Remarks 7.1 Limitations
A NF is a composition of diffeomorphisms—it is bijective and invertible.
Since we use the NF to transform the density on the parameter space of a NN, which has tens of millions dimension, the proposed reﬁnement framework is thus costly for all-layer BNNs or when the number of last-layer parameters is prohibitively large.
Interesting future works to alleviate this limitation are to reﬁne a posterior in the subspace of the parameter space [38] or in the “pruned” parameter space [50].
Moreover, by reﬁning a Gaussian with a NF, one loses the convenient properties associated with
Gaussians, making the resulting BNN unsuitable for cheaply addressing tasks such as continual learning [51] and model selection [52].
Nevertheless, this is not unique to our method—any samplebased method such as MCMC and ensembles are unsuitable for such tasks.
7.2 Conclusion
We have shown that MC integration tends to produce underperforming predictive distributions under a crude posterior approximation, even with a large number of samples.
While recently analytic approximations such as network linearization and the probit approximation have been extensively used in BNNs, we show that one must be careful with them since they are biased and do not have theoretical guarantees like MC integration.
Inspired by the high predictive performance of full-batch
HMC—even when a low number of samples are used—we conﬁrm the widely-held belief that a ﬁner-grained posterior approximation is the key to achieving better predictive performance in BNNs.
To that end, we proposed a post hoc method for last-layer parametric BNNs, based on normalizing ﬂows, that can cheaply reﬁne crude posterior approximations.
In conjunction with the post hoc last-layer Laplace approximation, the proposed method can thus give practitioners similar predictive performance to full-batch, expensive HMC in a cost-efﬁcient manner. Acknowledgments and Disclosure of Funding
The authors gratefully acknowledge ﬁnancial support by the European Research Council through
ERC StG Action 757275 / PANAMA; the DFG Cluster of Excellence “Machine Learning - New
Perspectives for Science”, EXC 2064/1, project number 390727645; the German Federal Ministry of Education and Research (BMBF) through the Tübingen AI Center (FKZ: 01IS18039A); and funds from the Ministry of Science, Research and Arts of the State of Baden-Württemberg. The authors are also grateful to Felix Dangel for feedback. AK is also grateful to the International Max Planck
Research School for Intelligent Systems (IMPRS-IS) for support.
Appendix A Derivation of the Multi-Class Probit Approximation
Remark. This derivation ﬁrst appeared in the ﬁrst author’s blog post [53]. Let p ( z ) : = N ( z | µ, Σ ) be the distribution a random variable z ∈ R c .
The multi-class probit approximation (MPA) is deﬁned by
[Equation]
where the vector division is deﬁned component-wise.
Its derivation, based on Lu et al. [54] is as follows. Notice that we can write the i -th component of softmax( z ) as 1 / (1 + (cid:80) j (cid:54)
=
i exp( − ( z i − z j ))) . So, for each i = 1 , . . , c , using z ij := z i − z j , we can write
[Equation]
Then, we use the following approximations: E ( f ( x )) ≈ f
( E ( x )) for a non-zero function f , 2. the mean-ﬁeld approximation N ( z | µ, Σ ) ≈ N ( z | µ, diag( Σ )) , and thus we have z ij := z i − z j ∼ N ( z ij | µ i − µ j , Σ ii + Σ jj ) , and 3. using the (binary) probit approximation (see Section 2.1.2), with a further approximation:
[Equation] we obtain (cid:90) [Equation] We identify the last equation above as the i -th component of softmax
[Equation]
We use the Pyro library [39] to implement the normalizing ﬂow (NF) used for the reﬁnement.
The
NF is trained by maximizing the evidence lower bound using the Adam optimizer [55] and the cosine learning rate decay [56] for 20 epochs, with an initial learning rate of 0 . Following [35], we do not use data augmentation. For the HMC baseline, we use the default implementation of NUTS in Pyro. We conﬁrm that the HMC used in our experiments are well-converged:
The average Gelman-Rubin (cid:98) R ’s are 0 . 998 , 0 . 999 , 0 . 997 , and 1 . 096 —below the standard threshold of 1 . 1 —for the last-layer F-MNIST, last-layer CIFAR-10, last-layer CIFAR-100, and all-layer F-MNIST experiments, respectively. For the MAP, VB, and CSGHMC baselines, we use the same settings as Daxberger et al.
[6]: We train them for 100 epochs with an initial learning rate of 0 . 1 , annealed via the cosine decay method [56]. The minibatch size is 128 , and data augmentation is employed. For MAP, we use weight decay of 5 × 10 − 4 . For VB and CSGHMC, we use the prior precision corresponding to the previous weight decay value.
For the LA baseline, we use the laplace-torch library [6]. The diagonal Hessian is used for CIFAR100 and all-layer F-MNIST, while the full Hessian is used for other cases. Following the current best-practice in LA, we tune the prior precision with post hoc marginal likelihood maximization [6]. Finally, for methods which require validation data, e.g. HMC (for ﬁnding the optimal prior precision), we obtain a validation test set of size 2000 by randomly splitting a test set.
Note that, these validation data are not used for testing. B.2 Datasets For the dataset-shift experiment, we use the following test sets: Rotated F-MNIST and Corrupted CIFAR-10 [57, 58].
Meanwhile, we use the following OOD test sets for each the in-distribution training set: • F-MNIST: MNIST, K-MNIST, E-MNIST. • CIFAR-10: LSUN, SVHN, CIFAR-100. • CIFAR-100: LSUN, SVHN, CIFAR-10.
To complement Table 3 in the main text, we present results for additional metrics (accuracy, Brier score, and ECE) in Table 5. We see that the trend Table 3 is also observable here. We also show that the reﬁnem
In Table 6, we observe that reﬁning an all-layer posterior improves its predictive quality further.
In Table 7, we present the detailed, non-averaged results to complement Table 4.
Moreover, we also present dataset-shift results on standard benchmark problems (Rotated F-MNIST and Corrupted CIFAR-10). In both cases, we observe that the performance of the reﬁned posterior approaches HMC’s.
C.2 Text classiﬁcation
We further validate the proposed method on text classiﬁcation problems.
We the benchmark in
Hendrycks et al. [59], Kristiadi et al. [60]: A GRU recurrent network architecture [61] is used and trained for 10 epochs under the 20NG, SST, and TREC datasets. The prior precision for HMC and the reﬁnement method is 40 .
We use the same normalizing ﬂow as in the image classiﬁcation experiment and follow a similar training procedure. The results are in Table 8. We observe similar results as in the previous experiment: Reﬁning a LA posterior, even with a very simple and cheap NF, makes the weight-space posterior closer to that of
HMC. Moreover, the calibration performance (in terms of NLL) also becomes better.
[Table] [Table]
We compare the theoretical costs of the reﬁnement method in Table 9. For comparisons, we include the costs of popular Bayesian baselines: deep ensemble [20], SWAG [13] [Table]
and MultiSWAG [43]. Our reﬁnement method introduces overheads over the very cheap Kronecker-factored last-layer LA [6,
Fig.
5, Tab. 5] due to the need of training the normalizing ﬂow and feeding posterior samples forward through it during prediction.
Combined with our ﬁndings in Section 6.4, our method only introduces a small overall overhead on top of the standard MAP network.
Finally, using synthetic datasets, we show the empirical costs of performing reﬁnement on different number of classes in Fig. We note that even with 4096 features (e
[Table]
g. in WideResNet-50-2) and 1000 classes (e
[Table]
g. in ImageNet), the reﬁnement method is cost-efﬁcient. C.4 Weight-space distributions obtained by reﬁnement To validate that the reﬁnement technique yields accurate posterior approximations
[Table]
we plot the empirical marginal densities q ( w i | D ) in Figs. 10 to 12.
We validate that the reﬁnement method makes the crude, base LA posteriors closer to HMC in the weight space.
[Table]
[Figure] [Figure] [Figure]
Corruption Intensity Corrupted CIFARMAP LA LA-Reﬁne-1 LA-ReﬁneLA-Reﬁne-10 LA-Reﬁne [Figure]
[Table]
[Table]
NM + C 3 + P
SCP
M + C 2 + P [Table] NM + C 3 + P 3 + [Table]
SFCP
M + C 2 + P 2 + FCP
[Table]
[Figure]
[Figure]
[Figure]
[Figure]