Refining the variational posterior through iterative optimization

1 INTRODUCTION . Uncertainty plays a crucial role in a multitude of machine learning applications , ranging from weather prediction to drug discovery . Poor predictive uncertainty risks potentially poor outcomes , especially in domains such as medical diagnosis or autonomous vehicles where some forms of high confidence errors may be especially costly ( Amodei et al. , 2016 ) . Thus , it is becoming increasingly important that the underlying model provides high quality uncertainty estimates along with its predictions . Yet , possibly the most widely used models , deep neural networks ( LeCun et al. , 2015 ) , are unable to accurately quantify model uncertainty . They are often overconfident in their predictions , even when their predictions are incorrect ( Guo et al. , 2017 ; Ovadia et al. , 2019 ) . By marginalizing over a posterior distribution over the parameters given the training data , Bayesian inference provides a principled approach to capturing uncertainty . In contrast , standard training of neural networks employs a point estimate of the parameters , which can not account for model uncertainty . Unfortunately , exact Bayesian inference is intractable in general for neural networks . To model epistemic uncertainty , variational inference ( VI ) instead approximates the true posterior with a simpler distribution . The most widely used one for neural networks is the mean-field approximation , where the posterior is represented using an independent Gaussian distribution over all the weights . Variational inference is appealing since it reduces the problem of inference to an optimization problem , minimizing the discrepancy between the true posterior and the variational posterior . The key challenge , however , is the task of training expressive posterior approximations that can capture the true posterior without significantly increasing the computational costs . This paper describes a novel method for training highly flexible posterior approximations . The idea is to start with a coarse , mean-field approximation q ( w ) and make iterative , local refinements to it . The regions of the local refinements are determined by sampling the values of additive auxiliary variables . The model parameters w are expressed using a number of auxiliary variables ak ( Figure 1 left ) for k = 1 , . . . , K that leave the marginal distribution unchanged . In each iteration , we sample the value of an auxiliary variable according to the current variational approximation q ( ak ) and refine the approximation by conditioning on the newly sampled value q ( w ) ≈ p ( w|x , y , a1 : k ) ( Figure 1 right illustrates the process ) . Each refinement step makes cheap , local adjustments to the variational posterior in the region of the sampled auxiliary variables . At the end , we draw one sample from the refined q ( w ) . The refinement iterations have to be repeated for each posterior sample . The algorithm results in samples from a highly complex distribution , starting from a simple mean-field approximation . While the distribution of the samples is difficult to quantify , we show that it is not limited to factorized , uni-modal forms , and that the procedure is guaranteed to improve the resulting ELBO without posing a significant computational overhead . Summary of contributions : • A novel algorithm for refining a variational distribution , increasing its flexibility . • We show that the refinement steps are guaranteed to improve the quality of the variational distribution under mild conditions . • We showcase the effectiveness of the method on Bayesian neural networks using a set of UCI regression and image classification benchmarks . We set a new state-of-the-art in uncertainty estimation using variational inference at ResNet scale ( ResNet-20 , ( He et al. , 2016 ) ) scale on CIFAR10 . 2 METHODS . In this section , we first describe standard variational inference ( VI ) , followed by a detailed description of the iterative refinement algorithm . While VI and our proposed methodology are generally applicable to latent variable models , in this work , we consider the application to Bayesian neural networks ( Figure 1 ) , where x are inputs , y outputs , and w the weights of network . 2.1 VARIATIONAL INFERENCE IN BAYESIAN NEURAL NETWORKS . Exact Bayesian inference in Bayesian neural networks ( BNN ) ( Figure 1 ) is often intractable and is NP-hard in the general case . Variational inference attempts to approximate the true posterior p ( w|x , y ) with an approximate posterior qφ ( w ) , typically from a simple family of distributions , for example independent Gaussians over the weights , i.e . the mean-field approximation . To ensure that the approximate posterior is close to the true posterior , the parameters of qφ ( w ) , φ are optimized to minimize their Kullback-Leibler divergence : DKL ( qφ ( w ) ‖ p ( w|x , y ) ) . At the limit of DKL ( qφ ( w ) ‖ p ( w|x , y ) ) = 0 , the approximate posterior exactly captures the true posterior , although this might not be achievable if p ( w|x , y ) is outside of the distribution family of qφ ( w ) . In order to minimize the KL-divergence , variational inference optimizes the Evidence Lower Bound ( ELBO ) w.r.t . φ ( denoted as L ( φ ) ) , which is a lower bound to the log marginal likelihood log p ( y|x ) . Since the the marginal likelihood can be expressed as the sum of the KL-divergence and the ELBO , maximizing the ELBO is equivalent to minimizing the KL divergence : log p ( y|x ) = DKL ( qφ ( w ) ∥∥ p ( w|x , y ) ) ︸ ︷︷ ︸ ≥0 +L ( φ ) ≥ L ( φ ) = Eqφ [ log p ( y|x , w ) ] −DKL ( qφ ( w ) ∥∥ p ( w ) ) , due to non-negativity of the KL-divergence . 1 : procedure REFINEANDSAMPLE ( φ ) 2 : for m = 1 , . . . , M do 3 : φ0 ← φ 4 : for k = 1 , . . . , K do 5 : ak ∼ qφk−1 ( ak ) . Sample ak 6 : qφk ( w ) ← qφk−1 ( w|ak ) . Initialize φk 7 : φk ← argmaxφk L|a1 : k ( φk ) . Optimize φk 8 : end for 9 : wm ∼ qφK ( w ) . Sample the refined posterior 10 : end for 11 : return w1 : M 12 : end procedure Algorithm 1 : Pseudocode for the iterative refinement of the variational posterior Following the optimization of φ , the model can be used to make predictions on unseen data . For a new input x′ , the predictive distribution p ( y′|x′ , y , x ) can be approximated by stochastically drawing a small number ( around M ≤ 10 ) of sample model parameters and averaging their prediction in an ensemble model : w1 : M ∼ qφ ( w ) , p ( y′|x′ , y , x ) ≈ 1 M M∑ i=1 p ( y′|x′ , wi ) . 2.2 REFINING THE POSTERIOR . The main issue with variational inference is the inflexibility of the posterior approximation . The most widely used variant of variational inference , mean-field variational inference , approximates the posterior with independent Gaussians across all dimensions . This approximation is too simplistic to capture the complexities of the posterior for complicated models such as BNNs . Our idea is to refine the posterior approximation through iterative optimization . Using the refinement procedure , it is feasible to train a detailed posterior in the regions of the posterior samples used for prediction while relying on a coarse-grained approximation further away from these samples . The graphical model is augmented with a finite number of auxiliary variables a1 : K as shown in Figure 1 . The constraints are that ( x , y ) must be conditionally independent of the auxiliary variables given w and that they must not affect the prior distribution p ( w ) . This is important in justifying the use of the initial variational approximation . A significant way in which we distinguish ourselves from hierarchical variational models ( Ranganath et al. , 2016 ) is that the model is unaffected by the presence of the auxiliary variables . Their purpose is solely to aid the inference procedure . Given a Gaussian prior N ( 0 , σ2w ) over w , we express w as a sum of independent auxiliary variables2 w = K∑ k=1 ak , with p ( ak ) = N ( 0 , σ2ak ) for k = 1 , . . . , K , while ensuring that ∑K k=1 σ 2 ak = σ2w so that the prior p ( w ) = N ( 0 , σ2w ) is unchanged . Locally refining the approximate posterior refers to iteratively sampling the values of the auxiliary variables a1 : K and then approximating the posterior conditional on the sampled values , i.e . qφk ( w ) approximates p ( w|x , y , a1 : k ) for iterations k = 1 , . . . , K ( Algorithm 1 ) . Starting from the initial mean field approximation qφ ( w ) , we sample the value of a1 from qφ ( a1 ) = ∫ p ( a1|w ) qφ ( w ) dw , 1The sudden drops after sampling are optimizer artefacts due to having to reset the parameters of Adam . 2While we are focusing on one specific definition of the auxiliary variables , additive auxiliary variables , note that all of our results straight-forwardly generalize to arbitrary joint distributions p ( w , a1 : K ) that meet the constraints . then optimize the approximation to the conditional posterior : qφ1 ( w ) ≈ p ( w|x , y , a1 ) . This procedure is then iteratively repeated for a2 : K . In iteration k , 1 ) ak ∼ ∫ p ( ak|a1 : k−1 , w ) qφk−1 ( w ) dw 2 ) φk = argminDKL ( qφk ( w ) ∥∥ p ( w|x , y , a1 : k ) ) . Analogously to variational inference , the KL divergence is minimized through the optimization of the conditional ELBO in each iteration : L|a1 : k ( φk ) = Eqφk [ log p ( y|x , w ) ] − DKL ( qφk ( w ) ‖ p ( w|a1 : k ) ) . In order to get independent samples from the variational posterior , we have to repeat the iterative refinement for each ensemble member w1 : M . Toy example We use a toy example to demonstrate the procedure . In this toy example , we have a single weight w with prior p ( w ) = N ( 0 , 1 ) and a complicated posterior with four modes ( synthetically generated data ) . We express w as the sum of two auxiliary variables w = a1 + a2 with p ( a1 ) = N ( 0 , 0.8 ) and p ( a2 ) = N ( 0 , 0.2 ) ( which recovers p ( w ) as per the constraint ) . As Figure 3b shows , a Gaussian approximation to the posterior fails to capture the multimodal nature of the true posterior . The first step of the refinement is to sample using qφ : a1 ∼ qφ ( a1 ) =∫ p ( a1|w ) qφ ( w ) dw . Next , we condition on the value of a1 and initialize the parameters of the vari- ational posterior accordingly : qφ1 ( w ) = N ( w|µφ1 , σ2φ1 ) ← qφ ( w|a1 ) . Both qφ ( a1 ) and qφ ( w|a1 ) can be computed analytically for Gaussian distributions . After optimizing φ1 , the approximate posterior qφ1 ( w ) is able to capture a good , local approximation to the posterior p ( w|a1 , x , y ) ( Figure 3c ) . In Figure 3d , we can see the histogram of the refined posterior , that is , the distribution we are generating samples from ( for each sample from w , we drew a sample from qφ ( a1 ) and optimized qφ1 ( w ) ) . Clearly , it is a much better approximation to the true posterior than the Gaussian approximation we started with , although it is important to note that the true posterior is not recovered exactly .

Summary: This paper describes a method for training flexible variational posterior distributions, which consists in making iterative locale refinements to an initial mean-field approximation, using auxiliary variables. The focus is on Gaussian latent variables, and the method is applied to Bayesian neural nets to perform variational inference (VI) over the weights. Empirical results show improvements upon the performance of the mean-field approach and some other baselines, on classification and regression tasks.