Collapsed Variational Bounds for Bayesian Neural Networks

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 . 1 Introduction . 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 ) . 2 Collapsed Variational Bounds for Bayesian Neural Networks . 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 . 2.1 Deriving collapsed variational bounds . 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 .

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