Collapsed Variational Bounds for Bayesian Neural Networks
Year: 2021

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.

text: 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.

heading: 1 Introduction
text: There has been a 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 on Variational Inference (VI) for BNNs has seen multiple advances enabling optimization of the Evidence Lower Bound (ELBO) [29, 65, 24, 7, 38], allowing for scaling to large neural networks and datasets [57, 63, 15, 16, 66]. However, the underconfidence of mean-field VI (MF-VI) in output-space has caused poor adoption of VI in applications, especially in sequential decision making [9, 60, 69]. Although a lot of work has improved performance by deriving better estimators of the gradient of the ELBO [29, 65, 24, 7, 38], this work has focused on the variational parameters of distributions over weights, with the prior over weights usually fixed to zero-mean isotropic Gaussians [55, 70, 74, 15, 20, 7, 40].
In this paper, we 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 [36, 43, 47, 64], which analytically solve the inference over prior parameters and allow us 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, 52, 50, 31], although there have also been arguments against such an approach [7]. Experiments using our derived ELBOs lead to 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-fits data in output space if the prior over weights is fixed and, (ii) the derived tighter ELBOs are good optimization targets to learn the hyperparameters of hierarchical Gaussian priors, even in large networks. We provide the code implementing the introduced algorithms at https://github.com/marctom/collapsed_bnns.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).

heading: 2 Collapsed Variational Bounds for Bayesian Neural Networks
text: We consider performing VI in a supervised learning scenario given N data points D = {(xi,yi)}Ni=1. We denote a neural network as fW , the observation model as p(y|fW (x)), where W denotes the vectorized weights and biases of the network, and the likelihood of the weights W as p(D|W ) =QN
i=1 p(yi|fW (xi)). We focus on Gaussian mean-field priors (with non-zero means) over vectorized network weights p(W |µp, 2p) = N (W |µp, 2p) meaning w ⇠ N (µp, 2p) for any coordinate w in W (both µp and 2p are also vectors in this notation). Similarly, we restrict our attention to mean-field Gaussian approximate posteriors q(W |µq, 2q) = N (W |µq, 2q). VI finds an approximate posterior q(W |µq, 2 q) by minimizing the KL divergence DKL(q(W |µq, 2q)||p(W |D,µp, 2p)) w.r.t. the variational parameters µq, 2q . This is done by optimizing the Evidence Lower Bound (ELBO) [28, 4]:
L(µq, 2 q;µp, 2 p) =
NX
i=1
E q(W |µq, 2q) log p yi|fW (xi)
DKL q(W |µq, 2 q)||p(W |µp, 2 p) .
(1) We denote the first term on RHS as Ldata(µq, 2q). As DKL(q(W |µq, 2q)||p(W |D,µp, 2p)) 0, the ELBO lower bounds log marginal likelihood log p(D|µp, 2p) L(µq, 2q;µp, 2p). 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 |µq, 2q) as p(y⇤|x⇤,D) ⇡ Eq⇤(W |µq, 2q)p(y ⇤ |fW (x⇤)).
For BNNs, prior means are most often set to zero µp = 0, and prior variances are set to a constant 2p = 1 (often scaled by the size of a previous layer [55, 13]). We consider a more general setup where the weights of a BNN are defined by a hierarchical model p(W |µp, 2p)p(µp)p( 2) and aim to perform inference over weights W and µp, 2p. For clarity, we will refer to µp, 2p as prior parameters, leaving the term hyperparameters for any other free variables.

