vadam
Code for ICML 2018 paper on "Fast and Scalable Bayesian Deep Learning by WeightPerturbation in Adam" by Khan, Nielsen, Tangkaratt, Lin, Gal, and Srivastava
view repo
Uncertainty computation in deep learning is essential to design robust and reliable systems. Variational inference (VI) is a promising approach for such computation, but requires more effort to implement and execute compared to maximumlikelihood methods. In this paper, we propose new naturalgradient algorithms to reduce such efforts for Gaussian meanfield VI. Our algorithms can be implemented within the Adam optimizer by perturbing the network weights during gradient evaluations, and uncertainty estimates can be cheaply obtained by using the vector that adapts the learning rate. This requires lower memory, computation, and implementation effort than existing VI methods, while obtaining uncertainty estimates of comparable quality. Our empirical results confirm this and further suggest that the weightperturbation in our algorithm could be useful for exploration in reinforcement learning and stochastic optimization.
READ FULL TEXT VIEW PDFCode for ICML 2018 paper on "Fast and Scalable Bayesian Deep Learning by WeightPerturbation in Adam" by Khan, Nielsen, Tangkaratt, Lin, Gal, and Srivastava
Deep learning methods have had enormous recent success in fields where prediction accuracy is important, e.g., computer vision and speech recognition. However, for these methods to be useful in fields such as robotics and medical diagnostics, we need to know the uncertainty of our predictions. For example, physicians might need such uncertainty estimates to choose a safe but effective treatment for their patients. Lack of such estimates might result in unreliable decisions which can sometime have disastrous consequences.
One of the goals of Bayesian inference is to provide uncertainty estimates by using the
posterior distributionobtained using Bayes’ rule. Unfortunately, this is infeasible in large models such as Bayesian neural networks. Traditional methods such as Markov Chain Monte Carlo (MCMC) methods converge slowly and might require a large memory
(Balan et al., 2015). In contrast, variational inference (VI) methods can scale to large models by using stochasticgradient (SG) methods, as recent work has shown (Graves, 2011; Blundell et al., 2015; Ranganath et al., 2014; Salimans et al., 2013). These works employ adaptive learningrate methods, such as RMSprop
(Tieleman & Hinton, 2012), Adam (Kingma & Ba, 2015) and AdaGrad (Duchi et al., 2011), for which easytouse implementations are available in existing codebases.Despite their simplicity, these VI methods require more computation, memory, and implementation effort compared to maximumlikelihood estimation (MLE). One reason for this is that the number of parameters in VI is usually much larger than in MLE, which increases the memory and computation costs. Another reason is that existing codebases are designed and optimized for tasks such as MLE, and their application to VI involves significant amount of modifications in the code. We ask the following question: is it possible to avoid these issues and make VI as easy as MLE?
In this paper, we propose to use naturalgradient methods to address these issues for Gaussian meanfield VI. By proposing a naturalmomentum method along with a series of approximations, we obtain algorithms that can be implemented with minimal changes to the existing codebases of adaptive learningrate methods. The main change involves perturbing the network weights during the gradient computation (see Fig. 1). An uncertainty estimate can be cheaply obtained by using the vector that adapts the learning rate. This requires lower memory, computation, and implementation efforts than existing methods for VI while obtaining uncertainty estimates of comparable quality. Our experimental results confirm this, and suggest that the estimated uncertainty could improve exploration in problems such as reinforcement learning and stochastic optimization.
Bayesian inference in models such as neural networks has a long history in machine learning
(MacKay, 2003; Bishop, 2006). Earlier work proposed a variety of algorithms such as MCMC methods (Neal, 1995), Laplace’s method (Denker & Lecun, 1991), and variational inference (Hinton & Van Camp, 1993; Barber & Bishop, 1998). The meanfield approximation has also been a popular tool from very early on (Saul et al., 1996; Anderson & Peterson, 1987). These previous works lay the foundation of methods now used for Bayesian deep learning (Gal, 2016).Recent approaches (Graves, 2011; Blundell et al., 2015) enable the application of Gaussian meanfield VI methods to large deeplearning problems. They do so by using gradientbased methods. In contrast, we propose to use naturalgradient methods which, as we show, lead to algorithms that are simpler to implement and require lower memory and computations than gradientbased methods. Natural gradients are also better suited for VI because they can improve convergence rates by exploiting the information geometry of posterior approximations (Khan et al., 2016). Some of our algorithms inherit these properties too.
A recent independent work on noisyAdam by Zhang et al. (2018) is algorithmically very similar to our Vadam method, however their derivation lacks a strong motivation for the use of momentum. In our derivation, we incorporate a naturalmomentum term based on Polyak’s heavyball method, which provides a theoretical justification for the use of momentum. In addition, we analyze the approximation error introduced in Vadam and discuss ways to reduce it.
Zhang et al. (2018) also propose an interesting extension by using KFAC, which could find better approximations than the meanfield method. The goal of this approach is similar to other approaches that employ structured approximations (Ritter et al., 2018; Louizos & Welling, 2016; Sun et al., 2017). Many other works have explored variety of approximation methods, e.g., Gal & Ghahramani (2016) use dropout for VI, HernandezLobato & Adams (2015); Hasenclever et al. (2017) use expectation propagation, Li et al. (2016); Balan et al. (2015) use stochasticgradient Langevin dynamics. Such approaches are viable alternatives to the meanfield VI approach we use.
Another related work by Mandt et al. (2017) views SG descent as VI but requires additional effort to obtain posterior approximations, while in our approach the approximation is automatically obtained within an adaptive method.
Our weightperturbed algorithms are also related to globaloptimization methods, e.g., Gaussianhomotopy continuation methods (Mobahi & Fisher III, 2015), smoothedoptimization method (Leordeanu & Hebert, 2008), graduated optimization method (Hazan et al., 2016), and stochastic search methods (Zhou & Hu, 2014). In particular, our algorithm is related to recent approaches in deep learning for exploration to avoid local minima, e.g., natural evolution strategy (Wierstra et al., 2014), entropySGD (Chaudhari et al., 2016), and noisy networks for reinforcement learning (Fortunato et al., 2018; Plappert et al., 2018). An earlier version of our work (Khan et al., 2017) focuses exclusively on this problem, and in this paper we modify it to be implemented within an adaptive algorithm like Adam.
We consider modeling of a dataset by using a deep neural network (DNN). We assume a probabilistic framework where each data example
is sampled independently from a probability distribution
parameterized by a DNN with weights , e.g., the distribution could be an exponentialfamily distribution whose mean parameter is the output of a DNN (Bishop, 2006).One of the most popular approaches to estimate given is maximumlikelihood estimation (MLE), where we maximize the loglikelihood: . This optimization problem can be efficiently solved by applying SG methods such as RMSProp, AdaGrad and Adam. For large problems, these methods are extremely popular, partly due to the simplicity and efficiency of their implementations (see Fig. 1 for Adam’s pseudocode).
One of the goals of Bayesian deep learning is to go beyond MLE and estimate the posterior distribution of to obtain an uncertainty estimate of the weights. Unfortunately, the computation of the posterior is challenging in deep models. The posterior is obtained by specifying a prior distribution and then using Bayes’ rule: . This requires computation of the normalization constant which is a very difficult task for DNNs. One source of the difficulty is the size of and which are usually very large in deep learning. Another source is the nonconjugacy of the likelihood and the prior , i.e., the two distributions do not take the same form with respect to (Bishop, 2006). As a result, the product does not take a form with which can be easily computed. Due to these issues, Bayesian inference in deep learning is computationally challenging.
Variational inference (VI) simplifies the problem by approximating with a distribution whose normalizing constant is relatively easier to compute. Following previous work (Ranganath et al., 2014; Blundell et al., 2015; Graves, 2011), we choose both and
to be Gaussian distributions with diagonal covariances:
(1) 
where is a known precision parameter with , and
are mean and standard deviation of
. The distribution is known as the Gaussian meanfield variational distribution and its parameters and can be obtained by maximizing the following variational objective:(2) 
A straightforward approach used in the previous work (Ranganath et al., 2014; Blundell et al., 2015; Graves, 2011) is to maximize by using an SG method, e.g., we can use the following update:
(3) 
where is the iteration number, denotes an unbiased SG estimate of at with respect to , and are learning rates which can be adapted using methods such as RMSprop or AdaGrad. These approaches make use of existing codebases for adaptive learningrate methods to perform VI, which can handle many network architectures and can scale well to large datasets.
Despite this, a direct application of adaptive learningrate methods for VI may result in algorithms that use more computation and memory than necessary, and also require more implementation effort. Compared to MLE, the memory and computation costs increase because the number of parameters to be optimized is doubled and we now have two vectors and to estimate. Using adaptive methods increases this cost further as these methods require storing the scaling vectors that adapt the learning rate for both and
. In addition, using existing codebases require several modifications as they are designed and optimized for MLE. For example, we need to make changes in the computation graph where the objective function is changed to the variational objective and network weights are replaced by random variables. Together, these small issues make VI more difficult to implement and execute than MLE.
The algorithms developed in this paper solve some of these issues and can be implemented within Adam with minimal changes to the code. We derive our algorithm by approximating a naturalgradient method and then using a naturalmomentum method. We now describe our method in detail.
In this section, we introduce a naturalgradient method to perform VI and then propose several approximations that enable implementation within Adam.
Naturalgradient VI methods exploit the Riemannian geometry of by scaling the gradient with the inverse of its Fisher information matrix (FIM). We build upon the naturalgradient method of Khan & Lin (2017), which simplifies the update by avoiding a direct computation of the FIM. The main idea is to use the expectation parameters of the exponentialfamily distribution to compute natural gradients in the naturalparameter space. We provide a brief description of their method in Appendix B.
For Gaussian meanfield VI, the method of Khan & Lin (2017) gives the following update:
(4)  
(5) 
where is a scalar learning rate and denotes the elementwise product between vectors and . We refer to this update as naturalgradient variational inference (NGVI). A detailed derivation is given in Appendix C.
The NGVI update differs from (3) in one major aspect: the learning rate in (4
) is adapted by the variance
. This plays a crucial role in reducing the NGVI update to an Adamlike update, as we show the next section. The update requires a constraint but, as we show in Section 3.2, we can eliminate this constraint using an approximation.We start by expressing the NGVI update in terms of the MLE objective, so that we can directly compute gradients on the MLE objective using backpropagation. We start by defining the MLE objective (denoted by
) and minibatch stochasticgradient estimates (denoted by ):(6) 
where is the negative loglikelihood of ’th data example, and the minibatch contains examples chosen uniformly at random. Similarly, we can obtain a minibatch stochasticapproximation of the Hessian which we denote by .
As we show in Appendix D, the NGVI update can be written in terms of the stochastic gradients and Hessian of :
VON :  (7)  
(8) 
where is an elementwise division operation between vectors and , and we have approximated the expectation with respect to using one MonteCarlo (MC) sample with and . The update can be easily modified when multiple samples are used. This update can leverage backpropagation to perform the gradient and Hessian computation. Since the scaling vector contains an online estimate of the diagonal of the Hessian, we call this the “variational onlineNewton” (VON) method. VON is expected to perform as well as NGVI, but does not require the gradients of the variational objective.
The Hessian can be computed by using methods such as automaticdifferentiation or the reparameterization trick. However, since is a nonconvex function, the Hessian can be negative which might make negative, in which case the method will break down. One could use a constrained optimization method to solve this issue, but this might be difficult to implement and execute (we discuss this briefly in Appendix D.1). In the next section, we propose a simple fix to this problem by using an approximation.
To avoid negative variances in the VON update, we propose to use the Generalized GaussNewton (GGN) approximation (Schraudolph, 2002; Martens, 2014; Graves, 2011):
(9) 
where is the ’th element of . This approximation will always be nonnegative, therefore if the initial at is positive, it will remain positive in the subsequent iterations. Using this approximation to update in (8) and denoting the vector of by , we get,
VOGN :  (10) 
Using this update in VON, we get the “variational online GaussNewton” (VOGN) algorithm.
The GGN approximation is proposed by Graves (2011) for meanfield Gaussian VI to derive a fast gradientbased method (see Eq. (17) in his paper^{1}^{1}1There is a discrepancy between Eq. (17) and (12) in Graves (2011), however the text below Eq. (12) mentions relationship to FIM, from which it is clear that the GGN approximation is used.). This approximation is very useful for our naturalgradient method since it eliminates the constraint on , giving VOGN an algorithmic advantage over VON.
How good is this approximation? For an MLE problem, the approximation error of the GGN in (9) decreases as the modelfit improves during training (Martens, 2014). For VI, we expect the same however, since are sampled from , the expectation of the error is unlikely to be zero. Therefore, the solutions found by VOGN will typically differ from those found by VON, but their performances are expected to be similar.
An issue with VOGN is that its implementation is not easy within existing deeplearning codebases. This is because these codebases are optimized to directly compute the sum of the gradients over minibatches, and do not support computation of individual gradients as required in (9). A solution for such computations is discussed by Goodfellow (2015), but this requires additional implementation effort. Instead, we address this issue by using another approximation in the next section.
To simplify the implementation of VOGN, we propose to approximate the Hessian by the gradient magnitude (GM) (Bottou et al., 2016):
(11) 
Compared to the GGN which computes the sum of squaredgradients, this approximation instead computes the square of the sum. This approximation is also used in RMSprop which uses the following update given weights :
RMSprop  (12)  
(13) 
where is the vector that adapts the learning rate and is a small positive scalar added to avoid dividing by zero. The update of uses the GM approximation to the Hessian (Bottou et al., 2016). Adam and AdaGrad also use this approximation.
Using the GM approximation and an additional modification in the VON update, we can make the VON update very similar to RMSprop. Our modification involves taking the squareroot over in (7) and then using the GM approximation for the Hessian. We also use different learning rates and to update and , respectively. The resulting update is very similar to the RMSprop update:
Vprop:  
(14) 
where with . We call this update “Variational RMSprop” or simply “Vprop”.
The Vprop update resembles RMSprop but with three differences (highlighted in red). First, the gradient in Vprop is evaluated at the weights sampled from . This is a weightperturbation where the variance of the perturbation is obtained from the vector that adapts the learning rate. The variance is also the uncertainty estimates. Therefore, VI can be performed simply by using an RMSprop update with a few simple changes. The second difference between Vprop and RMSprop is that Vprop has an extra term in the update of which is due to the Gaussian prior. Finally, the third difference is that the constant in RMSprop is replaced by .
It is clear that the GM approximation might not be the best approximation of the Hessian. Taking square of a sum leads to a sum with terms which, depending on the correlations between the individual gradients, would either shrink or expand the estimate. The following theorem formalizes this intuition. It states that, given a minibatch of size , the expectation of the GM approximation is somewhere between the GGN and square of the fullbatch gradient.
Denote the fullbatch gradient with respect to by and the corresponding fullbatch GGN approximation by . Suppose minibatches
are sampled from the uniform distribution
over all minibatches, and denote a minibatch gradient by , then the expected value of the GM approximation is the following,(15) 
where .
A proof is given in Appendix G. This result clearly shows the bias introduced in the GM approximation and also that the bias increases with the minibatch size. For a minibatch of size , we have
and the GM is an unbiased estimator of the GGN, but when
it is purely the magnitude of the gradient and does not contain any secondorder information.Therefore, if our focus is to obtain uncertainty estimates with good accuracy, VOGN with might be a good choice since it is as easy as Vprop to implement. However, this might require a small learningrate and converge slowly. Vprop with will converge fast and is much easier to implement than VOGN with , but might result in slightly worse estimates. Using Vprop with may not be as good because of the squareroot^{2}^{2}2Note that the squareroot does not affect a fixed point (see Appendix H) but it might still affect the steps taken by the algorithm. over .
We now propose a naturalmomentum method which will enable an Adamlike update.
Momentum methods generally take the following form that uses Polyak’s heavyball method:
(16) 
where is the function we want to maximize and the last term is the momentum term. We propose a naturalmomentum version of this algorithm which employs naturalgradients instead of the gradients. We assume to be an exponentialfamily distribution with naturalparameter . We propose the following naturalmomentum method in the naturalparameter space:
(17) 
where denotes the naturalgradients in the naturalparameter space, i.e., the gradient scaled by the Fisher information matrix of .
We show in Appendix E that, for Gaussian , we can express the above update as a VON update with momentum^{3}^{3}3This is not an exact update for (17). We make one approximation in Appendix E to derive it.:
(18)  
(19) 
where with . This update is similar to (17), but here the learning rates are adapted. An attractive feature of this update is that it is very similar to Adam. Specifically the Adam update shown in Fig. 1 can be expressed as the following adaptive version of (16) as shown in Wilson et al. (2017)^{4}^{4}4Wilson et al. (2017) do not use the constant , but in Adam a small constant is added for numerical stability.,
(20)  
(21) 
where are appropriately defined in terms of the Adam’s learning rate and : and .
Using a similar procedure as the derivation of Vprop, we can express the update as an Adamlike update, which we call “variational Adam” or simply “Vadam”. A pseudocode is given in Fig. 1, where we use learning rates of the Adam update insteof of choosing them accoring to and . A derivation is given in Appendix E.4.
Vprop and Vadam perform variational inference, but they can be modified to perform optimization instead of inference. We now derive such an algorithm which turns out to be a variational version of AdaGrad.
We follow Staines & Barber (2013) who consider minimization of blackbox functions via the variational optimization^{5}^{5}5The exact conditions on under which VO can be applied are also discussed by Staines & Barber (2013). (VO) framework. In this framework, instead of directly minimizing , we minimize its expectation under a distribution with respect to and . The main idea behind VO is that the expectation can be used as a surrogate to the original optimization problem since . The equality is attained when , i.e., all mass of is at the mode. The main advantage of VO is that is differentiable even when itself is nondifferentiable. This way we can use SG optimizers to solve such problems.
Similarly to Vprop, we can derive an algorithm for VO by noting that VO can be seen as a special case of the VI problem (2) where the KL term is absent and is the negative loglikelihood. With this in mind, we define the following variational objective with an additional parameter :
(22) 
The parameter
allows us to interpolate between inference and optimization. When
, the objective corresponds to VI with a negative loglikelihood , and when , it corresponds to VO. Similar objectives have been proposed in existing works (Blundell et al., 2015; Higgins et al., 2016) where is used to improve convergence.For twicedifferentiable , we can follow a similar derivation as Section 3, and obtain the following algorithm,
(23)  
(24) 
where with . This algorithm is identical to the VON algorithm when , but when , we perform VO with an algorithm which is a diagonal version of the Variational AdaptiveNewton (VAN) algorithm proposed in Khan et al. (2017). By setting the value of between 0 and 1, we can interpolate between VO and VI. When the function is not differentiable, we can still compute the derivative of by using methods such as REINFORCE (Williams, 1992).
When Hessian is difficult to compute, we can employ a GM approximation and take the squareroot as we did in Vprop. For , the updates turn out to be similar to AdaGrad, which we call “variational AdaGrad” or simply “VadaGrad”. The exact updates are given in Appendix F. Unlike Vprop and Vadam, the scaling vector in VadaGrad is a weighted sum of the past gradientmagnitudes. Therefore, the entries in never decrease, and the variance estimate of VadaGrad never expands. This implies that it is highly likely that will converge to a Dirac delta and therefore arrive at a minimum of .
Experiments on Bayesian logistic regression showing (a) posterior approximations on a toy example, (b) performance on ‘USPS3v5’ measuring negative ELBO, logloss, and the symmetric KL divergence of the posterior approximation to MFExact, (c) symmetric KL divergence of Vadam for various minibatch sizes on ‘BreastCancer’ compared to VOGN with a minibatch of size 1.
In this section, our goal is to show that the quality of the uncertainty approximations obtained using our algorithms are comparable to existing methods, and computation of uncertainty is scalable. We present results on Bayesian logistic regression for classification, Bayesian neural networks for regression, and deep reinforcement learning. An additional result illustrating avoidance of localminima using Vadam is in Appendix L. Another result showing benefits of weightperturbation in Vadam is in Appendix M. The code to reproduce our results is available at https://github.com/emtiyaz/vadam.
In this experiment, we compare the posterior approximations found with our algorithms to the optimal variational approximation that minimizes the variational objective. For Bayesian logistic regression we can compute the optimal meanfield Gaussian approximations using the method described in Marlin et al. (2011) (refer to as ‘MFExact’), and compare it to the following methods: VOGN with minibatch size and a momentum term (referred to as ‘VOGN1’), and Vadam with (referred to as ‘Vadam’). Since our goal is to compare the accuracy of posterior approximations and not the speed of convergence, we run both the methods for many iterations with a small learning rate to make sure that they converge. We use three datasets: a toy dataset (), USPS3vs5 () and BreastCancer (). Details are in Appendix I.
Fig. 2 visualizes the approximations on a twodimensional toy example from Murphy (2012). The true posterior distribution is shown with the contour in the background. Both, Vadam and VOGN1 find approximations that are different from MFExact, which is clearly due to differences in the type of Hessian approximations they use.
For real datasets, we compare performances using three metrics. First, the negative of the variational objective on the training data (the evidence lowerbound or ELBO), logloss on the test data, and the symmetric KL distance between MFExact and the approximation found by a method. Fig. 2 shows the results averaged over 20 random splits of the USPS3vs5 dataset. ELBO and logloss are comparable for all methods, although Vadam does slightly worse on ELBO and VOGN1 has slightly higher variance for logloss. However, performance on the KL distance clearly shows the difference in the quality of posterior approximations. VOGN1 performs quite well since it uses an unbiased approximation of the GNN. Vadam does worse due to the bias introduced in the GM approximation with minibatch , as indicated by Theorem 1.
Fig. 2 further shows the effect of where, for each , we plot results for 20 random initializations on one split of the BreastCancer dataset. As we decrease , Vadam’s performance gets better, as expected. For , it closely matches VOGN1. The results are still different because Vadam does not reduce to VOGN1, even when due to the use of the squareroot over .
Test RMSE  Test loglikelihood  

