Natural Langevin Dynamics Enhances Bayesian Neural Network Training

Bayesian methods help prevent overfitting by sampling model parameters from a posterior distribution rather than relying on a single maximum likelihood estimate. Stochastic Gradient Langevin Dynamics (SGLD) adds noise to SGD to approximate this Bayesian posterior for large neural networks and datasets.

The noise variance in SGLD is crucial: smaller noise is needed in sensitive parameter directions. Theory suggests using the inverse Fisher information matrix as noise covariance, matching the posterior variance. However, computing the full Fisher matrix is infeasible for large models.

This paper leverages efficient approximations of the Fisher matrix for deep nets, combining Amari's natural gradient with Fisher-preconditioned Langevin dynamics. This 'natural Langevin dynamics' adapts noise and gradients, improving convergence and regularization in neural network training.

Preconditioning SGLD with the inverse Fisher matrix shapes noise to reflect model sensitivity, reducing noise in critical directions. This also transforms the gradient step into a natural gradient descent, which is invariant to parameter reparameterizations and has optimal convergence properties.

For large neural networks, the authors use a quasi-diagonal Fisher matrix approximation that retains key invariance properties while being computationally efficient. This approximation stores only the diagonal and select off-diagonal terms per neuron, enabling scalable natural Langevin dynamics.

Experiments on MNIST compare four SGLD preconditioners: Euclidean (identity), RMSProp, Diagonal Outer Product, and Quasi-Diagonal Outer Product (QDOP). QDOP with Bayesian posterior ensembling outperforms other SGLD variants and approaches the performance of dropout regularization.

Using an ensemble of parameter samples from the SGLD trajectory better approximates the Bayesian posterior predictive distribution than using a single posterior mean. This ensemble approach yields improved test accuracy and negative log-likelihood, consistent with Bayesian theory.

The study highlights that naive SGLD or diagonal Fisher approximations offer limited gains. The quasi-diagonal Fisher approximation captures more invariances, leading to better noise shaping and gradient updates, thus enhancing regularization and generalization in neural networks.

In summary, natural Langevin dynamics with Fisher matrix preconditioning provides a principled, scalable way to incorporate Bayesian uncertainty in deep learning. It combines natural gradient benefits with noise adapted to model sensitivity, improving training and reducing overfitting.

The authors provide open-source code for their implementation, facilitating adoption and further research. This work bridges Bayesian theory, information geometry, and practical deep learning optimization, advancing robust neural network training methods.