heading: 2.1 Deriving collapsed variational bounds
text: In this section, we develop an efficient optimization scheme for variational BNNs with weights defined in a 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, 26] 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 a collapsed bound is simple: Suppose we maximize f(x, y) and we can derive the optimal x⇤(y) = argmaxxf(x, y). Substituting x⇤ into f results in f(x⇤(y), y), which only depends on y, is easier to optimize, and has the property that 8y f(x⇤(y), y) f(x, y). 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 2p, which we treat as latent variables and subsequently discuss collapsing the inference w.r.t. prior parameters µp and 2p. We use the factorization q(W ,µp, 2p|µq, 2q) = q(W |µq, 2 q)q(µp, 2 p), where we omit the variational parameters of q(µp, 2p) for concise notation. Learning q(µp, 2p) under this assumption does not complicate computing the BNN’s predictive distribution p(y|x,D), i.e. µq, 2q do not need to be sampled at the prediction time, since, p(y|x,D) =Ep(W ,µp, 2p|D)p(y|fW (x)) ⇡ Eq(W ,µp, 2p)p(y|fW (x))=Eq(W |µq, 2q)p(y|fW (x)). (2) To apply inference, we invoke the variational principle which lower bounds the log marginal likelihood log p(D) as:
log p(D) Eq(W |µq, 2q)q(µp, 2p) log p(D,W ,µp,
2 p)
q(W |µq, 2 q)q(µp, 2 p)
= Lq µq, 2 q, q(µp, 2 p) . (3)
The bound in Eq. (3) reduces to L(µq, 2q) in Eq. (1) when q(µp, 2p) = q(µp)q( 2p) and q(µp), q( 2p) are delta functions. If we are able to analytically derive the distribution q⇤(µp, 2p) maximizing
L q, we can substitute the resulting distribution into Eq. (3) to derive a tighter (collapsed) bound on the marginal likelihood, saving the need to perform coordinate ascent or gradient updates to learn the prior parameters:
log p(D) Eq(W |µq, 2q)q⇤(µp, 2p) log p(D,W ,µp,
2 p)
q(W |µq, 2 q)q ⇤(µp, 2 p)
=L⇤(µq, 2 q), (4)
where L⇤(µq, 2q) Lq(µq, 2q, q(µp, 2p)) and L⇤ depends only on the variational parameters µq, 2 q . 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 as a KL corrected bound [36] and marginal VI bound [43, 47].
Collapsing the variational bound for BNNs. We now provide a method to derive a collapsed variational bound by maximizing Lq(µq, 2q, q(µp, 2p)) w.r.t. the posterior over weights prior parameters q(µp, 2p) in BNNs. We use the property that the BNN likelihood term log p(y|x, fW (x)) does not depend on the prior parameters µp and 2p, i.e. p(D|W ,µp, 2p) = p(D|W ). The bound L q from Eq. (3) decomposes into three terms:
L q µq, 2 q, q(µp, 2 p) = Ldata(µq, 2 q) + µq, 2 q, q(µp, 2 p) +H[q(W |µq, 2 q)], (5)
where H[q(W |µq, 2q)] denotes the differential entropy, and we have defined
µq, 2 q, q(µp, 2 p) = E
q(µp, 2 p)
"
E q(W |µq, 2q) log p(W |µp, 2 p)
# DKL q(µp, 2 p)||p(µp)p( 2 p) .
(6) Only depends on the posterior q(µp, 2p), and therefore finding q⇤(µp, 2p) can be done by maximizing (µq, 2q, q(µp, 2p)).
We can significantly simplify this by noting that the objective in Eq. (6) is analogous to Eq. (1), hence optimizing takes the form of a nested VI problem: we want to infer q(µp, 2p) where the “data” distribution is replaced by the current variational posterior over weights q(W |µq, 2q). This implies the solutions q⇤(µp, 2p) are straightforward to derive in closed-form for priors/approximate posteriors from the exponential family [5, 6]. Specifically, we can write the maximizer q⇤(µp, 2p) as:
log q⇤(µp, 2 p) / log p(µp) + log p( 2 p) + Eq(W |µq, 2q) log p(W |µp, 2 p). (7)
Substituting the optimal variational posterior q⇤(µp, 2p) into Eq. (6) gives us ⇤, and we can substitute this into Eq. (5) to give our final collapsed bound L⇤(µq, 2q) = Ldata+ ⇤+H[q(W |µp, 2p)]. Furthermore, when calculating ⇤, the inferred prior parameters q⇤(µp, 2p) can be analytically integrated out in many cases, i.e. the outer expectation on the RHS in Eq. (6) can 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 q(W |µq, 2q) is given by L⇤(µq, 2q). Compared to gradient learning of µp and 2p, eliminating µp and 2p from L⇤ performs better (as we show in Section 4), and halves both the memory requirement and time taken per update. It is important to remember that optimizing L⇤ learns both the variational parameters µq, 2q and the posterior over prior parameters q⇤(µp, 2p), but the latter is implicit as we express it as a function of µq, 2 q 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 1. We next give two concrete examples of collapsed bounds.

