summaries: The paper proposes to improve standard variational inference by increasing the flexibility of the variational posterior by introducing a finite set of auxiliary variables. Motivated by the limited expressivity of mean field variational inference the author suggests to iteratively refine a ‘proposal’ posterior by conditioning on a sequence of auxiliary variables with decreasing variance. The key requirement to set the variance of the auxiliary variables such that the integrating over them leaves the original model unchanged. As noted by the authors this is a variant of auxiliary variables introduced by Barber & Agakov. The motivation and theoretical sections seems sound and the experimental results are encouraging, however maybe not completely supporting the claim of new ‘state of the art’ on uncertainty prediction.  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. The paper proposes a new way to improving the variational posterior. The q(w) is an integral over multiple auxiliary variables. These auxiliary variables can be specified at different scales that can be refined to match different scales of details of the true posterior. They show better performance regression and classification benchmark datasets. They also show that the training time is at a reasonable scale when being parallelized.
abstractText: Variational inference (VI) is a popular approach for approximate Bayesian inference that is particularly promising for highly parameterized models such as deep neural networks. A key challenge of variational inference is to approximate the posterior over model parameters with a distribution that is simpler and tractable yet sufficiently expressive. In this work, we propose a method for training highly flexible variational distributions by starting with a coarse approximation and iteratively refining it. Each refinement step makes cheap, local adjustments and only requires optimization of simple variational families. We demonstrate theoretically that our method always improves a bound on the approximation (the Evidence Lower BOund) and observe this empirically across a variety of benchmark tasks. In experiments, our method consistently outperforms recent variational inference methods for deep learning in terms of log-likelihood and the ELBO. We see that the gains are further amplified on larger scale models, significantly outperforming standard VI and deep ensembles on residual networks on CIFAR10.
references.authors: Felix V Agakov David Barber
references.title: An auxiliary variational method
references.venue: In International Conference on Neural Information Processing,
references.authors: Dario Amodei Chris Olah Jacob Steinhardt Paul Christiano John Schulman Dan Mané
references.title: Concrete problems in ai safety
references.venue: arXiv preprint arXiv:1606.06565,
references.authors: Charles Blundell Julien Cornebise Koray Kavukcuoglu Daan Wierstra
references.title: Weight uncertainty in neural network
references.venue: In International Conference on Machine Learning,
references.authors: Chris Cremer Xuechen Li David Duvenaud
references.title: Inference suboptimality in variational autoencoders
references.venue: arXiv preprint arXiv:1801.03558,
references.authors: Alex Graves
references.title: Practical variational inference for neural networks
references.venue: In Advances in neural information processing systems,
references.authors: Chuan Guo Geoff Pleiss Yu Sun Kilian Q. Weinberger
references.title: On calibration of modern neural networks
references.venue: In International Conference on Machine Learning,
references.authors: Fangjian Guo Xiangyu Wang Kai Fan Tamara Broderick David B Dunson
references.title: Boosting variational inference
references.venue: arXiv preprint arXiv:1611.05559,
references.authors: Kaiming He Xiangyu Zhang Shaoqing Ren Jian Sun
references.title: Deep residual learning for image recognition
references.venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,
references.authors: José Miguel Hernández-Lobato Ryan Adams
references.title: Probabilistic backpropagation for scalable learning of Bayesian neural networks
references.venue: In International Conference on Machine Learning,
references.authors: Geoffrey Hinton Drew Van Camp
references.title: Keeping neural networks simple by minimizing the description length of the weights
references.venue: In in Proc. of the 6th Ann. ACM Conf. on Computational Learning Theory. Citeseer,
references.authors: Devon Hjelm Ruslan R Salakhutdinov Kyunghyun Cho Nebojsa Jojic Vince Calhoun Junyoung Chung
references.title: Iterative refinement of the approximate posterior for directed belief networks
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Sergey Ioffe Christian Szegedy
references.title: Batch normalization: Accelerating deep network training by reducing internal covariate shift
references.venue: arXiv preprint arXiv:1502.03167,
references.authors: Yoon Kim Sam Wiseman Andrew C Miller David Sontag Alexander M Rush
references.title: Semi-amortized variational autoencoders
references.venue: arXiv preprint arXiv:1802.02550,
references.authors: Diederik P. Kingma Jimmy Ba
references.title: Adam: A method for stochastic optimization, 2014. cite arxiv:1412.6980Comment: Published as a conference paper at the 3rd International Conference for Learning Representations
references.venue: San Diego,
references.authors: Diederik P Kingma Max Welling
references.title: Auto-encoding variational bayes
references.venue: arXiv preprint arXiv:1312.6114,
references.authors: Durk P Kingma Tim Salimans Max Welling
references.title: Variational dropout and the local reparameterization trick
references.venue: In Advances in Neural Information Processing Systems
references.authors: Balaji Lakshminarayanan Alexander Pritzel Charles Blundell
references.title: Simple and scalable predictive uncertainty estimation using deep ensembles
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Yann LeCun Yoshua Bengio
references.title: Convolutional networks for images, speech, and time series
references.venue: The handbook of brain theory and neural networks,
references.authors: Francesco Locatello Gideon Dresdner Rajiv Khanna Isabel Valera Gunnar Raetsch
references.title: Boosting black box variational inference
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Christos Louizos Max Welling
references.title: Multiplicative normalizing flows for variational Bayesian neural networks
references.venue: In Proceedings of the 34th International Conference on Machine Learning-Volume
references.authors: Joseph Marino Yisong Yue Stephan Mandt
references.title: Iterative amortized inference
references.venue: arXiv preprint arXiv:1807.09356,
references.authors: Andrew C. Miller Nicholas J. Foti Ryan P. Adams
references.title: Variational boosting: Iteratively refining posterior approximations
references.venue: In International Conference on Machine Learning,
references.authors: Cuong V Nguyen Yingzhen Li Thang D Bui Richard E Turner
references.title: Variational continual learning
references.venue: arXiv preprint arXiv:1710.10628,
references.authors: Kazuki Osawa Siddharth Swaroop Anirudh Jain Runa Eschenhagen Richard E Turner Rio Yokota Mohammad Emtiyaz Khan
references.title: Practical deep learning with Bayesian principles
references.venue: Advances in Neural Information Processing Systems,
references.authors: Yaniv Ovadia Emily Fertig Jie Ren Zachary Nado D Sculley Sebastian Nowozin Joshua V Dillon Balaji Lakshminarayanan Jasper Snoek
references.title: Can you trust your model’s uncertainty? evaluating predictive uncertainty under dataset shift
references.venue: Advances in Neural Information Processing Systems,
references.authors: Carsten Peterson
references.title: A mean field theory learning algorithm for neural networks
references.venue: Complex systems,
references.authors: Rajesh Ranganath Dustin Tran David Blei
references.title: Hierarchical variational models
references.venue: In International Conference on Machine Learning,
references.authors: Carlos Riquelme George Tucker Jasper Roland Snoek
references.title: Deep Bayesian bandits showdown
references.venue: In International Conference on Representation Learning,
references.authors: Francisco Ruiz Michalis Titsias
references.title: A contrastive divergence for combining variational inference and mcmc
references.venue: In International Conference on Machine Learning,
references.authors: Tim Salimans Diederik Kingma Max Welling
references.title: Markov chain monte carlo and variational inference: Bridging the gap
references.venue: In International Conference on Machine Learning,
references.authors: Yeming Wen Paul Vicol Jimmy Ba Dustin Tran Roger Grosse
references.title: Flipout: Efficient pseudoindependent weight perturbations on mini-batches
references.venue: arXiv preprint arXiv:1803.04386,
references.authors: Cheng Zhang Judith Butepage Hedvig Kjellstrom Stephan Mandt
references.title: Advances in variational inference. IEEE transactions on pattern analysis and machine intelligence, 2018a
references.authors: Yichuan Zhang José Miguel Hernández-Lobato Zoubin Ghahramani
references.title: Ergodic measure preserving flows
references.venue: arXiv preprint arXiv:1805.10377,
sections.heading: 1 INTRODUCTION
sections.text: 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 cannot 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.
sections.heading: 2 METHODS
sections.text: 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.
sections.heading: 2.1 VARIATIONAL INFERENCE IN BAYESIAN NEURAL NETWORKS
sections.text: 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) .
sections.heading: 2.2 REFINING THE POSTERIOR
sections.text: 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.
sections.heading: 2.3 THEORETICAL JUSTIFICATION
sections.text: Our theoretical claims are twofold. Firstly, that through this procedure, we are optimizing a lower bound to the ELBO and secondly, that the refinement cannot result in a worse posterior approximation than the initial mean-field approximation that we start with (in the ELBO sense). That is,
  ELBOref ≥ ELBOaux ≥ ELBOinit ,
  where ELBOref denotes the ELBO of the refined posterior qref, ELBOaux refers to the objective that the refinement process is optimizing, and ELBOinit is the ELBO of the initial mean-field approximation.
  Lower bound to the ELBO Consider the case with two auxiliary variables a1 and a2. The initial training optimizes the ELBOinit = Eqφ [ log p(y|x,w) ] −DKL ( qφ(w) ∥∥ p(w)) and the refinement step optimizes the conditional ELBO, L|a1(φ1) = Eqφ1 [ log p(y|x,w) ] −DKL ( qφ1(w) ∥∥ p(w|a1)).
  The key observation is that we can define ELBOaux that is a lower bound to ELBOref and is increased both by the initial training and the refinement steps:
  ELBOaux = Eqφ [ Eqφ1 [ log p(y|x,w)− log qφ1(w)
  p(w|a1)
  ] − log qφ(a1)
  p(a1) ] = Eqref [ log p(y|x,w) ] −DKL ( qref(w, a1)
  ∥∥ p(w, a1)) ≤ Eqref [ log p(y|x,w) ] −DKL ( qref(w)
  ∥∥ p(w)) = ELBOref , since the KL divergence of the joint distribution is greater than or equal to that of the marginals.
  Guarantee of improvement Improvement in the ELBO (ELBOref ≥ ELBOinit) is guaranteed under two assumptions. First, that the conditional variational posterior, qφ(w|a1), is within the variational family of qφ1 . Second, that the process that optimizes φ1 does not make it worse than the value it was initialized with. The first assumption holds for Gaussian families: qφ(w|a1) is Gaussian and can be computed in closed form. The second assumption is reasonable to assume for most optimizers and, in addition, it can be ensured by comparing the initial value to the final value and choosing the one with the more desirable objective value.
  The argument goes as follows. By initializing qφ1 such that it coincides with qφ(w|a1), we can ensure that ELBOaux ≥ ELBOinit, since they are equal at initialization time and the optimization process does not decrease ELBOaux. From this combined with our previous result, it follows that ELBOref ≥ ELBOinit and therefore ensuring that the ELBO improves through the refining steps. Note that this also implies that it is not necessary to optimize until convergence: any amount of optimization increases the ELBO.
  Figure 2 shows that the ELBO improvement occurs on real world datasets. With the sampling of each auxiliary variable, the ELBO improves forming a staircase pattern. Table 1 and 2 serve as further empirical evidence that the ELBO improves as a result of the refinement steps.
  Extending to multiple auxiliary variables For simplicity, we stated the arguments with two auxiliary variables, but they straight-forwardly extend to any finite number of auxiliary variables. In this scenario, there are K ELBOaux-s, one for each auxiliary variable, upper bounded by ELBOref and lower bounded by ELBOinit.
