Deterministic Variational Inference for Robust Bayesian Neural Networks

1 INTRODUCTION . Bayesian approaches to neural network training marry the representational flexibility of deep neural networks with principled parameter estimation in probabilistic models . Compared to “ standard ” parameter estimation by maximum likelihood , the Bayesian framework promises to bring key advantages such as better uncertainty estimates on predictions and automatic model regularization ( MacKay , 1992 ; Graves , 2011 ) . These features are often crucial for informing downstream decision tasks and reducing overfitting , particularly on small datasets . However , despite potential advantages , such Bayesian neural networks ( BNNs ) are often overlooked due to two limitations : First , posterior inference in deep neural networks is analytically intractable and approximate inference with Monte Carlo ( MC ) techniques can suffer from crippling variance given only a reasonable computation budget ( Kingma et al. , 2015 ; Molchanov et al. , 2017 ; Miller et al. , 2017 ; Zhu et al. , 2018 ) . Second , performance of the Bayesian approach is sensitive to the choice of prior ( Neal , 1993 ) , and although we may have a priori knowledge concerning the function represented by a neural network , it is generally difficult to translate this into a meaningful prior on neural network weights . Sensitivity to priors and initialization makes BNNs non-robust and thus often irrelevant in practice . In this paper , we describe a novel approach for inference in feed-forward BNNs that is simple to implement and aims to solve these two limitations . We adopt the paradigm of variational Bayes ( VB ) for BNNs ( Hinton & van Camp , 1993 ; MacKay , 1995c ) which is normally deployed using Monte ∗Work done during an internship at Microsoft Research , Cambridge . †Work done while at Microsoft Research , Cambridge . Carlo variational inference ( MCVI ) ( Graves , 2011 ; Blundell et al. , 2015 ) . Within this paradigm we address the two shortcomings of current practice outlined above : First , we address the issue of high variance in MCVI , by reducing this variance to zero through novel deterministic approximations to variational inference in neural networks . Second , we derive a general and robust Empirical Bayes ( EB ) approach to prior choice using hierarchical priors . By exploiting conjugacy we derive data-adaptive closed-form variance priors for neural network weights , which we experimentally demonstrate to be remarkably effective . Combining these two novel ingredients gives us a performant and robust BNN inference scheme that we refer to as “ deterministic variational inference ” ( DVI ) . We demonstrate robustness and improved predictive performance in the context of non-linear regression models , deriving novel closed-form results for expected log-likelihoods in homoscedastic and heteroscedastic regression ( similar derivations for classification can be found in the appendix ) . Experiments on standard regression datasets from the UCI repository , ( Dheeru & Karra Taniskidou , 2017 ) , show that for identical models DVI converges to local optima with better predictive loglikelihoods than existing methods based on MCVI . In direct comparisons , we show that our Empirical Bayes formulation automatically provides better or comparable test performance than manual tuning of the prior and that heteroscedastic models consistently outperform the homoscedastic models . Concretely , our contributions are : • Development of a deterministic procedure for propagating uncertain activations through neural networks with uncertain weights and ReLU or Heaviside activation functions . • Development of an EB method for principled tuning of weight priors during BNN training . • Experimental results showing the accuracy and efficiency of our method and applicability to heteroscedastic and homoscedastic regression on real datasets . 2 VARIATIONAL INFERENCE IN BAYESIAN NEURAL NETWORKS . We start by describing the inference task that our method must solve to successfully train a BNN . Given a model M parameterized by weights w and a dataset D = ( x , y ) , the inference task is to discover the posterior distribution p ( w|x , y ) . A variational approach acknowledges that this posterior generally does not have an analytic form , and introduces a variational distribution q ( w ; θ ) parameterized by θ to approximate p ( w|x , y ) . The approximation is considered optimal within the variational family for θ∗ that minimizes the Kullback-Leibler ( KL ) divergence between q and the true posterior . θ∗ = argmin θ DKL [ q ( w ; θ ) ||p ( w|x , y ) ] . Introducing a prior p ( w ) and applying Bayes rule allows us to rewrite this as optimization of the quantity known as the evidence lower bound ( ELBO ) : θ∗ = argmax θ { Ew∼q [ log p ( y|w , x ) ] −DKL [ q ( w ; θ ) ||p ( w ) ] } . ( 1 ) Analytic results exist for the KL term in the ELBO for careful choice of prior and variational distributions ( e.g . Gaussian families ) . However , whenM is a non-linear neural network , the first term in equation 1 ( referred to as the reconstruction term ) can not be computed exactly : this is where MC approximations with finite sample size S are typically employed : Ew∼q [ log p ( y|w , x ) ] ≈ 1 S S∑ s=1 log p ( y|w ( s ) , x ) , w ( s ) ∼ q ( w ; θ ) . ( 2 ) Our goal in the next section is to develop an explicit and accurate approximation for this expectation , which provides a deterministic , closed-form expectation calculation , stabilizing BNN training by removing all stochasticity due to Monte Carlo sampling . 3 DETERMINISTIC VARIATIONAL APPROXIMATION . Figure 1 shows the architecture of the computation of Ew∼q [ log p ( D|w ) ] for a feed-forward neural network . The computation can be divided into two parts : first , propagation of activations though ℒ 𝔼𝒘∼𝑞 log 𝑝 𝒚 𝒘 , 𝒙 𝒂0 𝒉1 𝒂1 Layer 1 Layer 2 Layer 𝐿… ( a ) ( b ) 𝒂𝐿𝒙 𝒚 𝑞 ( 𝑾 0 ; 𝜽 ) 𝑞 ( 𝑾 1 ; 𝜽 ) 𝑞 ( 𝒂 0 ) 𝑞 ( 𝒂 1 ) 𝑞 ( 𝒂 𝐿 ) Figure 1 : Architecture of a Bayesian neural network . Computation is divided into ( a ) propagation of activations ( a ) from an input x and ( b ) computation of a log-likelihood function L for outputs y . Weights are represented as high dimensional variational distributions ( blue ) that induce distributions over activations ( yellow ) . MCVI computes using samples ( dots ) ; our method propagates a full distribution . parameterized layers and second , evaluation of an unparameterized log-likelihood function ( L ) . In this section , we describe how each of these stages is handled in our deterministic framework . 3.1 MOMENT PROPAGATION . We begin by considering activation propagation ( figure 1 ( a ) ) , with the aim of deriving the form of an approximation q̃ ( aL ) to the final layer activation distribution q ( aL ) that will be passed to the likelihood computation . We compute aL by sequentially computing the distributions for the activations in the preceding layers . Concretely , we define the action of the lth layer that maps a ( l−1 ) to al as follows : hl = f ( a ( l−1 ) ) , al = hlW l + bl , where f is a non-linearity and { W l , bl } ⊂ w are random variables representing the weights and biases of the lth layer that are assumed independent from weights in other layers . For notational clarity , in the following we will suppress the explicit layer index l , and use primed symbols to denote variables from the ( l − 1 ) th layer , e.g . a′ = a ( l−1 ) . Note that we have made the non-conventional choice to draw the boundaries of the layers such that the linear transform is applied after the nonlinearity . This is to emphasize that al is constructed by linear combination of many distinct elements of h′ , and in the limit of vanishing correlation between terms in this combination , we can appeal to the central limit theorem ( CLT ) . Under the CLT , for a large enough hidden dimension and for variational distributions with finite first and second moments , elements ai will be normally distributed regardless of the potentially complicated distribution for hj induced by f 1 . We empirically observe that this claim is approximately valid even when ( weak ) correlations appear between the elements of h during training ( see section 3.1.1 ) . Having argued that a adopts a Gaussian form , it remains to compute the first and second moments . In general , these can not be computed exactly , so we develop an approximate expression . An overview of this derivation is presented here with more details in appendix A . First , we modelW , b and h as independent random variables , allowing us to write : 〈ai〉 = 〈hj〉 〈Wji〉+ 〈bi〉 , Cov ( ai , ak ) = 〈hjhl〉Cov ( Wji , Wlk ) + 〈Wji〉Cov ( hj , hl ) 〈Wlk〉+ Cov ( bi , bk ) , ( 3 ) where we have employed the Einstein summation convention and used angle brackets to indicate expectation over q . If we choose a variational family with analytic forms for weight means and covariances ( e.g . Gaussian with variational parameters 〈Wji〉 and Cov ( Wji , Wlk ) ) , then the only difficult terms are the moments of h : 〈hj〉 ∝ ∫ f ( αj ) exp [ − ( αj−〈a ′ j〉 ) 2 2Σ′jj ] dαj , ( 4 ) 〈hjhl〉 ∝ ∫ f ( αj ) f ( αl ) exp [ −1 2 ( αj−〈a′j〉 αl−〈a′l〉 ) > ( Σ′jj Σ ′ jl Σ′lj Σ ′ ll ) −1 ( αj−〈a′j〉 αl−〈a′l〉 ) ] dαj dαl , ( 5 ) 1We are also required to choose a Gaussian variational approximation for b to preserve the Gaussian distribution of a. where we have used the Gaussian form of a′ parameterized by mean 〈a′〉 and covariance Σ′ , and for brevity we have omitted the normalizing constants . Closed form solutions for the integral in equation 4 exist for Heaviside or ReLU choices of non-linearity f ( see appendix A ) . Furthermore , for these non-linearities , the 〈 a′j 〉 → ±∞ and 〈a′l〉 → ±∞ asymptotes of the integral in equation 5 have closed form . Figure 2 shows schematically how these asymptotes can be used as a first approximation for equation 5 . This approximation is improved by considering that ( by definition ) the residual decays to zero far from the origin in the ( 〈 a′j 〉 , 〈a′l〉 ) plane , and so is well modelled by a decaying function exp [ −Q ( 〈 a′j 〉 , 〈a′l〉 , Σ′ ) ] , where Q is a polynomial in 〈a′〉 with a dominant positive even term . In practice we truncate Q at the quadratic term , and calculate the polynomial coefficients by matching the moments of the resulting Gaussian with the analytic moments of the residual . Specifically , using dimensionless variables µ′i = 〈a′i〉 / √ Σ′ii and ρ ′ ij = Σ ′ ij/ √ Σ′iiΣ ′ jj , this improved approximation takes the form 〈hjhl〉 = S′jl { A ( µ′j , µ ′ l , ρ ′ jl ) + exp [ −Q ( µ′j , µ′l , ρ′jl ) ] } , ( 6 ) where the expressions for the dimensionless asymptote A and quadratic Q are given in table table 1 and derived in appendix A.2.1 and A.2.2 . The dimensionful scale factor S′jl is 1 for a Heaviside non-linearity or Σ ′ jl/ρ ′ jl for ReLU . Using equation 6 in equation 3 gives a closed form approximation for the moments of a as a function of moments of a′ . Since a is approximately normally distributed by the CLT , this is sufficient information to sequentially propagate moments all the way through the network to compute the mean and covariances of q̃ ( aL ) , our explicit multivariate Gaussian approximation to q ( aL ) . Any deep learning framework supporting special functions arcsin and Φ will immediately support backpropagation through the deterministic expressions we have presented . Below we briefly empirically verify the presented approximation , and in section 3.2 we will show how it is used to compute an approximate log-likelihood and posterior predictive distribution for regression and classification tasks .

Summary: This work is tackling two difficulties in current VB applied to DNNs ("Bayes by backprop"). First, MC approximations of intractable expectations are replaced by deterministic approximations. While this has been done before, the solution here is new and very interesting. Second, a Gaussian prior with length scales is learned by VB empirical Bayes alongside the normal training, which is also very useful.