PaperlayerView exact revision on arXiv

Source-linked AI summary

Noisy Natural Gradient as Variational Inference

Guodong Zhang, Shengyang Sun, David Duvenaud, Roger Grosse

arXiv:1712.02390v2cs.LGstat.ML

TL;DR

Variational Bayesian neural networks face a tradeoff between simple, inexpensive posterior families and costly expressive inference. This paper shows that adaptive weight noise with natural-gradient ascent implicitly optimizes an ELBO, enabling scalable posterior fitting; its noisy K-FAC variant improves regression predictions and predictive-variance matching while supporting more efficient exploration.

  • Problem

    Variational Bayesian neural networks need expressive posterior families to capture weight correlations, but existing approaches can be expensive or impose restrictive factorization assumptions.

  • Method

    The paper reinterprets noisy natural-gradient optimization as variational inference and instantiates it with noisy natural gradient, Adam, and K-FAC for different Gaussian posterior structures.

  • Results

    Noisy K-FAC achieves the highest CIFAR10 accuracy in all configurations except with both data augmentation and batch normalization, including 85.52% versus 82.39% without extra regularization.

  • Takeaways & Limitations

    The resulting variational BNNs capture weight correlations efficiently, match HMC predictive variances better than fully factorized posteriors, and support more efficient exploration in active learning and reinforcement learning.

  • Takeaways & Limitations

    The method replaces difficult or unstable likelihood-Hessian updates with an NGPE Fisher approximation, while the ELBO regularization parameter need not equal the proper Bayesian value λ = 1 in practice.

Abstract

from arXiv · show

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.

1. Introduction

The paper connects natural-gradient optimization with variational inference and uses this connection to train more expressive Bayesian neural-network posteriors efficiently.

  • Natural-gradient procedures can be reinterpreted as optimizing a variational posterior with the exact Fisher matrix to maximize a unified ELBO.Both weight updates and Fisher estimation become natural-gradient ascent on the same objective.
  • The proposed noisy natural-gradient method provides an alternative training procedure for variational Bayesian neural networks.
  • For factorial Gaussian posteriors, noisy natural gradient matches Bayes By Backprop performance while converging faster.
  • Noisy K-FAC fits full matrix-variate Gaussian posteriors efficiently using correlated weight noise and a Kronecker-factored curvature approximation.

2. Background

The background introduces variational inference and natural-gradient optimization, then explains how Fisher approximations such as K-FAC make natural-gradient updates tractable for large neural networks.

  • Variational Inference for Bayesian Neural Networks: Variational Bayesian neural networks fit an approximate posterior q(w) over weights by maximizing an evidence lower bound.
  • Variational Inference for Bayesian Neural Networks: Proper Bayesian inference uses λ = 1, although other regularization values may perform better on some problems.
  • Gradient Estimators for Gaussian Distribution: Gaussian variational posteriors admit an alternative gradient estimator based on derivatives with respect to the mean and covariance.This estimator is presented alongside the pathwise derivative, or reparameterization trick.
  • Natural Gradient: Natural gradient is a second-order optimization method, with separate formulations for point estimation and variational inference.
  • Natural Gradient: For variational inference, the Fisher matrix is computed from q and serves as a metric on variational parameters rather than weights.Because q is tractable, this natural gradient can be computed exactly and may be simpler than the ordinary gradient.
  • Natural Gradient: NGPE and NGVI are closely related when the variational posterior is Gaussian.
  • Kronecker-Factored Approximate Curvature: K-FAC approximates the Fisher matrix with layerwise Kronecker factors, avoiding the quadratic storage cost of the exact Fisher.The resulting natural-gradient computation uses matrix transformations comparable in size to each layer’s weights.

3. Variational Inference using Noisy Natural Gradient

