[Paper Review] Practical Deep Learning with Bayesian Principles

Kazuki Osawa, Siddharth Swaroop+5 more|arXiv (Cornell University)|Jun 6, 2019

Domain Adaptation and Few-Shot Learning50 references50 citations

TL;DR

The paper demonstrates practical training of deep networks using Variational Inference with a natural-gradient method (VOGN), achieving competitive performance to Adam/SGD on CIFAR-10 and ImageNet while preserving Bayesian benefits like calibrated predictions and improved OOD uncertainty and continual learning.

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

Motivation & Objective

• Motivate and enable practical Bayesian deep learning by addressing scalability and performance gaps.
• Show that natural-gradient variational inference (VOGN) can train large networks efficiently with standard deep-learning tricks (batch norm, data augmentation, distributed training).
• Demonstrate preserved Bayesian benefits: calibrated predictive probabilities, improved out-of-distribution uncertainty, and improved continual-learning behavior.
• Provide empirical evidence across multiple architectures and datasets (CIFAR-10, ImageNet) showing competitive performance to non-Bayesian baselines.

Proposed method

• Formulate deep learning as Bayesian inference via variational inference with a Gaussian posterior q(w).
• Use natural-gradient updates for VI, yielding updates similar in form to SG/DL optimizers (VOGN).
• Adopt batch normalization, data augmentation, momentum, and distributed training to accelerate convergence.
• Employ Gauss-Newton-based variance updates (diagonal Sigma) to obtain a practical second-order VI method.
• Introduce data-augmentation scaling (rho) to compensate for effective dataset size in Bayesian training.
• Provide a distributed training scheme combining data and MC-sample parallelism to scale to ImageNet.

Experimental results

Research questions

• RQ1 Can natural-gradient variational inference (VOGN) train deep networks at scale with performance comparable to Adam/SGD on large datasets?
• RQ2 Do Bayesian posterior approximations via VOGN yield calibrated predictions and improved out-of-distribution uncertainty while preserving practical training dynamics?
• RQ3 What is the impact of Bayesian principles on continual learning and knowledge retention in sequential tasks?
• RQ4 How do standard deep-learning techniques (batch norm, data augmentation, distributed training) interact with VI to deliver practical Bayesian deep learning?
• RQ5 What are the trade-offs (speed, calibration, uncertainty quality) when using VOGN vs traditional optimizers and MC-dropout?

Key findings

• VOGN achieves similar convergence and performance to Adam/SGD on CIFAR-10 and ImageNet across multiple architectures.
• VOGN provides well-calibrated predictive probabilities and improved uncertainties on out-of-distribution data compared to non-Bayesian methods.
• VOGN with batch normalization and data augmentation matches the speed/epochs of standard optimizers on large-scale tasks, albeit with higher per-epoch cost due to VI computations.
• Compared to BBB and MC-dropout, VOGN often yields better calibration and lower overconfidence, particularly on ImageNet and ResNet-18.
• On continual-learning tasks, VOGN performs competitively with or better than existing Bayesian continual-learning methods (e.g., VCL) in accuracy, and trains faster per task in some settings.
• Table 1 shows VOGN achieving competitive or best-in-class metrics across CIFAR-10 (LeNet-5, AlexNet, ResNet-18) and ImageNet (ResNet-18) against Adam, SGD, MC-dropout, OGN, K-FAC, and Noisy K-FAC.