PaperlayerView exact revision on arXiv

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

arXiv:1906.02506v2stat.MLcs.LG

TL;DR

Deep-learning Bayesian methods have struggled with scalability and competitive training performance. This paper applies natural-gradient variational inference with deep-learning techniques, achieving comparable training performance while retaining Bayesian benefits such as calibrated uncertainty and improved continual-learning behavior. The main practical boundary is higher runtime than Adam.

  • Problem

    Principled Bayesian methods for deep learning remain difficult to scale and often fail to match standard methods’ performance on large problems.

  • Method

    The paper trains deep networks with natural-gradient variational inference using VOGN and borrows established deep-learning techniques such as normalization, augmentation, momentum, and distributed training.

  • Results

    VOGN achieves similar performance in about the same number of epochs as Adam across multiple architectures and datasets, while retaining calibrated probabilities, improved out-of-distribution uncertainty, and boosted continual-learning performance.

  • Takeaways & Limitations

    Practical deep learning with Bayesian methods is possible while preserving benefits of Bayesian principles for uncertainty and continual learning.

  • Takeaways & Limitations

    VOGN is slower than Adam or SGD because its Gauss-Newton update requires extra computation for individual gradients; its runtime is usually within 2-5 times Adam’s.

Abstract

from arXiv · show

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

Deep learning’s data, forgetting, confidence, and robustness problems motivate Bayesian principles, but scalable Bayesian inference remains difficult. The paper aims to make principled methods such as variational inference practical for deep learning.

  • Deep learning can overfit small datasets, forget past knowledge during sequential learning, and produce unreliable confidence estimates.
  • Bayesian principles represent uncertainty with posterior distributions, support sequential learning through Bayes’ rule, and reduce overfitting through model averaging.
  • Variational inference and MCMC have historically been difficult to scale to large datasets such as ImageNet.
  • The paper targets practical, principled Bayesian methods because MC-dropout has limitations for continual learning, sequential decision-making, and flexible posterior approximation.
  • Natural-gradient variational methods resemble Adam, allowing existing deep-learning techniques to support training across architectures and datasets.

2 Deep Learning with Bayesian Principles and Its Challenges

Standard stochastic-gradient training scales well, whereas Bayesian inference requires costly posterior computation. Variational inference offers a scalable approximation, but existing optimization methods often converge too slowly on large deep-learning problems.

  • Stochastic-gradient methods minimize training losses efficiently and scale to very large problems, especially with established deep-learning techniques.
  • The full Bayesian posterior is expensive because computing the marginal likelihood requires a difficult high-dimensional integral.
  • Variational inference approximates the posterior with a parametric distribution such as Gaussian q(w) := N(w|µ, Σ).
  • VI estimates variational parameters by maximizing the evidence lower bound, converting Bayesian inference into an optimization problem.
  • Existing direct applications of stochastic-gradient methods to VI often converge too slowly to achieve reasonable performance on large problems.

3 Practical Deep Learning with Natural-Gradient Variational Inference

The paper makes natural-gradient VI practical by implementing VOGN in a form compatible with deep-learning training techniques. Its design combines Bayesian updates with normalization, augmentation, scheduling, momentum, and distributed computation.

  • Natural-gradient VI uses updates that resemble stochastic-gradient methods, making deep-learning optimization techniques applicable.
  • VOGN estimates a Gaussian posterior with mean µ and diagonal covariance, updating its scale through a moving average.
  • VOGN’s Gauss-Newton update squares individual-example gradients before averaging, unlike Adam’s squared minibatch gradient.
  • Batch normalization, data augmentation, momentum, initialization, learning-rate scheduling, and distributed training are incorporated into the VOGN implementation.
  • The distributed algorithm parallelizes computation over data and Monte-Carlo samples, while one GPU aggregates Gauss-Newton information and updates posterior parameters.
  • VOGN requires extra computation for individual gradients, making it slower than Adam or SGD, although the implementation is more efficient than a naive approach.

4 Experiments