The paper links natural-gradient optimization of network weights to natural-gradient variational inference, showing that adaptive weight noise can approximate updates for Gaussian posteriors. This yields noisy optimization methods for diagonal, matrix-variate, and cross-layer covariance structures while addressing the computational cost of full covariance.

  • 3.1. Connection between NGPE and NGVI: Noisy Natural Gradient (NNG) approximates natural-gradient variational-inference updates using natural-gradient point estimation with adaptive weight noise.The connection supports training Gaussian variational posteriors with noisy versions of existing optimizers.
  • 3. Variational Inference using Noisy Natural Gradient: The mean update is preconditioned by Λ^-1, while the precision update behaves like an exponential moving average of the Hessian.The resulting mean step resembles a stochastic Newton update using Λ.
  • 3. Variational Inference using Noisy Natural Gradient: Approximating the negative log-likelihood Hessian with the NGPE Fisher avoids difficult Hessian computation and preserves positive semidefiniteness.The Hessian can be undefined for ReLU networks and indefinite for non-convex neural-network objectives.
  • 3.1. Simplifying the Update Rules: NNG includes intrinsic damping of strength λ/Nη and samples weights from a Gaussian posterior with covariance Σ.Additional extrinsic damping can be added for stability.
  • 3.3–3.5. Structured Covariance: Full covariance requires O((dim w)^2) parameters, motivating structured approximations such as diagonal, Kronecker-factored, and block-tridiagonal covariance.K-FAC’s Kronecker factors produce a matrix-variate Gaussian posterior, while block tridiagonal structure captures dependencies between adjacent layers.
  • 3.4. Fitting Matrix Variate Gaussian Posteriors with Noisy K-FAC: Noisy K-FAC fits a matrix-variate Gaussian posterior using correlated weight noise and a Kronecker-factored approximation to the Fisher.The factors are estimated online and are automatically positive semidefinite.

4. Related Work

Related work spans practical variational inference for factorized Gaussian posteriors, richer Gaussian posterior families, and tractable natural-gradient approximations such as K-FAC. The supplied material also situates the paper’s regression evaluation within standard RMSE and test log-likelihood reporting.

  • Variational inference: Graves proposed practical variational inference with fully factorized Gaussian posteriors using a simple but biased gradient estimator.Blundell et al. later proposed an unbiased estimator using the reparameterization trick, while local reparameterization reduces stochastic-gradient variance.
  • Structured posterior families: Prior work modeled weight correlations with matrix-variate Gaussian posteriors, unitary transformations, factorial Gaussians, and normalizing-flow-inspired constructions.These approaches target more expressive Gaussian variational families than fully factorized posteriors.
  • Natural-gradient optimization: Natural-gradient research developed tractable approximations including stochastic variational Bayes updates and K-FAC for efficient natural-gradient optimization in deep neural networks.K-FAC was subsequently adopted for many tasks and shown to support distributed computation.
  • Evaluation: Table 1 reports averaged test RMSE and log-likelihood for the regression benchmarks.These are described as standard metrics in the supplied evaluation context.
  • Concurrent work: Khan et al. independently related natural gradients to variational optimization and derived a diagonal stochastic Newton update, while concurrent work connected noisy RMSprop to variational inference.The latter method essentially resembles noisy Adam.

5. Experiments

Experiments show that noisy natural-gradient methods match or improve predictive performance, calibration, uncertainty estimation, and exploration across regression, CIFAR10 classification, active learning, and reinforcement learning.

  • Regression: NNG-FFG performed similarly to BBB, while NNG-MVG achieved substantially better RMSE and test log-likelihood than BBB and PBP on regression datasets.
  • Classification: Noisy K-FAC achieved the highest CIFAR10 accuracy in all configurations except with both data augmentation and Batch Normalization.
  • Classification: 3% gain: noisy K-FAC reached 85.52% versus K-FAC’s 82.39% without extra regularization.
  • Classification: Noisy K-FAC produced well-calibrated predictions, whereas point estimates tended to be poorly calibrated on CIFAR10.
  • Active Learning: NNG-MVG and NNG-BlkTri were evaluated against HMC predictive variances, while active-learning results showed that better variance estimates did not reliably improve performance.
  • Reinforcement Learning: Replacing BBB with NNG-MVG in VIME considerably improved exploration efficiency on all three reinforcement-learning tasks.

6. Conclusion

The conclusion presents noisy natural gradient as a bridge between point-estimation optimization and variational inference, enabling efficient structured Gaussian posteriors and improved uncertainty-driven exploration.

  • The paper identifies a relationship between natural-gradient ascent for neural-network weights and natural-gradient variational inference for Gaussian posteriors.
  • Noisy versions of widely used optimization algorithms provide simple procedures for fitting variational Bayesian neural networks with different posterior structures.
  • Matrix-variate Gaussian posteriors compactly capture correlations between different weights, while full-covariance representations require O((dim w)^2) parameters.
  • MVG posteriors matched HMC predictive variances better than fully factorized posteriors and led to more efficient exploration.

