Collapsed Variational Bounds for Bayesian Neural Networks
Authors: Marcin B. Tomczak
Siddharth Swaroop
Andrew Y. K. Foong
Richard E Turner
Venue: NeurIPS 2021 Poster (2021)
Keywords: Bayesian Neural Networks
BNNs
Bayesian Deep Learning
variational inference
VI
underfitting
collapsed ELBO
overpruning

Abstract: Recent interest in learning large variational Bayesian Neural Networks (BNNs) has been partly hampered by poor predictive performance caused by underfitting, and their performance is known to be very sensitive to the prior over weights. Current practice often fixes the prior parameters to standard values or tunes them using heuristics or cross-validation. In this paper, we treat prior parameters in a distributional way by extending the model and collapsing the variational bound with respect to their posteriors. This leads to novel and tighter Evidence Lower Bounds (ELBOs) for performing variational inference (VI) in BNNs. Our experiments show that the new bounds significantly improve the performance of Gaussian mean-field VI applied to BNNs on a variety of data sets, demonstrating that mean-field VI works well even in deep models. We also find that the tighter ELBOs can be good optimization targets for learning the hyperparameters of hierarchical priors.

50 3 although there Collapsed Variational Bounds for Bayesian Neural Networks
Marcin Tomczak, Siddharth Swaroop, Andrew Y. Foong, Richard E. Turner University of Cambridge Cambridge, UK
sss163,,221,,,yp ret26 ret.2)) ac.
Abstract
Recent interest in learning large variational Bayesian Neural Networks (BNNs) has been partly hampered by poor predictive performance caused by underfitting, and their performance is known to be very sensitive the prior over weights. Current practice often fixes the prior parameters to standard values or tunes them using heuristics or cross- validaani....... In this paper, we treat prior parameters in a distributional way by extending the model and collapsing the variational bound with respect to their posteriors. This leads to novel and tighter Evidence Lower Bounds (ELBOs) for performing variational inference (VI) in BNNs. Our experiments show that the new bounds significantly improve the performance of Gaussian mean-field VI applied to BNNs on variety of data sets, demonstrating that mean-field VI works well even in deep models. We also find that the tighter ELBOs can be good optimization targets for learning the hyperparameters of hierarchical priors.
Introduction
There has been lot of recent interest in developing methods for Bayesian Neural Networks (BNNs), and variational methods provide computationally cheap approximations to the posterior over weights when compared to alternatives like MCMC. Research Variational Inference (VI) for BNNs has
multiple advances enabling optimization of the Evidence Lower Bound (ELBO)
P
65
24
f B8)
allowing for scaling large neural networks and datasets (57
However, the under- confidence of mean-field VI (MF-VI) in output-space has caused poor adoption of VI in applications,
66
especially in sequential decision making [ 60 69 Although lot of work has improved performance
24! 7 [8], this work has focused on the variational parameters of distributions over weights, with the prior over weights usually fixed to cero-mean isotropic Gaussians (55 (70 774 1 20
by deriving better estimators of the gradient of the ELBO [29 65
In this paper, also apply inference to the prior parameters of Gaussian BNNs (e.g. means and variances of a Gaussian prior over weights). We do this by using collapsed VI bounds B6 [ 34 which analytically solve the inference over prior parameters and allow to derive tighter ELBOs for performing VI in BNNs. Many previous works argue that point estimates for parameters of Gaussian prior in BNNs can be learned reliably by maximizing the ELBO 72
have also been arguments against such an approach [7]. Experiments using our derived ELBOs lead two findings: (i) MF- VI performs well in deep models where the prior parameters of Gaussian distributions over weights have been learned, but severely under-fitt fits data in output space if the prior over weights is fixed and, (ii) the derived tighter ELBOs are good optimization targets learn the hyperparameters of hierarchical Gaussian priors, even in large networks. We provide the code implementing the introduced algorithms at
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

2 Collapsed Variational Bounds for Bayesian Neural Networks
We consider performing in supervised learning scenario given N data points D {(.,, y; } We denote a neural network as fw, the observation model as f(y] where W denotes the vectorized weights and biases of the network, and the likelihood of the weights W as p(D) W)
I p(y fw (zz:)) We focus on Gaussian mean -ield priors (with non-zero means) over vectorized network weights pp o)) = NW w of)) meaning w N (pp. oz) for any coordinate w in W (both Mp and o are also vectors in this notation). Similarly, we restrict our attention to mean-field Gaussian approximate posteriors g( ((Wiz Ie +2) = -V(WI (WW (.) VI finds an approximate posterior q( My: og)) by minimizing the KL divergence (9 Ipt, ( ( Hpp of)) w.r.t. the variational parameters Mg o. This is done by optimizing the Evidence Lower Bound (ELBO)
28
C(u oe up o
log logp(v.l/w(az)) - Dww.(gWWiil My oWWl Hpp o.)). (WI,,
E
We denote the first term on RHS as Ldata (My o.)). As DEL (g( W Mp. o. Mp. o 0.
the ELBO lower bounds log marginal likelihood log p(l x²) 2 L(u oz: Ug. o o Previous
work 7 38) introduces low variance updates for mean-field Vi by reparametrizing weights W. Bayesian predictions are approximated by integrating the learned variational posterior q" (W Mg. og) as p(y" x D) = (W ffw (w(xx))
For BNNs, prior means are most often set to zero H, 0, and prior variances set to a constant of = 71 (often scaled by the size of a previous layer [55 ) We consider a more general setup where the weights of BNN are defined by hierarchical model (WIp,,, o
())((,)))(()))) and aim
to perform inference over weights W and up. o, For clarity, we will refer to up o as prior parameters, leaving the term hyperparameters for any other free variables.
2.1 Deriving collapsed variational bounds
In this section, we develop an efficient optimization scheme for variational BNNs with weights defined in hierarchical fashion, where the inference over prior parameters is done analytically and prior parameters are subsequently marginalized out. This defines new learning objectives for learning variational BNNs, which we show by providing specific examples later in the text. More specifically, we apply collapsed variational bounds 36 43 47 20 to BNNs. and present a general, systematic way of deriving novel, tighter lower bounds on the model evidence which are useful for learning variational posteriors over weights. The idea behind collapsed bound is simple: Suppose we maximize f(r, y) and we can derive the optimal x: (v) = argmax, (xr, v). Substituting xx into results in (xx(()),)) y). which only depends on y. is easier to optimize, and has the property that Vy f(=(()), 2 f)z Collapsed variational bounds apply this reasoning to the ELBO with factorized variational posteriors taking the roles of x and y.
Inferring prior parameters in BNN. We consider inference over both network weights W and prior parameters μp and o which we treat as latent variables and subsequently discuss collapsing the inference w.r.t. prior parameters f and o² We use the factorization q My: Tyll, ? =
q( My a (p) of), where we omit the variational parameters of q( o²)) for concise notation. Learning pt o²) under this assumption does not complicate computing the BNN's predictive distribution p(y (z. D), i.e. My o² do not need to be sampled at the prediction time, since,
p(yla xx D) = E ", -, (HDD))(ll (z)) = E u.p oz) (yll (z)) = E,(Win, E, fw (x)).
(2)
To apply inference, we invoke the variational principle which lower bounds the log marginal likelihood log p( D) as:
p(D, W. ppp
log (Dg)(D)= 2 Ee(Wle,
log g(W M 22 22) oz)
L (Mep 2 q(u, o
(3)
The bound in Eq. B reduces to L(ug o²) in Eq. ① when (1p oB) =(M)) and g(pp))
(()r)) are delta functions. If we are able to analytically derive the distribution 9 (M,, o maximizing
2

---

L we can substitute the resulting distribution into Eq. (3) to derive a tighter (collapsed) bound on the marginal likelihood, saving the need to perform coordmatt ascent or gradient updates learn the prior parameters:
PD, W. up oz)
logp(D) 2 Eg WWI,,, 200 log (( B og 22)y (pp. o - (u::
(4)
where [ 0g) 2
q. Kpp R.)) and D depends only on the variational parameters Hy o². Collapsed variational bounds are desirable as they have been shown to make learning significantly more efficient [64]. The tighter bound in Eq. (4) has been also referred to KL
corrected bound [36 and marginal VI bound 43 47
Collapsing the variational bound for BNNs. We now provide method to derive a collapsed variational bound by maximizing (Hg, oy A(p o)) w.r.t. the posterior over weights prior parameters (Hp op))) in BNNs. We use the property that the BNN likelihood term log p(yl x fw (a))
does not depend on the prior parameters Hp and o i.. p(D)W, Hp o p(D())) The bound _a from Eq.
G decomposes into three terms:
L (Mg o Aup o})) Latta (1g 02) + (Mg og Mlp oB)) + H of where H[g(W Mn, x.) denotes the differential entropy, and we have defined
(5)
@ (pg o. app o()) = E ((11..0)))
E 2(Wi (3)
log p( W Ww o/) - -DKu(a(AA)
(6)
Only & depends on the posterior glppp p)), and therefore finding g" (mppp of)) can be done by maximizing ((uy og 1(pp: o.)).
We can significantly simplify this by noting that the objective in Eq. O6 is analogous to Eq. ① hence optimizing & takes the form of a nested VI problem: we want to infer q(t o where the "data" distribution is replaced by the current variational posterior over weights q ((Wl,, (Mg oz)). This implies the solutions qº (Mpp. o.) are straightforward to derive in closed-form for priors/approximate posteriors from the exponential family [5 Specifically, we can write the maximizer 9" (Mpp as:
log q" (M x x log (og (eg(W, log p W Mpp o
Substituting the optimal variational posterior q" (Mp, 2n) into Eq. tute this into Eq. 5 give our final collapsed bound L (p. 2 Ladaaa + +h + H[g(([[gg]))
6 gives Φ*, and can substi-
(o(Wlp ppp R))] Furthermore, when calculating Ø", the inferred prior parameters q" q"((,,.. ? can be analytically integrated out in many cases, i.e. the outer expectation on the RHS in Eq. 6
be solved in
closed-form (we provide examples in Section 2.2 leading to a concise optimization objective.
The final optimization target to learn the variational posterior g g( (W is given by (" a²) Compared to gradient learning of fe, and of eliminating M, and o from L performs better (as we show in Section 4. and halves both the memory requirement and time taken per update. It is important remember that optimizing L learns both the variational parameters My og and the posterior over prior parameters q (4pp: 2 but the latter is implicit as we express it as a function of My o given by Eq. 7
In summary, there are four steps to derive a collapsed bound L" which we directly use to learn the variational posterior of a BNN. These steps are given in Algorithm We next give two concrete examples of collapsed bounds.
2.2 Examples of tighter ELBOs
Learn prior means, fix prior variances. As easier example, first consider fixing the prior variance of weights, and learning prior means. We follow the four steps of Algorithm | to derive our collapsed bound. Detailed derivations are in Appendix B
Step This step defines model and the family of variational approximations. We choose p( (1,) WWWWWWWWW HMp, >1) with Gaussian prior p(u,la)) -N(4,,0, a1), where and are
3

---

Algorithm Deriving a collapsed bound L for BNN in four steps.
Step
Choose priors p(up).p(o.) p( oz),
and approximate posteriors g(ppp 2)), q (W og) over network weights.
Step II:
Calculate optimal prior parameters 9" (Mp: p: oi) using Eq.
Step III:
Form &' and solve (n, [ (?) log W Hpp p in Eq. 660
Step IV:
The collapsed bound is L Ldata + x* + H|g(WW Mg 2
hyperparameters. We also employ mean-field Gaussian posterior q( ( (Wlpg of)) and Gaussian posterior (114)) (1t,)
Step H: This step analytically finds the optimal variational distribution q" (kp). Substituting our distributions into Eq. for every coordinate Up gives us
(14g Up)
2
log (1ppp) x
(8)
27
2nx
In this case the variational posterior q" (1e matches the true posterior, as the above dependency is equivalent Bayesian inference over Mp given Gaussian likelihood/prior and observation Mg. so qº (ppp) = ((uplat at7 1).
Step III: This step forms &' by calculating the divergences between optimal posteriors q" and priors and simplifies the derived optimization objective by marginalizing out approximate prior parameters. The divergence between two Gaussians D (q (Ip(()) is straightforward to compute. Next
we solve the integral I = ((,)) [ (3) log W Mpp +
ooo2 q (My pp) up) 1
o
THE g
Hy
I A
E
N
1)
27
27
2( + 7)3
Step IV: This step forms collapsed bound that will be used to learn variational posteriors. Substituting the above into L Ldata + x0 + H[[g( W lpp اf)ل results in our new bound:
D
S (pg: 2 A Ldatt (1y 2²)) 2n oTo + areglyy o/2 Ms] + 1 log of :
log Oregs
2
where we have defined the hyperparameter Oreg 7/(7 + a)) E (0, 1), and for fixed there is 1-to-1 relation between a and Oregg
When a - (so Oreg > 1) we recover the standard expression for the ELBO given by Eq. ① with
q( Mg: o and prior (WW10, This corresponds to a dogmatic prior p( = NNu, (14,00 a1) delta spike 0, and no inference of Mp, When allow for inference, the regulariza- tion of the (approximate) posterior mean H4 decreases. When areg - 0 (uninformative prior) the modeling can be prone over-fitting due insufficient regularization in the model specification. Note that this is reflected by the term 7 log Oreg which can be interpreted as an Occam's 's Razor penalty [59]: this term penalizes small areg
converging
D
diverges to negative infinity when Oreg O. Since areg (0, 1), this enables finding well-performing values of aryeg by evaluating C (without having validation set) 30), and we show this in our experiments in Section $
We find that Oreg E (0.01, 0.1) significantly improves upon the predictions of MF-VI with fixed Gaus- sian priors. For existing implementations learning BNNs with Gaussian variational posterior/prior we suggest the default setting aregg = 0.05 in front of the mean regularization in the expression for KL divergence instantaneously improve the predictions (as opposed to down-weighting the whole KL divergence term). Larger Oreg increases the strength of the regularization of the model.
Learn both prior means and variances. We now discuss scheme to learn hyperparameters of both prior means and variances. We follow the same four steps from Algorithm | detailed derivations to Appendix B
We again defer the
Step I: p(W Mpp: Tp) NN(W up + and p(u,ItT) N,,, + p(7) (( We again consider posteriors q( W Mg o NNW My o2). ((p) NN(H,, and q(T,) G(Tp).

---

L
C C+
C
14
14
14-
13
13
1.5
95
10-
13
16
02
6.2
82
-027
0.21
180
1001
-100
-1.1
-1.5
100
1150
875
Hg
14
14
Figure 1: Comparison of regularizing terms in L. S and C. scaled to unity. In the absence of gradients from Ldata (approximate) posteriors weights converge to the blue region (red cross- mark). Optimizing prunes weights to prior Gaussians which still influence predictions causing under-fitiing (see Section A and is opposite optimizing fmm which prunes spiked Gaussians.
Step II: The optimal posteriors q' are given by 4² (1pp) N( (.))))) and (Tp)
G(Tpl(a + 1, 2[4 u -
Steps III and IV: The bound L (H1 o becomes.
1 111 Lyv (My: 22)) Ldutaa og)) (a 2
8
1
log 31 log
2
2
+
2
1" log o
(10)
where = 1/(111111 t). Eq. (10) recovers the case of setting prior mean Mp = 0 and learning only prior variances of when the prior precision lTp over H, goes to oo, i.e. / + 0c. Setting < in Eq. (10
weakens the regularization of posterior mean y. allowing it to vary more. The regularizer in Eq. TTO
is a decreasing function of Mu. i.e. without gradients from Lattaa the posterior mean M4 of a weight w converges to 0. In Section we show that C; outperforms standard MF-VI.
Comparison of regularizing terms. We compare the regularization terms on μ, and o² in the standard mean-field ELBO in Eq. TT and introduced ELBOs Eq. ( and Eq. (T)) Figure [ where we use prior the (()= = .V (w We consider what happens in the absence of gradients from the data term Ldata better understand the bounds. In Figure (left) we see that Gaussian posteriors (in the absence of data) optimized with Eq. ( converge to their prior N ((w00, 1) 7 68
In Figure (middle) we are down-weighting !2 M4 Eq. hence weakening the regularization of ng. This allows posterior means vary more, and can be interpreted as scaling the μo axis in Figure by vo. As show in Section this enables posterior means Hlg saturate activations and reduce the noise from pruned weights. Optimizing Eq. (10) is roughly opposite to Eq. ① and causes weights to be pruned (00 2,3) (Figure right, where 3 0.01). Weights pruned in this way do influence the predictions for small |3 and this fixes the excessive injection noise by pruned weights, as we
show in Section
3 Related work
511, but optimizing Performance of BNNs is known to be sensitive to prior hyperparameters 54. hyperparameters in BNNs has not been widely adopted, although there exceptions [50 72
Some authors heuristically propose using Gaussian centered on the MAP estimate as a prior for The Laplace approximation has been used for hyperparameter selection in BNNs but it does not use the EI ELBO. Hierarchical models in BNNs have been explored in
BNNs 8
39
31
61
The closest our approach learning hierarchical horseshoe prior over BNNs weights 23 However, their approach is computationally costly. does consider collapsed bounds derive tighter ELBOs, and does learn prior means. There is an active discussion in the community as to whether the true BNN posterior provides satisfactory predictions with work supporting the Bayesian approach [54 74 II and arguing against [70Í]. Similarly, there are conflicting views on the performance of VF-VInn in BNNs with fixed Gaussian priors as some authors claim can work well 177 and others arguing against
many works 2
58 B. 12
| A number of problems of mean-field VI BNNs have been exposed in shallow networks including over-pruning [ and poor uncertainty in output space 19
18 23 To resolve under-fitting of mean-field BNNs,
previous work has re-weighted terms arising from the KL divergence e.g. 45
5

---

concrete
winte
yacht
enorgy
power
kirStu
protein
-08
11111
20
2.8
-3.0
-12
2254
VI
dropout
boston
concuete
wine
yachi
poweer
kinffm
prokión
-25-
00
0.0
-220
-03
650
-25-
12-
-10
16
-100
-236
22%
1:00
20
-125
-24-
-13
+1.0
-2.0
-12
C3EMEN1
-20
hidden layers
hidden Jayers
hidden
hidden
bidden layers
hidden
hidden
biciden Jayers
Figure 2: Top: Comparison of collapsed bounds with other algorithms on UCI regression data sets. Bottom: CM-MF-VI (red) and standard MF-VI with default isotropic Gaussian prior (blue) for deeper networks, where CM-MF-VI avoids severe degradation of predictions.
4 Experiments
In this section we explore the predictive performance of the introduced variational bounds and
benchmark them together with other algorithms. We find that using the bounds in Eq. ? and Eq. (TC)
significantly improves the predictive performance of standard MF-VI BNNs (fixed Gaussian priors). We also find that learning the prior means according Eq. is more effective and gives sufficiently good predictive performance, whereas learning prior variances (Eq. (T0) mostly does not lead to
significant gains in the quality predictions. We defer the details of the experimental setup to Appendix C and report additional experimental results in Appendix D
The algorithms propose in this paper optimize various tighter ELBOs: CM-MF-VV VI optimizes
L. from Eq. ( CM-MF-VV VI OPT additionally optimizes Oreng using the tighter ELBO (Eq.
9):
CV-MF-VI optimizes for just prior variances, using the bound Lmo from Eq. IC when ō (prior
mean H4, fixed to 0); and CMV-MF-V VI optimizes both prior means and variances from Eq. (TO). We
compare the following baselines: MF-VI LRT learns 7ero-mean mean-field Gaussian priors with fixed variances 55 (which we call standard MF-VI) using local reparametrization B MF-VI FV learns a mean-field Gaussian posterior with fixed variances [68] MF-VI BD is mean-field Gaussian VI while doig-weighting KL penalty by 0.5; MF-VI EB learns prior means with ELBO gradient updates 51 / BBB is Bayes by Backprop using the ADAM optimizer (as opposed SGD in (7);
MF-VI ki-teed normal distribution
633; ELRG/SLANG learn low rank posteriors 66 53
(a.)
ww'++b x
(b.)
Removing pruned posterior weights
21
MF-VI
30
+
activation
last layer
Figure (a.) Pruned weights in standard MF-VI contribute noise to predictions as noise leaks through the activation function. (b.) Removing pruned weights from BNN can improve performance, as the removed weights contribute only to the predictive variance and cause under-fitting.

---

400 units
S00 units
algorithm
test NLL +
test ER +
test NLL 4
test ER 4
CM-MF-VI
0.047 + 0.006
1.34 + 0.24%
0.048 1: 0.005
1.42 = 0.26%
CV-MF-VI VI
0.068 1 0.007
2.13 0.34%
0.066 1 0.003
2.09 0.10%
CMV-MF-VI
0.049 1 0.005
1.25%
0.052 0.004
1.60 = 0.01%
CM-ME-VI (4000 batch)
0.041 + 0.002
1.31 0.04%
0.042 + 0.002
1.20 = 0.06%
MF-VI (8
0.094 + 0.002
2.42=0228%
0.099 + 0.001
2.58 = 0.07%
MF-VI BD
0.092 + 0.002
2.48
0.099 + 0.004
2.56 0..2%
MF-VI FV
0.052 + 0.001
0.13%
0.053 + 0.004
1.58 0.26%
BBB ADAM 7
0.095 + 0.008
3.4800.00 0.45%
0.097 0.007
2.49 = 0.33%
EIRG-VI K 5 66
0.053 + 0.006
1.54 + 0.18%
0.058 + 0.005
1.68 + 0.17%
K-TIED K- = 10
63
0.105 + 0.004
2.67 + 0.16%
0.108 + 0.004
2.61 + 0.17%
MF-VI SGD 7
1.82%
1.99%
SLANG K=32 53
1.72%
Table 1: Test NLL and error rate for vectorized MNIST classification with two hidden layer BNN. CM-MF-VI outperforms other algorithms and collapsed bounds improve upon standard MF-VI.
(2.)
wF
(c.)
pakeroor meep OMMMAA
peraakee thara CMVAMF-N
bis
www
www cod He
activation
(b))
!wix<b
(d.)
CMd MI-V1
EMV-ME--1
activation
Figure 4: Learning prior parameters in BNN fixes under-fitting caused by over-pruning, In (c.). and (f.) we show the distribution of statistics of variational posterior over weights in the first layer. The top row plots (c.) and (c.) the posterior means, where red is the means of the bias parameters, and black is the means of the weight parameters into hidden unit. The bottom row plots (d.) and (f.) show posterior log variances instead of posterior means. Most weights have been pruned. In CM-MF- VI, shown in (a.), units where posterior weights converge to prior have posterior biases set to low value, which saturates the activation stopping them from influencing the predictions. We see this in (c.). where the pruned units all have biases with negative posterior means. This saturation is made possible by weaker mean regularization in Eq. (9) governed by Oreg. In CMV-MF-VI shown in (b.), pruned units converge spiked Gaussians. We see this in the (f.), where the posterior log variances of pruned units are all very small.
UCI regression. We first consider 20 train-test splits for 8 UCI regression data sets [4] learn 2 hidden layer BNNs (results for 1 hidden layer in Appendix D with 50 units and ReLU activations [27 21 67 but we use a heteroscedastic observation model p(y) Sw (x), fv
We
Ny (x), expí. (x))). where fw (x), (x) are two heads of the network. We optimize the objectives for 200K steps with the ADAM optimizer (37 with default settings. We report comparison of test log-likelihood in Figure 2 (top). We observe that (i) the collapsed bounds CM-MF- VI, CV-MF-VI, CMV-MF-VI outperform MF-VI on all data (only CMV-MF-VI is worse on wine), (ii) bounds learning prior means CM-MF-VI, CMV-MF-VI outperform or match the bound learning only prior variances CV-MF-VI on all data sets, (ii) CM-MF-VI outperforms matches MC-dropout on all data sets except for boston and wine (the performance of dropout is unstable unless p is tuned per data set, e.g. it is poor on kin8nm and not shown). In addition, CM-MF-VI OPT matches the performance of CM-MF-VI showing that the ELBO can be used to learn hyperparameter areg- Next we report a comparison between CM-MF-VI and standard MF-VI as the depth of the network (bottom). We see standard MF-VI starts to rapidly
increases in Figure 2 sets (for larger data sets standard MF-VI performs well as it converges to MAP). Learning prior means/ verann/vriances by optimizing bounds in Eq. 9 Eq. 10 allows to mitigate rapid degradation of
most data
7

---

predictions. We now explain why standard MF-VI VI performs badly, finding this is due to over- ver-pruning BNNs. We investigate this effect and find that learning prior means and variances mitigates the effect of pruned units on the predictive distribution.
Test NLL
Prodictive entropy
Entropy in-distribution
Entropy out-of-distrbution
MF VI
10
MC des dropout
28
CM SIFIVV VI
18-1
MB-V1
84
MC covousl CMUAAE-T
82-
00
0.1
D6
08
00
02
0-1
06
05
000
92
93
04
000 0.5
20
Figure 5: Comparison of standard MF-VI and collapsed bounds (CM-MF-VI, CV-MF-VI) for in- distribution and out-o--distribution data. Left: replacing p% of MNIST image with random bits. Right: learning BNN on MNIST (in-distribution) and testing on fashionMNIST nout-of-distribution)
Over-pruning in BNNs. Applying mean-field VI to BNN prunes most of the posterior weights, i.e. approximate posteriors converge to prior 7
68 23 52. While this is beneficial for compressing the model, pruned weights contribute noise to predictions and make modeling data difficult. With I hidden layer networks and the regular ELBO, output layer weights are set as close to 0 as possible [8.. thereby reducing the noise from pruned weights in earlier layers. However, some noise This noise is one of the reasons for increasing under-confidence in output space as the size of the network increases [68] 16
still leaks through, as output layer weights are not quite delta-function
66 23 45 In Figure B we highlight this problem for 1 hidden layer MF-VI BNN with 50 hidden units, ReLU activations, zero-mean Gaussian prior over weights and heteroscedastic Gaussian likelihood learned with 200 random boston regression data points. In Figure 3 (right) we plot test NLL as pruned hidden layer's s units are removed in ascending order of average KL penalties of incoming weights (only last 4 units are modeling data). Removing pruned weights lowers test NLL, meaning that pruned weights contribute only noise to predictions. This is schematically explained in Figure 3 (a), which shows pruned posterior weights injecting noise. See Appendix E for further description of peer-pruning.
63
Fixing under-fitting caused by over-pruning in BNNs. We now explain why learning the hyperpa- rameters of the Gaussian prior in the BNN mitigates problems arising from over-pruning In short, learning the means and variances of Gaussian priors mitigates the influence of network units that have been pruned. These units otherwise cause under-fitting by contributing only to predictive variance. In Figure
3 we plot the distribution of learned mean-field Gaussian posterior means and log variances (middle and right) for units feeding into the output layer. We use a hidden layer heteroscedastic BNN with 50 hidden units and ReLU activations learned on 200 data points from boston regression data set. We see that CM-MF-VV VI sets means of posterior biases to lower values (%= -2) to saturate the activations for units where incoming weights converge to the prior. Saturated activations stop the pruned weights from contributing variance into the predictions and result in better modeling of data, schematically show in Figure g (a). For CMV-MF-V VI, posterior weights corresponding to inactive units converge to delta spikes with very slightly negative means, in line with observations made in Figure 1. hence they do not influence predictions, shown in Figure (b).
MNIST classification with MIP. Classifying vectorized MNIST images using two hidden layer
network with ReI.U activation is standard benchmark for BNNs 7 We demonstrate that collapsed bounds give large improvements in test NLL compared to using standard MF-VI and provides the best predictions across tested algorithms. We report test NLLs and test error rates (ER) averaged over random seeds in Table I For 400 hidden units, CM-MF-VIIL WLMM--/CCC-MM--VVV achieve test NLL of 0.047, compared to 0.094 for standard MF-VI. CM-MF-VI (4000 batch) uses batches of 4000 images for an optimization step and achieves the best performance across the tested models, showing the developed algorithms can leverage low variance updates.
Perturbed MNIST images. We now investigate if the in-domaín improvements of CM-MF- (OOD) performance. We find this is not the
VI/CMV-MFVV VI come at the cost of out-of-distribune case: OOD performance is as good as before. To show this, we learn a hidden layer BNN with ReLU activations on vectorized MNIST images as previously, but test it on the MNSH data
8

---

Test NLL 1/ERR
MNIST
K-MNIST
F-MNIST
SVHN
CIFARIO
CMV-MF-VI
0.021 + 0.001
0.152 0.006
0.253 0.006
0.313 + 0.006
0.807
0.005
CM-MF-VI
0.021 0.001
0.141 0.006
0.204+00003
=.05
0.809 0.009
CV-ME-VI
0.088 0.0003
0.239 = 0.005
0.317
0.321
0.821 0.005
CM-ME-VI OPT
0.001
0.158 0.004
1258
0.314
0.789 0.005
MF-VI
0.061 0.0001
0.319 = 0.006
0.371
0.340 0.001
0.848 0.009
MAP
0.048 0.007
0.402 0.027
0.775 0.012
1.134 0.065
MC dropout
0.001
0.222 0.011
0.326
0.400 0.009
1.018 0.017
ME-VI EB
0.060 0.001
0.319 = 0.003
0.372 0.003
0.340 - 0.005
0.843 = 0.010
CMV-ME-VI
0.73 0.02%
4.00 + 0.06%
8.90
0.22%
8.49 + 0.20%
27.76 + 0.37%
CM-ME-VT
0.67 + 0.04%
3.75 = 0.30%
9.05 + 0.31%
8.21 + 0.20%
27.63 + 0.48%
CV-MF-VI
1.24 0.12%
6.85 + 0.10%
0.13%
8.96 + 0.20%
2 0.14%
CM-MF-VI OPT
0.72 0.05%
4.23 323=0000 0.06%
0.07%
8.47 0.03%
26.80 1.401004499 0.41%
MF-VI
1,51 0.05%
9.17=02259
12.96 0.28%
9.20 0.15%
28.71 0.16%
MAP
1.16 0.11%
7.92 = 0.36%
1..229
12.78 + 0.12%
34.63 = 1.48%
MC dropout
0.81 + 0.05%
6.07 = 0.49%
12.01 + 0.17%
9.78 + 0.22%
33.86 = 0.42%
MF-VI EB
1.53 + 0.07%
9.18 = 0.19%
12.96 + 0.08%
9.31 = 0.28%
28.56 = 0.45%
Table 2: Image classification with LeNet CNN. Collapsed bounds visibly outperform standard MF-VI.
1 hidden layer
2 hidden layer
3 hidden layer
4 hiddien layer
012
010
010-
0.1-
0.00
1200
Test NLL
0.08
0.7-
0.0
0.5
0.00
404
0.4
0.9
0.8
1.2
10-
11-
lest NEL
neg ELBO
06-
36
Dryg
Brreg
Figure 6: Top: performance of CM-MF-VI for different settings of hyperparameter and different network architectures (a = is standard MF-VI). Setting a = 0.01 visibly improves predictions compared to standard MF-VI VI (fixed Gaussian prior). Bottom: ELBO Cr, correlates with test NLL.
set FX and randomly perturbed MNIST images where every pixel is flipped with probability p. We report test NLL,, predictive entropy for perturbed MNIST experiment in Figure (left). We see that CM-MF-VI gives lower test NI.L than MF-VI for any p. but it also maintains growing entropy when - i.c. when we transition to OOD data. In Figure (right) we plot histograms of predictive entropy for MNIST in-distribution data and OOD data fashionMNIST. CM-MF-VI significantly improves upon MF-VI's under-confidence without excessive reduction in uncertainty.
Limitations of the algorithm: adjusting hyperparameters The bounds in Eq. and Eq. 10 have hyperparameters, corresponding to the choice of hyper-prior. These additional parameters are the main limitation of the derived algorithms. For example, CM-MF-VI VI has one hyperparameter Oregg which allows us to control the strength of regularization of the model. analyze the test predictive performance for different settings of (rreg for CM-MF-VI in Figure (top) averaged over 5 random seeds. We found the setting Oreg = 0.075 works well across many different problems and robust to the choice of size/depth of the network. Increasing areeg to 0.2 provides strong regularization and

---

model
test NLLL/ER
CMVMMF-V
CM-MP-VI
CV-MF-VI
MF-VI
MC dropout
MAP
RESNETI8
STLI0
1.00
1.10 0.02
1.57 0.01
1.01
1.17 0.05
1.64 0.03
SVHN
0.15 4.000
0.14 2114000 0.00
0.20 1 0.01
0.22 0.00
0.18 4: 0.00
0.35 0.01
CIFARIOO
1.43
1.53 0.00
2.001
0.03
1.70
4.25 = 0.05
CIFARIO
0.41 0.000
0.39 0,00
0.59 1 0.00
0.68 1: 0.02
0.49 0.000
0.93 0.02
SHUPFLENET
STLI0
0.96 0.01
0.99 0.02
1.22 + 0.05
1.70 0.08
1.28 0.001
1.78 0.08
SVHN
0.27 0.01
0.26
0.31
0.01
3.08 1: +.11
0.32 0,00
CIFARIOO
2.03 0.001
1.99
2.21
2.28 0.002
4.48 0.03
CIFARIO
0.65 0.000
0.65 0.00
0.71
0.00
1.10 0.02
ALEXNET
STLIO
1.50 0.07
1.48 0.01
1,80 0.07
1.86 0.03
1.07 0.06
1.80 0.32
SVHN
0.30 1: 0.00
0.28 0.00
0.42 0.01
0.01
0.38 0.01
0.72 0.06
CIFARIOO
2.24 4- 0.02
2.16 0.06
2.62 0.02
2.95 4- 0.004
2.19 4: 0.01
7.03 0.14
CIFARIO
0.72 0.01
0.69 0.011
0.01 0.99
1.19 0.01
0.74 0.01
1.79 0.06
RESNETI8
STL10
37,69
0.25%
39. + 0.53%
64.8 + 0.54%
06.58
0.63%
29.98
1.17%
29 30 + 0.56%
SVHN
3.76 1: 0.02%
3.75 0.02%
5.26 0.21%
5.73
1: 0.09%
4.11 1: 0.15%
4.94 0.07%
CIFARIOD
39.41
0.39%
40.39 + 0.37%
53.78 0.54% 50.46
0.72%
0.36% 45.51
47.92 + 0.34%
CIFARIO
13.75 0.06%
13.34 0 0.24%
20.22
0.30%
1.14%
16.36 0.28%
15.31 + 0.35%
SHUFFLENET
STL10
34.18 + 0.18%
31.9 + 0.78%
16. 2.83%
73.74
5.67%
47.72 1.12%
41.41 + 1.13%
SVHN
7.85 + 0.32%
+.31++0.277)
8.90
8.87
0.21%
0.60% 26.
8,28 0. 1.%
CIFARIOO
51.01 0.07%
52.66 = 0.36%
58. 1.0%
00.81
0.51%
0.19% 73.96
63.52 + 0. 0.22%
CIFARIO
22.71 1: 0.61%
22
0.14%
24.66
0.53%
25.01
0.47%
0.08% 42.54
18.19- + 0.59%
ALEXNET
STLI0
56.59
4.82%
55.97
2.99%
74.38
1 5.70%
78.28
1.32%
37.33
0.99%
35.74 0.67%
SVHN
7,92 + 0.20%
7.41 + 0. 13%
11.61
0.19%
13.1
0.20%
8.87 + 0.19%
10.35 + 0.71%
CIFARIOO
58.11
0.83%
51.91
1.07%
0.61% 73.17 1.39%
39
23%
60.14 + 0.19%
CIFARIO
24.35 0.30%
23.40 40.40
0.20%
34.73 + 0.34% 41.63 0.31%
24.74 0.43%
26.04 + 0.66%
Table 3: Test NLL and error rates (ER) for the collapsed bounds on large scale CNN experiments. CMV-MF-VI, CM-MF-VI bounds provide visibly better predictions than MF-VI, which under-fits.
works well for data sets very prone to over-fitting. For models/data prone under-fitting, setting Oreg S 0.025 can give even better predictive performance. In Figure 6 (bottom) we show normalized values of ELBO to log p(D, (Ereg) with prior O'eeg ~ (5 where D is the number of network parameters, and test NLL for different values Oreg. This has an important practical implication: approximately tuning Oreg can be guided by the ELBO. Figure X (bottom) can be compared with Fig 6 in 35 showing the same property using approximation of log (D), as opposed to using ELBO,
Image classification LeNet. We now consider image classification with the LeNet architecture 44! on 6 data sets: MNIST, faShionMMISST, K-MNIST [00, CIFARIO 10, CIFARI 100 [4 and SVHN [56]. We optimize the objectives for 800 epochs (except MAP for 50 epochs and MC dropout for 100 epochs as they tend to overfit) using batch size 512 and ADAM optimizer with default parameters. We report test NLL and test error rates (ER) averaged over 3 random seeds and standard deviation error bars in Table 2 We again observe the introduced bounds CM-MF-VI, CV-MF-VI and CMV-MF-VI outperform standard MF-VI in both test NLL and test error rate. CM-MF-VI performs slightly better than CMV-MF-VI, II, but the differences are not statistically significant. Both CM-MF-VI and CMV-MF-VI VI outperform MAP, CV-MF-VI and MC dropout in test NLL and test ER.
Image classification large CNNs. We follow by experimenting with larger CNNs: ResNet18 25), ShuffleNet[48) and AlexNet We use CIFAR 10, CIFAR 100, STL10 (T and SVHN. We again compare CMV-MF-VV VI, CV-MF-VV VI and CM-MF-VI to MF-VI, MC dropout and MAP. We optimize the objectives for 800 epochs (MAP and MC dropout early stopped at 200) with the default ADAM optimizer and the same data augmentation as in (57]. and average results over 3 random seeds. We gather the results in Table 3 Experiments with large CNNs confirm our findings from previous experiments: (i) CM--FF-I/CMVV.A
always outperform standard MF-VI by a visible margin and result in good predictive performance (e.g. outperforming SOTA VOGN (57). (ii) learning prior means CM-ME-VI/CMV-MF- outperforms learning just prior variances (CV-MF-VI).
5 Conclusions
We developed a family of algorithms optimizing variational posteriors in BNNs based on collapsed variational bounds. We demonstrated that learning the prior parameters of BNN weights fixes their predictive under-confidence resulting in good empirical performance and robustness to fitt-fitting. The developed algorithms allowed us to demonstrate that the ELBO can be suitable optimization target for learning hyperparameters of BNNs. Importantly, the introduced algorithms do not incur additional computational cost compared to applying MF-VI to BNNs and can be readily applied to improve the predictive performance of existing implementations. We hope that our approach will enable the practical use of VI based approximate inference in large network architectures.
10

---

Acknowledgements
Marcin B. Tomezak funded by a Cambridge Trust scholarship. Siddharth Swaroop supported by an EPSRC DTP studentship and Microsoft Research EMEA PhD Award. Andrew Y. K. Foong supported by a Trinity Hall Research Studentship and the George and Lilian Schiff' Foundation. Richard E. Turner is supported by Google, Amazon, ARM, Improbable, EPSRC grants EP/M0269571 and EP/L000776/1, and the UKRI Centre for Doctoral Training the Application of Artificial Intelligence to the study of Environmental Risks (AI4ER). The authors thank David Burt for insightful discussions.
References
[1] Laurence Aitchison. A statistical theory of cold posteriors in deep neural networks, 2021.
[2] Javier Antoran, James Allingham, and José Miguel Hernández-Lobato. Depth uncertainty in neural networks. H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 10620-10634 Curran Associates, Inc., 2020.
[3] Andrei Atanov, Arsenii Ashukha, Kirill Struminsky, Dmitry Vetrov, and Max Welling. The deep weight prior, 2019.
[4] D. Barber and Christopher Bishop. Ensemble learning in bayesian neural networks. In Generalization in Neural Networks and Machine Learning, pages 215-237. Springer Verlag, January 1998.
[5] Christopher M. Bishop. Pattern Recognition and Machine Learning (Information Science and Statistics), Springer-Verlag, Berlin, Heidelberg, 2006.
[6] David M. Blei, Alp Kucukelbir, and Jon D. McAuliffe. Variational inference: review for statisticians. Journal of the American Statistical Association, 1 112(518):859-8777 Feb 2017.
[7] Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural networks. Proceedings of the 32nd International Conference on International Conference Machine Learning Volume 37, ICML'11, page 1613-1622. JMLR.org. 2015.
[8] Thang Bui, Daniel Hemnandez-Lobato Jose Hernandez- Yingzhen Li, and Richard Turner. Deep gaussian processes for regression using approximate expectation propagation. In Maria Florina Balcan and Kilian Q Weinberger, editors, Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pages 1472-1481, New York, New York, USA. 20-22 Jun 2016. PMLR.
[9] Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in a handful trials using probabilistic dynamics models. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems 31, pages 4754-4765
Curran Associates, Inc., 2018
[10] Tarin Clamuwat, Mikel Bober-Irizar, Asanobu Kitamoto, Alex Lamb, Kazuaki Yamamoto, and David Ha. Deep learning for classical japanese literature, 2018.
[11) Adam Coates, Andrew Ng. and Honglak Lee. An analysis of single-layer networks in unsu- pervised feature learning. In Geoffrey Gordon, David Dunson, and Miroslav Dudík, editors, Proceedings of the Fourteenth International Conference Artificial Intelligence and Statistics, volume 15 of Proceedings of Machine Learning Research, pages 215-223, Fort Lauderdale, FL, USA, 11-13 Apr 2011. PMLR.
[12] Beau Coker, Melanie Fernandez Pradier, and Finale Velez. Porb-nets: Poisson process radial basis function networks. In Jonas Peters and David Sontag, editors, Proceedings of the 36th Conference Uncertainty in Artificial Intelligence (UAI), volume 124 of Proceedings of Machine Learning Research, pages 1338-1347. PMLR, 03-06 Aug 2020.
[13] Alexander G. de Matthews, Mark Rowland, Jiri Hron, Richard E, Turner, and Zoubin Ghahramani. Gaussian process behaviour in wide deep neural networks, 2018.
11

---

[14] Dheeru Dua and Casey Graff. UCI machine learning repository, 2017.
[15]
Michael W. Dusenberry, Ghassen Jerfel, Yeming Wen, Yian Ma. Jasper Snoek, Katherine Heller, Balaji Lakshminarayanan, and Dustin Tran. Efficient and scalable bayesian neural nets with rank-1 1 factors. International Conference on Machine Learning (ICML), 2020.
[16] Sebastian Farqubar, Michael A. Osborne, and Yarin Gal. Radial bayesian neural networks: Beyond discrete support in large-scale bayesian deep learning. Silvia Chiappa and Roberto Calandra, editors, Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics, volume 108 of Proceedings of Machine Learning Research, pages 1352-1362. PMLR, 26-28 Aug 2020.
[17] Sebastian Farquhar, Lewis Smith, and Yarin Gal. Liberty or depth: Deep bayesian neural nets do not need complex weight posterior approximations. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 4346-4357. Curran Associates, Inc., 2020.
[18] Andrew Foong, David Burt, Yingzhen Li, and Richard Turner. On the expressiveness of approx- imate inference in bayesian neural networks. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 15897-15908. Curran Associates, Inc., 2020.
[19] Andrew Y. K. Foong, Yingzhen Li, José Miguel Hernández-Lobato, and Richard E, Turner. 'in-between' uncertainty in bayesian neural networks, 2019.
[20] Vincent Fortuin, Adrià Garriga-Alonso, Florian Wenzel, Gunnar Rätsch, Richard Turner, Mark van der Wilk, and Laurence Aitchison. Bayesian neural network priors revisited, 2021.
[21] Yarin Gal and Zoubin Ghahramani. Dropout as bayesian approximation: Representing model uncertainty in deep learning. In Maria Florina Balcan and Kilian Q. Weinberger, editors, Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pages 1050~1059, New York, New York, USA, 20-22 Jun 2016. PMLR.
[22] Soumya Ghosh, Jiayu Yao, and Finale Deshi-Velez. Structured variational learning of bayesian neural networks with horseshoe priors, 2018.
[23] Soumya Ghosh, Jiayu Yao, and Finale Doshi-Velez. Model selection selection in bayesian neural networks via horseshoe priors. Journal of Machine Learning Research, 20(182):1-44 2019.
[24] Alex Graves. Practical variational inference for neural networks. In J. Shawe-Taylor, R. S. Zemel, P. L. Bartlett, F. Pereira, and K. Q. Weinberger, editors, Advances in Neurul Information Processing Systems 24, pages 2348-2356. Curran Associates, Inc., 2011.
[25] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition, 2015.
[26] James Hensman, Magnus Rattray, and Neil D. Lawrence. Fast variational inference in the conjugate exponential family, 2012.
[27] Jose Miguel Hernandez- Lobato and Ryan Adams. Probabilistic backpropagation for scalable learning of bayesian neural networks. In Francis Bach and David Blei, editors, Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 1861-1869, Lille, France, 07-09 Jul 2015. PMLR.
[28) Geoffrey E. Hinton and Drew van Camp. Keeping the neural networks simple by minimizing the description length of the weights. In Proceedings of the Sixth Annual Conference on Computational Learning Theory, '9,, for Computing Machinery.
'93, page 5-13, New York, NY, USA, 1993. Association
[29] Matthew D. Hoffman, David M. Blei, Chong Wang, and John Paisley, Stochastic variational inference. J. Mach. Learn. Res., 14(1):1303-1347, May 2013.
12

---

[30] Alexander Immer, Matthias Bauer, Vincent Fortuin, Gunnar Rätsch, and Khan Mohammad Emtiyaz. Scalable marginal likelihood estimation for model selection in deep learning. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning. volume 139 of Proceedings of Machine Learning Research, pages 1563-4577 PMLR, 18-24 Jul 2021.
[31] Alexander Immer, Matthias Bauer, Vincent Fortuin, Gunnar Rätsch, and Mohammad Emtiyaz Khan. Scalable marginal likelihood estimation for model selection in deep learning, 2021.
[32] Pavel Izmailov, Sharad Vikram, Matthew D. Hoffman, and Andrew Gordon Wilson. What are bayesian neural network posteriors really like?, 2021.
[33] Theofanis Karaletsos and Thang D) Bui. Hierarchical gaussian process priors for bayesian neural network weights. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 17141-17152. Curran Associates, Inc., 2020.
[34] Theofanis Karaletsos, Peter Dayan, and Zoubin Ghahramani. Probabilistic mettarrepresentations of neural networks, 2018.
[35] Mohammad Emtiyaz E Khan, Alexander Immer, Ehsan Abedi, and Maciej Korzepa. Approx- imate inference turns deep networks into gaussian processes. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché Buché-BBe,, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[36] Nathaniel J. King and Neil D. Lawrence. Fast variational inference for gaussian process models through kl-correction. In Johannes Fürnkranz, Tobias Scheffer, and Myra Spiliopoulou, editors, Machine Learning: ECML 2006, pages 270-281, Berlin, Heidelberg, 2006. Springer Berlin Heidelberg.
[37] Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. International Conference on Learning Representations, 12 2014.
[38] Diederik P. Kingma, Tim Salimans, and Max Welling. Variational dropout and the local reparameterization trick. In Proceedings of the 28th International Conference on Neural Information Processing Systems Volume 2, NIPS 15, page 2575-2583, Cambridge, MA, USA, 2015. MIT Press.
[39] Ranganath Krishnan, Mahesh Subedar, and Omesh Tickoo. Specifying weight priors in bayesian deep neural networks with empirical bayes, 2019.
[40] Agustinus Kristiadi, Matthias Hein, and Philipp Hennig. Being bayesian, even just a bit, fixes overconfidence ReLU networks. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 5436-5446. PMLR, 13-18 Jul 2020.
[41] Alex Krizhevsky, Learning multiple layers of features from tiny images. University of Toronto, 05 2012.
[42] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E. Hinton. Imagenet classification with deep convolutional neural networks. In Proceedings of the 25th International Conference on Neural Information Processing Systems Volume 1, NIPS' 12, page 1097-1105, Red Hook, NY, USA, 2012. Curran Associates Inc.
[43] Miguel Lázaro-Grrdiille and Michalis K. Titsias. Variational heteroscedastic gaussian process regression. In Proceedings of the 28th International Conference on International Conference on Machine Learning, ICML' 11, page 841-848, Madison, WI, USA, 2011. Omnipress.
[44] Yann LeCun, Patrick Haffner, Léon Bottou, and Yoshua Bengio. Object recognition with gradient-based learning- In Shape, Contour and Grouping in Computer Vision, page 319, Berlin, Heidelberg, 1999. Springer-Verie Verlag.
[45] Noel Loo, Siddharth Swaroop, and Richard Turner. Generalized variational continual learning, 2020.
13

---

[46] Christos Louizos, Karen Ullrich, and Max Welling, Bayesian compression for deep learning. In
I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[47] Miguel Lázaro-Gredilla Steven Van Vaerenbergh, and Neil Lawrence. Overlapping mixtures of gaussian processes for the data association problem, 2011.
[48] Ningning Ma, Xiangyu Zhang, Hai-Tao Zheng, and Jian Sun. Shufflenet v2: Practical guidelines for efficient cnn architecture design, 2018.
[49] David J. C. MacKay. The evidence framework applied to classification networks. 4(5):720-736, September 1992.
[50] David J. C. MacKay. Probable networks and plausible predictions a review of practical bayesian methods for supervised neural networks. Network: Computation in Neural Systems, 6(3):469-505, 1995.
[51] David J. C. MacKay. Comparison of approximate methods for handling hyperparameters. Neural Compud., 115):1035-1006, July 1999.
[52] Chris J Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. 2016.
[53] Aaron Mishkin, Frederik Kunstner, Didrik Nielsen, Mark Schmidt, and Mohammad Emtiyaz Khan. Slang: Fast structured covariance approximations for bayesian deep learning with natural gradient. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Bianchi, and R. Gamett, editors, Advances in Neural Information Processing Systems 31, pages 6245-6255. Curran Associates, Inc., 2018
[54] Radford M. Neal. Bayesian learning via stochastic dynamics. In S. J. Hanson, J. D. Cowan, and C. L. Giles, editors, Advances in Neural Information Processing Systems 5, pages 475-482. Morgan- Kangan-Kaufman,, 1993.
[55] Radford M. Neal. Priors for Infinite Networks, pages 29-53. Springer New York, New York, NY, 1996.
[56) Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Ng. Reading digits natural images with unsupervised feature learning. NIPS, 01 2011.
[57] Kazuki Osawa, Siddharth Swaroop. Mohammad Emtiyaz E Khan, Anirudh Jain, Runa Eschen- hagen, Richard E Turner, and Rio Yokota. Practical deep learning with bayesian principles. In Advances in Neural Information Processing Systems 32, pages 4289 4301. Curran Associates, Inc., 2019.
[58] Hiske Overweg, Anna-Lena Popkes, Ari Ercole, Yingzhen Li, José Miguel Hernández- Yordan Zaykov, and Cheng Zhang. Interpretable outcome prediction with sparse bayesian neural networks in intensive care, 2019.
[59] Carl Edward Rasmussen and Zoubin Ghahramani. Occam's razor. In T. K. Leen, T. G. Dietterich, and V. Tresp, editors, Advances in Neural Information Processing Systems 13, pages 294-300. MIT Press, 2001.
[60] Carlos Riquelme, George Tucker, and Jasper Snoek. Deep bayesian bandits showdown: An empirical comparison of bayesian deep networks for thompson sampling, 2018.
[61] Hippolyt Ritter, Aleksandar Botev, and David Barber. A scalable laplace approximation for neural networks. In International Conference on Learning Representations, 2018
[62] Siddharth Swaroop, Cuong V. Nguyen, Thang D. Bui, and Richard Turner. Improving and understanding variational continual learning. arXiv preprint arXiv:1905.02099 2019.
14

---

[63] Jakub Swiatkowski, Kevin Roth, Bastiaan S. Veeling, Linh Tran, Joshua V. Dillon, Stephan Mandt, Jasper Snoek, Tim Salimans, Rodolphe Jenatton, and Sebastian Nowozin. The k-tied normal distribution: A compact parameterization of gaussian mean field posteriors in bayesian neural networks, 2020.
[64] Yee Teh, David Newman, and Max Welling. A collapsed variational bayesian inference algorithm for latent dirichlet allocation. In B. Schölkopf, J. Platt, and T. Hoffman, editors, Advances in Neural Information Processing Systems, volume 19. MIT Press, 2007.
[65] Michalis Titsias and Miguel Lázaro-Gredilla. Doubly stochastic variational bayes for non- conjugate inference. In Eric P. Xing and Tony Jebara, editors, Proceedings of the 31st Inter- national Conference on Machine Learning, volume 32 of Proceedings of Machine Learning Research, pages 1971-1979, Bejing, China, 22-24 Jun 2014. PMLR.
[66] Marcin Tomezak, Siddharth Swaroop, and Richard Turner. Efficient low rank gaussian varia- tional inference for neural networks. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 4610-4622 Curran Associates, Inc., 2020.
[67] Marcin B. Tomezak, Siddharth Swaroop, and Richard E. Turner. Neural network ensembles and variational inference revisited. In 1st Symposium on Advances in Approximate Bayesian Inference, pages 1-11, 2018.
[68] Brian Trippe and Richard Turner. Overpruning in variational bayesian neural networks, 2018.
[69] Richard Turner, Pietro Berkes, and Maneesh Sahani. Two problems with variational expectation maximisation for time-series models. Bayesian Time Series Models, 01 2008.
[70] Florian Wenzel, Kevin Roth, Bastiaan S. Veeling, Jakub Swiętkowski, Linh Tran, Stephan Mandt, Jasper Snoek, Tim Salimans, Rodolphe Jenatton, and Sebastian Nowozin. How good is the bayes posterior in deep neural networks really?, 2020.
[71] Andrew Gordon Wilson and Pavel Izmailov. Bayesian deep learning and a probabilistic perspective of generalization, 2020.
[72] Anqi Wu, Sebastian Nowozin, Edward Meeds, Richard E. Turner, Jose Miguel Hernandez- Lobato, and Alexander L. Gaunt. Deterministic variational inference for robust bayesian neural networks. In International Conference Learning Representations, 2019.
[73] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist. mnist: novel image dataset for benchmarking machine learning algorithms. arXiv preprint (XXi:17700.074777 2017.
[74] Ruqi Zhang, Chunyuan Li, Jianyi Zhang, Changyou Chen, and Andrew Gordon Wilson. Cyclical stochastic gradient meme for bayesian deep learning. International Conference on Learning Representations, 2020.
15


Reviews:
[{"review_id": "sr-HRdpKVi", "reviewer": "Reviewer_rzGC", "summary": "Although marginal likelihood can be used to select hyper-parameters, the ELBO is generally regarded as too weak a bound to do this effectively.\nThe authors propose applying collapsed variational bounds to variational inference of BNNs.\nThey show why this maintains the tightness of the bound for parameter inference while also providing a tightish bound for hyper-parameter optimization by exploiting the independence of the two.\nThey empirically show that their method performs well in a range of settings, including in tests of ood uncertainty.", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "# Overall review\nI found this paper difficult to assign a score to, and I remain quite uncertain.\nOn the one hand, I believe that you have proposed a sensible approach which fixes a moderately important problem.\nYou have also provided extensive experimental support for your hypotheses.\nOn the other hand, I found the paper quite hard to read and somewhat poorly organized.\nI feel that it would benefit a lot from being redrafted in a way that makes it clearer what its contributions are and how they are being examined.\nFor that reason I am currently recommending borderline rejection, but I could easily be persuaded upwards if other reviewers think particularly highly of the significance of the work.\n\n# Major comments\n\nThere is a lot of redundancy.\nThe abstract could probably be pruned by a third.\nMuch of the introduction is relatively boiler-plate and does not focus on the specific contribution that *you* are making.\nI did not quite understand what the role of the \"Preliminaries\" section was. Probably you could take a small part of this and place it at the start of what is now S3.\nSimilarly, I would reserve a discussion of over-pruning until the experiments section and make a note in the introduction that your method resolves an over-pruning issue in VI.\nWhen I read it I had no idea why this was being mentioned and Figure 1 is somewhat a distraction at this stage.\n\nI think some of the details of the derivations in S3 could be moved to an appendix, with more space devoted to an explanation of why you are doing things and what the intuitive purpose is and less on the details.\n\nThe captions on many figures seem incomplete, with only some subfigures discussed explicitly.\nThe figures are all far too small and difficult to read.\n\nI would be inclined to pick your experiments more carefully for the main body so that you can devote more time/space to explaining and interpreting each thoroughly.\nThey feel rushed and crowded.\n\n# Minor remarks\n\nRegarding the introduction: I'm easily persuaded that a tighter bound to the marginal likelihood would be good, so a lot of this background isn't super necessary.\nI would have liked more of a sketch of what your approach is going to be.\n\nPlease make your references hyperlinks.\nEspecially when you use a numerical citation style it is almost impossible to figure out what you are citing otherwise.\n\nThe spacing of the paper is very tight, perhaps some of the style-sheet was accidentally overridden?\nThis makes it much harder to read.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "These are adequately addressed.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "rDB448UR5Ho", "reviewer": "Reviewer_KgBQ", "summary": "This paper applies the idea of the collapsed variational bound studied by Teh et al., 2007 to BNNs. The bound is constructed by assuming that VB's Gaussian prior mean and variance are also part of the inference. (In contrast, one commonly sets them to some fixed values, e.g. $\\mathcal{N}(0, 1)$.) The collapsed bound seems to have a similar computational cost but is tighter than the standard ELBO.\n\nExtensive experiments show that the proposed objective works well, even on large networks, the regime which mean-field VB tends to underfit. It is also shown that the objective can mitigates the over-pruning effect that plagues VB.", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I think the proposed method is interesting and is backed up by extensive experiments. This could be the answer to the difficulty of training VB due to hyperparameter selection. Nevertheless, I think the experiments can be expanded more, esp. regarding uncertainty quantification, OOD detection. Currently, the authors only show MNIST-FMNIST OOD detection with MLP. I think this is inadequate and the authors should use the standard OOD detection suites (MNIST, FMNIST, SVHN, CIFAR10, CIFAR100 in-distributions against 3+ OOD test sets) with a more commonly used network instead (ResNets, Wide-ResNets).\n\nWith that being said, the major issue with this paper is that it is *not* well-written. There is no paragraph break *at all* in this paper, making it a continuous blob of text. (Instead, the authors used \\paragraph ``with its whitespace removed.) Maybe this due to some \"LaTex hacks\" that the authors used to make the text fit the page limit. If this is so, I suggest the authors not use them.\nAll in all, the presentation of this paper needs to be improved significantly and I urge the authors to do so meticulously. \n\n**Questions:**\n* Based on the discussion in L.80-84: How does the collapsed bound compare to EM?\n\n**Suggestions:**\n* Use paragraph break and use whitespace judiciously\n* Use display math judiciously, instead of putting long formulas inline\n* Put L.120-123 inside LaTex's algorithm env. Inline, it is very confusing to read\n* Figures are very hard to read. I suggest the author to use tikz\n\n**Minor suggestions:**\n* L.46: \"..., *the* ELBO lowers bound ...\", also add a period before \"Replacing ...\"\n* Fig. 3: Add y-label to the second row\n* Tab. 1: What is \"test ER\"? (It's discussed in the next page, but it should be much closer to Tab. 1)\n* Tab. 2, 3: Unclear which one is test-NLL, which one is test-ER.\n\n\n---\n\n**Post-Rebuttal**\n\nWhile I like the idea, I still think that it would be beneficial for the paper to be given more careful treatment in its writing and presentation---this will ultimately change the paper substantially and thus, I believe, is more suitable for resubmission. My updated score (weak reject) reflects this ambivalence.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors discussed the limitations of their method in the paper.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "TMNX0x3m_9Z", "reviewer": "Reviewer_PHT4", "summary": "This work introduces an improvement over the typical ELBO by deriving a tighter bound to the log marginal likelihood. It further shows that the ELBO can be used to learn appropriate hyperparameters which can help in situations where MF-VI would have underfit. A few variants of this method are derived where some combination of prior means/variances are learned. These methods were shown to be competitive with and at times outperform existing methods in experiments on standard benchmarking datasets. ", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Other work in this area has largely centered on improvements on gradients of the ELBO this work focuses on an improvement on the bound of the log marginal likelihood. These improved bounds are shown mathematically with sufficient detail. This work builds on existing ideas from Bayesian hierarchical models. The idea of learning hyperparameters seems to be related in spirit to empirical Bayes applied to neural networks. [1]\n\nIt is hard to read the algorithm as it is currently written in paragraph form; breaking this into its own section/some sort of indentation would help.. The provided example right after this was illustrative and helpful.\n\nJust above equation (2) there is a missing period “.”: “0, ELBO lower bounds log marginal likelihood $\\log p(D|µ_p,σ^2_p) ≥ L(µ_q,σ^2_q;µ_p,σ^2_p)$ Replacing” → “0, ELBO lower bounds log marginal likelihood $\\log p(D|µp,σ2 p) ≥ L(µ_q,σ^2_q;µ_p,σ^2_p)$. Replacing”\n\n\n[1] G. Casella, “An Introduction to Empirical Bayes Data Analysis,” The American Statistician, vol. 39, no. 2, pp. 83–87, 1985, doi: 10.2307/2682801.\n\n\n----------------------\nThanks for the reply. After further discussion and thought I have lowered my rating to marginally above the acceptance threshold. The idea is good but the structure and writing is weak.\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors have not addressed this, but this decision is appropriate for the content. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "47c_5bVYRXf", "reviewer": "Reviewer_A9Jx", "summary": "This paper defines bayesian neural network models with hierarchical prior distributions. These models are trained efficiently using collapsed variational bounds, in which some of the variables of the hierarchical ELBO are integrated out. The authors show that these models outperform their hierarchical counterparts in a number of experiments and standard benchmarks.", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I found this paper to be a very interesting read, and a solid contribution to the fundamental research question on how to best model uncertainty in deep neural networks.\n\nUsing hyperpriors makes a lot of sense to me, since as it is also shown in the experiments, fixing the values of the hyperparameters as normally done is not a robust solution.\nTo the best of my knowledge, this is the first work that extends MF-VI for bayesian deep learning using ideas on collapsed bounds previously introduced for hierarchical bayesian models.\nBy using distribution from the exponential family the authors show a scalable and effective way to learn such models.\n\nBecause of its nature, the paper is quite dense from the theoretical point of view. In most parts derivations are well explained, but there are some places that require clarification:\n\n1. line 87: \"where for clarity we omit. the parameterization..\" this actually caused me some confusion since I did not know why you would do this. You should say that you do it since it will be implicitly defined.\n\n2. line 107: \"straightforward to derive...\" missing a reference for readers not familiar with MF-VI (Bishop's PRML book could be a good one for example)\n\n3. line 146: why are there 2 different priors over W?\n\n4. line 148: distribution G in p($\\tau$) is not defined, i assume it is a Gamma distribution?\n\n5. line 215: what is p? is it the dropout probability?\n\nOverall, I believe this paper can have a big impact on the area of bayesian deep learning, so I vote for accepting it.\n\n\n___________________________________________\nReply to author's rebuttal.\nThanks for your reply, I will leave the score unchanged and argue for acceptance.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper tightens the standard variational bound optimized in Bayesian Deep Learning by drawing on collapsed variational inference. The authors consider the prior parameters as latent variables and derive a hierarchical variational inference procedure in which the top-level latent variables are marginalized out. The authors showed strong empirical performance of their method compared to a variety of baselines.

The paper provides a non-trivial methodological contribution to the Bayesian deep learning community. Its mathematical derivations are not simple but generally well-explained (the authors sometimes don’t clearly distinguish between latent variables and variational parameters at times).  The main point of criticism was suboptimal structuring and writing. I strongly encourage the authors to make their paper more accessible by following the detailed advice that the reviewers provided. Overall, this is a very good paper. 
Decision: Accept (Poster)