Practical Deep Learning with Bayesian Principles
Kazuki Osawa,1 Siddharth Swaroop,2,⇤ Anirudh Jain,3,⇤,† Runa Eschenhagen,4,†
Richard E. Turner,2 Rio Yokota,1 Mohammad Emtiyaz Khan5,‡ .
Tokyo Institute of Technology, Tokyo, Japan
University of Cambridge, Cambridge, UK
3
Indian Institute of Technology (ISM), Dhanbad, India
4
University of Osnabrück, Osnabrück, Germany
5
RIKEN Center for AI Project, Tokyo, Japan
1

2

Abstract
Bayesian methods promise to fix many shortcomings of deep learning, but they
are impractical and rarely match the performance of standard methods, let alone
improve them. In this paper, we demonstrate practical training of deep networks
with natural-gradient variational inference. By applying techniques such as batchnormalisation, data augmentation, and distributed training, we achieve similar
performance in about the same number of epochs as the Adam optimiser, even on
large datasets such as ImageNet. Importantly, the benefits of Bayesian principles
are preserved: predictive probabilities are well-calibrated, uncertainties on outof-distribution data are improved, and continual-learning performance is boosted.
This work enables practical deep learning while preserving benefits of Bayesianprinciples. A PyTorch implementation1 is available as a plug-and-play optimiser.

1

Introduction

Deep learning has been extremely successful in many fields such as computer vision [29], speechprocessing [17], and natural-language processing [39], but it is also plagued with several issues
that make its application difficult in many other fields. For example, it requires a large amount ofhigh-quality data and it can overfit when dataset size is small. Similarly, sequential learning can cause
forgetting of past knowledge [27], and lack of reliable confidence estimates and other robustness
issues can make it vulnerable to adversarial attacks [6]. Ultimately, due to such issues, application ofdeep learning remains challenging, especially for applications where human lives are at risk.
Bayesian principles have the potential to address such issues. For example, we can represent uncertainty using the posterior distribution, enable sequential learning using Bayes’ rule, and reduce
overfitting with Bayesian model averaging [19]. The use of such Bayesian principles for neuralnetworks has been advocated from very early on. Bayesian inference on neural networks were all proposed in the 90s, e.g., by using MCMC methods [41], Laplace’s method [35], and variational inference(VI) [18, 2, 49, 1]. Benefits of Bayesian principles are even discussed in machine-learning textbooks[36, 3]. Despite this, they are rarely employed in practice. This is mainly due to computational
concerns, unfortunately overshadowing their theoretical advantages.
The difficulty lies in the computation of the posterior distribution, which is especially challenging fordeep learning. Even approximation methods, such as VI and MCMC, have historically been difficult* These two authors contributed equally.
† This work is conducted during an internship at RIKEN Center for AI project.
‡ Corresponding author: emtiyaz.khan@riken.jp
1
The code is available at https://github.com/team-approx-bayes/dl-with-bayes.
33rd Conference on Neural Information Processing Systems (NeurIPS 2019), Vancouver, Canada.

Figure 1: Comparing VOGN [24], a natural-gradient VI method, to Adam and SGD, training ResNet18 on ImageNet. The two left plots show that VOGN and Adam have similar convergence behaviour
and achieve similar performance in about the same number of epochs. VOGN achieves 67.38% on
validation compared to 66.39% by Adam and 67.79% by SGD. Run-time of VOGN is 76 seconds perepoch compared to 44 seconds for Adam and SGD. The rightmost figure shows the calibration curve.
VOGN gives calibrated predictive probabilities (the diagonal represents perfect calibration).

to scale to large datasets such as ImageNet [47]. Due to this, it is common to use less principledapproximations, such as MC-dropout [9], even though they are not ideal when it comes to fixing theissues of deep learning. For example, MC-dropout is unsuitable for continual learning [27] since its
posterior approximation does not have mass over the whole weight space. It is also found to perform
poorly for sequential decision making [45]. The form of the approximation used by such methods
is usually rigid and cannot be easily improved, e.g., to other forms such as a mixture of Gaussians.
The goal of this paper is to make more principled Bayesian methods, such as VI, practical for deeplearning, thereby helping researchers tackle its key limitations.
We demonstrate practical training of deep networks by using recently proposed natural-gradient VI
methods. These methods resemble the Adam optimiser, enabling us to leverage existing techniques
for initialisation, momentum, batch normalisation, data augmentation, and distributed training. As aresult, we obtain similar performance in about the same number of epochs as Adam when training
many popular deep networks (e.g., LeNet, AlexNet, ResNet) on datasets such as CIFAR-10 and
ImageNet (see Fig. 1). The results show that, despite using an approximate posterior, the training
methods preserve the benefits coming from Bayesian principles. Compared to standard deep-learning
methods, the predictive probabilities are well-calibrated, uncertainties on out-of-distribution inputs
are improved, and performance for continual-learning tasks is boosted. Our work shows that practical
deep learning is possible with Bayesian methods and aims to support further research in this area.
Related work. Previous VI methods, notably by Graves [14] and Blundell et al. [4], require significant implementation and tuning effort to perform well, e.g., on convolution neural networks (CNN).
Slow convergence is found to be especially problematic for sequential problems [45]. There appears
to be no reported results with complex networks on large problems, such as ImageNet. Our work
solves these issues by applying deep-learning techniques to natural-gradient VI [24, 56].
In their paper, Zhang et al. [56] also employed data augmentation and batch normalisation for anatural-gradient method called Noisy K-FAC (see Appendix A) and showed results on VGG on
CIFAR-10. However, a mean-field method called Noisy Adam was found to be unstable with batchnormalisation. In contrast, we show that a similar method, called Variational Online Gauss-Newton(VOGN), proposed by Khan et al. [24], works well with such techniques. We show results for
distributed training with Noisy K-FAC on Imagenet, but do not provide extensive comparisons sincetuning it is time-consuming. Many of our techniques can speed-up Noisy K-FAC, which is promising.
Many other approaches have recently been proposed to compute posterior approximations by training
deterministic networks [46, 37, 38]. Similarly to MC-dropout, their posterior approximations are notflexible, making it difficult to improve the accuracy of their approximations. On the other hand, VI
offers a much more flexible alternative to apply Bayesian principles to deep learning.
2

2

Deep Learning with Bayesian Principles and Its Challenges

The success of deep learning is partly due to the availability of scalable and practical methods fortraining deep neural networks (DNNs). Network training is formulated as an optimisation problem
where a loss between the data and the DNN’s predictions is minimised. For example, in a supervised
learning task with a dataset D of N inputs xi and corresponding outputs yi of length K, we minimiseP¯¯a loss of the following form: `(w)
+ w> w, where `(w):= N1 i `(yi , f w (xi )), f w (x) 2 RK
denotes the DNN outputs with weights w, `(y, f ) denotes a differentiable loss function between an
output y and the function f , and > 0 is the L2 regulariser.2 Deep learning relies on stochasticgradient (SG) methods to minimise such loss functions. The most commonly used optimisers, suchas stochastic-gradient descent (SGD), RMSprop [53], and Adam [25], take the following form3 (alloperations below are element-wise):
wt+1

wt

ĝ(wt ) + wt
↵t p,
st+1 + ✏st+1(1t )st +t (ĝ(wt ) +2wt ) ,(1)