Dataset  N  D  MCDropout  BBVI  Vadam  MCDropout  BBVI  Vadam 
Boston  506  13  2.97 0.19  3.58 0.21  3.93 0.26  2.46 0.06  2.73 0.05  2.85 0.07 
Concrete  1030  8  5.23 0.12  6.14 0.13  6.85 0.09  3.04 0.02  3.24 0.02  3.39 0.02 
Energy  768  8  1.66 0.04  2.79 0.06  1.55 0.08  1.99 0.02  2.47 0.02  2.15 0.07 
Kin8nm  8192  8  0.10 0.00  0.09 0.00  0.10 0.00  0.95 0.01  0.95 0.01  0.76 0.00 
Naval  11934  16  0.01 0.00  0.00 0.00  0.00 0.00  3.80 0.01  4.46 0.03  4.72 0.22 
Power  9568  4  4.02 0.04  4.31 0.03  4.28 0.03  2.80 0.01  2.88 0.01  2.88 0.01 
Wine  1599  11  0.62 0.01  0.65 0.01  0.66 0.01  0.93 0.01  1.00 0.01  1.01 0.01 
Yacht  308  6  1.11 0.09  2.05 0.06  1.32 0.10  1.55 0.03  2.41 0.02  1.70 0.03 
Performance comparisons for BNN regression. The better method out of BBVI and Vadam is shown in boldface according to a paired ttest with
value. Both methods perform comparably but MCDropout outperforms them.