Experiments show that VOGN can train deep networks competitively with standard optimisers while retaining advantages in uncertainty estimation and continual learning. Its performance is supported by practical deep-learning techniques, although it incurs higher per-epoch runtime and Bayesian trade-offs.

  • Performance on CIFAR-10 and ImageNet: Momentum, batch normalisation, and data augmentation significantly improve VOGN’s convergence and performance.The improvements are reported for accuracy and log likelihoods, with data augmentation also improving Figure 4 results.
  • Performance on CIFAR-10 and ImageNet: VOGN achieves competitive training performance across CIFAR-10 and ImageNet experiments, with accuracy, log likelihoods, and convergence epochs comparable to standard methods.The comparisons include Adam, SGD, and MC-dropout across multiple architectures and datasets.
  • Performance on CIFAR-10 and ImageNet: VOGN is 2-5 times slower per epoch than Adam and SGD, mainly because it computes individual gradients.Despite this runtime cost, the reported number of convergence epochs remains comparable.
  • Continual learning: 100 versus 800 epochs per task: VOGN converges faster than VCL while achieving at least as high average accuracy over 10 Permuted MNIST tasks.The continual-learning comparison evaluates sequential learning on Permuted MNIST.
  • Quality of the Predictive Probabilities: VOGN performs best or tied best on 10 of 15 uncertainty metrics and second-best on the remaining 5.The metrics are NLL, ECE, and AUROC across five dataset/architecture combinations; lower is better for NLL and ECE, higher for AUROC.
  • Quality of the Predictive Probabilities: VOGN generally assigns low predictive entropy to in-distribution data and high entropy to out-of-distribution data, unlike the other evaluated methods.The out-of-distribution tests use CIFAR-10, SVHN, and LSUN inputs with AUROC and FPR at 95% TPR reported.

5 Conclusions

The paper demonstrates that VOGN and related natural-gradient VI methods can train deep networks at practical scale while retaining Bayesian benefits. Its structured covariance approximations support distributed training, though Noisy K-FAC updates are more involved than VOGN.

  • VOGN enables practical natural-gradient variational training across architectures and datasets, including ImageNet.
  • VOGN achieves accuracies and convergence rates comparable to SGD and Adam while retaining calibrated uncertainty and strong out-of-distribution performance.
  • Continual-learning results confirm that VOGN preserves Bayesian benefits beyond standard optimisation performance.
  • Noisy K-FAC approximates the full covariance structure, making its updates more involved than those of VOGN.
  • Noisy K-FAC estimates a matrix-variate Gaussian using covariance factors Σ2,t ⊗ Σ1,t for each fully-connected layer.

B Details on fast implementation of the Gauss-Newton approximation

The implementation makes Gauss–Newton variational training practical by computing per-example gradients, using layer-wise approximations, and distributing matrix calculations across GPUs. OGN provides a deterministic stepping stone before switching to VOGN.

  • Fast gradient computation: Per-example gradients are computed efficiently by modifying the backward pass and extending the method to convolutional and batch-normalisation layers.
  • Convolutional layer: For convolutional layers, unfolding the input tensor into M_A enables matrix multiplication for both outputs and per-example weight gradients.
  • Convolutional layer: PyTorch function hooks expose per-example inputs and output gradients, allowing weight gradients to be recovered after the backward pass.
  • Layer-wise block-diagonal approximation: The Gauss–Newton matrix is approximated as layer-wise block diagonal, ignoring correlations between weights in different layers.
  • Distributed training: Distributed training aggregates local Gauss–Newton matrices, parallelises block-diagonal inversion, and broadcasts updated posterior standard deviations.
  • Deterministic stepping stone: OGN removes sampling by evaluating gradients at the Gaussian mean, providing a deterministic stepping stone for applying Adam-like techniques before VOGN.

D.1 Bayes by Backprop for CIFAR-10/LeNet-5 training

The Bayes by Backprop comparison uses established CIFAR-10/LeNet-5 settings, Monte Carlo sampling, and repeated runs to evaluate optimiser performance under a controlled experimental setup.

  • Training setup: Bayes by Backprop uses the local reparameterisation trick and small initial means and variances.
  • Training setup: Training uses 10 Monte Carlo samples for linear layers and one for convolutional layers per minibatch.
  • Evaluation setup: Testing uses 100 Monte Carlo samples for linear layers and 10 for convolutional layers per minibatch.
  • Metrics: Accuracy is reported as higher-is-better, while negative log likelihood and expected calibration error are lower-is-better metrics.
  • Metrics: ImageNet accuracy and negative log likelihood are reported as median values from the final five epochs.
  • Experimental protocol: The experiments use two hidden layers with 100 ReLU units each and run VCL with matched settings from prior work.
  • Experimental protocol: VOGN runs for 100 epochs per task and averages results over 20 runs with different random seeds.

E Effect of prior variance and dataset size reweighting factor

Prior variance improves validation performance but also increases overfitting, while dataset reweighting has limited effect when its scale is correctly chosen. The study therefore prioritises tuning prior variance.

  • Prior variance: Increasing prior variance improves validation accuracy and log likelihood when other hyperparameters remain fixed.
  • Prior variance: 2.29 and 0.12 are the train-test accuracy and log-likelihood gaps at prior variance 7.5e-4 for VOGN on ResNet-18 ImageNet.
  • Prior variance: 6.38 and 0.34 are the respective train-test gaps when prior variance increases to 7.5e-3.
  • Prior variance: Higher prior variance makes VOGN and Noisy K-FAC converge toward non-Bayesian solutions where overfitting is an issue.
  • Dataset reweighting: When dataset reweighting ρ has the correct order of magnitude, it affects performance less than tuning the prior variance δ.