where t is the iteration, ↵t > 0 and 0 < t < 1 are learning rates,P✏ > 0 is a small scalar, and ĝ(w)1is the stochastic gradients at w defined as follows: ĝ(w) := M
i2Mt rw `(yi , f w (xi )) using aminibatch Mt of M data examples. This simple update scales extremely well and can be applied
to very large problems. With techniques such as initialisation protocols, momentum, weight-decay,
batch normalisation, and data augmentation, it also achieves good performance for many problems.
In contrast, the full Bayesian approach to deep learning is computationally very expensive. The¯
posterior distribution can be obtained using Bayes’ rule: p(w|D) = exp N `(w)/⌧p(w)/p(D)where 0 < ⌧  1.4 This is costly due to the computation of the marginal likelihood p(D), a
high-dimensional integral that is difficult to compute for large networks. Variational inference (VI)is a principled approach to more scalably estimate an approximation to p(w|D). The main idea
is to employ a parametric approximation, e.g., a Gaussian q(w) := N (w|µ, ⌃) with mean µ andcovariance ⌃. The parameters µ and ⌃ can then be estimated by maximising the evidence lower
bound (ELBO):⇥⇤¯
ELBO: L(µ, ⌃) := N Eq `(w)
⌧ DKL [q(w) k p(w)],(2)
where DKL [·] denotes the Kullback-Leibler divergence. By using more complex approximations, we
can further reduce the approximation error, but at a computational cost. By formulating Bayesianinference as an optimisation problem, VI enables a practical application of Bayesian principles.

Despite this, VI has remained impractical for training large deep networks on large datasets. Existing
methods, such as Graves [14] and Blundell et al. [4], directly apply popular SG methods to optimise
the variational parameters in the ELBO, yet they fail to get a reasonable performance on large problems, usually converging very slowly. The failure of such direct applications of deep-learning methods
to VI is not surprising. The techniques used in one field may not directly lead to improvements in theother, but it will be useful if they do, e.g., if we can optimise the ELBO in a way that allows us to
exploit the tricks and techniques of deep learning and boost the performance of VI. The goal of thiswork is to do just that. We now describe our methods in detail.

3

Practical Deep Learning with Natural-Gradient Variational Inference

In this paper, we propose natural-gradient VI methods for practical deep learning with Bayesianprinciples. The natural-gradient update takes a simple form when estimating exponential-family
approximations [23, 22]. When p(w) := N (w|0, I/ ), the update of the natural-parameter isperformed by using the stochastic gradient of the expected regularised-loss:⇥⇤¯⌧ ⇢) t ⇢rµ Eq `(w)
+ 12 ⌧ w> w ,(3)
t+1 = (12This regulariser is sometimes set to 0 or a very small value.
Alternate versions with weight-decay and momentum differ from this update [34]. We present a form usefulto establish the connection between SG methods and natural-gradient VI.
4
This is a tempered posterior [54] setup where ⌧ is set 6= 1 when we expect model misspecification and/or
adversarial examples [10]. Setting ⌧ = 1 recovers standard Bayesian inference.
33

where ⇢ > 0 is the learning rate, and we note that the stochastic gradients are computed with respect
to µ, the expectation parameters of q. The moving average above helps to deal with the stochasticityof the gradient estimates, and is very similar to the moving average used in deep learning (see (1)).
When ⌧ is set to 0, the update essentially minimises the regularised loss (see Section 5 in Khan et al.
[24]). These properties of natural-gradient VI makes it an ideal candidate for deep learning.
Recent work by Khan et al. [24] and Zhang et al. [56] further show that, when q is Gaussian, theupdate (3) assumes a form that is strikingly similar to the update (1). For example, the Variational
Online Gauss-Newton (VOGN) method of Khan et al. [24] estimates a Gaussian with mean µt and a
diagonal covariance matrix ⌃t using the following update:ĝ(wt ) + ˜µt
1 X
2
µt+1
µt ↵t
, st+1(1 ⌧ t )st + t(gi (wt )) ,(4)
M
st+1 + ˜
i2Mt

where gi (wt ) := rw `(yi , fwt (xi )), wt ⇠ N (w|µt , ⌃t ) with ⌃t := diag(1/(N (st + ˜))), ˜ :=⌧ /N , and ↵t , t > 0 are learning rates. Operations are performed element-wise. Similarly to (1),the vector st adapts the learning rate and is updated using a moving average.
A major difference inPVOGN is that the update of st is now based on a Gauss-Newton approximation
12[14] which uses M
fundamentally different from the SG update in (1)i2Mt (gi (wt )) . This isP
1which instead uses the gradient-magnitude ( M i2Mt gi (wt ) + wt )2 [5]. The first approach uses
the sum outside the square while the second approach uses it inside. VOGN is therefore a secondorder method and, similarly to Newton’s method, does not need a square-root over st . Implementationof this step requires an additional calculation (see Appendix B) which makes VOGN a bit slower thanAdam, but VOGN is expected to give better variance estimates (see Theorem 1 in Khan et al. [24]).
The main contribution of this paper is to demonstrate practical training of deep networks using
VOGN. Since VOGN takes a similar form to SG methods, we can easily borrow existing deeplearning techniques to improve performance. We will now describe these techniques in detail.
Pseudo-code for VOGN is shown in Algorithm 1.
Batch normalisation: Batch normalisation [20] has been found to significantly speed up and stabilise
training of neural networks, and is widely used in deep learning. BatchNorm layers are inserted
between neural network layers. They help stabilise each layer’s input distribution by normalising therunning average of the inputs’ mean and variance. In our VOGN implementation, we simply use theexisting implementation with default hyperparameter settings. We do not apply L2 regularisation andweight decay to BatchNorm parameters, like in Goyal et al. [13], or maintain uncertainty over the
BatchNorm parameters. This straightforward application of batch normalisation works for VOGN.
Data Augmentation: When training on image datasets, data augmentation (DA) techniques canimprove performance drastically [13]. We consider two common real-time data augmentation
techniques: random cropping and horizontal flipping. After randomly selecting a minibatch at each
iteration, we use a randomly selected cropped version of all images. Each image in the minibatch has
a 50% chance of being horizontally flipped.
We find that directly applying DA gives slightly worse performance than expected, and also affects
the calibration of the resulting uncertainty. However, DA increases the effective sample size. We
therefore modify it to be ⇢N where ⇢ 1, improving performance (see step 2 in Algorithm 1). The
reason for this performance boost might be due to the complex relationship between the regularisation¯
and N . For the regularised loss `(w)
+ w> w, the two are unidentifiable, i.e., we can multiply
by a constant and reduce N by the same constant without changing the minimum. However, in a
Bayesian setting (like in (2)), the two quantities are separate, and therefore changing the data mightalso change the optimal prior variance hyperparameter in a complicated way. This needs further
theoretical investigations, but our simple fix of scaling N seems to work well in the experiments.
We set ⇢ by considering the specific DA techniques used. When training on CIFAR-10, the random
cropping DA step involves first padding the 32x32 images to become of size 40x40, and then taking
randomly selected 28x28 cropped images. We consider this as effectively increasing the dataset sizeby a factor of 5 (4 images for each corner, and one central image). The horizontal flipping DA step
doubles the dataset size (one dataset of unflipped images, one for flipped images). Combined, thisgives ⇢ = 10. Similar arguments for ImageNet DA techniques give ⇢ = 5. Even though ⇢ is anotherhyperparameter to set, we find that its precise value does not matter much. Typically, after setting an
estimate for ⇢, tuning a little seems to work well (see Appendix E).
4

Algorithm 1: Variational Online Gauss Newton (VOGN)
1: Initialise µ0 , s0 , m0 .
2: N
⇢N , ˜⌧ /N .
3: repeat4:
Sample a minibatch M of size M .
5:Split M into each GPU (local minibatch Mlocal ).
6:for each GPU in parallel do
7:
for k = 1, 2, . . . , K do
8:
Sample ✏ ⇠ N (0, I).
9:w(k)µ + ✏ with
(1/(N (s + ˜ + )))1/2 .
(k)10:Compute gi
rw `(yi , f w(k) (xi )), 8i 2 Mlocal
using the method described in Appendix B.
P
(k)
1
11:
ĝk
i2Mlocal gi .
M
P
(k) 2
1
12:
ĥk
i2Mlocal (gi ) .
M
13:
end for P
PK
K
1
1
14:
ĝ
k=1 ĝk and ĥ
k=1 ĥk .
K
K15:
end for
16:
AllReduce ĝ, ĥ.
˜
17:m
1 m + (ĝ + µ).
18:s(1 ⌧ 2 )s + 2 ĥ.
19:µµ ↵m/(s + ˜ + ).
20: until stopping criterion is met

w(i) ⇠ q(w)<latexit sha1_base64="Dv2ZWT7aVb6yCMx0AIW4B6zAHBY=">AAACC3icbVDLSsNAFJ3UV42vqEs3Q4vQbkqigi6LblxWsA9oYplMJ+3QySTOTJQSunfjr7hxoYhbf8Cdf+OkDaitBy4czrmXe+/xY0alsu0vo7C0vLK6Vlw3Nza3tnes3b2WjBKBSRNHLBIdH0nCKCdNRRUjnVgQFPqMtP3RRea374iQNOLXahwTL0QDTgOKkdJSzyq5IVJDP0jvJzdphVYnrqSheVv5kas9q2zX7CngInFyUgY5Gj3r0+1HOAkJV5ghKbuOHSsvRUJRzMjEdBNJYoRHaEC6mnIUEuml018m8FArfRhEQhdXcKr+nkhRKOU49HVndqKc9zLxP6+bqODMSymPE0U4ni0KEgZVBLNgYJ8KghUba4KwoPpWiIdIIKx0fKYOwZl/eZG0jmrOcc2+OinXz/M4iuAAlEAFOOAU1MElaIAmwOABPIEX8Go8Gs/Gm/E+ay0Y+cw++APj4xv4CZr8</latexit>M<latexit sha1_base64="uZvzERQPKMMR9NiLGAcS/T5qG+0=">AAAB8nicbVBNS8NAFHypX7V+VT16WSyCp5KooMeiFy9CBWsLaSib7bZdutmE3RehhP4MLx4U8eqv8ea/cdPmoK0DC8PMe+y8CRMpDLrut1NaWV1b3yhvVra2d3b3qvsHjyZONeMtFstYd0JquBSKt1Cg5J1EcxqFkrfD8U3ut5+4NiJWDzhJeBDRoRIDwShaye9GFEeMyuxu2qvW3Lo7A1kmXkFqUKDZq351+zFLI66QSWqM77kJBhnVKJjk00o3NTyhbEyH3LdU0YibIJtFnpITq/TJINb2KSQz9fdGRiNjJlFoJ/OIZtHLxf88P8XBVZAJlaTIFZt/NEglwZjk95O+0JyhnFhCmRY2K2EjqilD21LFluAtnrxMHs/q3nndvb+oNa6LOspwBMdwCh5cQgNuoQktYBDDM7zCm4POi/PufMxHS06xcwh/4Hz+AIM0kWU=</latexit>

Mlocal Mlocal Mlocal Mlocal<latexit sha1_base64="ZwiUQGlfNIiAPSgLdNtgjp+cW94=">AAAB/HicbVDLSsNAFL3xWesr2qWbwSK4KokKuiy6cSNUsA9oQ5hMJ+3QySTMTIQQ6q+4caGIWz/EnX/jpM1CWw8MHM65l3vmBAlnSjvOt7Wyura+sVnZqm7v7O7t2weHHRWnktA2iXksewFWlDNB25ppTnuJpDgKOO0Gk5vC7z5SqVgsHnSWUC/CI8FCRrA2km/XBhHWY4J5fjf1cx4bNvXtutNwZkDLxC1JHUq0fPtrMIxJGlGhCcdK9V0n0V6OpWaE02l1kCqaYDLBI9o3VOCIKi+fhZ+iE6MMURhL84RGM/X3Ro4jpbIoMJNFVLXoFeJ/Xj/V4ZWXM5GkmgoyPxSmHOkYFU2gIZOUaJ4ZgolkJisiYywx0aavqinBXfzyMumcNdzzhnN/UW9el3VU4AiO4RRcuIQm3EIL2kAgg2d4hTfryXqx3q2P+eiKVe7U4A+szx9z0JVI</latexit><latexit sha1_base64="ZwiUQGlfNIiAPSgLdNtgjp+cW94=">AAAB/HicbVDLSsNAFL3xWesr2qWbwSK4KokKuiy6cSNUsA9oQ5hMJ+3QySTMTIQQ6q+4caGIWz/EnX/jpM1CWw8MHM65l3vmBAlnSjvOt7Wyura+sVnZqm7v7O7t2weHHRWnktA2iXksewFWlDNB25ppTnuJpDgKOO0Gk5vC7z5SqVgsHnSWUC/CI8FCRrA2km/XBhHWY4J5fjf1cx4bNvXtutNwZkDLxC1JHUq0fPtrMIxJGlGhCcdK9V0n0V6OpWaE02l1kCqaYDLBI9o3VOCIKi+fhZ+iE6MMURhL84RGM/X3Ro4jpbIoMJNFVLXoFeJ/Xj/V4ZWXM5GkmgoyPxSmHOkYFU2gIZOUaJ4ZgolkJisiYywx0aavqinBXfzyMumcNdzzhnN/UW9el3VU4AiO4RRcuIQm3EIL2kAgg2d4hTfryXqx3q2P+eiKVe7U4A+szx9z0JVI</latexit><latexit sha1_base64="ZwiUQGlfNIiAPSgLdNtgjp+cW94=">AAAB/HicbVDLSsNAFL3xWesr2qWbwSK4KokKuiy6cSNUsA9oQ5hMJ+3QySTMTIQQ6q+4caGIWz/EnX/jpM1CWw8MHM65l3vmBAlnSjvOt7Wyura+sVnZqm7v7O7t2weHHRWnktA2iXksewFWlDNB25ppTnuJpDgKOO0Gk5vC7z5SqVgsHnSWUC/CI8FCRrA2km/XBhHWY4J5fjf1cx4bNvXtutNwZkDLxC1JHUq0fPtrMIxJGlGhCcdK9V0n0V6OpWaE02l1kCqaYDLBI9o3VOCIKi+fhZ+iE6MMURhL84RGM/X3Ro4jpbIoMJNFVLXoFeJ/Xj/V4ZWXM5GkmgoyPxSmHOkYFU2gIZOUaJ4ZgolkJisiYywx0aavqinBXfzyMumcNdzzhnN/UW9el3VU4AiO4RRcuIQm3EIL2kAgg2d4hTfryXqx3q2P+eiKVe7U4A+szx9z0JVI</latexit><latexit sha1_base64="ZwiUQGlfNIiAPSgLdNtgjp+cW94=">AAAB/HicbVDLSsNAFL3xWesr2qWbwSK4KokKuiy6cSNUsA9oQ5hMJ+3QySTMTIQQ6q+4caGIWz/EnX/jpM1CWw8MHM65l3vmBAlnSjvOt7Wyura+sVnZqm7v7O7t2weHHRWnktA2iXksewFWlDNB25ppTnuJpDgKOO0Gk5vC7z5SqVgsHnSWUC/CI8FCRrA2km/XBhHWY4J5fjf1cx4bNvXtutNwZkDLxC1JHUq0fPtrMIxJGlGhCcdK9V0n0V6OpWaE02l1kCqaYDLBI9o3VOCIKi+fhZ+iE6MMURhL84RGM/X3Ro4jpbIoMJNFVLXoFeJ/Xj/V4ZWXM5GkmgoyPxSmHOkYFU2gIZOUaJ4ZgolkJisiYywx0aavqinBXfzyMumcNdzzhnN/UW9el3VU4AiO4RRcuIQm3EIL2kAgg2d4hTfryXqx3q2P+eiKVe7U4A+szx9z0JVI</latexit>

w(1) w(3) w(5)w(2) w(4) w(6)

w(7)
w(8)

ĝ ĥ

ĝ ĥ

<latexit sha1_base64="DZVriWHhAABp/VIrxpVe/AkNe74=">AAAB+XicbVDLSsNAFL3xWesr6tLNYBHqpiQq6LLoxmUF+4A2lsl00g6dTMLMpFJC/sSNC0Xc+ifu/BsnbRbaemDgcM693DPHjzlT2nG+rZXVtfWNzdJWeXtnd2/fPjhsqSiRhDZJxCPZ8bGinAna1Exz2oklxaHPadsf3+Z+e0KlYpF40NOYeiEeChYwgrWR+rbdC7Ee+UH6lD2mVfcs69sVp+bMgJaJW5AKFGj07a/eICJJSIUmHCvVdZ1YeymWmhFOs3IvUTTGZIyHtGuowCFVXjpLnqFTowxQEEnzhEYz9fdGikOlpqFvJvOcatHLxf+8bqKDay9lIk40FWR+KEg40hHKa0ADJinRfGoIJpKZrIiMsMREm7LKpgR38cvLpHVecy9qzv1lpX5T1FGCYziBKrhwBXW4gwY0gcAEnuEV3qzUerHerY/56IpV7BzBH1ifPzvyk2I=</latexit><latexit sha1_base64="hvhnf9uGBxMR9D/7hkTLQSq9t6g=">AAAB+XicbVDLSsNAFL2pr1pfUZduBotQNyWxgi6LblxWsA9oa5lMJ+3QySTMTCol5E/cuFDErX/izr9x0mahrQcGDufcyz1zvIgzpR3n2yqsrW9sbhW3Szu7e/sH9uFRS4WxJLRJQh7KjocV5UzQpmaa004kKQ48Ttve5Dbz21MqFQvFg55FtB/gkWA+I1gbaWDbvQDrsecnT+ljUqmdpwO77FSdOdAqcXNShhyNgf3VG4YkDqjQhGOluq4T6X6CpWaE07TUixWNMJngEe0aKnBAVT+ZJ0/RmVGGyA+leUKjufp7I8GBUrPAM5NZTrXsZeJ/XjfW/nU/YSKKNRVkcciPOdIhympAQyYp0XxmCCaSmayIjLHERJuySqYEd/nLq6R1UXVrVef+sly/yesowgmcQgVcuII63EEDmkBgCs/wCm9WYr1Y79bHYrRg5TvH8AfW5w8+/pNk</latexit><latexit sha1_base64="B2LnCCBsDGx8M7KwkiP54CtpI2E=">AAAB+XicbVDLSsNAFL2pr1pfUZduBotQNyXxgS6LblxWsA9oa5lMJ+3QySTMTCol5E/cuFDErX/izr9x0mahrQcGDufcyz1zvIgzpR3n2yqsrK6tbxQ3S1vbO7t79v5BU4WxJLRBQh7KtocV5UzQhmaa03YkKQ48Tlve+DbzWxMqFQvFg55GtBfgoWA+I1gbqW/b3QDrkecnT+ljUrk8Tft22ak6M6Bl4uakDDnqffurOwhJHFChCcdKdVwn0r0ES80Ip2mpGysaYTLGQ9oxVOCAql4yS56iE6MMkB9K84RGM/X3RoIDpaaBZyaznGrRy8T/vE6s/etewkQUayrI/JAfc6RDlNWABkxSovnUEEwkM1kRGWGJiTZllUwJ7uKXl0nzrOqeV537i3LtJq+jCEdwDBVw4QpqcAd1aACBCTzDK7xZifVivVsf89GCle8cwh9Ynz9CCpNm</latexit><latexit sha1_base64="peZMHpnhsnVKmu0FyhJJG0tzlO4=">AAAB+XicbVDLSsNAFL2pr1pfUZduBotQNyWpgi6LblxWsA9oa5lMJ+3QySTMTCol5E/cuFDErX/izr9x0mahrQcGDufcyz1zvIgzpR3n2yqsrW9sbhW3Szu7e/sH9uFRS4WxJLRJQh7KjocV5UzQpmaa004kKQ48Ttve5Dbz21MqFQvFg55FtB/gkWA+I1gbaWDbvQDrsecnT+ljUqmdpwO77FSdOdAqcXNShhyNgf3VG4YkDqjQhGOluq4T6X6CpWaE07TUixWNMJngEe0aKnBAVT+ZJ0/RmVGGyA+leUKjufp7I8GBUrPAM5NZTrXsZeJ/XjfW/nU/YSKKNRVkcciPOdIhympAQyYp0XxmCCaSmayIjLHERJuySqYEd/nLq6RVq7oXVef+sly/yesowgmcQgVcuII63EEDmkBgCs/wCm9WYr1Y79bHYrRg5TvH8AfW5w89eJNj</latexit><latexit sha1_base64="Jjh+roV7RXCtHv15jiKBS8Q2kV8=">AAAB+XicbVDLSsNAFL2pr1pfUZduBotQNyXRgi6LblxWsA9oa5lMJ+3QySTMTCol5E/cuFDErX/izr9x0mahrQcGDufcyz1zvIgzpR3n2yqsrW9sbhW3Szu7e/sH9uFRS4WxJLRJQh7KjocV5UzQpmaa004kKQ48Ttve5Dbz21MqFQvFg55FtB/gkWA+I1gbaWDbvQDrsecnT+ljUqmdpwO77FSdOdAqcXNShhyNgf3VG4YkDqjQhGOluq4T6X6CpWaE07TUixWNMJngEe0aKnBAVT+ZJ0/RmVGGyA+leUKjufp7I8GBUrPAM5NZTrXsZeJ/XjfW/nU/YSKKNRVkcciPOdIhympAQyYp0XxmCCaSmayIjLHERJuySqYEd/nLq6R1UXUvq859rVy/yesowgmcQgVcuII63EEDmkBgCs/wCm9WYr1Y79bHYrRg5TvH8AfW5w9AhJNl</latexit><latexit sha1_base64="ZRrNh8WvBse2CToET/IVGFUdSXg=">AAAB+XicbVDLSsNAFL2pr1pfUZduBotQNyVRUZdFNy4r2Ae0tUymk3boZBJmJpUS8iduXCji1j9x5984abPQ1gMDh3Pu5Z45XsSZ0o7zbRVWVtfWN4qbpa3tnd09e/+gqcJYEtogIQ9l28OKciZoQzPNaTuSFAcepy1vfJv5rQmVioXiQU8j2gvwUDCfEayN1LftboD1yPOTp/QxqVyepn277FSdGdAycXNShhz1vv3VHYQkDqjQhGOlOq4T6V6CpWaE07TUjRWNMBnjIe0YKnBAVS+ZJU/RiVEGyA+leUKjmfp7I8GBUtPAM5NZTrXoZeJ/XifW/nUvYSKKNRVkfsiPOdIhympAAyYp0XxqCCaSmayIjLDERJuySqYEd/HLy6R5VnXPq879Rbl2k9dRhCM4hgq4cAU1uIM6NIDABJ7hFd6sxHqx3q2P+WjByncO4Q+szx9DkJNn</latexit>ĝ ĥ<latexit sha1_base64="MlFgSWDuhHul1vp77q5wm5HaDKU=">AAAB+XicbVDLSsNAFL2pr1pfUZduBovgqiQq6LLoxmUF+4AmlMl00g6dPJi5KZTQP3HjQhG3/ok7/8ZJm4VWDwwczrmXe+YEqRQaHefLqqytb2xuVbdrO7t7+wf24VFHJ5livM0SmaheQDWXIuZtFCh5L1WcRoHk3WByV/jdKVdaJPEjzlLuR3QUi1AwikYa2LY3pph7EcVxEOaj+Xxg152GswD5S9yS1KFEa2B/esOEZRGPkUmqdd91UvRzqlAwyec1L9M8pWxCR7xvaEwjrv18kXxOzowyJGGizIuRLNSfGzmNtJ5FgZksIupVrxD/8/oZhjd+LuI0Qx6z5aEwkwQTUtRAhkJxhnJmCGVKmKyEjamiDE1ZNVOCu/rlv6Rz0XAvG87DVb15W9ZRhRM4hXNw4RqacA8taAODKTzBC7xaufVsvVnvy9GKVe4cwy9YH988yZQL</latexit><latexit sha1_base64="gooRaEzmAkh1JPQL65zWVoIqQzM=">AAAB+XicbVDLSsNAFJ3UV62vqEs3g0VwVRIVdFl047KCfUATymQ6aYZOJmHmplBC/sSNC0Xc+ifu/BsnbRZaPTBwOOde7pkTpIJrcJwvq7a2vrG5Vd9u7Ozu7R/Yh0c9nWSKsi5NRKIGAdFMcMm6wEGwQaoYiQPB+sH0rvT7M6Y0T+QjzFPmx2QiecgpASONbNuLCOReTCAKwjwqipHddFrOAvgvcSvSRBU6I/vTGyc0i5kEKojWQ9dJwc+JAk4FKxpepllK6JRM2NBQSWKm/XyRvMBnRhnjMFHmScAL9edGTmKt53FgJsuIetUrxf+8YQbhjZ9zmWbAJF0eCjOBIcFlDXjMFaMg5oYQqrjJimlEFKFgymqYEtzVL/8lvYuWe9lyHq6a7duqjjo6QafoHLnoGrXRPeqgLqJohp7QC3q1cuvZerPel6M1q9o5Rr9gfXwDPk+UDA==</latexit>ĝ ĥ<latexit sha1_base64="MlFgSWDuhHul1vp77q5wm5HaDKU=">AAAB+XicbVDLSsNAFL2pr1pfUZduBovgqiQq6LLoxmUF+4AmlMl00g6dPJi5KZTQP3HjQhG3/ok7/8ZJm4VWDwwczrmXe+YEqRQaHefLqqytb2xuVbdrO7t7+wf24VFHJ5livM0SmaheQDWXIuZtFCh5L1WcRoHk3WByV/jdKVdaJPEjzlLuR3QUi1AwikYa2LY3pph7EcVxEOaj+Xxg152GswD5S9yS1KFEa2B/esOEZRGPkUmqdd91UvRzqlAwyec1L9M8pWxCR7xvaEwjrv18kXxOzowyJGGizIuRLNSfGzmNtJ5FgZksIupVrxD/8/oZhjd+LuI0Qx6z5aEwkwQTUtRAhkJxhnJmCGVKmKyEjamiDE1ZNVOCu/rlv6Rz0XAvG87DVb15W9ZRhRM4hXNw4RqacA8taAODKTzBC7xaufVsvVnvy9GKVe4cwy9YH988yZQL</latexit><latexit sha1_base64="gooRaEzmAkh1JPQL65zWVoIqQzM=">AAAB+XicbVDLSsNAFJ3UV62vqEs3g0VwVRIVdFl047KCfUATymQ6aYZOJmHmplBC/sSNC0Xc+ifu/BsnbRZaPTBwOOde7pkTpIJrcJwvq7a2vrG5Vd9u7Ozu7R/Yh0c9nWSKsi5NRKIGAdFMcMm6wEGwQaoYiQPB+sH0rvT7M6Y0T+QjzFPmx2QiecgpASONbNuLCOReTCAKwjwqipHddFrOAvgvcSvSRBU6I/vTGyc0i5kEKojWQ9dJwc+JAk4FKxpepllK6JRM2NBQSWKm/XyRvMBnRhnjMFHmScAL9edGTmKt53FgJsuIetUrxf+8YQbhjZ9zmWbAJF0eCjOBIcFlDXjMFaMg5oYQqrjJimlEFKFgymqYEtzVL/8lvYuWe9lyHq6a7duqjjo6QafoHLnoGrXRPeqgLqJohp7QC3q1cuvZerPel6M1q9o5Rr9gfXwDPk+UDA==</latexit><latexit sha1_base64="MlFgSWDuhHul1vp77q5wm5HaDKU=">AAAB+XicbVDLSsNAFL2pr1pfUZduBovgqiQq6LLoxmUF+4AmlMl00g6dPJi5KZTQP3HjQhG3/ok7/8ZJm4VWDwwczrmXe+YEqRQaHefLqqytb2xuVbdrO7t7+wf24VFHJ5livM0SmaheQDWXIuZtFCh5L1WcRoHk3WByV/jdKVdaJPEjzlLuR3QUi1AwikYa2LY3pph7EcVxEOaj+Xxg152GswD5S9yS1KFEa2B/esOEZRGPkUmqdd91UvRzqlAwyec1L9M8pWxCR7xvaEwjrv18kXxOzowyJGGizIuRLNSfGzmNtJ5FgZksIupVrxD/8/oZhjd+LuI0Qx6z5aEwkwQTUtRAhkJxhnJmCGVKmKyEjamiDE1ZNVOCu/rlv6Rz0XAvG87DVb15W9ZRhRM4hXNw4RqacA8taAODKTzBC7xaufVsvVnvy9GKVe4cwy9YH988yZQL</latexit><latexit sha1_base64="gooRaEzmAkh1JPQL65zWVoIqQzM=">AAAB+XicbVDLSsNAFJ3UV62vqEs3g0VwVRIVdFl047KCfUATymQ6aYZOJmHmplBC/sSNC0Xc+ifu/BsnbRZaPTBwOOde7pkTpIJrcJwvq7a2vrG5Vd9u7Ozu7R/Yh0c9nWSKsi5NRKIGAdFMcMm6wEGwQaoYiQPB+sH0rvT7M6Y0T+QjzFPmx2QiecgpASONbNuLCOReTCAKwjwqipHddFrOAvgvcSvSRBU6I/vTGyc0i5kEKojWQ9dJwc+JAk4FKxpepllK6JRM2NBQSWKm/XyRvMBnRhnjMFHmScAL9edGTmKt53FgJsuIetUrxf+8YQbhjZ9zmWbAJF0eCjOBIcFlDXjMFaMg5oYQqrjJimlEFKFgymqYEtzVL/8lvYuWe9lyHq6a7duqjjo6QafoHLnoGrXRPeqgLqJohp7QC3q1cuvZerPel6M1q9o5Rr9gfXwDPk+UDA==</latexit><latexit sha1_base64="wdyEarqCEeVHbD9YSbbId9H4C4Y=">AAAB+XicbVDLSsNAFL2pr1pfUZduBotQNyVRoS6LblxWsA9oa5lMJ+3QySTMTCol5E/cuFDErX/izr9x0mahrQcGDufcyz1zvIgzpR3n2yqsrW9sbhW3Szu7e/sH9uFRS4WxJLRJQh7KjocV5UzQpmaa004kKQ48Ttve5Dbz21MqFQvFg55FtB/gkWA+I1gbaWDbvQDrsecnT+ljUqmdpwO77FSdOdAqcXNShhyNgf3VG4YkDqjQhGOluq4T6X6CpWaE07TUixWNMJngEe0aKnBAVT+ZJ0/RmVGGyA+leUKjufp7I8GBUrPAM5NZTrXsZeJ/XjfW/nU/YSKKNRVkcciPOdIhympAQyYp0XxmCCaSmayIjLHERJuySqYEd/nLq6R1UXUvq879Vbl+k9dRhBM4hQq4UIM63EEDmkBgCs/wCm9WYr1Y79bHYrRg5TvH8AfW5w9FFpNo</latexit><latexit sha1_base64="0AX6CIJdpG7lxM3dZNYUgXYXyYA=">AAAB+XicbVDLSsNAFL2pr1pfUZduBotQNyVRwS6LblxWsA9oa5lMJ+3QySTMTCol5E/cuFDErX/izr9x0mahrQcGDufcyz1zvIgzpR3n2yqsrW9sbhW3Szu7e/sH9uFRS4WxJLRJQh7KjocV5UzQpmaa004kKQ48Ttve5Dbz21MqFQvFg55FtB/gkWA+I1gbaWDbvQDrsecnT+ljUqmdpwO77FSdOdAqcXNShhyNgf3VG4YkDqjQhGOluq4T6X6CpWaE07TUixWNMJngEe0aKnBAVT+ZJ0/RmVGGyA+leUKjufp7I8GBUrPAM5NZTrXsZeJ/XjfWfq2fMBHFmgqyOOTHHOkQZTWgIZOUaD4zBBPJTFZExlhiok1ZJVOCu/zlVdK6qLqXVef+qly/yesowgmcQgVcuIY63EEDmkBgCs/wCm9WYr1Y79bHYrRg5TvH8AfW5w9GnJNp</latexit><latexit sha1_base64="gooRaEzmAkh1JPQL65zWVoIqQzM=">AAAB+XicbVDLSsNAFJ3UV62vqEs3g0VwVRIVdFl047KCfUATymQ6aYZOJmHmplBC/sSNC0Xc+ifu/BsnbRZaPTBwOOde7pkTpIJrcJwvq7a2vrG5Vd9u7Ozu7R/Yh0c9nWSKsi5NRKIGAdFMcMm6wEGwQaoYiQPB+sH0rvT7M6Y0T+QjzFPmx2QiecgpASONbNuLCOReTCAKwjwqipHddFrOAvgvcSvSRBU6I/vTGyc0i5kEKojWQ9dJwc+JAk4FKxpepllK6JRM2NBQSWKm/XyRvMBnRhnjMFHmScAL9edGTmKt53FgJsuIetUrxf+8YQbhjZ9zmWbAJF0eCjOBIcFlDXjMFaMg5oYQqrjJimlEFKFgymqYEtzVL/8lvYuWe9lyHq6a7duqjjo6QafoHLnoGrXRPeqgLqJohp7QC3q1cuvZerPel6M1q9o5Rr9gfXwDPk+UDA==</latexit><latexit sha1_base64="MlFgSWDuhHul1vp77q5wm5HaDKU=">AAAB+XicbVDLSsNAFL2pr1pfUZduBovgqiQq6LLoxmUF+4AmlMl00g6dPJi5KZTQP3HjQhG3/ok7/8ZJm4VWDwwczrmXe+YEqRQaHefLqqytb2xuVbdrO7t7+wf24VFHJ5livM0SmaheQDWXIuZtFCh5L1WcRoHk3WByV/jdKVdaJPEjzlLuR3QUi1AwikYa2LY3pph7EcVxEOaj+Xxg152GswD5S9yS1KFEa2B/esOEZRGPkUmqdd91UvRzqlAwyec1L9M8pWxCR7xvaEwjrv18kXxOzowyJGGizIuRLNSfGzmNtJ5FgZksIupVrxD/8/oZhjd+LuI0Qx6z5aEwkwQTUtRAhkJxhnJmCGVKmKyEjamiDE1ZNVOCu/rlv6Rz0XAvG87DVb15W9ZRhRM4hXNw4RqacA8taAODKTzBC7xaufVsvVnvy9GKVe4cwy9YH988yZQL</latexit>ĝ ĥ<latexit sha1_base64="MlFgSWDuhHul1vp77q5wm5HaDKU=">AAAB+XicbVDLSsNAFL2pr1pfUZduBovgqiQq6LLoxmUF+4AmlMl00g6dPJi5KZTQP3HjQhG3/ok7/8ZJm4VWDwwczrmXe+YEqRQaHefLqqytb2xuVbdrO7t7+wf24VFHJ5livM0SmaheQDWXIuZtFCh5L1WcRoHk3WByV/jdKVdaJPEjzlLuR3QUi1AwikYa2LY3pph7EcVxEOaj+Xxg152GswD5S9yS1KFEa2B/esOEZRGPkUmqdd91UvRzqlAwyec1L9M8pWxCR7xvaEwjrv18kXxOzowyJGGizIuRLNSfGzmNtJ5FgZksIupVrxD/8/oZhjd+LuI0Qx6z5aEwkwQTUtRAhkJxhnJmCGVKmKyEjamiDE1ZNVOCu/rlv6Rz0XAvG87DVb15W9ZRhRM4hXNw4RqacA8taAODKTzBC7xaufVsvVnvy9GKVe4cwy9YH988yZQL</latexit><latexit sha1_base64="gooRaEzmAkh1JPQL65zWVoIqQzM=">AAAB+XicbVDLSsNAFJ3UV62vqEs3g0VwVRIVdFl047KCfUATymQ6aYZOJmHmplBC/sSNC0Xc+ifu/BsnbRZaPTBwOOde7pkTpIJrcJwvq7a2vrG5Vd9u7Ozu7R/Yh0c9nWSKsi5NRKIGAdFMcMm6wEGwQaoYiQPB+sH0rvT7M6Y0T+QjzFPmx2QiecgpASONbNuLCOReTCAKwjwqipHddFrOAvgvcSvSRBU6I/vTGyc0i5kEKojWQ9dJwc+JAk4FKxpepllK6JRM2NBQSWKm/XyRvMBnRhnjMFHmScAL9edGTmKt53FgJsuIetUrxf+8YQbhjZ9zmWbAJF0eCjOBIcFlDXjMFaMg5oYQqrjJimlEFKFgymqYEtzVL/8lvYuWe9lyHq6a7duqjjo6QafoHLnoGrXRPeqgLqJohp7QC3q1cuvZerPel6M1q9o5Rr9gfXwDPk+UDA==</latexit>

Learning rate
Momentum rate
Exp. moving average rate
Prior precision
External damping factor
Tempering parameter# MC samples for training
Data augmentation factor

↵
1
2

⌧
K
⇢

Figure 2: A pseudo-code for our distributed VOGN algorithm is shown in Algorithm 1, and the
distributed scheme is shown in the right figure. The computation in line 10 requires an extra
calculation (see Appendix B), making VOGN slower than Adam. The bottom table gives a list ofalgorithmic hyperparameters needed for VOGN.
Momentum and initialisation: It is well known that both momentum and good initialisation can
improve the speed of convergence for SG methods in deep learning [51]. Since VOGN is similar
to Adam, we can implement momentum in a similar way. This is shown in step 17 of Algorithm 1,
where 1 is the momentum rate. We initialise the mean µ in the same way the weights are initialised
in Adam (we use init.xavier_normal in PyTorch [11]). For the momentum term m, we use thesame initialisation as Adam (initialised to 0). VOGN requires an additional initialisation for thevariance 2 . For this, we first run a forward pass through the first minibatch, calculate the average ofthe squared gradients and initialise the scale s0 with it (see step 1 in Algorithm 1). This implies thatthe variance is initialised to 20 = ⌧ /(N (s0 + ˜)). For the tempering parameter ⌧ , we use a schedule
where it is increased from a small value (e.g., 0.1) to 1. With these initialisation protocols, VOGN isable to mimic the convergence behaviour of Adam in the beginning.
Learning rate scheduling: A common approach to quickly achieve high validation accuracies is touse a specific learning rate schedule [13]. The learning rate (denoted by ↵ in Algorithm 1) is regularly
decayed by a factor (typically a factor of 10). The frequency and timings of this decay are usually
pre-specified. In VOGN, we use the same schedule used for Adam, which works well.
Distributed training: We also employ distributed training for VOGN to perform large experiments
quickly. We can parallelise computation both over data and Monte-Carlo (MC) samples. Data
parallelism is useful to split up large minibatch sizes. This is followed by averaging over multiple
MC samples and their losses on a single GPU. MC sample parallelism is useful when minibatch size
is small, and we can copy the entire minibatch and process it on a single GPU. Algorithm 1 andFigure 2 illustrate our distributed scheme. We use a combination of these two parallelism techniqueswith different MC samples for different inputs. This theoretically reduces the variance during training(see Equation 5 in Kingma et al. [26]), but sometimes requires averaging over multiple MC samples
to get a sufficiently low variance in the early iterations. Overall, we find that this type of distributed
training is essential for fast training on large problems such as ImageNet.
Implementation of the Gauss-Newton update in VOGN: As discussed earlier, VOGN uses theGauss-Newton approximation, which is fundamentally different from Adam. In this approximation,
the gradients on individual data examples are first squared and then averaged afterwards (see step5

12 in Algorithm 1 which implements the update for st shown in (4)). We need extra computation
to get access to individual gradients, due to which, VOGN is slower Adam or SGD (e.g., in Fig. 1).
However, this is not a theoretical limitation and this can be improved if a framework enables an easycomputation of the individual gradients. Details of our implementation are described in Appendix
B. This implementation is much more efficient than a naive one where gradients over examples are
stored and the sum over the square is computed sequentially. Our implementation usually brings therunning time of VOGN to within 2-5 times of the time that Adam takes.
Tuning VOGN: Currently, there is no common recipe for tuning the algorithmic hyperparametersfor VI, especially for large-scale tasks like ImageNet classification. One key idea we use in our
experiments is to start with Adam hyperparameters and then make sure that VOGN training closely
follows an Adam-like trajectory in the beginning of training. To achieve this, we divide the tuning
into an optimisation part and a regularisation part. In the optimisation part, we first tune thehyperparameters of a deterministic version of VOGN, called the online Gauss-Newton (OGN)
method. This method, described in Appendix C, is more stable than VOGN since it does not require
MC sampling, and can be used as a stepping stone when moving from Adam/SGD to VOGN. Afterreaching a competitive performance to Adam/SGD by OGN, we move to the regularisation part,
where we tune the prior precision , the tempering parameter ⌧ , and the number of MC samples K for
VOGN. We initialise our search by setting the prior precision using the L2-regularisation parameterused for OGN, as well as the dataset size N . Another technique is to warm-up the parameter ⌧towards ⌧ = 1 (also see the “momentum and initialisation" part). Setting ⌧ to smaller values usuallystabilises the training, and increasing it slowly also helps during tuning. We also add an external
damping factor > 0 to the moving average st . This increases the lower bound of the eigenvalues ofthe diagonal covariance ⌃t and prevents the noise and the step size from becoming too large. We
find that a mix of these techniques works well for the problems we considered.

4

Experiments

In this section, we present experiments on fitting several deep networks on CIFAR-10 and ImageNet. Our experiments demonstrate practical training using VOGN on these benchmarks and show
performance that is competitive with Adam and SGD. We also assess the quality of the posterior
approximation, finding that benefits of Bayesian principles are preserved.
CIFAR-10 [28] contains 10 classes with 50,000 images for training and 10,000 images for validation.
For ImageNet, we train with 1.28 million training examples and validate on 50,000 examples,classifying between 1,000 classes. We used a large minibatch size M = 4, 096 and parallelise themacross 128 GPUs (NVIDIA Tesla P100). We compare the following methods on CIFAR-10: Adam,
MC-dropout [9]. For ImageNet, we also compare to SGD, K-FAC, and Noisy K-FAC. We do notconsider Noisy K-FAC for other comparisons since tuning is difficult. We compare 3 architectures:
LeNet-5, AlexNet, ResNet-18. We only compare to Bayes by Backprop (BBB) [4] for CIFAR-10
with LeNet-5 since it is very slow to converge for larger-scale experiments. We carefully set thehyperparameters of all methods, following the best practice of large distributed training [13] as theinitial point of our hyperparameter tuning. The full set of hyperparameters is in Appendix D.
4.1 Performance on CIFAR-10 and ImageNet
We start by showing the effectiveness of momentum and batch normalisation for boosting theperformance of VOGN. Figure 3a shows that these methods significantly speed up convergence andperformance (in terms of both accuracy and log likelihoods).
Figures 1 and 4 compare the convergence of VOGN to Adam (for all experiments), SGD (onImageNet), and MC-dropout (on the rest). VOGN shows similar convergence and its performanceis competitive with these methods. We also try BBB on LeNet-5, where it converges prohibitively
slowly, performing very poorly. We are not able to successfully train other architectures using thisapproach. We found it far simpler to tune VOGN because we can borrow all the techniques used forAdam. Figure 4 also shows the importance of DA in improving performance.
Table 1 gives a final comparison of train/validation accuracies, negative log likelihoods, epochs
required for convergence, and run-time per epoch. We can see that the accuracy, log likelihoods,
and the number of epochs are comparable. VOGN is 2-5 times slower than Adam and SGD. This6

(a) Effect of momentum and batch normalisation.

(b) Continual Learning

Figure 3: Figure (a) shows that momentum and batch normalisation improve the performance of
VOGN. The results are for training ResNet-18 on CIFAR-10. Figure (b) shows comparison for acontinual-learning task on the Permuted MNIST dataset. VOGN performs at least as well (average
accuracy) as VCL over 10 tasks. We also find that, for each task, VOGN converges much faster,taking only 100 epochs per task as opposed to 800 epochs taken by VCL (plots not shown).

Figure 4: Validation accuracy for various architectures trained on CIFAR-10 (DA: Data Augmentation). VOGN’s convergence and validation accuracies are comparable to Adam and MC-dropout.
is mainly due to the computation of individual gradients required in VOGN (see the discussion in
Section 3). We clearly see that by using deep-learning techniques on VOGN, we can perform practical
deep learning. This is not possible with methods such as BBB.
Due to the Bayesian nature of VOGN, there are some trade-offs to consider. Reducing the priorprecision ( in Algorithm 1) results in higher validation accuracy, but also larger train-test gap (moreoverfitting). This is shown in Appendix E for VOGN on ResNet-18 on ImageNet. As expected,
when the prior precision is small, performance is similar to non-Bayesian methods. We also show
the effect of changing the effective dataset size ⇢ in Appendix E: note that, since we are going totune the prior variance anyway, it is sufficient to set ⇢ to its correct order of magnitude. Anothertrade-off concerns the number of Monte-Carlo (MC) samples, shown in Appendix F. Increasing thenumber of training MC samples (up to a limit) improves VOGN’s convergence rate and stability,
but also increases the computation. Increasing the number of MC samples during testing improves
generalisation, as expected due to averaging.
Finally, a few comments on the performance of the other methods. Adam regularly overfits thetraining set in most settings, with large train-test differences in both validation accuracy and log
likelihood. One exception is LeNet-5, which is most likely due to the small architecture which resultsin underfitting (this is consistent with the low validation accuracies obtained). In contrast to Adam,
MC-dropout has small train-test gap, usually smaller than VOGN’s. However, we will see in Section
4.2 that this is because of underfitting. Moreover, the performance of MC-dropout is highly sensitive
to the dropout rate (see Appendix G for a comparison of different dropout rates). On ImageNet, Noisy
K-FAC performs well too. It is slower than VOGN, but it takes fewer epochs. Overall, wall clock
time is about the same as VOGN.
4.2 Quality of the Predictive Probabilities
In this section, we compare the quality of the predictive probabilities for various methods. For
Bayesian methods, we compute these probabilities by averaging over the samples from the posterior
approximations (see Appendix H for details). For non-Bayesian methods, these are obtained using the7Dataset/
Architecture

Optimiser

Train/Validation
Accuracy (%)

Validation
NLL

EpochsTime/epoch (s)

ECE

AUROC

CIFAR-10/
LeNet-5(no DA)Adam
BBB
MC-dropout
VOGN

71.98 / 67.67
66.84 / 64.61
68.41 / 67.65
70.79 / 67.32

0.937
1.018
0.99
0.938

210
800
210
210

6.96
11.43†
6.95
18.33

0.021
0.045
0.087
0.046

0.794
0.784
0.797
0.8

CIFAR-10/
AlexNet(no DA)Adam
MC-dropout
VOGN

100.0 / 67.94
97.56 / 72.20
79.07 / 69.03

2.83
1.077
0.93

161
160
160

3.12
3.25
9.98

0.262
0.140
0.024

0.793
0.818
0.796

CIFAR-10/
AlexNetAdam
MC-dropout
VOGN

97.92 / 73.59
80.65 / 77.04
81.15 / 75.48

1.480
0.667
0.703

161
160
160

3.08
3.20
10.02

0.262
0.114
0.016

0.793
0.828
0.832

CIFAR-10/
ResNet-18

Adam
MC-dropout
VOGN

97.74 / 86.00
88.23 / 82.85
91.62 / 84.27

0.55
0.51
0.477

160
161
161

11.97
12.51
53.14

0.082
0.166
0.040

0.877
0.768
0.876

SGD
Adam
MC-dropout
OGN
VOGN
K-FAC
Noisy K-FAC

82.63 / 67.79
80.96 / 66.39
72.96 / 65.64
85.33 / 65.76
73.87 / 67.38
83.73 / 66.58
72.28 / 66.44

1.38
1.44
1.43
1.60
1.37
1.493
1.44

90
90
90
90
90
60
60

44.13
44.40
45.86
63.13
76.04
133.69179.27

0.067
0.064
0.012
0.128
0.029
0.158
0.080

0.856
0.855
0.856
0.854
0.854
0.842
0.852

ImageNet/ResNet-18

Table 1: Performance comparisons on different dataset/architecture combinations. Out of the 15
metrics (NLL, ECE, and AUROC on 5 dataset/architecture combinations), VOGN performs the best
or tied best on 10 ,and is second-best on the other 5. Here DA means ‘Data Augmentation’, NLL
refers to ‘Negative Log Likelihood’ (lower is better), ECE refers to ‘Expected Calibration Error’(lower is better), AUROC refers to ‘Area Under ROC curve’ (higher is better). BBB is the Bayes
By Backprop method. For ImageNet, the reported accuracy and negative log likelihood are themedian value from the final 5 epochs. All hyperparameter settings are in Appendix D. See Table 3 forstandard deviations. † BBB is not parallelised (other methods have 4 processes), with 1 MC sample
used for the convolutional layers (VOGN uses 6 samples per process).

point estimate of the weights. We compare the probabilities using the following metrics: validation
negative log-likelihood (NLL), area under ROC (AUROC) and expected calibration curves (ECE)[40, 15]. For the first and third metric, a lower number is better, while for the second, a higher
number is better. See Appendix H for an explanation of these metrics. Results are summarised inTable 1. VOGN’s uncertainty performance is more consistent and marginally better than the othermethods, as expected from a more principled Bayesian method. Out of the 15 metrics (NLL, ECE
and AUROC on 5 dataset/architecture combinations), VOGN performs the best or tied best on 10,
and is second-best on the other 5. In contrast, both MC-dropout’s and Adam’s performance varies
significantly, sometimes performing poorly, sometimes performing decently. MC-dropout is best on 4,
and Adam is best on 1 (on LeNet-5; as argued earlier, the small architecture may result in underfitting).
We also show calibration curves [7] in Figures 1 and 14. Adam is consistently over-confident, with
its calibration curve below the diagonal. Conversely, MC-dropout is usually under-confident. On
ImageNet, MC-dropout performs well on ECE (all methods are very similar on AUROC), but thisrequired an excessively tuned dropout rate (see Appendix G).
We also compare performance on out-of-distribution datasets. When testing on datasets that are
different from the training datasets, predictions should be more uncertain. We use experimental
protocol from the literature [16, 31, 8, 32] to compare VOGN, Adam and MC-dropout on CIFAR-10.
We also borrow metrics from other works [16, 30], showing predictive entropy histograms and also
reporting AUROC and FPR at 95% TPR. See Appendix I for further details on the datasets and metrics.
Ideally, we want predictive entropy to be high on out-of-distribution data and low on in-distribution
data. Our results are summarised in Figure 5 and Appendix I. On ResNet-18 and AlexNet, VOGN’s
predictive entropy histograms show the desired behaviour: a spread of entropies for the in-distribution
data, and high entropies for out-of-distribution data. Adam has many predictive entropies at zero,8

Figure 5: Histograms of predictive entropy for out-of-distribution tests for ResNet-18 trained on
CIFAR-10. Going from left to right, the inputs are: the in-distribution dataset (CIFAR-10), followedby out-of-distribution data: SVHN, LSUN (crop), LSUN (resize). Also shown are the FPR at 95%
TPR metric (lower is better) and the AUROC metric (higher is better), averaged over 3 runs. We
clearly see that VOGN’s predictive entropy is generally low for in-distribution and high for out-ofdistribution data, but this is not the case for other methods. Solid vertical lines indicate the mean
predictive entropy. The standard deviations are small and therefore not reported.
indicating Adam tends to classify out-of-distribution data too confidently. Conversely, MC-dropout’s
predictive entropies are generally high (particularly in-distribution), indicating MC-dropout has toomuch noise. On LeNet-5, we observe the same result as before: Adam and MC-dropout both perform
well. The metrics (AUROC and FPR at 95% TPR) do not provide a clear story across architectures.
4.2.1

Performance on a Continual-learning task

The goal of continual learning is to avoid forgetting of old tasks while sequentially observing newtasks. The past tasks are never visited again, making it difficult to remember them. The fieldof continual learning has recently grown, with many approaches proposed to tackle this problem[27, 33, 43, 48, 50]. Most approaches consider a simple setting where the tasks (such as classifying a
subset of classes) arrive sequentially, and all the data from that task is available. We consider thesame setup in our experiments.
We compare to Elastic Weight Consolidation (EWC) [27] and a VI-based approach called Variational
Continual Learning (VCL) [43]. VCL employs BBB for each task, and we expect to boost its
performance by replacing BBB by VOGN. Figure 3b shows results on a common benchmark called
Permuted MNIST. We use the same experimental setup as in Swaroop et al. [52]. In Permuted
MNIST, each task consists of the entire MNIST dataset (10-way classification) with a different fixed
random permutation applied to the input images’ pixels. We run each method 20 times, with different
random seeds for both the benchmark’s permutations and model training. See Appendix D.2 for
hyperparameter settings and further details. We see that VOGN performs at least as well as VCL,
and far better than a popular approach called EWC [27]. Additionally, as found in the batch learning
setting, VOGN is much quicker than BBB: we run VOGN for only 100 epochs per task, whereas
VCL requires 800 epochs per task to achieve best results [52].

5Conclusions

We successfully train deep networks with a natural-gradient variational inference method, VOGN,
on a variety of architectures and datasets, even scaling up to ImageNet. This is made possible dueto the similarity of VOGN to Adam, enabling us to boost performance by borrowing deep-learningtechniques. Our accuracies and convergence rates are comparable to SGD and Adam. Unlike them,
however, VOGN retains the benefits of Bayesian principles, with well-calibrated uncertainty andgood performance on out-of-distribution data. Better uncertainty estimates open up a whole rangeof potential future experiments, for example, small data experiments, active learning, adversarial
experiments, and sequential decision making. Our results on a continual-learning task confirm this.
Another potential avenue for research is to consider structured covariance approximations.

9Acknowledgements
We would like to thank Hikaru Nakata (Tokyo Institute of Technology) and Ikuro Sato (Denso IT
Laboratory, Inc.) for their help on the PyTorch implementation. We are also thankful for the RAIDEN
computing system and its support team at the RIKEN Center for AI Project which we used extensively
for our experiments. This research used computational resources of the HPCI system provided by
Tokyo Institute of Technology (TSUBAME3.0) through the HPCI System Research Project (Project
ID:hp190122). K. O. is a Research Fellow of JSPS and is supported by JSPS KAKENHI Grant
Number JP19J13477.

References[1] James R Anderson and Carsten Peterson. A mean field theory learning algorithm for neuralnetworks. Complex Systems, 1:995–1019, 1987.
[2] David Barber and Christopher M Bishop. Ensemble learning in Bayesian neural networks.
Generalization in Neural Networks and Machine Learning, 168:215–238, 1998.
[3] Christopher M. Bishop. Pattern Recognition and Machine Learning (Information Science andStatistics). Springer-Verlag, Berlin, Heidelberg, 2006. ISBN 0387310738.
[4] Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty
in neural networks. In International Conference on Machine Learning, pages 1613–1622, 2015.
[5] Léon Bottou, Frank E Curtis, and Jorge Nocedal. Optimization methods for large-scale machinelearning. arXiv preprint arXiv:1606.04838, 2016.
[6] John Bradshaw, Alexander G de G Matthews, and Zoubin Ghahramani. Adversarial examples,
uncertainty, and transfer testing robustness in Gaussian process hybrid deep networks. arXiv
preprint arXiv:1707.02476, 2017.
[7] Morris H. DeGroot and Stephen E. Fienberg. The comparison and evaluation of forecasters.
The Statistician: Journal of the Institute of Statisticians, 32:12–22, 1983.
[8] Terrance DeVries and Graham W. Taylor. Learning confidence for out-of-distribution detection
in neural networks. arXiv preprint arXiv:1802.04865, 2018.
[9] Yarin Gal and Zoubin Ghahramani. Dropout as a Bayesian approximation: Representing
model uncertainty in deep learning. In International Conference on Machine Learning, pages
1050–1059, 2016.
[10] S. Ghosal and A. Van der Vaart. Fundamentals of nonparametric Bayesian inference, volume 44.
Cambridge University Press, 2017.
[11] Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial
intelligence and statistics, pages 249–256, 2010.
[12] Ian Goodfellow. Efficient Per-Example Gradient Computations. ArXiv e-prints, October 2015.
[13] Priya Goyal, Piotr Dollár, Ross B. Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo
Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch SGD:
training imagenet in 1 hour. CoRR, abs/1706.02677, 2017.
[14] Alex Graves. Practical variational inference for neural networks. In Advances in Neural
Information Processing Systems, pages 2348–2356, 2011.
[15] Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration of modern neuralnetworks. In Proceedings of the 34th International Conference on Machine Learning-Volume
70, pages 1321–1330. JMLR. org, 2017.
[16] Dan Hendrycks and Kevin Gimpel. A baseline for detecting misclassified and out-of-distribution
examples in neural networks. In International Conference on Learning Representations, 2017.
10

[17] Geoffrey Hinton, Li Deng, Dong Yu, George Dahl, Abdel-rahman Mohamed, Navdeep Jaitly,
Andrew Senior, Vincent Vanhoucke, Patrick Nguyen, Brian Kingsbury, et al. Deep neuralnetworks for acoustic modeling in speech recognition. IEEE Signal processing magazine, 29,
2012.
[18] Geoffrey E Hinton and Drew Van Camp. Keeping the neural networks simple by minimizing thedescription length of the weights. In Annual Conference on Computational Learning Theory,
pages 5–13, 1993.
[19] Jennifer A Hoeting, David Madigan, Adrian E Raftery, and Chris T Volinsky. Bayesian model
averaging: a tutorial. Statistical science, pages 382–401, 1999.
[20] Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training
by reducing internal covariate shift. CoRR, abs/1502.03167, 2015. URL http://arxiv.org/
abs/1502.03167.
[21] Mohammad Khan. Variational learning for latent Gaussian model of discrete data. PhD thesis,
University of British Columbia, 2012.
[22] Mohammad Emtiyaz Khan and Wu Lin. Conjugate-computation variational inference: converting variational inference in non-conjugate models to inferences in conjugate models. In
International Conference on Artificial Intelligence and Statistics, pages 878–887, 2017.
[23] Mohammad Emtiyaz Khan and Didrik Nielsen. Fast yet simple natural-gradient descent forvariational inference in complex models. In 2018 International Symposium on Information
Theory and Its Applications (ISITA), pages 31–35. IEEE, 2018.
[24] Mohammad Emtiyaz Khan, Didrik Nielsen, Voot Tangkaratt, Wu Lin, Yarin Gal, and Akash
Srivastava. Fast and scalable Bayesian deep learning by weight-perturbation in Adam. In
International Conference on Machine Learning, pages 2616–2625, 2018.
[25] Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International
Conference on Learning Representations, 2015.
[26] Diederik P Kingma, Tim Salimans, and Max Welling. Variational dropout and the local
reparameterization trick. In Advances in Neural Information Processing Systems, pages 2575–
2583, 2015.
[27] James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins,
Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska-Barwinska, et al.
Overcoming catastrophic forgetting in neural networks. Proceedings of the national academy ofsciences, 114(13):3521–3526, 2017.
[28] Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images.
Technical report, Citeseer, 2009.
[29] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep
convolutional neural networks. In Advances in neural information processing systems, pages
1097–1105, 2012.
[30] Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable
predictive uncertainty estimation using deep ensembles. In Advances in Neural Information
Processing Systems 30, pages 6402–6413. Curran Associates, Inc., 2017.
[31] Kimin Lee, Honglak Lee, Kibok Lee, and Jinwoo Shin. Training confidence-calibrated classifiers for detecting out-of-distribution samples. In International Conference on Learning
Representations, 2018.
[32] Shiyu Liang, Yixuan Li, and R. Srikant. Enhancing the reliability of out-of-distribution image
detection in neural networks. In International Conference on Learning Representations, 2018.
[33] David Lopez-Paz and Marc Aurelio Ranzato. Gradient episodic memory for continual learning.
In NIPS, 2017.
11

[34] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International
Conference on Learning Representations, 2019.
[35] David Mackay. Bayesian Methods for Adaptive Models. PhD thesis, California Institute ofTechnology, 1991.
[36] David JC MacKay. Information theory, inference and learning algorithms. Cambridge university
press, 2003.
[37] Wesley Maddox, Timur Garipov, Pavel Izmailov, Dmitry Vetrov, and Andrew Gordon Wilson.
A simple baseline for Bayesian uncertainty in deep learning. arXiv preprint arXiv:1902.02476,
2019.
[38] Stephan Mandt, Matthew D Hoffman, and David M Blei. Stochastic gradient descent as
approximate Bayesian inference. Journal of Machine Learning Research, 18:1–35, 2017.
[39] Tomas Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean. Efficient estimation of word
representations in vector space. arXiv preprint arXiv:1301.3781, 2013.
[40] Mahdi Pakdaman Naeini, Gregory F. Cooper, and Milos Hauskrecht. Obtaining well calibrated
probabilities using Bayesian binning. In Proceedings of the Twenty-Ninth AAAI Conference on
Artificial Intelligence, AAAI’15, pages 2901–2907. AAAI Press, 2015.
[41] Redford M Neal. Bayesian learning for neural networks. PhD thesis, University of Toronto,
1995.
[42] Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y. Ng.
Reading digits in natural images with unsupervised feature learning. In NIPS Workshop on
Deep Learning and Unsupervised Feature Learning 2011, 2011.
[43] Cuong V Nguyen, Yingzhen Li, Thang D Bui, and Richard E Turner. Variational continual
learning. arXiv preprint arXiv:1710.10628, 2017.
[44] Kazuki Osawa, Yohei Tsuji, Yuichiro Ueno, Akira Naruse, Rio Yokota, and Satoshi Matsuoka.
Second-order optimization method for large mini-batch: Training resnet-50 on imagenet in 35
epochs. CoRR, abs/1811.12019, 2018.
[45] Carlos Riquelme, George Tucker, and Jasper Snoek. Deep Bayesian bandits showdown: An
empirical comparison of Bayesian deep networks for Thompson sampling. arXiv preprint
arXiv:1802.09127, 2018.
[46] Hippolyt Ritter, Aleksandar Botev, and David Barber. A scalable Laplace approximation forneural networks. In International Conference on Learning Representations, 2018.
[47] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng
Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. Imagenet large scale visual
recognition challenge. International journal of computer vision, 115(3):211–252, 2015.
[48] Andrei A. Rusu, Neil C. Rabinowitz, Guillaume Desjardins, Hubert Soyer, James Kirkpatrick,Koray Kavukcuoglu, Razvan Pascanu, and Raia Hadsell. Progressive neural networks. arXiv
preprint arXiv:1606.04671, 2016.
[49] Lawrence K Saul, Tommi Jaakkola, and Michael I Jordan. Mean field theory for sigmoid belief
networks. Journal of Artificial Intelligence Research, 4:61–76, 1996.
[50] Jonathan Schwarz, Jelena Luketina, Wojciech M. Czarnecki, Agnieszka Grabska-Barwinska,
Yee Whye Teh, Razvan Pascanu, and Raia Hadsell. Progress & compress: A scalable framework
for continual learning. In International Conference on Machine Learning, 2018.
[51] Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance ofinitialization and momentum in deep learning. In International conference on machine learning,
pages 1139–1147, 2013.
[52] Siddharth Swaroop, Cuong V. Nguyen, Thang D. Bui, and Richard E. Turner. Improving andunderstanding variational continual learning. arXiv preprint arXiv:1905.02099, 2019.
12

[53] Tijmen Tieleman and Geoffrey Hinton. Lecture 6.5-RMSprop: Divide the gradient by a runningaverage of its recent magnitude. COURSERA: Neural Networks for Machine Learning 4, 2012.
[54] V. G. Vovk. Aggregating strategies. In Proceedings of the Third Annual Workshop on Computational Learning Theory, COLT ’90, pages 371–386, San Francisco, CA, USA, 1990. Morgan
Kaufmann Publishers Inc. ISBN 1-55860-146-5.
[55] Fisher Yu, Yinda Zhang, Shuran Song, Ari Seff, and Jianxiong Xiao. LSUN: construction of alarge-scale image dataset using deep learning with humans in the loop. CoRR, abs/1506.03365,
2015.
[56] Guodong Zhang, Shengyang Sun, David K. Duvenaud, and Roger B. Grosse. Noisy natural
gradient as variational inference. arXiv preprint arXiv:1712.02390, 2018.

13