Stochastic Gradient Langevin dynamics that Exploit Neural Network Structure

Zachary Nado, Jasper Snoek, Roger B. Grosse, David Duvenaud, Bowen Xu, James Martens

Abstract: Tractable approximate Bayesian inference for deep neural networks remains challenging. Stochastic Gradient Langevin Dynamics (SGLD) offers a tractable approximation to the gold standard of Hamiltonian Monte Carlo. We improve on existing methods for SGLD by incorporating a recently-developed tractable approximation of the Fisher information, known as K-FAC, as a preconditioner.