Refining the variational posterior through iterative optimization
Authors: Marton Havasi
Jasper Snoek
Dustin Tran
Jonathan Gordon
José Miguel Hernández-Lobato
Venue:  (2020)
Keywords: uncertainty estimation
variational inference
auxiliary variables
Bayesian neural networks

Abstract: 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.

Under review as a conference paper at ICLR 2020
REFINING THE VARIATIONAL POSTERIOR THROUGH ITERATIVE OPTIMIZATION
Anonymous authors Paper under double-blind review
ABSTRACT
Variational inference (VT) is a popular approach for approximate Bayesian infer- ence that particularly promising for highly parameterized models such deep neural networks. A key challenge of variational inference is to approximate the posterior over model parameters with a distribution that 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 iter- atively refining it. Each refinement step makes cheap, local adjustments and only requires optimization simple variational families. We demonstrate theoretically that our method always improves 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 CIFARIO.
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 al.] (016). Thus, is becoming increasingly important that the underlying model provides high quality uncertainty estimates along with its pre- dictions. 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 posterior distribution over the parameters given the training data, Bayesian inference provides a principled approach to capturing uncertainty. contrast, standard training of neural networks employs a point estimate of the parameters, which cannot account for model un- certainty. Unfortunately, exact Bayesian inference is intractable in general for neural networks. To model epistemic uncertainty, variational inference (VT) instead approximates the true posterior with a simpler distribution. The most widely used one for neural networks is the mean-field approxi- mation, where the posterior represented using an independent Gaussian distribution over all the weights. Variational inference is appealing since it reduces the problem of inference to an optimiza- tion 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 novel method for training highly flexible posterior approximations. The idea to start with a coarse, mean-field approximation qww) and make iterative, local refinmeenss 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 aę (Figure [ left) for k = .,.., K that leave the marginal distribution unchanged. In each iteration, the value of an auxiliary variable according to the current variational approximation q(ax) and refine the approximation by conditioning 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 sample from
sample
the newly sampled value q(w) p(w|z y, a1:k) (Figure

---

Under review as a conference paper at ICLR 2020
1a,
Figure I: (Left) The supervised learning model and augmented model respectively where w is ex- pressed sum of independent auxiliary variables. (Right) Illustration of the refining algorithm. In each iteration the value of an auxiliary variable fixed and the posterior is locally adjusted. In the final iteration, a sample is drawn from w. Through the iterations, the variational distribution is able to approximate well the true posterior in a small region.
the refined qw)). The refinement iterations have to be repeated for each posterior sample. The algorithm results in samples from a highly complex distribution, starting from simple mean-field approximation. While the distribution of the samples is difficult to quantify, show that it is not limited to factorized, uni-modal forms, and that the procedure is guaranteed to improve the resulting ELBO without posing significant computational overhead.
Summary of contributions:
A novel algorithm for refining 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 Bayesian neural networks using of UCI regression and image classification benchmarks. We set new state-of-the in uncertainty estimation using variational inference ResNet scale (ResNet-20,, (He al., 2016)) scale on CIFARIO.
2 METHODS
In this section, we first describe standard variational inference (VI), followed by detailed descrip- tion of the iterative refinement algorithm. While VI and proposed methodology generally applicable to latent variable models, in this work, we consider the application to Bayesian neural networks (Figure D. where x are inputs, y outputs, and the weights of network.
2.1 VARIATIONAL INFERENCE BAYESIAN NEURAL NETWORKS
Exact Bayesian inference in Bayesian neural networks (BNN) (Figure I is often intractable and NP-hard in the general case. Variational inference attempts to approximate the true posterior p[w|z, y) with approximate posterior go (w),, typically from simple family of distributions, for example independent Gaussians over weights, i.e. the mean-field approximation. sure that the approximate posterior is close to the true posterior, the parameters of (w),, & are optimized to minimize their Kullback-Leible divergence: DKL(qs (w) || p(w[x, y)). At the limit of DKL (w) p(w|,, zz v)) O, the approximate posterior exactly captures the true posterior, although this might not be achievable (pw/[e (z, y) outside of the distribution family of qo (w).
en-
In order to minimize the KL-divergence, variational inference optimizes the Evidence Lower Bound (ELBO) w.r.t. & (denoted as L())), which a lower bound to the log marginal likelihood log p(y (r)) Since the the marginal likelihood can be expressed the sum of the KL-divergence and the ELBO. maximizing the ELBO is equivalent to minimizing the KL divergence:
log.(ylæz) DKL. (go(w) p( +L(0) 2 ((Φ) Egg [log p(yle, w)] DRI. (us(w) || p(w)),
due non-negativity of the Kl-ddivegeecce.
2

---

Under review as a conference paper at ICLR 2020
-50000
1: procedure REFINEANDSAAPLECC)
2: for 171 I,..., , M do
-52000
3:
@0 +ー Φ for k = 1,...,, K do
4:
-54000
5:
aE ~ 9 (ax)
Sample dg
6:
ad, (w) s 9du-1 (wlar)
Initialize oR
560000
7:
Øk + arg maxos Liama (Øu)
> Optimize on
-8000
8: end for
9: Wm 9ox (w)
P Sample the refined posterior
10000
Starting ELBO
10: end for
11: return WI.M
500
1000
1500
2000
12: end procedure
Figure 2: The ELBO staircase: ELBO aux is increasing as we sam- ple the auxiliary variables (sin- gle sample Monte Carlo estimate, LeNet-5/CIFARIO]
Algorithm 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 new imput.z' x², the predictive distribution p(y'lz' la' y, x) can be approximated by stochastically drawing small number (around M < 10) of sample model parameters and averaging their prediction in an ensemble model:
M
WIM ~ (w)), ply/o'
1 M plu/lz',, w))
Σ
2.2 REFINING THE POSTERIOR
The main issue with variational inference 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 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 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 finite number of auxiliary variables 0ı. as shown in Figure The constraints are that (r. 3) must be conditionally independent of the auxiliary variables given and that they must not affect the prior distribution p(w). This important in justifying the use of the initial variational approximation. A significant way in which we distinguish ourselves from hierarchical variational models (Ranganath al. 2016) is that the model unaffected by the presence of the auxiliary variables. Their purpose is solely to aid the inference procedure. Given Gaussian prior A (0, o
over w, we express w as sum of independent auxiliary variables
K
with plax) = (O, ozr) for h: -1...,,K
Ok:
k=1
while ensuring that EK 22 o so that the prior p(w) NN(0, is unchanged.
Locally refining the approximate posterior refers to iteratively sampling the values of the auxiliary variables 0 and then approximating the posterior conditional the sampled values, i.e. qax (w) approximates p(w|r, (1 for iterations & = 1,..., K (Algorithm I. Starting from the initial mean field approximation (w), we sample the value of a₁ from gs (.))
plagl//)ga (w) dw,
1 'Tee sudden drops after sampling are optimizer artefacts due to having reset the parameters of Adam. While we are focusing on one specific definition of the auxiliary variables. additive auxiliary variables. note that all of our results trraightfforwadd generalize to arbitrary joint distributions p(w, (1xx) that meet the constraints.
3

---

Under review as a conference paper at ICLR 2020
The tru passerorr
The aspronimets pastartor
Vartztoonl posterir conditona
Midogrum of the peight
vi6) ViNG of
X6-4e1
He wol
0.4
0.2
(a)
(b)
(c)
(d)
Figure 3: A simple multi--oddl example demonstrating how our method can capture a more com- plex distribution by refining simple mean-field posterior. In (a) the true posterior is too com- plex to be well approximated by a Gaussian. (b) The Gaussian approximate posterior after opti- mizing the ELBO (ELBO -1.79). (c) After sampling 01, we optimize the conditional ELBO w.r.t. Φ1 [1. Wm is drawn from gø, an (w). (d) Samples from the refined posterior approximation. ELBO 1.45.
then optimize the approximation to the conditional posterior: qor (w) Re p(ularyy, This proce- dure is then iteratively repeated for 02.K. In iteration k,
1) ag ~ f plaxlerkki,
w)). pl (w) dw 2) Φx arg DKL (yon (w) y
Analogously to variational inference, the KL divergence is minimized through the opti- mization of the conditional ELBO in each iteration: Lia.. = E [logpyyla ( w)] (go (w) || p(w]a1:))) In order to get independent samples from the variational posterior, have to repeat the iterative refinement for each ensemble member WI.M-
Toy example We toy example to demonstrate the procedure. In this toy example, we have a single weight w with prior p(w) = (o 1) and a complicated posterior with four modes (syn- thetically generated data). We express as the sum of two auxiliary variables w = ai a2 with p(a1)= (O, 0.8) and p(az) (0, 0.2) (which recovers p(w) as per the constraint).
As Figure shows, a Gaussian approximation to the posterior fails to capture the multimodal na- ture of the true posterior. The first step of the refinement is to sample using qa= a1 ~ (o(a))= = (paal//)pr (w)) dw. Next, we condition on the value of an and initialize the parameters of the vari- ational posterior accordingly: Go, (w) (w)ps, 22,) < qo (w)(a)). Both (a)) and g ( can be computed analytically for Gaussian distributions. After optimizing Φ1, the approximate pos- terior 4s, (w) is able to capture a good, local approximation to the posterior p(w)a [1l 2, y) (Figure 3c). In Figure [d] we can see the histogram of the refined posterior, that is, the distribution we are generating samples from (for each sample from we drew a sample from (a1) and optimized qo. (w)). Clearly, it is a much better approximation to the true posterior than the Gaussian approx- imation we started with, although it is important to note that the true posterior is not recovered exactly.
2.3 THEORETICAL JUSTIFICATION
Our theoretical claims are twofold. Firstly, that through this procedure, we are optimizing lower bound to the ELBO and secondly, that the refinement cannot result a worse posterior approxima- tion than the initial mean-field approximation that we start with (in the ELBO sense). That is,
EL 2 ELBO ELBO mit
where EL denotes the ELBO of the refined posterior gref. ELBO refers the objective that the refinement process is optimizing, and ELBO... is the ELBO of the initial mean-fiel field approxi- mation.
Lower bound to the ELBO Consider the case with two auxiliary variables a and 0g. The initial training optimizes the ELBO in Eg, log p(y|2x w)] DRL (ww p(w)) and the refinement step optimizes the conditional BLBO, L/a, ( Bgg, [log Pogg(la,,, (,, w)] DRL. (gs, (w) (())))))-
4

---

Under review as a conference paper at ICLR 2020
The key observation is that we can define ELBOaux that is a lower bound to ELBO and is increased both by the initial training and the refinement steps:
p(wla1) log {s(a1)
Pon (w)
ELBO aux = Egg [Eg. log z, w) - log
p(a1)
= Equ log |r, w)] DKL (Ber( (. a1) p(w, a1))
S gut
log (x, w)] DKL (Ter((w) ) || p(w)) = ELBO,
since the KL divergence of the joint distribution is greater than equal that of the marginals.
Guarantee of improvement Improvement in the ELBO (ELBO,er 2 ELBO is guaranteed under two assumptions. First, that the conditional variational posterior, go (w(w()), is within the variational family of 4á.- Second, that the process that optimizes Φ1 does not make worse than the value it was initialized with. The first assumption holds for Gaussian families: 4 (o(w)a)) is Gaussian and can computed closed form. The second assumption 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 qe, such that it coincides with go (w)a),, we can ensure that ELBOmM 2 ELBOmnt. 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 ELBO 2 ELBO and therefore ensuring that the ELBO improves through the refining steps. Note that this also implies that it is not necessary optimize until convergence: any amount of optimization increases the ELBO.
Figure 2 shows that the ELBO improvement occurs real world datasets. With the sampling of each auxiliary variable, the ELBO improves forming a staircase pattern. Table | and 2 serve as further empirical evidence that the ELBO improves result of the refinement steps.
Extending to multiple auxiliary variables For simplicity, we stated the arguments with two aux- iliary variables, but they straight-forwardly extend to any finite number of auxiliary variables. In this scenario, there are K ELBOaax -S, one for each auxiliary variable, upper bounded by ELBO,e and lower bounded by ELBOm
RELATED WORKS
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 approx- imate inference methods. There a wealth of research on approximate inference methods, here we focus works closely related to this paper.
Variational inference (Hinton & Van Camp 1993) tries to approximate the true posterior distribution over parameters with variational posterior from 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 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, [oouizos & Welling (2017) showed that normalizing flows can be employed to further increase the flexibility of the variational posterior. Zhang et al. (201 8a) provide a detailed survey of recent advances VI.
Our method is 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 aux- iliary variables. The key distinction, however, is that instead of trying to train 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 O(MK) (K is the number of aux- iliary variables and M 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
5

---

Under review as a conference paper at ICLR 2020
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 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 large scale and iteratively move towards smaller scale refinements whereas these methods only refine the posterior at the scale of the MCMC steps. [uo et al. (2016); Miller et al. (2017) Locatelo al. (2018) used boosting to refine the variational posterior, iteratively adding parameters such as mixture components to min- imize 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 mod- els (Hjelm et al. 2016); Cremer et al. 20118] Kim et al.) 20188 Marino et al. 2018). These methods, however, focus on reducing the amortization gap and they do not increase the flexibility the vari- ational approximation.
Lastly, there are non- non-Bayesian strategies for estimating epistemic uncertainty in deep learning. Boot- strapping (Breiman 1996) and deep ensembles (Lakshminarayanan et al.) 2017) are perhaps the most promising. Deep ensembles, in particular, have been demonstrated to achieve strong perfor- mance benchmark regression and classification problems and uncertainty benchmarks includ- ing -ut-of-distribution detection (Lakshminarayanan et al. 2017) and prediction under distribution shift (Ovadia et al.) 2019). Both methods rely on constructing a of independently trained models to estimate the uncertainty. Intuitively, the amount of disagreement across models reflects the uncer- tainty 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.
4 EXPERIMENTS
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 bench- marks using Bayesian neural networks as the underlying model. We looked at the marginal log- likelihood of the predictions as well as accuracy in classification tasks.
Refinement (Refined VI) After training the initial mean-field approximation, refine M 10 ensemble members, each with K = 5 auxiliary variables. The means on their prior distributions are fixed 0., and their variances form geometric series (each auxiliary variable has variance 0.3 times the previous one, which roughly halves the standard deviation of the prior each time): o², = 0.70 o2 = 0.216² w. o² = 0.063o2 o² = 0.018962 and a² = 0.0081 wo 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 A: Ir=0.3 0.001.
4.1 BASELINES
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 state-of-the- art non- Bayesian approach. For all methods we used batch size of 256 and the Adam optimizer with the default learning rate of 0.00
The hyperparameters of each haseline were tuned using Bayesian optimization package, however found batch size and learning be consistent across methods.

