Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Eigenvalue Corrected Noisy Natural Gradient

Juhan Bae, Guodong Zhang, Roger Grosse

TL;DR

The paper tackles uncertainty estimation in neural networks by extending Noisy K-FAC to a more expressive posterior. It introduces eigenvalue corrected matrix-variate Gaussian (EMVG) posteriors via Noisy EK-FAC, computing diagonal variance in the EK-FAC eigenbasis to increase posterior flexibility without sacrificing tractability. The method combines EK-FAC's diagonal correction with the noisy natural gradient framework, yielding improved ELBO and predictive performance on regression and classification benchmarks. This approach enhances practical Bayesian deep learning by delivering scalable, better-calibrated uncertainty estimates in large models.

Abstract

Variational Bayesian neural networks combine the flexibility of deep learning with Bayesian uncertainty estimation. However, inference procedures for flexible variational posteriors are computationally expensive. A recently proposed method, noisy natural gradient, is a surprisingly simple method to fit expressive posteriors by adding weight noise to regular natural gradient updates. Noisy K-FAC is an instance of noisy natural gradient that fits a matrix-variate Gaussian posterior with minor changes to ordinary K-FAC. Nevertheless, a matrix-variate Gaussian posterior does not capture an accurate diagonal variance. In this work, we extend on noisy K-FAC to obtain a more flexible posterior distribution called eigenvalue corrected matrix-variate Gaussian. The proposed method computes the full diagonal re-scaling factor in Kronecker-factored eigenbasis. Empirically, our approach consistently outperforms existing algorithms (e.g., noisy K-FAC) on regression and classification tasks.

Eigenvalue Corrected Noisy Natural Gradient

TL;DR

The paper tackles uncertainty estimation in neural networks by extending Noisy K-FAC to a more expressive posterior. It introduces eigenvalue corrected matrix-variate Gaussian (EMVG) posteriors via Noisy EK-FAC, computing diagonal variance in the EK-FAC eigenbasis to increase posterior flexibility without sacrificing tractability. The method combines EK-FAC's diagonal correction with the noisy natural gradient framework, yielding improved ELBO and predictive performance on regression and classification benchmarks. This approach enhances practical Bayesian deep learning by delivering scalable, better-calibrated uncertainty estimates in large models.

Abstract

Variational Bayesian neural networks combine the flexibility of deep learning with Bayesian uncertainty estimation. However, inference procedures for flexible variational posteriors are computationally expensive. A recently proposed method, noisy natural gradient, is a surprisingly simple method to fit expressive posteriors by adding weight noise to regular natural gradient updates. Noisy K-FAC is an instance of noisy natural gradient that fits a matrix-variate Gaussian posterior with minor changes to ordinary K-FAC. Nevertheless, a matrix-variate Gaussian posterior does not capture an accurate diagonal variance. In this work, we extend on noisy K-FAC to obtain a more flexible posterior distribution called eigenvalue corrected matrix-variate Gaussian. The proposed method computes the full diagonal re-scaling factor in Kronecker-factored eigenbasis. Empirically, our approach consistently outperforms existing algorithms (e.g., noisy K-FAC) on regression and classification tasks.

Paper Structure

This paper contains 18 sections, 21 equations, 3 figures, 2 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Natural Gradient
  3. Kronecker-Factored Approximate Curvature
  4. An Alternative Interpretation of Natural Gradient
  5. Eigenvalue Corrected Kronecker-Factored Approximate Curvature
  6. Variational Bayesian Neural Networks
  7. Noisy Natural Gradient
  8. Fitting Matrix-Variate Gaussian Posteriors with Noisy K-FAC
  9. Methods
  10. Related Work
  11. Experiments
  12. Regression
  13. Classification
  14. Conclusion
  15. Matrix-Variate Gaussian
  16. ...and 3 more sections

Figures (3)

  • Figure 1: A cartoon illustration to describe the relationships of FFG, MVG, and EMVG.
  • Figure 2: The diagonal re-scaling factor in K-FAC has Kronecker structure with $n + p$ degrees of freedom. The diagonal re-scaling matrix in EK-FAC is the second moment of the gradient vector with $n \times p$ degrees of freedom.
  • Figure 3: Training curve for all three methods. Note that BBB, noisy K-FAC, and noisy EK-FAC use FFG, MVG, EMVG accordingly. EMVG has the most flexible variational posterior distribution.