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