Posterior Reﬁnement Improves Sample Efﬁciency in Bayesian Neural Networks
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Monte Carlo (MC) integration is the de facto method for approximating the predictive distribution of Bayesian neural networks (BNNs).
But, even with many MCsamples, Gaussian-based BNNs could still yield bad predictive performance dueto the posterior approximation’s error.
Meanwhile, alternatives to MC integrationtend to be more expensive or biased.
In this work, we experimentally show that thekey to good MC-approximated predictive distributions is the quality of the approximate posterior itself.
However, previous methods for obtaining accurate posteriorapproximations are expensive and non-trivial to implement.
We, therefore, proposeto reﬁne Gaussian approximate posteriors with normalizing ﬂows.
When appliedto last-layer BNNs, it yields a simple post hoc method for improving pre-existingparametric approximations.
We show that the resulting posterior approximation iscompetitive 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, resultingin Bayesian neural networks (BNNs), is a principled way of mitigating this issue [4].
Indeed, evencrude approximations of BNNs’ posteriors, such as the Laplace approximation [5], can lead to goodpredictive uncertainty quantiﬁcation (UQ) performance [6], provided the predictive distributions arecorrectly computed [7].
A prediction in a BNN amounts to an integration of the likelihood w.r.t.
the (approximate) posteriormeasure.
Due to the non-linearity of NNs, no analytic solution to the integral exists, even whenthe likelihood and the approximate posterior are both Gaussian.
A low-cost, unbiased, stochasticapproximation can be obtained via Monte Carlo (MC) integration: obtain S samples from theapproximate posterior and then compute the empirical expectation of the likelihood w.r.t.
thesesamples.
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 errorscaling of Θ(1 / √ S ) , intuitively, MC integration with a small S is inadequate for an accurateprediction—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], itis 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 thatfew-sample MC integration is inaccurate for even an “easy” integral such as when the domain ofintegration is the output space of a binary-classiﬁcation BNNs, cf.
Fig.
1 (left).
Further, we showthat its alternative, the network linearization, disagrees with large-sample MC integration, makingit unsuitable as a general-purpose predictive approximation method, i.e. when the Gauss-Newton
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 ) might still be useful for making predictions, provided posterior measure
“closeenough” implies improving
accuracy approximations.
Nevertheless, prior methods for obtaining expressive posteriors [e.g., 10, 21] require either signiﬁcantmodiﬁcation to the NN or storing many copies of the parameters.
Moreover, they require training fromscratch 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 byleveraging normalizing ﬂows [22].
Contrary to the existing normalizing ﬂow methods with a prioribase 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 conjunctionwith 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 competi-
tive with the much more expensive full-batch Hamiltonian Monte Carlo.
Preliminaries2.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 randomparameter θ ∼ p ( θ ) on R d and a likelihood p ( y | f θ ( x )) , along with the associated posteriorp ( θ | D ) ∝ p ( θ ) (cid:81)
m i
=1 p ( y i | f θ ( x i )) .
However, the exact posterior p ( θ | D ) is intractable in generaldue to its normalization constant.
Approximate BNNs , which approximate p ( θ | D ) with only itssamples 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.
Themain idea is to construct a local Gaussian approximation to p ( θ | D ) , centered at a maximum aposteriori (MAP) estimate θ MAP :
= arg max θ p ( θ | D ) , with the covariance matrix given by thenegative inverse-Hessian ( −∇ 2 θ log p ( θ | D ) | θ MAP ) − 1 .
Since the MAP estimation is the standardtraining procedure for NNs, the LA can be efﬁciently applied on top of pre-trained NNs [6], especiallydue to recent efﬁcient second-order optimization libraries [23, 24].
Variational Bayes [VB, 25] assumes a family of simple approximate distributions over the parameterspace, e.g. M :=
{ q ( θ ) :
=
N ( θ | µ, Σ ) : ( µ, Σ ) ∈ R d + d 2 } , and ﬁnds a q ∈ M that is the closest top ( θ | 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 obtainparametric approximations to the posterior.
Instead, they collect samples from asymptotically thetrue posterior and use them to approximate integrals under the posterior measure.
A popular MCMCmethod is the Hamiltonian Monte Carlo (HMC) method [18, 19], where sampling processes arecasted as Hamiltonian dynamics.
Nevertheless, MCMC methods are generally expensive since theyrequire 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 ∗ isdistributed as p
[Equation]
|
[Equation]
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 approximateintegrals 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 witherror 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 Sincreases.
In the realm of BNNs, however, S is often chosen to be a small number, e.g. S = 20 , dueto the computational cost of evaluating f θ .
Analytic approximations Since a Gaussian N ( θ | µ, Σ ) is the de facto choice for q ( θ ) , an analyticapproximation 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 ) (cid:62) ( θ − µ )where J f ( x ) ∈ R d × c is the Jacobian matrix of the network output w.r.t.
θ at µ .
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 forbinary classiﬁcation where p ( y | f ∗ ) =
σ ( f ∗ ) for the logistic function σ , one can use the so-calledprobit approximation [27, 15]: p ( y = 1 | x ∗ , D )
≈ σ ( f µ ( x ∗ ) / (cid:112)
+ π/ 8 S ( x ∗ )) .
Moreover, forthe multi-class case, we can use the multi-class probit approximation [MPA, 28]:
[Equation]
where the vector division is taken component-wise.
[Figure] [Figure]
Probit 0.
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 .
Ifp 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 constructinga sophisticated diffeomorphism F , by composing several simple parametric ones.
The resultingdiffeomorphism 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 aknown 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 closeto p ∗ .
In Bayesian inference, the reverse KL-divergence D KL ( (cid:101) p φ , p ∗ ) is often used.
Pitfalls of BNNs’ Approximate Predictive Distributions
3.1 Accurate MC integration requires many samples
[Figure] [Figure]
[Figure] [Figure]
f
[Figure]
We begin our analysis with the simplest yet practicallyrelevant case.
Let x ∗ ∈ R n and f ∗ :
= f ( x ∗ ) , with p ( f ∗ | 126 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. 2)
I ( m,
[Equation]
) :=
[Equation]
Note that this integral is prevalent in practice, e.g. in linearized or last-layer classiﬁcation BNNs [4, 7, 17, 32].
We compare a MC integration with S = 100 against the probit approximation, using a trapezoidquadrature with 20000 evaluation points to represent a gold-standard baseline.
That is, we computethe discrepancy | (cid:101) I ( m, s ) − I ∗ ( m, s ) | where (cid:101) I ( m, s ) is I ( m, s ) computed either with MC integrationor 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 —MCintegration 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 probitapproximation, as to be expected given its theoretical guarantees.
This highlights the ﬂaw of fewsample MC integration—accurate MC integration generally requires large number of samples.
-    -   (a) Predictive distributions.
(b) Densities of pred. means.
MC
Lin.
(c) Densities of pred.
std. dev.
3.2 Many-sample MC integration is not sufﬁcient
[Table]
However, even with a large S , MC integration can still failto yield good predictive performance in BNNs.
This canhappen when the approximation q ( θ ) used in (1) is an inaccurate approximation of p ( θ | D ) —virtually the case forevery parametric BNN.
Thankfully, there is some evidencethat the error of MC integration might be relatively smallin comparison to the error generated by crude posteriorapproximations.
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 performa last-layer Laplace approximation with the exact Hessian and a full-batch NUTS-HMC [19], underthe 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 posteriorapproximations are most important for BNNs.
Furthermore, this also shows that with a ﬁne-grainedposterior approximation, the predictive performance of a BNN is less sensitive to the number of MCsamples, leading to better test-time efﬁciency.
3.3 Analytic alternatives to MC integration are not the deﬁnitive answers
[Table]
Of course, ﬁne-grained posterior approximations are expensive.
So, a natural question is whether one can keep acheap but crude posterior approximation by replacing MCintegration.
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 perexample Jacobian J f ( x ∗ ) , where for each test point x ∗ ,one must store the associated d × c matrix, where d couldeasily be tens of millions (e.g. in ResNets) and c could bein the order of thousands (e.g. ImageNet).
Second, while
Immer et al.
[7] argued that network linearization is thecorrect way to make predictions in Gauss-Newton-based
Gaussian approximations, it is not generally applicable.
We show this in Fig.
Everything elsebeing equal, MC integration ( S = 10000 ) and linearization yield different results especially interms of predictive uncertainty.
Since one should prefer MC integration in this many-sample regime,linearization is thus not accurate for the general cases.
Moreover, we show that the multi-class probit approximation (MPA),which is often used on top of a linearized output distributionp ( f ( x ∗ ) | x ∗ , D ) to obtain the predictive distribution p ( y ∗ | x ∗ , D )[6, 17, 32], can also obscure the predictive performance, albeit oftenfor the better [6, 32].
To that end, we use the same experimentsetting 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 ofp ( f ( x ∗ ) | x ∗ , D ) , as shown in (2).
That is, the MPA “biases” thepredictive distribution p ( y | x ∗ , D ) since it generally assumes that f ∗has lower uncertainty than it actually has.
And since p ( f ∗ | x ∗ , D ) isusually induced by the LA or VB which are often underconﬁdent onlarge networks [7, 33, 34], the bias of MPA towards overconﬁdencethus counterbalances the underconﬁdence of p ( f ∗ | x ∗ , D ) .
Therefore, in this case, the MPA canyield better-calibrated predictive distributions than MC integration [32, 6].
Nevertheless, it can alsofail even in simple cases such as in Fig.
5, 6 where the MPA yields underconﬁdent predictions evennear the training data.
Moreover, the structured nature of the error of the MPA (cf.
Fig.
3 for thebinary case) might also contribute to the cases where the MPA differs from MC integration.
Bothexamples above thus highlight the need of careful consideration when analytic alternatives to MCintegration is employed.
Reﬁning Gaussian Approximate Posteriors
The previous analysis indicates that accurate approximate posteriors q ( θ ) are most important for
BNNs’ predictive distributions.
While one can obtain them via MCMC methods such as HMC[18, 19], in practice, those methods are very expensive since they require a full-batch of data intheir updates [35].
While mini-batch versions of MCMC methods exist, they do not seem to yieldas good of results as their full-batch counterparts—indeed, Daxberger et al.
[6] even showed thata well-tuned last-layer LA can outperform a state-of-the-art all-layer stochastic-gradient MCMCmethod [21].
Furthermore, these sample-based methods—both the full- and mini-batch versions,along with deep ensembles and their variants—are costly in terms of storage since one effectivelymust store S copies of the network.
We, therefore, propose a simple post hoc technique for “reﬁning”
Gaussian approximations using normalizing ﬂows.
The resulting method is thus analytic yet canproduce high-quality samples.
Let f θ be a NN equipped with a Gaussian approximate posterior q ( θ ) =
N ( θ | µ, Σ ) (e.g., via a
LA, VB, or SWAG [13]) on the parameter space under a dataset D .
Given a NF F φ of length (cid:96) withparameter φ , 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 MCintegration:
[Equation]
Due to the expressiveness of NFs [36], we can expect based on the previous analysis that a large S isnot necessary here to obtain good predictive performance.
We shall validate this in Section 6.
Lastbut not least, this reﬁnement technique is especially useful for last-layer BNNs since their parameterspaces typically have manageable dimensions, cf.
Section 6.4.
Related Work
Normalizing ﬂows have previously been used for approximate Bayesian inference in BNNs.
Anobvious way to do so, based on the ﬂexibility of NFs in approximating any density [36], would be toapply 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 posteriorinstead of an a priori distribution—is more cost-effective.
In a more sophisticated model, Louizosand Welling [10] combine VB with NF by assuming a compound distribution on each NN’s weightmatrix and use the NF to obtain an expressive mixing distribution.
However, their method requirestraining 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 impliesthat they still assume a Gaussian posterior.
To obtain a non-Gaussian posterior, Miller et al.
[41] forma mixture-of-Gaussians approximation by iteratively adding component distributions.
But, at everyiteration, 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 methodshave high storage costs since they must store many high-dimensional Gaussians.
By contrast, ourmethod only does an ELBO optimization once and only requires storage of the base Gaussian andthe parameters of a NF.
In a similar vein, Havasi et al.
[42] perform an iterative ELBO optimizationfor reﬁning a sample of a variational approximation.
But, this inner-loop optimization needs to beconducted each time a sample is drawn, e.g., during MC integration at test time.
Our method, on theother hand, only requires sampling from a Gaussian and evaluations of a NF, cf.
Experiments6.1 Setups
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 threestandard OOD test sets for each in-distribution dataset, see Section A.2 for the full list.
Finally, forthe 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 [43].
Meanwhile, for the CIFAR experiments, we use the WideResNet architecture with a depth of 16and widen factor of 4 [WRN-16-4, 37].
For the NF, we use the radial ﬂow [22] which is among thesimplest and cheapest non-trivial NF architectures.
Baselines We focus on last-layer Bayesian methods to validate the reﬁnement technique.
Weuse 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-standardbaseline—all HMC results presented in this paper are well-converged in term of the Gelman-Rubindiagnostic [44], see Appendix A. Furthermore, we compare our method against recent, all-layer BNNbaselines: variational Bayes with the Flipout estimator [VB, 45] and cyclical stochastic-gradient HMC[CSGHMC, 21].
For all methods, we use MC integration with 20 samples to obtain the predictivedistribution, except for HMC and CSGHMC where we use S = 600 and S = 12 , respectively.
Finally, 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.
More implementation details are in Appendix A.
Metrics We use standard metrics to measure both calibration and OOD-detection performance.
Forthe former, we use the negative log-likelihood (NLL) and the expected calibration error [ECE, 46].
For the latter, we employ the false-positive rate at 95% true-positive rate (FPR95).
Finally, to measurethe closeness between an approximate posterior to the true posterior, we use the maximum-mean
[Figure]
[Table]
discrepancy [MMD, 47] 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.
The LA matches the weight-space posterior mean, but inaccurate the further away from it.
It evenassigns probability mass on what are supposed to be low-density regions.
While VB yields a moreaccurate result than the LA, it still assigns some probability mass on low-density regions due to thesymmetry of the Gaussian approximation.
Furthermore, it is unable to match the posterior’s modewell.
The proposed reﬁnement method, on the other hand, is able to make the LA more accurate—ityields a skewed, non-Gaussian approximation, similar to HMC.
We further quantify the previous observation using the MMD distance between each approximation’ssamples and HMC’s samples.
The LA, as expected, obtain the worst weight-space MMD.
While VBis better than the LA with an MMD, the reﬁned LA achieve even better MMD.
This quantiﬁes theprevious visual observation.
6.3 Image classiﬁcation
[Table]
We present the calibration results in Table 3 using the LAand HMC as baselines, which represents a two “extremes”in the continuum of posterior approximations.
As haspreviously shown by e.g. Guo et al.
[48], the vanilla MAPapproximation yields uncalibrated, low-quality predictivedistributions in terms of NLL and ECE.
While the LA isa cheap way to improve MAP predictive, its predictiveperformance is still lagging behind HMC.
By reﬁning itwith a NF, the LA becomes even better and closer to the
HMC predictive.
We also note that one does not needa complicated nor long (thus expensive) NF to achievethese 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 fewernumber 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 LAbaseline can already be better than all-layer baselines—asalso observed by Daxberger et al.
[6]—reﬁning it can yieldbetter results: Even with a small NF, e.g., (cid:96) = 5 , the OODdetection performance of the reﬁned LA is close to that ofthe gold-standard HMC.
Finally, we show that it is indeed desirable to do reﬁnement ,i.e. using an approximate posterior as the base distributionof the NF, instead of starting from scratch, i.e. startingfrom a data-independent distribution such as N (0 , I ) .
Asshown in Fig.
7, starting from an approximate posterioryields 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:
Wesee a 53% increase in training time from (cid:96) =
5 to (cid:96) = 10 —the latter is required for the a priori NFapproximation 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.
Supposeone 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 ourexperiments, we found that a low number of epochs (we use 20 ) is already sufﬁcient for improving acrude approximate posterior.
Thus, the entire reﬁnement process is quick, especially when comparedto MAP estimation, ELBO optimization, or HMC sampling.
Concluding Remarks7.1 Limitations
A NF is a composition of diffeomorphisms—it is bijective and invertible.
Since we use the NF totransform the density on the parameter space of a NN, which has tens of millions dimension, theproposed reﬁnement framework is thus costly for all-layer BNNs or when the number of last-layerparameters is prohibitively large.
Interesting future works to alleviate this limitation are to reﬁne aposterior in the subspace of the parameter space [38] or in the “pruned” parameter space [49].
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 continuallearning [50] and model selection [51].
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 undera crude posterior approximation, even with a large number of samples.
While recently analyticapproximations such as network linearization and the probit approximation have been extensivelyused in BNNs, we show that one must be careful with them since they are biased and do not havetheoretical 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 hoclast-layer Laplace approximation, the proposed method can thus give practitioners similar predictiveperformance to full-batch, expensive HMC in a cost-efﬁcient manner. Checklist For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] Section 7.1.
(c) Did you discuss any potential negative societal impacts of your work?
[N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A]
(b) Did you include complete proofs of all theoretical results?
[N/A]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes]
(b) Did you specify all the training details (e
g., data splits, hyperparameters,
how they
were chosen)?
[Yes] Appendix A.
(c) Did you report error bars (e
g., with respect to the random seed after running experiments multiple times)?
[Yes]
(d) Did you include the total amount of compute and the type of resources used (e
g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] Section 6.4.
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]