Structured Stochastic Gradient MCMC

Stochastic gradient Markov Chain Monte Carlo (SGMCMC) is a scalable algorithm for asymptotically exact Bayesian inference in parameterrich models, such as Bayesian neural networks. However, since mixing can be slow in high dimensions, practitioners often resort to variational inference (VI). Unfortunately, VI makes strong assumptions on both the factorization and functional form of the posterior. To relax these assumptions, this work proposes a new non-parametric variational inference scheme that combines ideas from both SGMCMC and coordinate-ascent VI. The approach relies on a new Langevin-type algorithm that operates on a self-averaged posterior energy function, where parts of the latent variables are averaged over samples from earlier iterations of the Markov chain. This way, statistical dependencies between coordinates can be broken in a controlled way, allowing the chain to mix faster. This scheme can be further modified in a dropout manner, leading to even more scalability. We test our scheme for ResNet-20 on CIFAR-10, SVHN, and FMNIST. In all cases, we find improvements in convergence speed and/or final accuracy compared to SGMCMC and parametric VI.

Automated summary

Researchers propose a new non-parametric variational inference scheme that combines ideas from SGMCMC and coordinate-ascent VI, aiming to relax the assumptions on the posterior distribution. They introduce a new Langevin-type algorithm that operates on a self-averaged posterior energy function to break the statistical dependencies between coordinates, allowing for faster mixing. The scheme is tested on various tasks and shows improvements in convergence speed and/or final accuracy compared to SGMCMC and parametric VI.