We show results on the standard UCI benchmark. We repeat the experimental setup used in Gal & Ghahramani (2016)
. Following their work, we use a neural network with one hidden layer, 50 hidden units, and ReLU activation functions. We use the 20 splits of the data provided by
Gal & Ghahramani (2016) for training and testing. We use Bayesian optimization to select the prior precision and noise precision of the Gaussian likelihood. Further details of the experiments are given in Appendix J.We compare Vadam to MCDropout (Gal & Ghahramani, 2016) using the results reported in Gal & Ghahramani (2016). We also compare to an SG method using the reparameterization trick and the Adam optimizer (referred to as ‘BBVI’). For a fair comparison, the Adam optimizer is run with the same learning rates as Vadam, although these can be tuned further to get better performance.
Table 1 shows the performance in terms of the test RMSE and the test loglikelihood. The better method out of BBVI and Vadam is shown in boldface found using a paired ttest with . Both methods perform comparably, which supports our conclusion, however, MCDropout outperforms both the methods. We also find that VOGN shows similar results to Vadam and BBVI (we omit the results due to lack of space). The convergence plots for the final runs is given in Appendix J.
A good exploration strategy is crucial in reinforcement learning (RL) since the data is sequentially collected. We show that weightperturbation in Vadam improves exploration in RL. Due to space constraints, we only provide a brief summary of our results, and give details in Appendix K.
We consider the deep deterministic policy gradient (DDPG) method for the HalfCheetah task using a twolayer neural networks with and ReLU hidden units (Lillicrap et al., 2015). We compare Vadam and VadaGrad to two SGD methods, one of which does exploration (referred to as ‘SGDExplore’), and the other does not (referred to as ‘SGDplain’). The rightmost plot in Figure 3 shows the cumulative rewards (higher is better) of each method against training iterations. VadaGrad and Vadam clearly learn faster than both SGDPlain and SGDExplore. We also compare the performances against Adam variants of SGDPlain and SGDExplore. Their results, given in the Appendix K, show that Vadam and VadaGrad still learn faster, but only in the beginning and Adam based methods can catch up quickly. This suggests that the exploration strategy has a high impact on the early learning performance in the HalfCheetah task, and the effect of good exploration decreases over time as the agent collect more informative training samples.
In this paper, we present new VI algorithms which are as simple to implement and execute as algorithms for MLE. We obtain them by using a series of approximations and a natural momentum method for a naturalgradient VI method. The resulting algorithms can be implemented within Adam with minimal changes. Our empirical findings confirm that our proposed algorithms obtain comparable uncertainty estimates to existing VI methods, but require less computational and implementation effort^{6}^{6}6We made many new changes in this cameraready version of the paper. A list of the changes is given in Appendix A..
An interesting direction we hope to pursue in the future is to generalize our naturalgradient approach to other types of approximation, e.g., exponetialfamily distributions and their mixtures. We would also like to further explore the application to areas such as RL and stochastic optimization.
We thank the anonymous reviewers for their feedback. We greatly appreciate many insightful discussions with Aaron Mishkin (UBC) and Frederik Kunstner (EPFL), and also thank them for their help on carrying out experiments and reviewing the manuscript. We would also like to thank Roger Grosse and David Duvenaud from the University of Toronto for useful discussions. We would like to thank Zuozhu Liu (SUTD, Singapore) for his help with the experiment on deep RL and logistic regression. Finally, we are thankful for the RAIDEN computing system at the RIKEN Center for AI Project, which we extensively used for our experiments.
Annual Conference on Computational Learning Theory
, pp. 5–13, 1993.International Conference on Artificial Intelligence and Statistics
, pp. 878–887, 2017.Fixedform variational posterior approximation through stochastic linear regression.
Bayesian Analysis, 8(4):837–882, 2013.Taking reviewer’s suggestions into account, we changed the title of our paper. The title of our submitted version was “Vadam: Fast and Scalable Variational Inference by Perturbing Adam”.
In the submitted version, we motivated our approach based on its ease of implementation. In the new version, we changed the motivation to make VI as easy to implement and execute as MLE.
In the new version, we have added a separate section on related work.
We improved the discussion of our approximation methods, and added an error analysis.
Overall conclusions of our paper have also slightly changed in the new version. The new conclusions suggest that there is a tradeoff between the easeofimplementation and quality of uncertainty approximation.
As per reviewers suggestions, we also made major improvements in our experiment results.
We added test loglikelihood in the BNN results. We changed the hyperparameter selection from grid search to Bayesian optimization. We removed two methods from the table, namely PBP and VIG, since they use different splits compared to our setting. We improved the performance of BBVI by using better initialization and learning rates. We corrected a scaling problem with our Vadam method. The new method gives a slightly worse performance than the results present in the submitted version.
We added a logistic regression experiment where we evaluate the quality of uncertainty estimates.
We added the details of the RL experiments which we forgot to add in the submitted version. We also added a comparison to Adambased methods in the appendix for the RL experiment.
We removed an unclear result about reducing overfitting.
We added an additional result comparing VOGN with Vadam and BBVI on Bayesian neural network.
Khan & Lin (2017) propose a naturalgradient method for variational inference. In this section, we briefly discuss this method.
Denote the variational objective by for the variational distribution which takes an exponentialfamily form with naturalparameter . The objective is given as follows:
(25) 
We assume that the exponentialfamily is in minimal representation, which ensures that there is a onetoone mapping between the natural parameter and the expectation parameter, denoted by . Therefore, it is possible to express in terms of . We denote this new objective by . We can also reparameterize in terms of and denote it by .
Naturalgradient methods exploit the Riemannian geometry of by scaling the gradient by the inverse of the Fisher information matrix. The method of Khan & Lin (2017) simplifies the update by avoiding a direct computation of the FIM. This is made possible due to a relationship between the natural parameter and the expectation parameter of an exponentialfamily distribution. The relationship dictates that the natural gradient with respect to is equal to the gradient with respect to . This is stated below where FIM is denoted by ,
(26) 
This relationship has been discussed in the earlier work of Hensman et al. (2012) and can also be found in Amari (2016).
The method of Khan & Lin (2017) exploits this result within a mirror descent framework. They propose to use a mirrordescent update in the expectation parameter space which is equivalent to the naturalgradient update in the naturalparameter space. Therefore, the naturalgradient update can be performed by using the gradient with respect to the expectation parameter. We give a formal statement below.
Consider the following mirrordescent step:
(27) 
where
is the KullbackLeibler divergence and
is the learning rate in iteration . Each step of this mirror descent update is equivalent to the following naturalgradient descent in the naturalparameter space:(28) 
A formal proof of this statement can be found in Raskutti & Mukherjee (2015).
Using (26), the naturalgradient update above can be simply written as the following:
(29) 
which involves computing the gradient with respect to but taking a step in the naturalparameter space.
As we show in the next section, the above relationship enables us to derive a simple naturalgradient update because, for a Gaussian distribution, the gradient with respect to leads to a simple update.
In this section, we derive the naturalgradient update for the Gaussian approximation with mean and covariance matrix . In the end, we will make the meanfield approximation: to get the final update.
We start by defining the natural and expectation parameters of a Gaussian:
(30)  
(31) 
Now we will express the gradient with respect to these expectation parameters in terms of the gradients with respect to and
using the chain rule (see Appendix B.1 in
Khan & Lin (2017) for a derivation):(32)  
(33) 
Next, using the definition of natural parameters, we can rewrite (29) in terms of and (here implies that it is gradient of the variational objective with respect to a variable at ):
(34)  
(35)  
(36)  
(37)  
(38)  
(39) 
In summary, the naturalgradient update is
(40)  
(41) 
By considering a Gaussian meanfield VI with a diagonal covariance: , we obtain
(42)  
(43) 
In update (5), we use stochastic gradients instead of exact gradients.
Note that there is an explicit constraint in the above update, i.e., the precision needs to be positive at every step. The learning rate can be adapted to make sure that the constraint is always satisfied. We discuss this method in Appendix D.1. Another option is to make an approximation, such as a GaussNewton approximation, to make sure that this constraint is always satisfied. We make this assumption in one of our methods called Variational Online GaussNewton Method.
In this section, we derive the variational onlineNewton (VON) method proposed in Section 3. We will modify the NGVI update in (41).
The variational lowerbound in (25) can be reexpressed as
(44) 
where . To derive VON, we use the Bonnet’s and Price’s theorems (Opper & Archambeau, 2009; Rezende et al., 2014) to express the gradients of the expectation of with respect to and in terms of the gradient and Hessian of , i.e.,
(45)  
(46) 
where and denote the gradient and Hessian of , respectively. Using these, we can rewrite the gradients of required in the NGVI update in (41) as
(47)  
(48)  
(49)  
(50)  
(51) 
By substituting these into the NGVI update of (41) and then approximating the expectation by one MonteCarlo sample , we get the following update:
(52)  
(53) 
By defining a matrix , we get the following:
VON (Fullcovariance):  