Scalable Bayesian Learning with posteriors

TL;DR

This paper introduces 'posteriors', a PyTorch library that makes scalable Bayesian learning feasible for large models and datasets, combining stochastic gradient MCMC, optimization, and ensemble methods.

Contribution

It presents a new library, 'posteriors', and a tempered stochastic gradient MCMC framework that unifies Bayesian inference, optimization, and ensemble methods for large-scale models.

Findings

The library enables scalable Bayesian inference in large models.

The tempered SGMCMC approach bridges Bayesian sampling and optimization.

Experiments demonstrate improved Bayesian approximation and insights into the cold posterior effect.

Abstract

Although theoretically compelling, Bayesian learning with modern machine learning models is computationally challenging since it requires approximating a high dimensional posterior distribution. In this work, we (i) introduce posteriors, an easily extensible PyTorch library hosting general-purpose implementations making Bayesian learning accessible and scalable to large data and parameter regimes; (ii) present a tempered framing of stochastic gradient Markov chain Monte Carlo, as implemented in posteriors, that transitions seamlessly into optimization and unveils a minor modification to deep ensembles to ensure they are asymptotically unbiased for the Bayesian posterior, and (iii) demonstrate and compare the utility of Bayesian approximations through experiments including an investigation into the cold posterior effect and applications with large language models.