sections.heading: 3 RELATED WORKS
sections.text: While in theory, the Bayesian approach can accurately capture uncertainty, in practice we find that exact inference is computationally infeasible in most scenarios and thus we have to resort to approximate inference methods. There is a wealth of research on approximate inference methods, here we focus on works closely related to this paper.
  Variational inference (Hinton & Van Camp, 1993) tries to approximate the true posterior distribution over parameters with a variational posterior from a simple family of distributions. Mean-field VI, which for neural networks traces back to Peterson (1987), uses independent Gaussian distributions over the parameters to try to capture the posterior. The advantage of the mean-field approximation is that the network can be efficiently trained using the reparameterization trick (Kingma & Welling, 2013) and the variational posterior has a proper density over the parameter space which then can be used across tasks like continual learning (Osawa et al., 2019; Nguyen et al., 2017) and contextual bandits (Riquelme et al., 2018). Recently, Louizos & Welling (2017) showed that normalizing flows can be employed to further increase the flexibility of the variational posterior. Zhang et al. (2018a) provide a detailed survey of recent advances in VI.
  Our method is a novel variant of the auxiliary variable approaches to VI (Agakov & Barber, 2004; Ranganath et al., 2016) that increase the flexibility of the variational posterior through the use of auxiliary variables. The key distinction, however, is that instead of trying to train a complex variational approximation over the joint distribution, we iteratively train simple, mean-field approximations at the sampled values of the auxiliary variables. While this poses anO(MK) (K is the number of auxiliary variables andM is the number of posterior samples) overhead over mean-field VI, the training itself is kept straightforward and efficient. The introduction of every new auxiliary variable increases
  the flexibility of the posterior approximation. In contrast to MCMC methods, it is unclear whether the algorithm approaches the true posterior in the limit of infinitely many auxiliary variables.
  There are also numerous methods that start with an initial variational approximation and refine it through a few MCMC steps (Salimans et al., 2015; Zhang et al., 2018b; Ruiz & Titsias, 2019). The distinction from our algorithm is that we refine the posterior starting at large scale and iteratively move towards smaller scale refinements whereas these methods only refine the posterior at the scale of the MCMC steps. Guo et al. (2016); Miller et al. (2017) and Locatello et al. (2018) used boosting to refine the variational posterior, iteratively adding parameters such as mixture components to minimize the residual of the ELBO. Our method does not add parameters at training time but instead iteratively refines the samples through the introduction of auxiliary variables. We did not include these models among our baselines because they have yet to be applied to Bayesian multi-layer neural networks.
  Further related works include methods that iteratively refine the posterior in latent variable models (Hjelm et al., 2016; Cremer et al., 2018; Kim et al., 2018; Marino et al., 2018). These methods, however, focus on reducing the amortization gap and they do not increase the flexibility of the variational approximation.
  Lastly, there are non-Bayesian strategies for estimating epistemic uncertainty in deep learning. Bootstrapping (Breiman, 1996) and deep ensembles (Lakshminarayanan et al., 2017) are perhaps the most promising. Deep ensembles, in particular, have been demonstrated to achieve strong performance on benchmark regression and classification problems and uncertainty benchmarks including out-of-distribution detection (Lakshminarayanan et al., 2017) and prediction under distribution shift (Ovadia et al., 2019). Both methods rely on constructing a set of independently trained models to estimate the uncertainty. Intuitively, the amount of disagreement across models reflects the uncertainty in the ensemble prediction. To induce diversity among the ensemble members, bootstrapping subsamples the training set for each member while deep ensembles use the randomness in weight initialization and mini-batch sampling.
