sgmcmc: An R Package for Stochastic Gradient Markov Chain Monte Carlo

Abstract

This paper introduces the R package sgmcmc; which can be used for Bayesian inference on problems with large datasets using stochastic gradient Markov chain Monte Carlo (SGMCMC). Traditional Markov chain Monte Carlo (MCMC) methods, such as Metropolis-Hastings, are known to run prohibitively slowly as the dataset size increases. SGMCMC solves this issue by only using a subset of data at each iteration. SGMCMC requires calculating gradients of the log likelihood and log priors, which can be time consuming and error prone to perform by hand. The sgmcmc package calculates these gradients itself using automatic differentiation, making the implementation of these methods much easier. To do this, the package uses the software library TensorFlow, which has a variety of statistical distributions and mathematical operations as standard, meaning a wide class of models can be built using this framework. SGMCMC has become widely adopted in the machine learning literature, but less so in the statistics community. We believe this may be partly due to lack of software; this package aims to bridge this gap.of the major advantages of sgmcmc is that gradients are calculated within the package using automatic differentiation (Griewank and Walther, 2008). This means that users need only specify the log likelihood function and log prior for their model. The package calculates the gradients using TensorFlow (TensorFlow Development Team, 2015), which has recently been made available for R (Allaire et al., 2016). TensorFlow is an efficient library for numerical computation which can take advantage of a wide variety of architectures, as such, sgmcmc keeps much of this efficiency. Both sgmcmc and TensorFlow are available on CRAN, so sgmcmc can be installed by using the standard install.packages function. Though after the TensorFlow package has been installed, the extra install tensorflow() function needs to be run, which installs the required Python implementation of TensorFlow. 1 The sgmcmc package also has a website with vignettes, tutorials and an API reference. 2 SGMCMC methods have become popular in the machine learning literature but less so in the statistics community. We partly attribute this to the lack of available software. To the best of our knowledge, there are currently no R packages available for SGMCMC, probably the most popular programming language within the statistics community. The only package we are aware of which implements scalable MCMC is the Python package edward (Tran et al., 2016). This package implements both SGLD and SGHMC, but does not implement SGNHT or any of the control variate methods.Section 2 introduces MCMC and discusses the software currently available for implementing MCMC algorithms, we discuss the scenarios where sgmcmc is designed to be used. In Section 3 we review the methodology behind the SGMCMC methods implemented in sgmcmc. Section 4 provides a brief introduction to TensorFlow. Section 5 overviews the structure of the package, as well as details of how the algorithms are implemented. Section 6...

Stochastic Gradient Markov Chain Monte Carlo

“…This simple argument suggests that, for the same level of accuracy, we can reduce the computational cost of SGLD by O(N) if we use control variates. This is supported by a number of theoretical results (e.g., Nagapetyan et al 2017;Brosse, Durmus, and Moulines 2018;Baker et al 2019a) which show that, if we ignore the preprocessing cost of findingθ , the computational cost per-effective sample size of SGLD with control variates has a computational cost that is O(1), rather than the O(N) for SGLD with the simple gradient estimator (4).…”

“…The intuition behind this idea is that if each u i (θ) ≈ ∇U i (θ ), then this estimator can have a much smaller variance. Recent works-for example, Baker et al (2019a) and Huggins and Zou (2017) (see Bardenet, Doucet, and Holmes 2017;Bierkens, Fearnhead, and Roberts 2019;Pollock et al 2020, for similar ideas used in different Monte Carlo procedures)-have implemented this control variate technique with each u i (θ ) set as a constant. These approaches propose (i) using SGD to find an approximation to the mode of the distribution we are sampling from, which we denote asθ ; and (ii) set u i (θ ) = ∇U i (θ ).…”

“…These results give theoretical justification for using SGLD, and show we can sample from an arbitrarily good approximation to our posterior distribution if we choose K large enough, and h small enough. They have also been used to show the benefits of using control variates when estimating the gradient, which results in a computational cost that is O(1), rather than O(N), per effective sample size (Chatterji et al 2018;Baker et al 2019a). Perhaps the main benefit of results such as ( 6) is that they enable us to compare the properties of the different variants of SGLD that we will introduce in Section 3, and in particular how different algorithms scale with dimension, d (see Section 3 for details).…”

“…A drawback of these algorithms is that, while producing consistent estimates (Teh, Thiery, and Vollmer 2016), they converge at a slower rate than traditional MCMC algorithms. In recent years, SGMCMC algorithms have become a popular tool for scalable Bayesian inference, particularly in the machine learning community, and there have been numerous methodological (Chen, Fox, and Guestrin 2014;Ma, Chen, and Fox 2015;Dubey et al 2016;Baker et al 2019a) and theoretical developments (Teh, Thiery, and Vollmer 2016;Vollmer, Zygalakis, and Teh 2016;Durmus and Moulines 2017;Dalalyan and Karagulyan 2019) along with new application areas for these algorithms (Balan et al 2015;Gan et al 2015;Wang, Fienberg, and Smola 2015). This article presents a review of some of the key developments in SGMCMC and highlights some of the opportunities for future research.…”

Control variates for stochastic gradient MCMC

“…There is a variety of software available implementing these methods (Tran et al, 2016;Baker et al, 2016). In particular Baker et al (2016) implements the control variate methodology we discuss in this article. This paper investigates stochastic gradient Langevin dynamics (SGLD), a popular SGMCMC algorithm that discretises the Langevin diffusion.…”