A Complete Recipe for Stochastic Gradient MCMC
Many recent Markov chain Monte Carlo (MCMC) samplers leverage continuous dynamics to define a transition kernel that efficiently explores a target distribution. In tandem, a focus has been on devising scalable variants that subsample the data and use stochastic gradients in place of full-data gradients in the dynamic simulations. However, such stochastic gradient MCMC samplers have lagged behind their full-data counterparts in terms of the complexity of dynamics considered since proving convergence in the presence of the stochastic gradient noise is nontrivial. Even with simple dynamics, significant physical intuition is often required to modify the dynamical system to account for the stochastic gradient noise. In this paper, we provide a general recipe for constructing MCMC samplers-including stochastic gradient versions-based on continuous Markov processes specified via two matrices. We constructively prove that the framework is complete. That is, any continuous Markov process that provides samples from the target distribution can be written in our framework. We show how previous continuous-dynamic samplers can be trivially "reinvented" in our framework, avoiding the complicated sampler-specific proofs. We likewise use our recipe to straightforwardly propose a new state-adaptive sampler: stochastic gradient Riemann Hamiltonian Monte Carlo (SGRHMC). Our experiments on simulated data and a streaming Wikipedia analysis demonstrate that the proposed SGRHMC sampler inherits the benefits of Riemann HMC, with the scalability of stochastic gradient methods.
The authors propose a general framework for designing new MCMC samplers, including methods that use stochastic gradients. Their approach is to define a stochastic dynamical system whose stationary distribution is the target distribution from which we want to sample from. The stochastic dynamical system is represented through a stochastic differential equation that is simulated through an epsilon-discretization approach. As the step-size parameter epsilon goes to zero, the bias in the simulation vanishes. The proposed approach can handle stochastic approximations obtained by sub-sampling the data in mini-batches. For this, the additional noise produced by the sub-sampling process is removed from the noise used to simulate the SDE. The proposed framework constructs the SDE using two state dependent matrices D(z) and Q(z) and by modifying these matrices the authors can explore the space of MCMC methods. The authors proof that their approach samples from the correct stationary distribution as the step size epsilon tends to zero. The authors also proof that any continuous Markov process with a particular stationary distribution can be represented in their framework and consequently, such framework is complete. After this, the authors show how several already existing samplers can be represented in their framework: Hamiltonian Monte Carlo, Stochastic Gradient Hamiltonian Monte Carlo, Stochastic Gradient Langevin Dynamics, Stochastic Gradient Riemannian Langevin Dynamics and Stochastic Nose-Hoover Thermostast. Finally, the authors proposed based on their framework a new sampler technique called Stochastic Gradient Riemann Hamiltonian Monte Carlo (SGRHMC). This method uses the Ficsher information matrix to precondition the dynamics through the matrices D(z) and Q(z). The SGRHMC method is shown to provide very good results on experiments with simulated data and streeming data extracted from Wikipedia.Quality:The paper is technically very sound, with theorems showing the applicability of the framework to any target distribution and its completeness in terms of how many continuous Markov processes it can represent.Clarity:The paper is clearly written and well organized. A few things could be improved for clarity: How do you construct in practice the matrix \hat{B}_t? When you substract \hat{B}_t to eliminate the stochastic gradient noise term, what happens if you get a covariance matrix that is not possitive definite? Perhaps this is not possible if \epsilon is small. In line 232 the minus sign in Q should probably be changed from the upper right identity matrix to the lower left one.Originality:The proposed framework is original. It allows the authors to re-design existing MCMC samples and propose new ones with better properties. The author's contribution is significantly different from previous work.Significance:This work is highly significant, specifying a principled way to design new scalable MCMC samples that work by sub-sampling the data and are able to sample from the correct target distribution in the limit of the step size tending to zero.