Noisy Natural Gradient as Variational Inference

TL;DR

The paper reveals a fundamental connection between natural gradient methods and variational inference for Bayesian neural networks, enabling Noisy Natural Gradient (NNG) as a practical training paradigm. By reinterpreting NG updates as variational updates with adaptive weight noise, it supports full-covariance, diagonal, matrix-variate Gaussian posteriors and scalable training via noisy Adam and noisy K-FAC. Empirical results show improved predictive uncertainty, calibration, and exploration across regression, classification, active learning, and reinforcement learning. This approach provides a scalable path to expressive variational posteriors in large neural networks, with tangible benefits for uncertainty-driven decision making.

Abstract

Variational Bayesian neural nets combine the flexibility of deep learning with Bayesian uncertainty estimation. Unfortunately, there is a tradeoff between cheap but simple variational families (e.g.~fully factorized) or expensive and complicated inference procedures. We show that natural gradient ascent with adaptive weight noise implicitly fits a variational posterior to maximize the evidence lower bound (ELBO). This insight allows us to train full-covariance, fully factorized, or matrix-variate Gaussian variational posteriors using noisy versions of natural gradient, Adam, and K-FAC, respectively, making it possible to scale up to modern-size ConvNets. On standard regression benchmarks, our noisy K-FAC algorithm makes better predictions and matches Hamiltonian Monte Carlo's predictive variances better than existing methods. Its improved uncertainty estimates lead to more efficient exploration in active learning, and intrinsic motivation for reinforcement learning.

Figures (4)

• Figure 1: Normalized precision matrices for Gaussian variational posteriors trained using noisy natural gradient. We used a network with 2 hidden layers of 15 units each, trained on the Boston housing dataset.
• Figure 2: Reliability diagrams niculescu2005predictingguo2017calibration for K-FAC (left) and noisy K-FAC (right) on CIFAR10. Reliability diagrams show accuracy as a function of confidence. Models trained without BN (top) and with BN (bottom). ECE = Expected Calibration Error guo2017calibration; smaller is better.
• Figure 3: Performance of [TRPO] TRPO baseline with Gaussian control noise, [TRPO+BBB] VIME baseline with BBB dynamics network, and [TRPO+NNG-MVG] VIME with NNG-MVG dynamics network (ours). The darker-colored lines represent the median performance in 10 different random seeds while the shaded area show the interquartile range.
• Figure 4: Training curves for all three methods. For each method, we tuned the learning rate for updating the posterior mean. Note that BBB and NNG-FFG use the same form of $q$, while NNG-MVG uses a more flexible $q$ distribution.