Mixture Distributions for Scalable Bayesian Inference

1 INTRODUCTION . Neural Networks models have been applied in diverse fields from weather forecasting , to autonomous vehicle driving , to online advertisement and many more ( Goodfellow et al . ( 2016 ) ) . However , vanilla feed-forward NNs are susceptible to the problem of over-fitting . In addition to this , NNs trained using Maximum Likelihood Estimation ( MLE ) , or Maximum A Posteriori ( MAP ) can not provide an estimate of the uncertainty in predicted value and tend to produce overconfident results on out-of-distribution test data . Thus we aim to build deep learning frameworks which are more robust , secure and reliable - specifically , in the face of uncertainty we would like the model to be able to say ” I Don ’ t Know ! ” A principled approach to build such models is through Bayesian inference , where instead of a point estimate for the network parameters , we infer the posterior distribution of the weights given the data . The predictive distribution of unseen data ( x∗ , y∗ ) is given by q ( y∗|x∗ ) =∫ p ( y∗|x∗ , W ) p ( W |D ) dW , where p ( W |D ) is the true posterior computed from Bayes rule . Un- fortunately , since modern NNs have an exponentially large number of parameters and do not lend themselves to being integrated into mathematical equations , exact inference remains intractable . Several methods for approximate Bayesian inference have been investigated . Until now , Hamiltonian Monte Carlo ( HMC ) ( Neal ( 2012 ) ) has been considered the gold standard for approximate Bayesian inference . However , HMC requires explicit storage of samples from the posterior and is typically not scalable to larger networks and datasets . Thus methods based on Variational Inference ( VI ) have gained popularity recently for the task of approximate inference . VI relies on using a surrogate posterior ( qθ ( W ) ) as an approximation for the true posterior . One of the first application of VI in NNs was by Hinton & Van Camp ( 1993 ) but the optimization remained intractable for most Bayesian Neural Networks ( BNNs ) . Graves ( 2011 ) revisited similar ideas and proposed a simple but biased estimator for performing VI with a fully factorized posterior ( mean field assumption ) . Recently , Bayes-by-Backprop algorithm ( BBB ) also employed VI using an unbiased estimator for Variational loss , a fully factorized Gaussian posterior and a non-Gaussian prior Blundell et al . ( 2015 ) . Interestingly , Gal & Ghahramani ( 2016 ) showed links between Bernoulli Dropout ( Srivastava et al . ( 2014 ) , Hinton et al . ( 2012 ) ) and approximate inference in Deep Gaussian Process ( Damianou & Lawrence ( 2013 ) ) , thus allowing for the extraction of uncertainty estimates of a NN in a principled way . Louizos & Welling ( 2016 ) arrived at the same conclusion through structured posterior approximations via matrix variate Gaussian Posterior ( Gupta & Nagar ( 2018 ) ) and local reparameterization trick ( Kingma et al . ( 2015 ) ) . Bayesian methods are far more computationally expensive , as compared to standard NNs , making Deep Ensembles a very attractive alternative as they scale well to large datasets and models . The intuition behind Deep Ensembles is straightforward , NNs initialized and trained independently can be expected to output similar predictions on training data , while disagreeing on points away from training set . Despite their empirical success in uncertainty estimation and adversarial robustness ( Lakshminarayanan et al . ( 2017 ) , Strauss et al . ( 2018 ) ) , Deep Ensembles have lacked theoretic support from the Bayesian perspective . Thus , it is interesting to see how ensembles of deterministic and stochastic NNs can be placed into the VI framework . Another drawback of Deep Ensembles is their heavy parametric overhead . We modify the BBB algorithm to learn mixture distributions rather than the unimodal Gaussian distribution , which reduces this burden . We further show that this way , Deep Ensembles can be linked with the Bayesian framework . In essence , using our methods , we are able to achieve accuracy and uncertainty estimates ( with lower parametric cost ) which are never worse than Deep Ensembles , BBB , Concrete Dropout Contributions : We design an efficient and simple way of sampling from mixture distributions which incorporates the ” local reparameterization trick ” ( Kingma et al . ( 2015 ) ) to obtain low variance gradients for tuning of parameters in mixture distributions . This is also of independent interest for designing more complex posteriors using any mixture distributions besides the ones used in this paper . We also fit ensembles of deterministic and Bayesian NNs into VI framework with no changes made to the original algorithm . We improve upon the uncertainty estimates of Deep Ensembles by using mixture distribution with relatively less number of parameters . Distribution of the paper - In section 2 , we show how the original BBB algorithm can be modified for generalized mixture distributions while remaining back-prop compatible . For simplicity , we stick to a mixture of Gaussian ’ s . In section 3 , we provide links between Ensembles and VI with mixture distributions . Section 4 contains experiments that demonstrate that our methods generate superior uncertainty estimates as compared to unimodal distributions or Deep Ensembles . 2 CONCRETE MIXTURE OF GAUSSIAN . A feed-forward network trained with gradient descent will arrive at point estimates . However , in the case of BNNs , the weights are not point estimates but a probability distribution . Let us revisit the problem of finding a functional relationship y = fW ( x ) between input x and output y , given a labelled training dataset D = { ( x1 , y1 ) , ( x2 , y2 ) , . . . } from a Bayesian standpoint . Our task is to find a distribution over the parameters p ( W |D ) , given the input data . Given this posterior , we can predict test output y∗ given a new test input x∗ by marginalizing the likelihood over the space of parameters W . However , even for modest sized NNs , the number of uncertain parameters prohibit this calculation analytically . We require approximate inference methods in such cases . We define an approximating variational distribution qθ ( W ) with parameters θ . We would minimize the Kullback-Leibler ( KL ) Divergence with respect to the parameters θ to maximize the similarity between proposed posterior and the true posterior . Note that the problem of minimizing KL Divergence can be reformulated to minimizing of another term usually called the Variational free energy in literature Friston et al . ( 2007 ) , Blundell et al . ( 2015 ) . Evaluating the objective of maximizing this is computationally expensive for large datasets therefore . Thus we split the entire data into M mini batchesD1 , D2 , . . . , DM . The unbiased Monte Carlo ( MC ) approximation to the minibatch cost it given as : W j ∼ qθ ( W ) , L̂V I−MC ( θ ) = − ∑ i∈Dj log ( p ( yi|fW j ( xi ) ) ) + 1 M KL ( qθ ( W ) ||p ( W ) ) ( 1 ) In order to use stochastic backpropogation , we need the MC sample W j to be reparameterized as g ( θ , ) . We shall see how this is possible for Mixture of Gaussians ( MoG ) next . 2.1 SAMPLING FROM CONCRETE MIXTURE OF GAUSSIAN . Blundell et al . ( 2015 ) proposed the Bayes-by-Backprop algorithm for tuning of weight posterior ’ s parameters . For simplicity , the posterior was assumed to be fully factorized Gaussian . The assumption of fully factorized posterior is also known as the Mean Field assumption in the literature . qθ ( W ) = N∏ i=1 N ( wi ; µi , σ 2 i ) θ = ( µi , σ 2 i ) ( i=1 , ... , N ) N = Total no . of weights ( 2 ) Louizos & Welling ( 2016 ) used a matrix variate Gaussian posterior Gupta & Nagar ( 2018 ) to capture correlations between weights of same layer . Both of the above assumption do not account for multiple modes in true posterior . A natural choice for posterior that can capture multiple modes is a mixture distribution , where each mixture component would correspond to a separate mode of the true posterior . We show method to sample from for mixture of Gaussian ( MoG ) for independent weights and mixture of matrix variate Gaussian ( MoMG ) for correlated weights , while still being able to use the local reparameterization trick . This method can be extended for any mixture of distributions in general . To sample weight w from a mixture of K Gaussians i.e . w ∼ qθ ( w ) = ∑K j=1 pjN ( w ; µj , σ 2 j ) , ∑K j=1 pj = 1 , we assume a latent random variable which governs the component of the mixture that generates the value of w. We assume the latent random variable to be uniform random variable u i.e . u ∼ Uniform [ 0,1 ] . The latent random variable u has a pj probability of being in the range [ pj−1 , pj + pj−1 ) , similarly it has pj+1 probability of being in between [ pj + pj−1 , pj + pj−1 + pj+1 ) and so on , thus if u falls in the range [ pj−1 , pj + pj−1 ) we sample from the j ’ th distribution N ( w ; µj , σ 2 j ) . For example when sampling from mixture of three Gaussian , we will have to first sample two random variables zp1 and zp1+p2 , which can be seen as correlated Bernoulli random variables . u ∼ Uniform [ 0 , 1 ] , zp1 = 1u < p1 , zp1+p2 = 1u < p1+p2 and ∼ N ( 0 , 1 ) w = zp1 ( µ1 + σ1 ) + ( 1− zp1 ) ( zp1+p2 ) ( µ2 + σ2 ) + ( 1− zp1+p2 ) ( µ3 + σ3 ) ( 3 ) Unfortunately , in the above case where we sample w from a multimodal distribution , local reparametrization trick can not be applied . This is because the Bernoulli random variables zp1 and zp1+p2 can ’ t be parameterized in the form of g ( θ , ) . So , we replace them with their continuous relaxation , the Concrete distribution Gal et al . ( 2017 ) , Maddison et al . ( 2016 ) . This distribution can be viewed as softmax relaxation of the max function used in Gumbel-max trick Gal et al . ( 2017 ) . This allows us to represent the discrete random variable zp1 as ẑp1 = g ( θ , ) . Thus we replace zp1 and zp1+p2 in equation 4 with : zp1 ≈ ẑp1 = sigmoid ( 1 t . ( log ( p1 ) − log ( 1− p1 ) + log ( u ) − log ( 1− u ) ) zp1+p2 ≈ ẑp1+p2 = sigmoid ( 1 t . ( log ( p1 + p2 ) − log ( 1− p1 − p2 ) + log ( u ) − log ( 1− u ) ) ( 4 ) where , t is the temperature term and is usually kept as 0.1 or 0.01 ( < < 1 to guarantee most of the mass of the pdf is around 0 and 1 ) . So , we have now successfully reparameterized w as g ( θ , ) where θ = ( µ1 , σ1 , µ2 , σ2 , µ3 , σ3 , p1 , p2 ) . To prevent parametric explosion , we assume p1 , p2 are same for the entire layer . This is similar to assuming that Dropout parameter p is the same for entire layer . Thus now , with mean field assumption we can now define the posterior of the weight matrixW from mixture of K Gaussians as : qθ ( W ) = ∏ i ∑ j pijN ( wi ; µij , σ 2 ij ) , K∑ j=1 pij = 1 , θ = ( ( pij , µij , σ 2 ij ) j=1 , ... , K ) i=1 , ... , N ( 5 ) To model the correlations among the weights it is possible for us the treat the entire weight matrix W of a layer , as a MoMG Gupta & Nagar ( 2018 ) i.e . p ( W ) = ∑ i piMN ( Mi , Ui , Vi ) , where p ( W ′ ) = MN ( Mi , Ui , Vi ) = exp ( − 12 tr [ V −1 i ( W ′ −Mi ) TU−1i ( W ′ −Mi ) ] ) ( 2π ) rc/2|Vi|c/2|Ui|r/2 ( 6 ) Mi is a r× c matrix and is the mean of the distribution , Ui is a r× r matrix that provides covariance of the rows and Vi is a c × c matrix that governs the covariance of the columns of the matrix . It has been shown by Gupta et al . in Gupta & Nagar ( 2018 ) that matrix variate distributions are essentially multivariate Gaussian . p ( vec ( W ′ ) ) = N ( vec ( Mi ) , Vi ⊗ Ui ) , ⊗ is the Kronecker product ( 7 ) Thus , now instead of treating each weight independently we are now able to sample directly from its joint distribution which is assumed to be a MoMG . For sampling from a mixture of 2 matrix variate Gaussian , W = zp ( M1+U 1 2 1 EV 1 2 1 ) + ( 1− zp ) ( M2+U 1 2 2 EV 1 2 2 ) ) where Ei , j ∼ N ( 0 , 1 ) . We replace zp with its continuous relaxed version ẑp , thus making it possible to use the local reparameterization trick . We assume Ui = diag ( σ2ri ) and Vi = diag ( σ 2 ci ) i.e . independent rows and columns in this paper . This allows us to reduce the number of parameters as compared to MoG , while still modelling correlations between weights . Now that we are able to sample from MoG and MoMG we need to be able to approximate DKL ( qθ ( W ) ||p ( W ) ) , where we assume a Gaussian prior i.e . p ( W ) = N ( 0 , l−2I ) . We use Monte Carlo approximations similar to the one used in Blundell et al . ( 2015 ) for optimising the DKL ( qθ ( W ) ||p ( W ) ) . For all following sections , we interchangeably use the name Concrete MoG ( CMoG ) and Concrete MoMG ( CMoMG ) for MoG and MoMG respectively .

Summary: This paper proposes to use either relaxed mixture distributions or relaxed mixtures of matrix Gaussians as the approximate posterior for Bayesian neural networks. Naturally, taking the mixture variance to zero allows a stretched interpretation of ensembling to become Bayesian as well. Experiments are performed on small neural networks on regression tasks, as well as uncertainty estimation experiments on deep networks. Finally, a downstream task of bandit optimization (the mushroom experiment) is performed.