heading: 2.2 Examples of tighter ELBOs
text: Learn prior means, fix prior variances. As an easier example, we first consider fixing the prior variance of weights, and learning prior means. We follow the four steps of Algorithm 1 to derive our collapsed bound. Detailed derivations are in Appendix B.
Step I: This step defines a model and the family of variational approximations. We choose p(W |µp) = N (W |µp, 1) with a Gaussian prior p(µp|↵) = N (µp|0,↵1), where ↵ and are
Algorithm 1 Deriving a collapsed bound L⇤ for BNN in four steps. Step I: Choose priors p(µp), p( 2p), p(W |µp, 2p),
and approximate posteriors q(µp, 2p), q(W |µq, 2q) over network weights. Step II: Calculate optimal prior parameters q⇤(µp, 2p) using Eq. (7).
Step III: Form ⇤ and solve Eq⇤(µp, 2p) h Eq(W |µq, 2q) log p(W |µp, 2 p) i
in Eq. (6). Step IV: The collapsed bound is L⇤ = Ldata + ⇤ +H[q(W |µq, 2q)].
hyperparameters. We also employ a mean-field Gaussian posterior q(W |µq, 2q) and a Gaussian posterior q(µp).
Step II: This step analytically finds the optimal variational distribution q⇤(µp). Substituting our distributions into Eq. (7) for every coordinate µp gives us
log q⇤(µp) / (µq µp)2
2
µ2p 2↵ . (8)
In this case the variational posterior q⇤(µp) matches the true posterior, as the above dependency is equivalent Bayesian inference over µp given Gaussian likelihood/prior and observation µq, so q⇤(µp) = N (µp| ↵ ↵+ µq, ↵ ↵+ 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 DKL(q⇤(µp)||p(µp)) is straightforward to compute. Next we solve the integral I = Eq⇤(µp)[Eq(W |µq, 2q) log p(W |µp, 2 p)]:
I , E N (µp| ↵↵+ µq, ↵ ↵+ 1)
 1T 2q + (µq µp) T (µq µp)
2
, 1T 2q 2
µTq µq
2(↵+ )2 .
Step IV: This step forms a collapsed bound that will be used to learn variational posteriors. Substituting the above into L⇤ = Ldata + ⇤ +H[q(W |µp, 2p)] results in our new bound:
L ⇤ m(µq, 2 q) , Ldata(µq, 2q)
1
2
h 1T 2q + ↵regµ T q µq i + 1
2 1T log 2q +
D 2 log↵reg, (9)
where we have defined the hyperparameter ↵reg = /( + ↵) 2 (0, 1), and for fixed there is a 1-to-1 relation between ↵ and ↵reg .
When ↵ ! 0 (so ↵reg ! 1) we recover the standard expression for the ELBO given by Eq. (1) with q(W |µq, 2 q) and prior N (W |0, 1). This corresponds to a dogmatic prior p(µp) = N (µp|0,↵1) converging to a delta spike at 0, and no inference of µp. When we allow for inference, the regularization of the (approximate) posterior mean µq decreases. When ↵reg ! 0 (uninformative prior) the modeling can be prone to over-fitting due to insufficient regularization in the model specification. Note that this is reflected by the term D2 log↵reg which can be interpreted as an Occam’s Razor penalty [59]: this term penalizes small ↵reg as it diverges to negative infinity when ↵reg ! 0. Since ↵reg 2 (0, 1), this enables finding well-performing values of ↵reg by evaluating L⇤m (without having a validation set) [30], and we show this in our experiments in Section 4.
We find that ↵reg 2 (0.01, 0.1) significantly improves upon the predictions of MF-VI with fixed Gaussian priors. For existing implementations learning BNNs with Gaussian variational posterior/prior we suggest the default setting ↵reg = 0.05 in front of the mean regularization in the expression for KL divergence to instantaneously improve the predictions (as opposed to down-weighting the whole KL divergence term). Larger ↵reg increases the strength of the regularization of the model.
Learn both prior means and variances. We now discuss a scheme to learn hyperparameters of both prior means and variances. We follow the same four steps from Algorithm 1. We again defer the detailed derivations to Appendix B.
Step I: p(W |µp, ⌧ p) = N (W |µp, 1 ⌧p ) and p(µp|t⌧ ) = N (µp|0, 1 t⌧ ), p(⌧ ) = G(⌧ |↵, ). We again consider posteriors q(W |µq, 2q) = N (W |µq, 2q), q(µp|⌧ p) = N (µp) and q(⌧ p) = G(⌧ p).
Step II: The optimal posteriors q⇤ are given by q⇤(µp) = N (µp| 1 1+tµq, 1 (1+t)⌧p ) and q⇤(⌧ p) = G(⌧ p|(↵+ 1 2 )1, + t 2(1+t)µ 2 q + 1 2 2 q).
Steps III and IV: The bound L⇤(µq, 2q) becomes,
L ⇤ mv(µq, 2 q) , Ldata(µq, 2q) (↵+
1 2 )1T log
h 1+
2 µ2q +
1 2 2q
i + 1
2 1T log 2q, (10)
where = t/(1 + t). Eq. (10) recovers the case of setting prior mean µp = 0 and learning only prior variances 2p when the prior precision t⌧ p over µp goes to 1, i.e. t ! 1. Setting < 1 in Eq. (10) weakens the regularization of posterior mean µq , allowing it to vary more. The regularizer in Eq. (10) is a decreasing function of µq , i.e. without gradients from Ldata, the posterior mean µq of a weight w converges to 0. In Section 4 we show that L⇤mv outperforms standard MF-VI.
Comparison of regularizing terms. We compare the regularization terms on µq and 2q in the standard mean-field ELBO in Eq. (1) and introduced ELBOs in Eq. (9) and Eq. (10) in Figure 1, where we use prior the p(w) = N (w|0, 1). We consider what happens in the absence of gradients from the data term Ldata to better understand the bounds. In Figure 1 (left) we see that Gaussian posteriors (in the absence of data) optimized with Eq. (1) converge to their prior N (w|0, 1) [7, 68]. In Figure 1 (middle) we are down-weighting µTq µq in Eq. (9), hence weakening the regularization of µq . This allows posterior means to vary more, and can be interpreted as scaling the µp axis in Figure 1 by p ↵. As we show in Section 4, this enables posterior means µq to saturate activations and reduce the noise from pruned weights. Optimizing Eq. (10) is roughly opposite to Eq. (1) and causes weights to be pruned to N (0, 2 ) (Figure 1 right, where = 0.01). Weights pruned in this way do influence the predictions for small and this fixes the excessive injection of noise by pruned weights, as we show in Section 4.

heading: 3 Related work
text: Performance of BNNs is known to be sensitive to prior hyperparameters [54, 71, 31], but optimizing hyperparameters in BNNs has not been widely adopted, although there are exceptions [50, 72]. Some authors heuristically propose using a Gaussian centered on the MAP estimate as a prior for BNNs [8, 39]. The Laplace approximation has been used for hyperparameter selection in BNNs [49, 31, 61], but it does not use the ELBO. Hierarchical models in BNNs have been explored in many works [23, 22, 46, 33, 34, 58, 3, 12, 2]. The closest to our approach is learning a hierarchical horseshoe prior over BNNs weights [23]. However, their approach is computationally costly, does not consider collapsed bounds to derive tighter ELBOs, and does not 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, 32, 1] and arguing against [70]. Similarly, there are conflicting views on the performance of MF-VI in BNNs with fixed Gaussian priors as some authors claim it can work well [17] and others arguing against [63, 19, 18]. A number of problems of mean-field VI BNNs have been exposed in shallow networks including over-pruning [68] 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].