---

Under review as a conference paper at ICLR 2020
Naturally, investigate the improvement over variational inference with a mean-field Gaussian posterior approximation. We do inference over all weights and biases with Gaussian prior centered O. The variance of the prior is tuned through empirical Bayes. This model is trained for 30000 iterations.
Variational inference (VI) (Graves
2011;
Blundell et al. 2015)
Multiplicative Normalizing Flows (MNF), (Louizos & Welling 2017) To measure the perfor- mance against a more flexible class of posterior approximations, we look at Multiplicatiiv Nor- malizing 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 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 (c = 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.
4.2 REGRESSION BENCHMARKS
Following Hernández-Loball Adams 2015), we evaluate the methods on a set of UCI regression benchmarks (Table The datasets used random 80-20 split for training and testing. The model used for these datasets a feed forward neural network with a single hidden layer containing 50 units with ReLU activation function. As it is common practice, utilize the local reparameterization trick (Kingma et al. 2015).
On these benchmarks, are able to confirm that the refinement step consistently improves both the ELBO and the uncertainty estimates over VI. On out of the 9 datasets, Refined VI is one of the best performing approaches.
Doep Ensemble MLL
Relined VI
work)
MLL
MLL
ELBO
MLL
ELBO
boston housing
9.13611..719
1.920:1.113
2.877400.151
-668.272.7.6611 6.67
-2.51100.185
-6630.37917.716
concrete strength
4.062+0.130
-3.002+0.005
3.138=0.063
.3248. 1344-137-66..5
10771.144:64..00
naval. pwwllppopulaoee
3 3.905-00.013
3.473=200077
5.969=02245
13440.701::300
340
6.128:00171
54882 656ch 12:566::1222... 361
energy .efiiiencc
-0.66610.058
-0.75620.054
-0.74920.068
12296.211166,310
-0.70720.094
-1192.337.16..06
pach.hydrodynmmii.
0.984110.104
-1.339-10.170
-1.74910.232
.928.758
-1.626::0.231
2--99.0552.84.7
kin8am
1.135+00.12
1.125+0.112
1.066+0.019
6071 20711268++11.10 (788
1.099-0.D18
2 61727 (1172.009+-67..0 65669
power-plant
3.935220.140
-2.835=0.0333
-2.8262000200
1519665791100.44 130.4877
2.820=00.024
2--22368.
2--2268.965-1888 965:1555:55888
protein strucctrr
3.687.10.013
.2.028-00.007
-2.92610.010
-108806.00
174.52
2-108597.99211155 19989977992-15844422
wine
4.968+00.079
-..663+0.027
-0.973-0.054
-1346. 13446110++18.. 18.018
-0.968-0.056
2 2-1311 11311.898+-13 8984+17.487
Table 1: Results on the UCI regression benchmarks with single hidden layer containing 50 units.
Metrics: marginal log-likelihooa (MLL, higher is better), and the evidence lower bound (ELBO higher better). The mean values and standard deviations are shown in the table.
4.3 CLASSIFICATION BENCHMARKS
We examine the performance on commonly used image classification benchmarks (Table 2). The architecture used for this experiment is the LeNet5 (LeCCnn et al. 1995) architecture containing three convolutional layers with 6, 16 and 120 channels respectively followed by 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, does outperform them in MLL on the largest dataset, CIFAR10.
7

---

Under review as a conference paper at ICLR 2020
Deep Ensemble MLL & Acc
MNF
Refined VI (This wark)
MLL Acc
MLL Acc
ELBO
MLL & Acc
ELBO
-0.017-0.000 -9.4%=0.0
0.034++0.022
0.032-+0.001
76188533-47.555
-00225+0,001
2 -63310.824-42.357
mnist
39.18:0.11
99.19++0.1
99.2%%=+0.0
-0.201=0.002
-0.255=0.004
-0.255=:00003
-2269.330-232.654
-0,241=0.004
-20438.955279.672
fashion mnist
93. 03.1%_0.1
90.7%
90.7%:10..1
91.22%
cifar10
0.79110.000 0.009
0.7955
0.8115400.04
-57257.887- 2999.770
-0.768 0.0007
50989.217--28.977
76.
72.6%
72.3%-+0.5
73.5%-++0.5
Table 2: Results on image classification benchmarks with the LeNet-5 architecture, without data augmentation. Metrics: marginal ligllikeihhod (MLL, higher is better), accuracy (Acc, higher better), and the evidence lower bound (ELBO higher is better). The mean values and standard deviations are shown in the table.
4.4 LARGE SCALE MODELS
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 's ResNet implementation). We look at two models. Firstly, model where 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 a regular layer. For this model, used 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 investigate the effect of Batch Normalization (loff & Szegedy) 2015). While it is difficult to incorporate batch normalization into the Bayesian framework, its positive effects are undeniable.
Regarding batch normalization, 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 MIL.. outperforms Osawa et af. (2019) in both MLL and accuracy.
confirm the findings Osawa et al. (2019), that it provides
Deep Ensemble Acc
VI
Refined VI (This work)
MLL
MLL
Acc
MLL
Acc
ResNet
-0.698
82.7%
-0.795
72.6%
-0.696
75.5%
ResNet + BatchNorm
-0.561
83.6%
-0.672
77.6%
-0.593
79.7%
ResNet Hybrid
-0.698
82.7%
-0.465
84.2%
-0.432
85.8%
ResNet Hybrid + BatchNorm
-0.561
83.6%
-0.465
84.0%
-0.423
85.6%
Table 3: Results on CIFARIO with the ResNet architecture, without data augmentation. Mctries: marginal ligeliill (MLL, higher is better), accuracy (Acc, higher is better), and the evidence lower bound (ELBO higher is better). Note that the non-hybrid and the hybrid models are equivalent when trained deterministically.
4.5
COMPUTATIONAL COSTS
When introducing novel algorithm for variational inference, have to discuss the computational costs. Table 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 can- not be parallelized, but the following refinement iterations can be foraight-forwardly split into M threads.
No parallelism
Maximum parallelism
Deep Ensemble
433.7
43.4 s
MNF
990.9
990.9
VI Refined VI
53
531.5
708.6
566.9 s
Table 4: The training time of each method (LeNet-5/CIFARIO) on P100 using Tensorflow.
8

---

Under review as a conference paper at ICLR 2020
4.6 CONCLUSIONS
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 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 require- ments. It sets new state-of-the-art in uncertainty estimation using variational inference at ResNet scale (ResNet-20) on CIFAR1O.
REFERENCES
Felix V Agakov and David Barber. An auxiliary variational method. In International Conference on Neural Information Processing, pp. 561-566. Springer, 2004.
Dario Amodei, Chris Olah, Jacob Steinhardt, Paul Christiano, John Schulman, and Dan Mané. Con- crete problems in ai safety. arXiv preprint 1666v:66060065566 2016.
Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In International Conference Machine Learning, pp. 1613-1622, 2015.
Leo Breiman. Bagging predictors. Machine learning, 24(2):123-140, 1996.
Chris Cremer, Xuechen Li, and David Duvenaud. Inference suboptimality in variational autoen- coders. arXiv preprint /80i 033558, 2018.
Alex Graves. Practical variational inference for neural networks. In Advances in neural information processing systems, pp. 2348-2356 2011.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q. Weinberger. On calibration of modern neural networks. In International Conference Machine Learning, 2017.
Fangjian Guo, Xiangyu Wang, Kai Fan, Tamara Broderick, and David B Dunson. Boosting varia- tional inference. arXiv preprint arXiv:161 1.Iv::611..05555 2016
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp- 770-778, 2016.
José Miguel Herbandez-Lobato and Ryan Adams. Probabilistic backpropagation for scalable learn- ing of Bayesian neural networks. International Conference on Machine Learning, pp. 1861- 1869, 2015.
Geoffrey Hinton and Drew Van Camp. Keeping neural networks simple by minimizing the descrip- tion length of the weights. In in Proc. of the 6th Ann. ACM Conf. Computational Learning Theory. Citesecr, 1993.
Devon Hjelm, Ruslan R Salakhutdinov, Kyunghyun Cho, Nebojsa Jojic, Vince Calhoun, and Juny- oung Chung. Iterative refinement of the approximate posterior for directed belief networks. In Advances in Neural Information Processing Systems, pp. 4691-4699, 2016.
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. arXiv preprint arXiv:/502.000 1502.03 2015.
Yoon Kim, Sam Wiseman, Andrew C Miller, David Sontag, and Alexander M Rush. Semi-amortized variational autoencoders. arXiv preprint
2018.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2014. cite 2ediv11422.9800oommet Published as a conference paper at the 3rd International Conference for Learning Representations, San Diego, 2015.
Diederik Kingma and Max Welling. Auto-encoding variational bayes. 1312:11312.61 2013.
arXiv preprint

---

Under review as a conference paper at ICLR 2020
Durk P Kingma, Tim Salimans, and Max Welling. Variational dropout and the local reparameteri- zation trick. In Advances in Neural Information Processing Systems. 2015.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In Advances in Neural Information Processing Systems, pp. 6402-6413, 2017.
Yann LeCun, Yoshua Bengio, et al. Convolutional networks for images, speech, and time series. The handbook of brain theory and neural networks, 3361(10):1995, 1995.
Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. nature, 521(7553):436, 2015.
Francesco Locatello, Gideon Dresdner, Rajiv Khanna, Isabel Valera, and Gunnar Raetsch. Boosting black box variational inference. In Advances in Neural Information Processing Systems, 2018.
Christos Louizos and Max Welling, Multiplicative normalizing flows for variational Bayesian neural networks. In Proceedings of the 34th International Conference on Machine Cetrning-Voluumee 70, pp. 2218-2227. JMLR. org, 2017.
Joseph Marino, Yisong Yue, and Stephan Mandt. Iterative amortized inference. arXiv preprint arXiv:1807.09356. 2018.
Andrew C. Miller, Nicholas J. Foti, and Ryan P. Adams. Variational boosting: Iteratively refining posterior approximations. In International Conference on Machine Learning, 2017.
Cuong V Nguyen, Yingzhen Li. Thang D Bui, and Richard E Turner. Variational continual learning. arXiv preprint Drxii:/7110 1710.1068, 2017.
Kazuki Osawa, Siddharth Swaroop, Anirudh Jain, Runa Eschenhagen, Richard Turner, Rio Yokota, and Mohammad Emtiyaz Khan. Practical deep learning with Bayesian principles. Ad- vances Neural Information Processing Systems, 2019.
Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, D Sculley, Sebastian Nowozin, Joshua V Dil- lon, Balaji Lakshminarayanan, and Jasper Snoek. Can you trust your model's uncertainty? eval- uating predictive uncertainty under dataset shift. Advances in Neural Information Processing Systems, 2019.
Carsten Peterson. A mean field theory learning algorithm for neural networks. Complex systems, 1: 995-1019, 1987.
Rajesh Ranganath, Dustin Tran, and David Blei. Hierarchical variational models. In International Conference on Machine Learning, pp. 324-333, 2016.
Carlos Riquelme, George Tucker, and Jasper Roland Snoek. Deep Bayesian bandits showdown. In International Conference on Representation Learning, 2018.
Francisco Ruiz and Michalis Titsias. contrastive divergence for combining variational inference and meme. In International Conference on Machine Learning, pp. 5537-5545, 2019.
Tim Salimans, Diederik Kingma, and Max Welling. Markov chain monte carlo and variational inference: Bridging the gap. In International Conference on Machine Learning, pp. 1218-1226, 2015.
Yeming Wen, Paul Vicol, Jimmy Ba, Dustin Tran, and Roger Grosse. Flipout: Efficient pseudo- independent weight perturbations on
arXiv preprint
2018.
Cheng Zhang, Judith Butepage, Hedvig Kjellstrom, and Stephan Mandt. Advances in variational inference. IEEE transactions on pattern analysis and machine intelligence, 2018a.
Yichuan Zhang, José Miguel Hernández-Lobato, and Zoubin Ghahramani. Ergodic measure pre- serving flows. arXiv preprint arXiv: 2018b.
10

---

Under review as a conference paper at ICLR 2020
A CLOSED FORMS OF THE SAMPLING DISTRIBUTIONS
In Section 2.2 we claim that (0())) (used to sample a1) and go (ww(a)) (used to initialize do, have analytic solutions for additive Gaussians. For p(w) = NO0 ))), q(w) = NiHos oa), p(a1) = AO o and p(az) = NN
{o(ai) = / p(a1lw)as(w) dw = N ub
o2 o2
61
d
qo. (w) = (M se.) (o(wla1) . 0 u
u
0
0
0
a2
m2
+ to 0
0 o
These are derived by applying Bayes rule to the Gaussian probability density functions.
B NOTE ON ELBOAUS AUX 2 ELBO INIT
We formally show that ELBO 2 ELBO under the assumption that the conditional variational posterior, (w(aa)). is within the variational family of qa, for all values of a
For given a1, let olala be a set variational parameters such that Vw, g0ip a1 (w) = qo (w a1). Such 0igla, must exist as a result of the initial assumption. The formula for Gaussian distributions given in Appendix[A) A
Set $1j0
such that
Doipinn (16)
oppling (w)
if if Eeggging
log p(yla, w) log p(ulas)
E Eusrope
logp(ylæ, w) - log
p(wla1)
@lopt/as otherwise
(2)
where @loptlas is the result of the stochastic optimizer that attempts
maximize
Borion (ar)
Egoiies plyl(,,, w) - log p(w]es)
By construction, for all a1.
Donias (w)
Egerien
logp(ylz, w) - -log p(w(a1)
2
0dnpla (w)
Bggiales
logp(ula, w) log p(wla1)
(3)
1s (w [rlaa)] log logp(ul2, w) - -log p(w)a1)
E [o
Substituting this into ELBO... mux gives
ELBO aux
Odriag (w)
log ga(a1)
Egg Egosjen
log plylp(l pr, w) - llog
2
p(wlas (1))
p(a1)
go (wlaa))
go (a1)
Egg Egs(m)la))
log bogp(ulo,
w) log p(wla))
log p(a1)
(4)
Egg logp(y/ar, log go (w()
(w)
p(w)
ELBO jnic
concluding the proof.
"We the notation d instead of Φ1 emphasize the dependence the value of aj.
11


Reviews:
[{"review_id": "ryeMUJ66tS", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary.\n\nThis 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.\n\nMain comments.\n \nOverall, this paper is well written and easy to follow. It tackles an important topic in VI and proposes an interesting idea to improve the flexibility of the approximate distribution. I have the following comments/questions.\n\n- On the guarantee of improvement. I still have some doubts regarding the inequality “ELBO_aux >=  ELBO_init”. Can you please elaborate more on this and provide a detailed formal proof? Figure 2 shows that ELBO_aux can go below ELBO_init.\n- The focus of the paper is on Gaussian variables and a configuration where some key distributions, q(a_1) and q(w|a_1), are accessible in closed from. The generalization of the proposed method beyond these settings should be discussed and explored in experiments. \n- Important baselines are missing in the experiments. I would recommend including at least the other VI techniques relying on auxiliary variables to build flexible variational families [1,2]. This would help to better assess the impact/importance of the proposed method.\n\n[1] Ranganath, Rajesh, Dustin Tran, and David Blei. \"Hierarchical variational models.\" ICML. (2016).    \n[2] Maaløe, Lars, et al. \"Auxiliary deep generative models.\" ICML (2016).\n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "r1gLsnqAKS", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nI think the idea is quite interesting. They did a good illustration of the difference between their model and related works. They also compare with the state-of-art variational approximation methods. \n\nOne concern I have is the complexity. I don't think it's just O(MK) since it has to optimize for each phi_k for each posterior sample w. This could be quite large depending on problems. \n\nAlso is the refining only run for once or run after each update of the mean field q(w)? If it's the latter, the overhead would be much larger.\n\nWhen w is very high-dimensional, the number of auxiliary variables should be exponentially larger. Is that true? Or it's actually invariant to the dimensionality of the posterior distribution?\n\nThe paper proves that the refined ELBO is larger than the auxiliary ELBO which is larger than the initial mean field. But the initial ELBO should be a tight lower bound of the true log likelihood. Would that be a problem that ELBO_ref actually spill over the true log likelihood which is a dangerous sign?\n\nOverall I think the paper is well written. The experiments are carefully designed. The idea is interesting and useful.", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "BylUth6tqS", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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. \n\nOverall i find the motivation and theoretical contribution interesting. However I do not find the experimental section completely comprehensive why I currently think the paper is borderline acceptance. \n\nComments\n1) The computational demand using the method seems quite large by adding O(NumSamples * NumAuxiliary) additional computational cost on top of the standard VI. Here each factor M is quite large e.g. 200 epochs for CIFAR10 (if i understand the algorithm correctly?)\n2) For the UCI experiments the comparison is only made against DeepEnsembels or other VI methods, however to the best of my knowledge MCMC methods are superior in this setting given the small dataset size? \n3) The results on CIFAR10 do seem to demonstrate that the proposed method is superior to DeepEnsembles and standard VI in one particular setting where VI is only performed over a small subset of layers in a ResNet (why doesn’t it work for when doing VI on all the parameters?). However generally looking at the best obtained results of ~86% acc this is quite far from current best probabilistic models (see e.g. Heek2019 that gets 94% acc). Some of this can probably be attributed to differences in data-augmentation and model architecture however in general it makes it very hard to compare with other methods when the baselines are not competitive. \n\nMinor Comments:\n In relation to comment 3) above I think you should reword the sentence “It sets a new state-of-the-art in uncertainty estimation at ResNet scale on CIFAR10” in the conclusion.\n\n\n“In  order  to  get  independent  samples  from  the  variational  posterior,we have to repeat the iterative refinement for each ensemble member”: Does this imply that if we want M samples we first have to optimize using the standard VI and then to M optimizations to get q_k(w)?\n\n\nHow sensitive is the method to sequence of variances for a?\n\n[Heek2019]: Bayesian Inference for Large Scale Image Classification\n", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "BkgJiljxjB", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper presented an iteratively refined variational inference for Gaussian latent variables. The intuition is straightforward and makes sense to me. However, I have some concerns.\n\nDetailed comments:\n1. In theoretical justification, only K=2 is discussed. My intuition is that as K increases, the approximation of the true posterior should be closer. The summation of multiple Gaussian distributions can arbitrarily approximate any distribution given enough base distributions. I would like to see some theoretical discussion about K. At least in the experiment, the author should provide the performance of different Ks.\n2. The toy example in the paper is simply 1D Gaussian. I want to see more discussion for high dimensional latent variables. So in the experiments, how you parameterized the distribution for each weight? Totally independent? or allowing structural correlations? I am not sure the details of the implementation in this paper, but I also have a naive question for high dimensional Gaussian. Does it require to compute the matrix inverse when sampling a_k?\n3. Another related paper \"Guo, Fangjian, et al. \"Boosting variational inference.\" arXiv preprint arXiv:1611.05559 (2016).\" should be discussed as well.", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: In this paper a method for refining the variational approximation is proposed.

The reviewers liked the contribution but a number reservations such as missing reference made the paper drop below the acceptance threshold. The authors are encouraged to modify paper and send to next conference.

Reject. 
Decision: Reject