F Effect of number of Monte Carlo samples on ImageNet

On ImageNet, VOGN performance improves when more Monte Carlo samples are used during training or validation, although additional samples increase computation and eventually provide diminishing returns.

  • Training samples: 1 to 2 training MC samples per process speeds up VOGN training on ImageNet while maintaining stable training.ImageNet’s large minibatches and many parallel processes make one sample per process stable; increasing to two reduces training noise and speeds training.
  • Testing samples: 10 to 100 testing MC samples improve validation accuracy and log likelihood without changing training accuracy or log likelihood.The higher validation-sample count also increases runtime because validation is performed each epoch.
  • Practical implication: Additional compute can improve VOGN by refining the Monte Carlo approximation during training, testing, or both, up to a limit.

G MC-dropout’s sensitivity to dropout rate

MC-dropout is sensitive to its dropout rate, with the best setting depending on the architecture and dataset, especially for ImageNet.

  • Chosen rates: p = 0.1 works best for MC-dropout across architectures trained on CIFAR-10, while p = 0.05 is used for ResNet-18 on ImageNet.On CIFAR-10 LeNet-5, p = 0.1 yields accuracy and log-likelihood train-test gaps of 1.4% and 0.04.
  • ImageNet sensitivity: For ResNet-18 on ImageNet, MC-dropout is extremely sensitive to dropout rate, and even p = 0.1 performs badly.The authors identify this high sensitivity as an issue with MC-dropout.
  • CIFAR-10 sensitivity: On LeNet-5 with CIFAR-10, p = 0.01 produces accuracy and log-likelihood train-test gaps of 10.3% and 0.34.
  • CIFAR-10 sensitivity: On LeNet-5 with CIFAR-10, p = 0.2 produces accuracy and log-likelihood gaps of -7.71% and -0.02.

H Uncertainty metrics

The paper evaluates uncertainty using predictive probabilities and compares validation log loss, AUROC, calibration curves, and out-of-distribution behavior across optimisers.

  • Predictive probabilities: Bayesian predictive probabilities are estimated by averaging predictions from weight samples drawn from a variational Gaussian approximation.VOGN and MC-dropout use 10 validation-time MC samples, increasing testing cost relative to Adam or SGD.
  • Metrics: The uncertainty evaluation uses validation log loss, AUROC, calibration curves, and model outputs on out-of-distribution data.
  • Validation log likelihood: VOGN performs very well on validation negative log likelihood and usually matches the best-performing optimiser across datasets and architectures.The exception is CIFAR-10 with AlexNet, with or without data augmentation, where MC-dropout performs best.
  • AUROC: AUROC measures uncertainty-based separation of correctly and incorrectly classified examples after thresholding predicted class confidence.
  • Calibration: VOGN usually has better calibration curves and ECE than competing optimisers, except for LeNet-5.Adam is consistently over-confident, while MC-dropout is usually under-confident; MC-dropout performs well on ImageNet after extensive dropout-rate tuning.

I Out-of-distribution experimental setup and additional results

The paper evaluates out-of-distribution uncertainty on CIFAR-10-trained networks using entropy histograms and threshold-based detection metrics, with additional figures covering calibration and training choices.

  • Experimental setup: Out-of-distribution tests compare trained LeNet-5, AlexNet, and ResNet-18 models on CIFAR-10 validation data, SVHN, and LSUN crop and resize datasets.CIFAR-10 validation data is the in-distribution reference; SVHN and LSUN provide out-of-distribution inputs.
  • Uncertainty measures: Predictive-entropy histograms compare in-distribution and out-of-distribution inputs, alongside AUROC and FPR at 95% TPR.Ideally, out-of-distribution inputs have high entropy, while many in-distribution examples have low entropy.
  • Reported comparisons: Out-of-distribution results are reported for AlexNet with and without data augmentation, LeNet-5 without augmentation, and ResNet-18 in Figure 5.
  • Additional analyses: Appendix figures examine prior variance, dataset-size reweighting, Monte Carlo sample counts, and dropout-rate effects on ImageNet and CIFAR-10 training.
  • Calibration results: Figure 14 compares VOGN, Adam, and MC-dropout calibration curves on CIFAR-10, while the ResNet-18 ImageNet curve appears in Figure 1.VOGN is described as extremely well-calibrated except for LeNet-5, where all optimisers perform well.
Loading 1906.02506v2…