heading: 4 Experiments
text: 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. (9) and Eq. (10) significantly improves the predictive performance of standard MF-VI BNNs (fixed Gaussian priors). We also find that learning the prior means according to Eq. (9) is more effective and gives sufficiently good predictive performance, whereas learning prior variances (Eq. (10)) mostly does not lead to significant gains in the quality of predictions. We defer the details of the experimental setup to Appendix C and report additional experimental results in Appendix D.
The algorithms we propose in this paper optimize various tighter ELBOs: CM-MF-VI optimizes Lm from Eq. (9); CM-MF-VI OPT additionally optimizes ↵reg using the tighter ELBO (Eq. (9)); CV-MF-VI optimizes for just prior variances, using the bound Lmv from Eq. (10) when ! 1 (prior mean µp fixed to 0); and CMV-MF-VI optimizes both prior means and variances from Eq. (10). We compare to the following baselines: MF-VI LRT learns zero-mean mean-field Gaussian priors with fixed variances [55] (which we call standard MF-VI) using local reparametrization [38]; MF-VI FV learns a mean-field Gaussian posterior with fixed variances [68]; MF-VI BD is mean-field Gaussian VI while down-weighting KL penalty by 0.5; MF-VI EB learns prior means with ELBO gradient updates [51, 72]; BBB is Bayes by Backprop using the ADAM optimizer (as opposed to SGD in [7]); MF-VI k-tied normal distribution [63]; ELRG/SLANG learn low rank posteriors [66, 53].
UCI regression. We first consider 20 train-test splits for 8 UCI regression data sets [14]. We 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|f1W (x), f 2 W (x)) = N (y|f1W (x), exp(f 2 W (x))), where f 1 W (x), f 2 W (x) are two heads of the network. We optimize the objectives for 200K steps with the ADAM optimizer [37] with default settings. We report a comparison of test log-likelihood in Figure 2 (top). We observe that (i) the collapsed bounds CM-MFVI, 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, (iii) CM-MF-VI outperforms or 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 ↵reg. Next we report a comparison between CM-MF-VI and standard MF-VI as the depth of the network increases in Figure 2 (bottom). We see standard MF-VI starts to rapidly under-perform on most data sets (for larger data sets standard MF-VI performs well as it converges to MAP). Learning prior means/variances by optimizing bounds in Eq. (9) or Eq. (10) allows to mitigate rapid degradation of
predictions. We now explain why standard MF-VI performs badly, finding this is due to over-pruning in BNNs. We investigate this effect and find that learning prior means and variances mitigates the effect of pruned units on the predictive 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, 62]. While this is beneficial for compressing the model, pruned weights contribute noise to predictions and make modeling data difficult. With 1 hidden layer networks and the regular ELBO, output layer weights are set as close to 0 as possible [68], thereby reducing the noise from pruned weights in earlier layers. However, some noise still leaks through, as output layer weights are not quite delta-functions. This noise is one of the reasons for increasing under-confidence in output space as the size of the network increases [68, 16, 63, 66, 23, 45]. In Figure 3 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 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 over-pruning.
Fixing under-fitting caused by over-pruning in BNNs. We now explain why learning the hyperparameters 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 4 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 1 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-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, as we schematically show in Figure 4 (a). For CMV-MF-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 4 (b).
MNIST classification with MLP. Classifying vectorized MNIST images using two hidden layer network with ReLU activation is a 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 5 random seeds in Table 1. For 400 hidden units, CM-MF-VI/CMV-MF-VI 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-domain improvements of CM-MFVI/CMV-MF-VI come at the cost of out-of-distribution (OOD) performance. We find this is not the case: OOD performance is as good as before. To show this, we learn a 2 hidden layer BNN with ReLU activations on vectorized MNIST images as previously, but test it on the fashionMNIST data
set [73] 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 5 (left). We see that CM-MF-VI gives lower test NLL than MF-VI for any p, but it also maintains growing entropy when p ! 1 i.e. when we transition to OOD data. In Figure 5 (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. (9) 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 has one hyperparameter ↵reg which allows us to control the strength of regularization of the model. We analyze the test predictive performance for different settings of ↵reg for CM-MF-VI in Figure 6 (top) averaged over 5 random seeds. We found the setting ↵reg = 0.075 works well across many different problems and is robust to the choice of size/depth of the network. Increasing ↵reg to 0.2 provides strong regularization and
works well for data sets very prone to over-fitting. For models/data prone to under-fitting, setting ↵reg  0.025 can give even better predictive performance. In Figure 6 (bottom) we show normalized values of ELBO to log p(D,↵reg) with prior ↵reg ⇠ Exp(5D), where D is the number of network parameters, and test NLL for different values ↵reg. This has an important practical implication: approximately tuning ↵reg can be guided by the ELBO. Figure 6 (bottom) can be compared with Fig 6 in [35] showing the same property using approximation of log p(D), as opposed to using ELBO.
Image classification LeNet. We now consider image classification with the LeNet architecture [44] on 6 data sets: MNIST, fashionMNIST, K-MNIST [10], CIFAR10, CIFAR100 [41] 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, but the differences are not statistically significant. Both CM-MF-VI and CMV-MF-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[42]. We use CIFAR10, CIFAR100, STL10 [11] and SVHN. We again compare CMV-MF-VI, CV-MF-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-MF-VI/CMV-MF-VI 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-MF-VI/CMV-MF-VI outperforms learning just prior variances (CV-MF-VI).

heading: 5 Conclusions
text: 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 over-fitting. The developed algorithms allowed us to demonstrate that the ELBO can be a 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.

heading: Acknowledgements
text: Marcin B. Tomczak is funded by a Cambridge Trust scholarship. Siddharth Swaroop is supported by an EPSRC DTP studentship and a Microsoft Research EMEA PhD Award. Andrew Y. K. Foong is 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 in the Application of Artificial Intelligence to the study of Environmental Risks (AI4ER). The authors thank David Burt for insightful discussions.

Review summaries:
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. 

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.  Extensive 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.

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.

Although marginal likelihood can be used to select hyper-parameters, the ELBO is generally regarded as too weak a bound to do this effectively. The authors propose applying collapsed variational bounds to variational inference of BNNs. They 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. They empirically show that their method performs well in a range of settings, including in tests of ood uncertainty.