sections.heading: 4 EXPERIMENTS
sections.text: The goal of the experiments is twofold. First, we empirically confirm the improvement in the ELBO as claimed in Section 2.3. Second, we quantify the improvement in the uncertainty estimates due to the refinement. We conducted experiments on a selection of regression and classification benchmarks using Bayesian neural networks as the underlying model. We looked at the marginal loglikelihood of the predictions as well as accuracy in classification tasks.
  Refinement (Refined VI) After training the initial mean-field approximation, we refine M = 10 ensemble members, each with K = 5 auxiliary variables. The means on their prior distributions are fixed at 0., and their variances form a geometric series (each auxiliary variable has variance 0.3 times the previous one, which roughly halves the standard deviation of the prior each time): σ2a1 = 0.7σ 2 w, σ2a2 = 0.21σ 2 w, σ 2 a3 = 0.063σ 2 w, σ 2 a4 = 0.0189σ 2 w, and σ 2 a5 = 0.0081σ 2 w. In each refinement iteration, we optimized the posterior with Adam (Kingma & Ba, 2014) for 200 iterations. To keep the training stable, we scaled the learning rate according to the standard deviation of the conditional posterior: in iteration k, lr=0.3 k 2 0.001.
sections.heading: 4.1 BASELINES
sections.text: We used three baselines. First, mean-field variational inference in order to quantify the improvement provided by refining. Second, multiplicative normalizing flows (MNF) to compare against a more flexible posterior approximation and third, deep ensemble models to compare against a state-of-theart non-Bayesian approach. For all methods we used a batch size of 256 and the Adam optimizer with the default learning rate of 0.0013.
  3The hyperparameters of each baseline were tuned using a Bayesian optimization package, however we found batch size and learning rate to be consistent across methods.
  Variational inference (VI) (Graves, 2011; Blundell et al., 2015) Naturally, we investigate the improvement over variational inference with a mean-field Gaussian posterior approximation. We do inference over all weights and biases with a Gaussian prior centered at 0. The variance of the prior is tuned through empirical Bayes. This model is trained for 30000 iterations.
  Multiplicative Normalizing Flows (MNF), (Louizos & Welling, 2017) To measure the performance against a more flexible class of posterior approximations, we look at Multiplicative Normalizing Flows. In this work, the posterior means are augmented with a multiplier from a flexible distribution parameterized by the masked RealNVP. This model is trained with the default flow parameters for 30000 iterations.
  Deep ensemble models, (Lakshminarayanan et al., 2017) Deep ensemble models are shown to be surprisingly effective at quantifying uncertainty. While they are less principled than Bayesian methods, they are certainly a competitive baseline. For the regression datasets, we used adversarial training ( = 0.01) whereas in classification we did not use adversarial training (adversarial training did not give an improvement on the classification benchmarks). For each dataset, we trained 10 ensemble members for 5000 iterations each.