C.1. Regression Implementation Details

Regression implementation used repeated train-test splits, dataset-dependent network sizes, normalized inputs and targets during training, and scheduled optimization settings.

  • Datasets used 90% for training and 10% for testing, with repeated splits except for the two largest datasets.
  • Networks used 50 hidden units for most datasets and 100 hidden units for Year and Protein.
  • ELBO gradients with respect to variational parameters α and β were computed using automatic differentiation.
  • Input features and training targets were normalized to zero mean and unit variance during training, with target normalization removed at test time.

C.2. Classification Implementation Details

Classification experiments used a reduced-filter VGG16 and controlled optimization through Fisher-norm trust-region steps, update schedules, curvature approximations, and regularization settings.

  • Classification experiments used VGG16 with the number of filters in each convolutional layer reduced by half.
  • Step sizes were selected so the squared Fisher norm of each update stayed below a tunable trust-region value.
  • The trust-region bound followed an exponential decay schedule, with separate initial settings for noisy K-FAC under different Batch Normalization conditions and for noisy Adam.
  • Curvature statistics and inversions were updated less frequently to reduce K-FAC overhead, making the statistics somewhat stale without significantly affecting per-iteration optimization performance.
  • Regularization parameters were tuned by data-augmentation and Batch-Normalization settings, including λ = 0.1 with augmentation and λ = 0.5 otherwise.

C.3. Active Learning Implementation Details

The active-learning experiments used small one-hidden-layer networks, evaluated predictive-variance agreement with HMC, and iteratively selected points with highest predictive variance.

  • Protocol: Datasets were split into 20 training points, 100 test points, and pool sets for active-learning experiments.The protocol followed PBP, with all remaining data assigned to pools.
  • Uncertainty evaluation: Uncertainty quality was measured by Pearson correlation between each model’s predictive variance and HMC’s predictive variance.Higher correlation indicates stronger agreement between the two variance estimates.
  • Uncertainty evaluation: The comparison included PBP, NNG-FFG, NNG-MVG, and NNG-BlkTri.Predictive variances were computed from 1000 sampled weights on the test set.

C.4. Reinforcement Learning Implementation Details

The reinforcement-learning experiments used Bayesian neural networks in three tasks and derived layerwise intrinsic rewards for NNG-MVG using its matrix-variate posterior structure.

  • Experimental setup: CartPoleSwingup, MountainCar, and DoublePendulum used one-layer Bayesian neural networks with 32 hidden units for BBB and NNG-MVG.The models used ReLU activations and 10 variational-posterior samples during training.
  • Experimental setup: The policy network was a one-layer network with 32 tanh units, while TRPO used batch size 5000 and a fixed replay-pool size.The passage specifies the policy architecture and TRPO batch setting.
  • Reward design: Rewards were sparsified with task-specific success conditions: pole uprightness, escaping the MountainCar valley, or DoublePendulum distance below 0.1.CartPoleSwingup rewarded cos(θ) > 0.8, while DoublePendulum rewarded D < 0.1.
  • Intrinsic motivation: Because NNG-MVG assumes layerwise independence, intrinsic rewards were derived by analyzing one layer at a time.The intrinsic-reward expression uses the ELBO with q(φ) as the prior.
  • Matrix-variate computation: The matrix-variate posterior represents covariance as a Kronecker product, allowing covariance-related computations to use smaller matrices.The method decomposes the covariance gradient and ignores off-diagonal Fisher blocks so each term can be computed efficiently.

D. Additional Results

Additional results compare training convergence and regression performance across Bayesian neural-network methods, including the more flexible posterior used by NNG-MVG.

  • Regression benchmarks: Additional regression comparisons included PBP MV and VMG alongside the paper’s methods.VMG introduced pseudo input-output pairs to increase posterior flexibility.
  • Regression benchmarks: The regression benchmarks reported averaged test RMSE and log-likelihood.Table 5 presents these metrics for the regression benchmarks.
  • Convergence: NNG-FFG converged in fewer iterations than BBB while reaching similar ELBO values.The comparison used training curves with tuned posterior-mean learning rates.
  • Convergence: NNG-MVG converged in roughly 2 times fewer iterations than BBB and surpassed it substantially in ELBO despite greater flexibility and more parameters.NNG-MVG uses a more flexible q distribution than BBB and NNG-FFG.
Loading 1712.02390v2…