On the Parameterization of Gaussian Mean Field Posteriors in Bayesian Neural Networks

1 INTRODUCTION . Bayesian Neural Networks explicitly represent their parameter-uncertainty by forming a posterior distribution over model parameters , instead of relying on a single point estimate for making predictions , as is done in traditional deep learning . For neural network weights w , features x and labels y , the posterior distribution p ( w|x , y ) is computed using Bayes ’ rule , which multiplies the prior distribution p ( w ) and data likelihood p ( y|w , x ) and renormalizes . When predicting with Bayesian neural networks , we form an average over model predictions where each prediction is generated using a set of parameters that is randomly sampled from the posterior distribution . Bayesian neural networks are thus a type of ensembling , of which various types have proven highly effective in deep learning ( see e.g . Goodfellow et al. , 2016 , sec 7.11 ) . Besides offering improved predictive performance over single models , Bayesian ensembles are also more robust because ensemble members will tend to make different predictions on hard examples ( Raftery et al. , 2005 ) . In addition , the diversity of the ensemble represents predictive uncertainty and can be used for out-of-domain detection or other risk-sensitive applications ( Ovadia et al. , 2019 ) . Variational inference is a popular class of methods for approximating the posterior distribution p ( w|x , y ) , since the exact Bayes ’ rule is often intractable to compute for models of practical interest . This class of methods specifies a distribution qθ ( w ) of given parametric or functional form as the posterior approximation , and optimizes the approximation by solving an optimization problem . In particular , we minimize the Kullback-Leibler ( KL ) divergence DKL between the variational distribution qθ ( w ) and the true posterior distribution p ( w|x , y ) , which is given by DKL [ qθ ( w ) ||p ( w|x , y ) ] = Eq [ log qθ ( w ) p ( w|x , y ) ] = Eq [ log qθ ( w ) p ( w ) p ( y|w , x ) /p ( y|x ) ] . Here , we do not know the normalizing constant of the exact posterior p ( y|x ) , but since this term does not depend on w , we may ignore it for the purpose of optimizing our approximation q . We are then left with what is called the negative Evidence Lower Bound ( negative ELBO ) : Lq = DKL [ qθ ( w ) ||p ( w ) ] − Eq [ p ( y|w , x ) ] . In practice , the expectation of the likelihood p ( y|w , x ) with respect to q is usually not analytically tractable and instead is estimated using Monte Carlo sampling : Eq [ log p ( y|w , x ) ] ≈ 1 S S∑ s=1 log p ( y|w ( s ) , x ) , w ( s ) ∼ qθ ( w ) , where the ELBO can then be optimized by differentiating this stochastic approximation with respect to the variational parameters θ ( Salimans et al. , 2013 ; Kingma & Welling , 2013 ) . In Gaussian Mean Field Variational Inference ( GMFVI ) , we choose the variational approximation to be a fully factorized Gaussian distribution q = N ( µq , Σq ) with wlij ∼ N ( µlij , σ2lij ) , where l is a layer number , and i and j are the row and column indices in the layer ’ s weight matrix . Although this type of approximation is considered to be one of the simplest types of variational approximations , it already doubles the parameter count compared to deterministic neural networks . In addition , Bayesian neural networks with mean-field Gaussian posterior approximations often are harder to train than deterministic neural networks because they suffer from increased noise in stochastic gradient estimates . Beyond mean-field variational inference , recent work on approximate Bayesian inference has explored ever richer parameterizations of the approximate posterior in the hope of improving the performance of Bayesian neural networks ( see Figure 1 ) . In contrast , here we study a simpler , more compactly parameterized variational approximation , which we show can also work well for a variety of models . In particular we find that : • Converged posterior standard deviations under GMFVI consistently display strong lowrank structure . This means that by decomposing these variational parameters into a lowrank factorization , we can make our variational approximation more compact without decreasing our model ’ s performance . • Factorized parameterizations of posterior standard deviations are easier to train since they improve the signal-to-noise ratio of stochastic gradient estimates of the variational lower bound . Using a more limited parameterization of the approximate posterior can thus lead to faster convergence and a reduction in number of parameters compared to standard GMFVI . 2 MEAN FIELD POSTERIOR STANDARD DEVIATIONS NATURALLY HAVE LOW-RANK STRUCTURE . In this section we show that the converged posterior standard deviations of Bayesian neural networks trained using standard GMFVI consistently display strong low-rank structure . We also show that it is possible to compress the learned posterior standard deviation matrix using a low-rank approximation without decreasing the network ’ s performance . We first briefly introduce the mathematical notation for our GMFVI setting and the low-rank approximation that we explore . We then provide experimental results that support the two main claims of this section . To avoid any confusion among the readers , we would like to clarify that we use the terminology “ low-rank ” in a particular context . While variational inference typically makes use of low-rank decompositions to compactly represent the dense covariance of a Gaussian variational distribution ( see numerous references in Section 4 ) , we investigate instead underlying low-rank structures within the already diagonal covariance of a Gaussian fully-factorized variational distribution . We will make this explanation more formal in the next section . 2.1 METHODOLOGY . To introduce the notation we consider layers that consist of a linear transformation followed by a non-linearity f , al = hlWl + bl , hl+1 = f ( al ) , ( 1 ) where Wl ∈ Rm×n , hl ∈ R1×m and bl , al , hl+1 ∈ R1×n . To simplify the notation in the following , we drop the subscript l such that W = Wl , µq = µql , Σq = Σql and we focus on the kernel matrix W for a single layer . In GMFVI , ( Blundell et al. , 2015 ) , we model the variational posterior as q ( W ) = N ( µq , Σq ) = m∏ i=1 n∏ j=1 q ( wij ) , with q ( wij ) = N ( µij , σ2ij ) , ( 2 ) where µq ∈ Rmn×1 is the posterior mean vector , Σq ∈ Rmn×mn+ is the diagonal posterior covariance matrix . The weights are then usually sampled using a reparametrization trick ( Kingma & Welling , 2013 ) , i.e , for the s-th sample , we have w ( s ) ij = µij + σij ( s ) , ∼ N ( 0 , 1 ) . ( 3 ) In practice , we often represent the posterior standard deviation parameters σij in the form of a matrix A ∈ Rm×n+ . Note that we have the relationship Σq = diag ( vec ( A2 ) ) where the elementwise-squared A is vectorized by stacking its columns , and then expanded as a diagonal matrix into Rmn×mn+ . In the sequel , we start by empirically studying the properties of the spectrum of matrices A post training , while using standard Gaussian mean-field variational distributions . Interestingly , we observe that those matrices naturally exhibit a low-rank structure ( see Section 2.3 for the corresponding experiments ) , i.e , A ≈ UVT ( 4 ) for some U ∈ Rm×k , V ∈ Rn×k and k a small value ( e.g. , 2 or 3 ) . This observation motivates the introduction of the following variational family , which we name k-tied Normal : k-tied -N ( W ; µq , U , V ) = N ( µq , diag ( vec ( ( UVT ) 2 ) ) ) , ( 5 ) where the squaring of the matrix UVT is applied elementwise . Due to the tied parametrization of the diagonal covariance matrix , we emphasize that this variational family is smaller—i.e. , included in—the standard Gaussian mean-field variational distribution family . As formally discussed in Appendix B , the matrix variate Gaussian distribution ( Gupta & Nagar , 2018 ) , referred to as MN and already used for variational inference by Louizos & Welling ( 2016 ) and Sun et al . ( 2017 ) , is related to our k-tied Normal distribution with k = 1 whenMN uses diagonal row and column covariances . Interestingly , we prove that for k ≥ 2 , our k-tied Normal distribution can not be represented by any MN distribution . This illustrates the main difference of our approach from the most closely related previous work of Louizos & Welling ( 2016 ) ( see also Figure 1 ) . Furthermore , notice that diagonal covariance Σq repeatedly reuses the same elements of U and V , which results in parameter sharing across different weights . The total number of the standard deviation parameters in our method is k ( m + n ) from U and V , compared to mn in the standard GMFVI parametrization from A . Given that in our experiments the k is very low ( e.g . k = 2 ) this reduces the number of parameters from quadratic to linear in the dimensions of the layer , see Table 1 . More importantly , such parameter sharing across the weights leads to higher signal-to-noise ratio during training and thus faster convergence . We demonstrate this phenomena in the next section . In the rest of this section , we will first demonstrate that the standard GMFVI methods already learn a low-rank structure in the posterior standard deviation matrix A . Furthermore , we will show that replacing the full matrix A with its low-rank approximation does not reduce the predictive performance . 2.2 EXPERIMENTAL SETTING . Before providing the experimental results we briefly explain the key properties of the experimental setting . We analyse three types of GMFVI Bayesian neural network models : • Multilayer Perceptron ( MLP ) : a network of 3 dense layers and ReLu activations that we train on the MNIST dataset ( LeCun & Cortes , 2010 ) . We use the last 10,000 examples of the training set as a validation set . • Convolutional Neural Network ( CNN ) : a LeNet architecture ( LeCun et al. , 1998 ) with 2 convolutional layers and 2 dense layers that we train on the CIFAR-100 dataset ( Krizhevsky et al. , 2009b ) . We use the last 10,000 examples of the training set as a validation set . • Long Short-Term Memory ( LSTM ) : a model that consists of an embedding and an LSTM cell ( Hochreiter & Schmidhuber , 1997 ) , followed by a single unit dense layer . We train it on an IMBD dataset ( Maas et al. , 2011 ) , in which we use the last 5,000 examples of the training set as a validation set . • Residual Convolutional Neural Network ( ResNet ) : a ResNet-181 architecture ( He et al. , 2016 ) trained on all 50,000 training examples of the CIFAR-10 dataset ( Krizhevsky et al. , 2009a ) . In each of the four models we use a mean-field Normal posterior and a Normal prior with a single scalar standard deviation hyper-parameter for all the layers . We optimize the variational parameters using an Adam optimizer ( Kingma & Ba ( 2014 ) ) . For a more comprehensive explanation of the experimental setup used in this section please refer to Appendix A.1 . Finally , we highlight that our experiments focus primarily on the comparison across a broad range of model types rather than competing with the state-of-the-art results over the specifically used datasets . Therefore , we use mainly small to medium models ( MLP , CNN , LSTM ) that are known to train well using the standard GMFVI approach explored in this paper . Encouraged by the reviewers , we also briefly show that our results can extend to larger models such as the ResNet-18 model . However , scaling GMFVI to such larger model sizes is still a challenging research problem ( Osawa et al. , 2019 ) .

Summary: The paper considers variational Bayesian inference (learning) for neural networks assuming that both the prior distributions and the posterior distributions of the network weights are factorising Gaussians. It is well known that the respective optimisation task for the parameters of the posterior weight distributions (ELBO) is tractable by stochastic gradient ascent.  The authors propose to simplify (restrict) the model even further, by assuming that the posterior variances for fully connected layers, (when seen as a matrix) have low rank. It is pretty obvious that the ELBO optimisation task remains tractable in this case.