sections.heading: 4.2 REGRESSION BENCHMARKS
sections.text: Following Hernández-Lobato & Adams (2015), we evaluate the methods on a set of UCI regression benchmarks (Table 1). The datasets used a random 80-20 split for training and testing. The model used for these datasets is a feed forward neural network with a single hidden layer containing 50 units with a ReLU activation function. As it is common practice, we utilize the local reparameterization trick (Kingma et al., 2015).
  On these benchmarks, we are able to confirm that the refinement step consistently improves both the ELBO and the uncertainty estimates over VI. On 7 out of the 9 datasets, Refined VI is one of the best performing approaches.
sections.heading: 4.3 CLASSIFICATION BENCHMARKS
sections.text: We examine the performance on commonly used image classification benchmarks (Table 2). The architecture used for this experiment is the LeNet5 (LeCun et al., 1995) architecture containing three convolutional layers with 6, 16 and 120 channels respectively followed by a feed-forward layer with 84 units and an output layer with ReLu activations throughout the network. We use the local reparameterization trick (Kingma et al., 2015) for the dense layers and Flipout (Wen et al., 2018) for the convolutional layers to reduce the gradient noise.
  On the classification benchmarks, we again are able to confirm that the refinement step consistently improves both the ELBO and the uncertainty estimates over VI. While Refined VI is unable to outperform Deep Ensembles in classification accuracy, it does outperform them in MLL on the largest dataset, CIFAR10.
