Source-linked AI summary
Practical Deep Learning with Bayesian Principles
Kazuki Osawa, Siddharth Swaroop, Anirudh Jain, Runa Eschenhagen, Richard E. Turner, Rio Yokota, Mohammad Emtiyaz Khan
TL;DR
Principled Bayesian variational inference has been impractical for large deep networks. This paper trains networks with natural-gradient variational inference, achieving Adam-like performance in about the same number of epochs while retaining calibrated probabilities, improved out-of-distribution uncertainty, and boosted continual-learning performance.
Problem
Method
Results
Takeaways & Limitations
The improved uncertainty estimates motivate further experiments in small-data settings, active learning, adversarial research, and sequential decision making.
Source: §5 · p. 9VOGN runs slower than Adam or SGD because it requires extra computation for individual gradients, though its implementation typically takes 2–5× Adam’s runtime.
Source: §3 · p. 6
Abstract
from arXiv · showhide
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 batch normalisation, 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 out-of-distribution data are improved, and continual-learning performance is boosted. This work enables practical deep learning while preserving benefits of Bayesian principles. A PyTorch implementation is available as a plug-and-play optimiser.
1 Introduction
The paper argues that natural-gradient variational inference can make Bayesian deep learning practical while retaining Bayesian benefits. It applies standard deep-learning techniques to train networks at a pace comparable to Adam and reports improved calibration, out-of-distribution uncertainty, and continual learning.
- Deep learning can require abundant high-quality data, overfit small datasets, forget past knowledge during sequential learning, and lack reliable confidence estimates, making high-stakes applications challenging.Source: §1 · p. 1
- Bayesian inference can represent uncertainty, support sequential learning through Bayes’ rule, and reduce overfitting through model averaging, but posterior computation has hindered practical use.Less principled approximations such as MC-dropout are limited for continual learning and sequential decision making, and their rigid posterior forms are difficult to improve.Source: §1 · pp. 1–2
- VOGN reaches 67.38% validation accuracy versus 66.39% for Adam and 67.79% for SGD, with similar convergence in about the same epochs; its runtime is 76 seconds per epoch versus 44 seconds for Adam and SGD.These ImageNet results are for ResNet-18.Source: Figure 1 · p. 2, §1 · p. 2
- Compared with standard deep-learning methods, the approximate-posterior training preserves Bayesian benefits: well-calibrated predictive probabilities, improved uncertainty on out-of-distribution inputs, and boosted continual-learning performance.Source: §1 · p. 2
- Natural-gradient VI resembles Adam, allowing the use of familiar initialization, momentum, batch normalization, data augmentation, and distributed training techniques.The paper reports similar performance in about the same number of epochs as Adam across LeNet, AlexNet, and ResNet on datasets including CIFAR-10 and ImageNet.Source: §1 · p. 2
Related work
- Earlier VI methods by Graves and Blundell et al. [4] required substantial implementation and tuning; Zhang et al. [56] showed Noisy K-FAC on VGG/CIFAR-10, while Noisy Adam was unstable with batch normalization.The paper reports that VOGN [24] works with such techniques; it also shows distributed Noisy K-FAC training on ImageNet but provides no extensive comparisons because tuning is time-consuming. Other deterministic-network posterior approximations [37] [38] are described as less flexible than VI.Source: §1 · p. 2
2 Deep Learning with Bayesian Principles and Its Challenges
Standard deep learning succeeds partly because scalable stochastic-gradient optimization can be combined with practical techniques. Bayesian inference is computationally expensive, and although variational inference makes it an optimization problem, direct methods have struggled on large problems.
- Stochastic-gradient methods scale to large problems, and techniques such as batch normalization and data augmentation help them achieve good performance.Common optimizers include SGD, RMSprop, and Adam [25].Source: §2 · p. 3
- The full Bayesian posterior is costly to compute because its marginal likelihood requires a difficult high-dimensional integral for large networks.Source: §2 · p. 3
- Variational inference approximates the posterior with a parameterized distribution, such as a Gaussian, whose parameters are estimated by maximizing the evidence lower bound.More complex approximations can reduce approximation error, but at additional computational cost.Source: §2 · p. 3
- Directly applying stochastic-gradient methods to variational parameters has failed to achieve reasonable performance on large problems, often converging very slowly, as in Graves and Blundell et al. [4].Source: §2 · p. 3
- This work aims to optimize the ELBO in a way that exploits deep-learning techniques to improve variational inference.Source: §2 · p. 3
3 Practical Deep Learning with Natural-Gradient Variational Inference
The section presents VOGN as a practical natural-gradient variational-inference method for deep learning, using familiar deep-learning techniques to support training at scale.
Natural-gradient VI and VOGN
- For Gaussian approximations, VOGN gives a natural-gradient update resembling stochastic-gradient methods, with moving averages that help manage noisy gradient estimates.The update uses stochastic gradients with respect to the expectation parameters of q; the paper presents natural-gradient VI as a basis for practical deep learning.Source: §3 · pp. 3–4
Gauss-Newton update
- VOGN squares individual-example gradients before averaging them, unlike the cited stochastic-gradient update; its implementation runs within 2-5 times of Adam’s running time.This Gauss-Newton step makes VOGN a second-order method and requires extra computation, though the paper says the overhead is not a theoretical limitation.Source: §3 · pp. 4–6
Batch normalisation
- VOGN uses standard BatchNorm layers with default settings, without L2 regularisation, weight decay, or uncertainty over BatchNorm parameters.Batch normalization stabilizes each layer’s input distribution by normalizing running averages of input means and variances.Source: §3 · p. 4
Data Augmentation
- Random cropping and horizontal flipping directly applied to VOGN slightly worsen performance and uncertainty calibration, so the method scales N to ρN to improve performance.The passages identify ρ ≥ 1 as accounting for augmentation’s increase in effective sample size; the reason for the boost may involve the relationship between regularisation δ and N.Source: §3 · p. 4
Momentum and initialisation; Learning rate scheduling
- VOGN adopts Adam-like momentum, weight and momentum initialisation, and learning-rate scheduling; its variance initialisation and gradual τ increase help mimic Adam’s early convergence.The learning-rate schedule is the same as Adam’s, while τ starts small (for example, 0.1) and increases toward 1.Source: §3 · p. 5
Distributed training
- Distributed data and Monte-Carlo-sample parallelism reduces training variance in theory and is essential for fast training on large problems such as ImageNet.The scheme combines data parallelism with MC samples assigned to different inputs, and may average multiple samples early in training to reduce variance sufficiently.Source: §3 · p. 5, Figure 2 · p. 5
Tuning VOGN
- VOGN tuning proceeds from Adam/SGD hyperparameters through deterministic OGN optimisation, then tunes δ, τ, and MC-sample count K for regularisation.The paper also recommends warming τ toward 1 and adding damping γ to the moving average to bound covariance eigenvalues and prevent excessive noise or step sizes.Source: §3 · p. 6
4 Experiments
Experiments show that VOGN can be trained practically on CIFAR-10 and ImageNet with performance comparable to standard optimisers, while retaining Bayesian predictive-probability benefits and performing at least as well as VCL on a continual-learning task.
4.1 Performance on CIFAR-10 and ImageNet
- Momentum and batch normalisation significantly speed VOGN convergence and improve performance on ResNet-18 trained on CIFAR-10; data augmentation improves validation accuracy across architectures trained on CIFAR-10.On CIFAR-10, VOGN’s convergence and validation accuracies are comparable to Adam and MC-dropout; BBB converges prohibitively slowly on LeNet-5 and could not be trained successfully on other architectures.Source: §4.1 · p. 6, Figure 3 · p. 7, Figure 4 · p. 7
- 2–5 times slower per epoch than Adam and SGD, VOGN achieves comparable accuracy, log likelihoods, and convergence epochs.Its added runtime is mainly due to computing individual gradients.Source: §4.1 · pp. 6–7
4.2.1 Performance on a Continual-learning task
- On Permuted MNIST, VOGN performs at least as well as VCL and far better than EWC, using 100 epochs per task versus VCL’s 800.The comparison averages accuracy over 10 tasks in Figure 3b.Source: Figure 3 · p. 7, §4.2.1 · p. 9
4.1 Performance on CIFAR-10 and ImageNet
- Lower prior precision raises validation accuracy but also enlarges the train-test gap, while more training MC samples can improve convergence and stability at added computational cost.More test-time MC samples improve generalisation through averaging; setting effective dataset size ρ to its correct order of magnitude is sufficient when tuning prior variance δ.Source: §4.1 · p. 7
- Adam often overfits, MC-dropout’s small train-test gap reflects underfitting, and Noisy K-FAC is slower than VOGN but takes fewer epochs.Noisy K-FAC and VOGN have about the same wall-clock time; MC-dropout performance is highly sensitive to dropout rate.Source: §4.1 · p. 7
4.2 Quality of the Predictive Probabilities
- On 10 of 15 NLL, ECE, and AUROC metrics, VOGN is best or tied best, and it is second-best on the other 5.MC-dropout is best on 4 metrics and Adam on 1; Adam is consistently over-confident, while MC-dropout is usually under-confident.Source: Table 1 · p. 8, §4.2 · p. 8
- VOGN shows low predictive entropy in-distribution and high entropy out-of-distribution, unlike Adam’s frequent zero entropies and MC-dropout’s generally high entropies.Across architectures, AUROC and FPR at 95% TPR do not provide a clear overall story.Source: §4.2 · pp. 8–9, Figure 5 · p. 9
5 Conclusions · A Noisy K-FAC algorithm
The conclusion reports that VOGN achieves accuracy and convergence comparable to SGD and Adam while retaining Bayesian benefits. The appendix describes Noisy K-FAC’s covariance approximation and updates.
5 Conclusions
- VOGN achieves accuracies and convergence rates comparable to SGD and Adam across architectures and datasets, including ImageNet.The authors attribute practical scaling to VOGN’s similarity to Adam and the use of deep-learning techniques.Source: §5 · p. 9
- Unlike SGD and Adam, VOGN retains well-calibrated uncertainty and good performance on out-of-distribution data.Source: §5 · p. 9
- VOGN’s continual-learning results confirm its promise; future experiments could examine small data, active learning, adversarial settings, sequential decisions, and structured covariance approximations.Source: §5 · p. 9
A Noisy K-FAC algorithm
- Noisy K-FAC [56] approximates the full covariance structure with a matrix-variate Gaussian over layer weights, making its updates more involved than VOGN’s.Source: A Noisy K-FAC algorithm · p. 14
- The Noisy K-FAC update uses sampled weights and minibatch loss gradients, with the gradient’s input tied to the activation vector of the previous layer.Source: A Noisy K-FAC algorithm · p. 14
- Noisy K-FAC scales gradients with matrices A_t and S_t related to the approximation’s precision matrix; π_t is the average eigenvalue of A_t divided by that of S_t.Source: A Noisy K-FAC algorithm · p. 14
B Details on fast implementation of the Gauss-Newton approximation … D Hyperparameter settings
The appendix describes per-example Gauss–Newton gradient computation for convolutional and batch-normalisation layers, tractability and distributed-training approximations, a deterministic OGN variant, and where hyperparameter settings are reported.
B Details on fast implementation of the Gauss-Newton approximation
- The backward pass is modified to compute per-example minibatch gradients efficiently, extending Goodfellow [12] to convolutional and batch-normalisation layers.This enables efficient computation of the Gauss–Newton approximation.Source: B Details on fast implementation of… · p. 14
B.1 Convolutional layer
- For convolutions, PyTorch unfold forms an extended input matrix, while function hooks expose inputs and per-example output gradients for computing per-example weight gradients after backpropagation.The method applies for any stride and padding.Source: §B.1 · p. 14
B.2 Batch normalisation layer
- For batch normalisation after a fully connected layer, per-example gradients for scale γ and shift β are obtained using normalised inputs and graph-accessible activations and output gradients.The appendix describes accessing the per-example input and output gradient through the computational graph as for convolutional layers.Source: §B.2 · pp. 14–15
B.3 Layer-wise block-diagonal Gauss-Newton approximation
- To make Gauss–Newton tractable, the method ignores correlations between layers with a block-diagonal matrix and approximates each layer block with either Kronecker-factored or diagonal curvature.The diagonal option corresponds to a mean-field approximation in that layer.Source: §B.3 · p. 15
- In distributed training, each GPU computes a Gauss–Newton matrix for its local minibatch; one GPU aggregates and inverts them, then sends σ to processes for independent posterior draws.The inversion can also be parallelised after the block-diagonal approximation.Source: §B.3 · p. 15
- Batch-normalisation layers use diagonal blocks with K-FAC and Noisy K-FAC, matching Osawa et al., who reported no loss of accuracy; Table 2 lists the layer-type approximations.Noisy K-FAC uses a distributed scheme in which each process has both S and A.Source: §B.3 · p. 16, Table 2 · p. 16
C OGN: A deterministic version of VOGN
- OGN uses gradients at the Gaussian mean instead of Monte Carlo samples, enabling a deterministic update and making Adam techniques easier to apply when transitioning to VOGN.Like stochastic-gradient methods it converges to a loss local minimum, giving a Laplace rather than variational approximation; its uncertainty estimates are slightly worse and it preserves Bayesian benefits less effectively.Source: C OGN: A deterministic version of VOGN · p. 16
D Hyperparameter settings
- Hyperparameters for Table 1 results are reported in Table 4, distributed variational-inference settings in Table 5, with best-practice guidance cited to Goyal et al. [13] and Osawa et al..Source: D Hyperparameter settings · p. 16
D.1 Bayes by Backprop for CIFAR-10/LeNet-5 training · D.2 Continual learning experiment
These sections specify Bayes by Backprop’s CIFAR-10/LeNet-5 training choices and the setup and repeated-run protocol for the continual-learning experiment.
D.1 Bayes by Backprop for CIFAR-10/LeNet-5 training
- BBB follows Swaroop et al. [52]’s suggested procedure, using the local reparameterisation trick and small initial means and variances.The settings are for Bayes by Backprop [4].Source: §D.1 · p. 16
- BBB uses 10 training MC samples per minibatch for linear layers and 1 for convolutional layers, then 100 and 10, respectively, during testing.Source: §D.1 · p. 16
- BBB optimises twice as many parameters as Adam or SGD; fewer convolutional-layer MC samples speed epochs without empirically reducing convergence rate.Source: §D.1 · p. 16
- Table 3 compares optimisers across dataset/architecture combinations using accuracy, NLL, ECE, and AUROC, with means and standard deviations over three runs.For ImageNet, accuracy and NLL are medians from the final five epochs.Source: Table 3 · p. 17
D.2 Continual learning experiment
- Continual-learning models use two 100-unit ReLU hidden layers; VCL follows Swaroop et al. [52], while EWC uses the grid-searched λ = 100 setting also found by Nguyen et al. [43].Source: §D.2 · p. 18
- VOGN runs 100 epochs per task with default PyTorch linear-layer initialisation, initial precision 1e6, a standard normal prior, and α = 1e −3, M = 256, β1 = 0, β2 = 1e −3.It uses 10 MC samples during training and 100 during testing; between tasks, mean and precision are initialised as for the first task.Source: §D.2 · p. 18
- Each continual-learning method is run 20 times, varying random seeds for both benchmark permutation and model training.Source: §D.2 · p. 18
E Effect of prior variance and dataset size reweighting factor
Increasing prior variance improves validation accuracy and log likelihood but also increases the train-test gap, while pushing VOGN and Noisy K-FAC toward converged solutions more like their non-Bayesian counterparts. When the dataset reweighting factor ρ is set to the correct order of magnitude, tuning prior variance δ matters more for performance.
E Effect of prior variance and dataset size reweighting factor
- Increasing prior variance improves validation accuracy and log likelihood but always increases the train-test gap when other hyperparameters are held constant.The reported trend has no exceptions under the stated condition.Source: E Effect of prior variance and dataset… · p. 18
- 2.29 and 0.12 were the train-test accuracy and log likelihood gaps at prior variance 7.5e−4, versus 6.38 and 0.34 at 7.5e−3.The example is VOGN on ResNet-18 on ImageNet; validation accuracy and log likelihood also increase at the higher prior variance.Source: E Effect of prior variance and dataset… · p. 19
- With increased prior variance, VOGN and Noisy K-FAC converge to solutions more like their non-Bayesian counterparts, where overfitting is an issue.The passage characterizes this as expected from Bayesian principles.Source: E Effect of prior variance and dataset… · p. 19
- When ρ is set to the correct order of magnitude, it has little effect on performance, so the methodology is to tune δ instead.For ImageNet, ρ is set smaller than for CIFAR-10 because cropping uses a higher portion of the initial image: 224x224 crops from 256x256 images.Source: E Effect of prior variance and dataset… · p. 19
F Effect of number of Monte Carlo samples on ImageNet
On ImageNet, increasing Monte Carlo samples at training or validation time improves VOGN’s performance, but increases computation time.
F Effect of number of Monte Carlo samples on ImageNet
- More training or validation MC samples improve ImageNet performance, with gains limited by increased computation time.The experiments trained ResNet-18 on ImageNet using 128 GPUs, with 1 training MC sample per process and 10 validation samples as the reported setup.Source: F Effect of number of Monte Carlo… · p. 19
- Increasing training MC samples per process from 1 to 2 speeds convergence per epoch and yields a better converged solution, while raising epoch time by approximately 1.5×.Increasing training MC samples to 3 does not significantly improve final test performance.Source: F Effect of number of Monte Carlo… · p. 19
- Increasing validation MC samples from 10 to 100 improves validation accuracy and log likelihood, while leaving training accuracy and log likelihood unchanged.Because validation runs every epoch, using more validation samples increases runtime.Source: F Effect of number of Monte Carlo… · p. 19
- With more compute, users can improve VOGN by increasing training or test-time MC samples, or both, up to a limit.Source: F Effect of number of Monte Carlo… · p. 19
G MC-dropout’s sensitivity to dropout rate
MC-dropout is highly sensitive to dropout rate: p = 0.1 works best across architectures on CIFAR-10, but performs badly on ImageNet ResNet-18, where p = 0.05 is used. The LeNet-5 results illustrate how the train-test gaps change with p.
G MC-dropout’s sensitivity to dropout rate
- MC-dropout is highly sensitive to dropout rate: p = 0.1 works best for all architectures on CIFAR-10, while even p = 0.1 performs badly on ImageNet ResNet-18, leading to use of p = 0.05 there.The paper identifies this high sensitivity as an issue with MC-dropout.Source: G MC-dropout’s sensitivity to dropout… · p. 19, Figure 13 · p. 20
- At p = 0.01, LeNet-5’s train-test gaps are 10.3% in accuracy and 0.34 in log likelihood; at p = 0.1, they are 1.4% and 0.04, respectively.At p = 0.2, the gaps are -7.71% in accuracy and -0.02 in log likelihood; p = 0.1 is chosen for its high accuracy and log likelihood and small train-test gap.Source: Figure 12 · p. 20
H Uncertainty metrics
VOGN generally matches the best-performing optimisers on validation log likelihood, performs similarly on AUROC, and usually provides better calibration, with dataset- and architecture-specific exceptions.
H Uncertainty metrics
- VOGN and MC-dropout use 10 Monte Carlo validation samples from the variational Gaussian approximation, increasing test-time computation relative to Adam or SGD.Appendix F examines the effect of changing the number of validation samples.Source: H Uncertainty metrics · p. 20
- Uncertainty comparisons use predictive probabilities to evaluate validation log loss, AUROC, calibration curves, and model outputs on out-of-distribution data.The same approach is used across optimisers, with hyperparameters tuned for validation accuracy before uncertainty testing.Source: H Uncertainty metrics · pp. 20–21
- Final validation log likelihoods are very similar, with VOGN usually matching the best-performing optimisers; MC-dropout performs best on CIFAR-10/AlexNet with and without data augmentation.Source: H Uncertainty metrics · p. 21
- AUROC results are similar, especially on ImageNet, where MC-dropout is marginally better; VOGN is best or tied best in all but one CIFAR-10 experiment, while Adam is worst but performs surprisingly well on CIFAR-10/ResNet-18.Source: H Uncertainty metrics · p. 21
- Across datasets and architectures except underfitting LeNet-5, VOGN usually has better calibration curves and ECE; Adam is consistently over-confident, while MC-dropout is usually under-confident, except on ImageNet after extensive rate tuning.Calibration curves compare true and predicted accuracy, and ECE [15] averages their error over 20 bins; calibration curves.Source: H Uncertainty metrics · p. 21
I Out-of-distribution experimental setup and additional results · J Author contributions statement · K Changes in the camera-ready version compared to the submitted version
The supplementary sections define out-of-distribution tests and metrics, document author contributions, and list changes made for the camera-ready version.
I Out-of-distribution experimental setup and additional results
- VOGN is compared with Adam and MC-dropout on CIFAR-10-trained LeNet-5, AlexNet, and ResNet-18, using CIFAR-10 validation data in-distribution and SVHN and two LSUN variants out-of-distribution.The tests use CIFAR-10 validation images, the full SVHN training set, and LSUN images randomly cropped or down-sampled; the OOD datasets have no similar classes to CIFAR-10.Source: I Out-of-distribution experimental… · p. 21
- OOD performance is assessed with predictive-entropy histograms, AUROC, and FPR at 95% TPR, with high OOD entropy and low in-distribution entropy indicating desirable uncertainty.AUROC and FPR at 95% TPR are computed by thresholding the most likely class’s softmax output to distinguish high-uncertainty images as unknown.Source: I Out-of-distribution experimental… · pp. 21–22
- Figures 15–17 show OOD predictive-entropy histograms for CIFAR-10-trained AlexNet with and without data augmentation and LeNet-5 without augmentation, with AUROC and FPR at 95% TPR averaged over 3 runs.Standard deviations are described as very small and not reported.Source: Figure 15 · p. 26, Figure 16 · p. 27, Figure 17 · p. 27
J Author contributions statement
- M.E.K., A.J., and R.E. conceived the original idea; M.E.K. led the project, and all authors discussed and agreed on its main focus and ideas.R.Y. provided computing resources and access to HPCI systems.Source: J Author contributions statement · p. 22, K Changes in the camera-ready version… · p. 22
- A.J. began the proof of concept, K.O. built the main PyTorch and distributed implementations and led most large experiments, while S.S. wrote the BBB implementation.R.E. fixed implementation issues and identified a data-augmentation issue addressed jointly by S.S., A.J., K.O., and R.E.; the authors also shared experiment and writing contributions.Source: J Author contributions statement · p. 22
K Changes in the camera-ready version compared to the submitted version
- The camera-ready version added a continual-learning experiment, damping-factor details, an OGN method for tuning VOGN hyperparameters, and guidance on tuning VOGN.It also revised Noisy K-FAC notation and the Gauss-Newton approximation description, and added OGN results to Tables 1 and 3.Source: K Changes in the camera-ready version… · p. 22
- The camera-ready version added training curves and calibration curves for K-FAC, Noisy K-FAC, and OGN, and included additional prior-variance and Monte Carlo-sample analyses.Figure 7 collects the curves for Table 1 results; Figures 8–11 examine prior variance and dataset-size reweighting, and Figure 11 also examines Monte Carlo sample counts.Source: K Changes in the camera-ready version… · p. 22, Figure 7 · p. 23, Figure 8 · p. 24
Related papers
This paper cites
- Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift2015
- Efficient Estimation of Word Representations in Vector Space2013
- On Calibration of Modern Neural Networks2017
- Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning2015
- Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour2017
- A Baseline for Detecting Misclassified and Out-of-Distribution Examples in Neural Networks2016