sections.heading: 4.4 LARGE SCALE MODELS
sections.text: To demonstrate the performance on larger scale models, we apply the refining algorithm to Residual Networks (He et al., 2016) with 20 layers (based on Keras’s ResNet implementation). We look at two models. Firstly, a model where we do inference over all of the residual blocks and secondly, following Ovadia et al. (2019), a hybrid model (ResNet Hybrid) where inference is only done over the final layer of each residual block and every other layer is treated as a regular layer. For this model, we used a batch-size of 256 and we decayed the learning rate starting from 0.001 over 200 epochs. We used 10 auxiliary variables each reducing the prior variance by a factor of 0.5. Furthermore we investigate the effect of Batch Normalization (Ioffe & Szegedy, 2015). While it is difficult to incorporate batch normalization into the Bayesian framework, its positive effects are undeniable.
  Regarding batch normalization, we can confirm the findings of Osawa et al. (2019), that it provides a substantial improvement for VI, although interestingly, this improvement disappears in the case of the hybrid model. To our knowledge, the refined hybrid model is state-of-the-art in terms of MLL. It outperforms Osawa et al. (2019) in both MLL and accuracy.
sections.heading: 4.5 COMPUTATIONAL COSTS
sections.text: When introducing a novel algorithm for variational inference, we have to discuss the computational costs. Table 4 shows the wall-time required to train each model including the possibility of parallel training. Deep ensembles require fewer training iterations and parallelize very well, although it is important to note that the system we tested on is heavily optimized towards training these models both in hardware and software. For Refined VI, training the initial mean-field approximation cannot be parallelized, but the following refinement iterations can be straight-forwardly split into M threads.
sections.heading: 4.6 CONCLUSIONS
sections.text: In this paper, we describe a novel algorithm for refining a coarse variational approximation to the Bayesian posterior. We show, both theoretically and empirically, that the refined posterior is a better approximation to the posterior than the initial variational distribution. Our method outperforms the baseline variational approximations in both uncertainty estimation as well as computational requirements. It sets a new state-of-the-art in uncertainty estimation using variational inference at ResNet scale (ResNet-20) on CIFAR10.
sections.heading: A CLOSED FORMS OF THE SAMPLING DISTRIBUTIONS
sections.text: In Section 2.2, we claim that qφ(a1) (used to sample a1) and qφ(w|a1) (used to initialize qφ1 ) have analytic solutions for additive Gaussians. For p(w) = N (0, σ2w), q(w) = N (µφ, σ2φ), p(a1) = N (0, σ2a1) and p(a2) = N (0, σ 2 a2),
  qφ(a1) = ∫ p(a1|w)qφ(w) dw = N ( µφ σ2a1 σ2w , σ2φσ 4 a1 σ4w + σ2a1σ 2 a2 σ2w ) qφ1(w) = N ( µφ1 , σ 2 φ1 ) ← qφ(w|a1) = N (a1σ2φσ2w + µφσ2a2σ2w σ2φσ 2 a1 + σ 2 wσ 2 a2 , σ2φσ 2 wσ 2 a2 σ2a1σ 2 φ + σ 2 wσ 2 a2 ) . (1)
  These are derived by applying Bayes rule to the Gaussian probability density functions.
  B NOTE ON ELBOAUX ≥ ELBOINIT
  We formally show that ELBOaux ≥ ELBOinit under the assumption that the conditional variational posterior, qφ(w|a1), is within the variational family of qφ1 for all values of a1. For a given a1, let φ10|a1 be a set of variational parameters such that ∀w, qφ10|a1 (w) = qφ(w|a1). Such φ10|a1 must exist as a result of the initial assumption. The formula for Gaussian distributions is given in Appendix A.
  Set φ1|a1 4 such that
  φ1|a1 = φ10|a1 if Eqφ10|a1 [ log p(y|x,w)− log qφ10|a1 (w) p(w|a1) ] ≥ Eqφ1opt|a1 [ log p(y|x,w)− log qφ1opt|a1 (w) p(w|a1) ] φ1opt|a1 otherwise (2) where φ1opt|a1 is the result of the stochastic optimizer that attempts to maximize
  Eqφ1|a1
  [ log p(y|x,w)− log qφ1|a1 (w)
  p(w|a1)
  ] .
  By construction, for all a1,
  Eqφ1|a1
  [ log p(y|x,w)− log qφ1|a1 (w)
  p(w|a1)
  ] ≥
  Eqφ10|a1
  [ log p(y|x,w)− log qφ10|a1 (w)
  p(w|a1)
  ] =
  Eqφ(w|a1) [ log p(y|x,w)− log qφ(w|a1)
  p(w|a1)
  ] .
  (3)
  Substituting this into ELBOaux gives
  ELBOaux = Eqφ [ Eqφ1|a1 [ log p(y|x,w)− log qφ1|a1 (w)
  p(w|a1)
  ] − log qφ(a1)
  p(a1)
  ] ≥
  Eqφ [ Eqφ(w|a1) [ log p(y|x,w)− log qφ(w|a1)
  p(w|a1)
  ] − log qφ(a1)
  p(a1)
  ] =
  Eqφ [ log p(y|x,w)− log qφ(w)
  p(w)
  ] =
  ELBOinit
  (4)
  concluding the proof.
  4We use the notation φ1|a1 instead of φ1 to emphasize the dependence on the value of a1.