Stochastic gradient Markov chain Monte Carlo

Christopher Nemeth; Paul Fearnhead

Source-linked AI summary

Stochastic gradient Markov chain Monte Carlo

Christopher Nemeth, Paul Fearnhead

arXiv:1907.06986v1Jul 16, 2019stat.COstat.ML

TL;DR

Large datasets make exact MCMC expensive because each iteration generally processes all data. This paper reviews SGMCMC methods that use subsampled stochastic gradients, develops their theoretical basis, and compares them with MCMC; the reported benchmarks show substantial computational savings with comparable accuracy, subject to known biases and scope limitations.

Problem
Traditional MCMC can be computationally prohibitive for large datasets because exact inference generally requires processing all data at every iteration.
Method
The paper reviews SGMCMC algorithms based on diffusion approximations, including methods that replace full-data gradients with stochastic estimates from data subsamples.
Results
SGMCMC algorithms are reported to achieve parameter-inference accuracy comparable to traditional MCMC while significantly reducing computational cost.
Takeaways & Limitations
SGMCMC provides a scalable class of Monte Carlo methods with significant computational savings when the algorithms are well-tuned.
Takeaways & Limitations
For dependent data such as network, time-series, and spatial data, naive random subsampling can produce biased estimates of the potential and its gradient.

Abstract

from arXiv · show

Markov chain Monte Carlo (MCMC) algorithms are generally regarded as the gold standard technique for Bayesian inference. They are theoretically well-understood and conceptually simple to apply in practice. The drawback of MCMC is that in general performing exact inference requires all of the data to be processed at each iteration of the algorithm. For large data sets, the computational cost of MCMC can be prohibitive, which has led to recent developments in scalable Monte Carlo algorithms that have a significantly lower computational cost than standard MCMC. In this paper, we focus on a particular class of scalable Monte Carlo algorithms, stochastic gradient Markov chain Monte Carlo (SGMCMC) which utilises data subsampling techniques to reduce the per-iteration cost of MCMC. We provide an introduction to some popular SGMCMC algorithms and review the supporting theoretical results, as well as comparing the efficiency of SGMCMC algorithms against MCMC on benchmark examples. The supporting R code is available online.

1 Introduction

MCMC provides theoretically exact Bayesian inference but becomes computationally impractical as datasets and models grow. SGMCMC addresses this bottleneck by using subsampled stochastic gradients, trading exactness for substantially lower computational cost.

MCMC approximates posterior distributions with samples from a Markov chain whose invariant distribution is the posterior.More efficient variants use gradients to explore the posterior faster.
Growing datasets and increasingly complex, high-dimensional models can make traditional MCMC computationally prohibitive.This limits its use in fields including population genetics, brain imaging, and natural language processing.
Scalable alternatives use parallel partial posteriors or data subsampling, but introduce trade-offs between computational speed-up and accuracy.Parallel methods must merge samples from partial posteriors, while subsampling approximates full-data calculations.
SGMCMC methods derive from discrete-time approximations to continuous-time diffusion processes and use subsampled stochastic gradients.SGLD replaces the full-data gradient in unadjusted Langevin dynamics with a stochastic approximation computed from a subsample.
The paper introduces SGMCMC algorithms, reviews their theoretical results, and compares their efficiency with MCMC on benchmark examples.It also highlights opportunities for future research.

2 Langevin-based Stochastic Gradient MCMC

Langevin-based SGMCMC uses discretized diffusion dynamics and unbiased subsampled gradients to reduce per-iteration cost while approximating the posterior. Its accuracy depends on step size, gradient-estimator variance, initialization, and regularity conditions, with theory characterizing both convergence and distortion.

2.1 The Langevin Diffusion: The Langevin diffusion has the posterior π as its stationary distribution, providing the foundation for Langevin-based MCMC.Its continuous-time dynamics are approximated numerically because general Langevin dynamics are intractable.
2.2 Approximate MCMC using the Langevin Diffusion: ULA omits the Metropolis correction, producing a biased approximation of π despite its simpler iteration.Its full-gradient cost remains O(N), and its smaller required step size can necessitate more iterations.
2.3 Stochastic Gradient Langevin Dynamics: SGLD replaces the full gradient with an unbiased estimate from a random subsample n << N, greatly reducing per-iteration cost.With decreasing step size, the algorithm was informally justified as converging to the true Langevin dynamics.
2.3 Stochastic Gradient Langevin Dynamics: Control variates can reduce gradient-estimator variance from O(N^2/n) to O(N/n), suggesting an O(N) computational-cost reduction for comparable accuracy.The variance reduction is expected when θ is close to the approximate mode θ̂, while performance can worsen when it is far away.
2.4 Theory for SGLD: SGLD error bounds separate initialization bias, Euler-discretization error, and unbiased-gradient error, and support arbitrarily accurate posterior approximation with sufficiently large K and small h.The bounds require technical conditions, including assumptions that promote geometric ergodicity; they do not directly specify K, h, or subsample size in practice.

3 A General Framework for Stochastic Gradient MCMC

The paper develops a general stochastic-gradient MCMC framework by constructing diffusions with the desired stationary distribution and replacing exact gradients with unbiased subsampled estimates. The framework includes popular algorithms, whose accuracy and stability depend on diffusion choices, gradient-noise variance, and dimensionality.

General framework: The framework extends Langevin-based SGMCMC to a wider class of diffusion processes, including stochastic-gradient versions of Hamiltonian Monte Carlo.These algorithms arise by choosing drift, diffusion, and curl components so the diffusion has a specified stationary distribution.
General framework: The general state can include parameters and auxiliary variables such as velocity, allowing Hamiltonian-style dynamics while preserving the target as the parameter marginal.When auxiliary variables are included, the joint stationary density can be constructed so that the parameter and auxiliary-variable components are independent.
Stochastic gradients: SGMCMC replaces the exact gradient with an unbiased data-subsampled estimate and may correct its variance to avoid inflating the approximate target distribution.The variance correction requires D(ζt) − h B̂(ζt) to be positive semi-definite; otherwise the resulting algorithm may be unstable.
Algorithm design: Diffusion and curl matrices affect convergence: larger diffusion promotes escape from local modes, whereas smaller diffusion improves accuracy within a local mode.The matrix choices also determine the noise level and influence the sampler’s convergence rate.
Algorithm variants: SGLD and SG-HMC are common instances of the framework, while SG-RLD, SG-RHMC, and SG-NHT modify geometry, curvature adaptation, or gradient-noise control.SG-HMC adds velocity to improve mixing, whereas Riemannian variants require accurate and computationally manageable curvature estimates.
Theory: O(d2/ϵ2) iterations for SGLD versus O(d/ϵ) for SG-HMC are the stated Wasserstein-distance bounds under smooth, strongly log-concave targets.The SG-HMC advantage depends on sufficiently small gradient-estimate variance; excessive noise can make it behave like SGLD.

4 Diagnostic Tests

The section develops Stein-discrepancy diagnostics for assessing convergence and tuning stochastic-gradient MCMC, whose samples are asymptotically biased. A Gaussian example shows that kernel Stein discrepancy can identify a useful step size when standard effective sample size can prefer a biased setting.

Motivation: Standard MCMC diagnostics are inappropriate for SGMCMC because its samples target an asymptotically biased approximation rather than the exact stationary posterior.The section motivates Stein-based methods as diagnostics that assess how accurately samples approximate a distribution.
Discrepancy measures: A discrepancy compares empirical and target expectations over a function class, with convergence of the discrepancy corresponding to weak convergence under suitable choices.A small discrepancy also supports accurate posterior-expectation estimates when the functions of interest belong to the chosen class.
Stein discrepancy: Stein discrepancy avoids the unknown target expectations by using functions whose expectations under the target are zero, constructed from processes having that target as invariant distribution.The kernel Stein approach uses target gradients and can be computed from kernel evaluations over all pairs of sample points.
Kernel choice: The inverse multi-quadratic kernel detects non-convergence when c > 0 and β ∈ (−1, 0), whereas some Gaussian and Matern kernels may fail in dimensions d ≥ 3.Kernel choice is therefore important when using kernel Stein discrepancy as a convergence diagnostic.
Practical considerations: Kernel Stein discrepancy is more computationally expensive than standard MCMC metrics because its cost scales quadratically with sample size, although the computation can be parallelised.The example evaluates it alongside effective sample size over multiple Langevin step sizes.
Illustration: In the Gaussian example, h = 0.1 produced the closest-looking samples and the lowest kernel Stein discrepancy, whereas h = 1 maximised effective sample size.Large h produced over-dispersed samples, while small h limited exploration, supporting kernel Stein discrepancy over standard MCMC metrics for this tuning task.

5 Extending the SGMCMC framework

The SGMCMC framework extends beyond unconstrained Langevin dynamics to improve mixing and accommodate constrained parameters, structured models, and dependent data. These extensions rely on transformations or alternative stochastic processes, while unbiased gradient estimation remains difficult for many dependent-data models.

SGMCMC extensions aim to improve Markov-chain mixing and broaden applicability beyond standard Langevin dynamics.
Unconstrained Euclidean parameters and additive potentials restrict SGMCMC for constrained spaces and many time-series or spatial models.
Constrained parameters can violate their boundaries under Langevin dynamics, while transformations may mix poorly near boundaries.
Mirrored Langevin transforms constrained sampling to an unconstrained space with convergence rates comparable to unconstrained SGLD.
For simplex-constrained LDA parameters, specialized methods include stochastic gradient Riemannian Langevin dynamics and CIR-based sampling, though transformations can create boundary bias.
Independent-data subsampling yields unbiased gradients, but common random subsampling is biased for many network, time-series, and spatial models; state-space bias can decay geometrically with buffer size.Specialized methods use model structure or buffering to address some dependent-data settings.

6 Simulation Study

The study compares SGMCMC algorithms with STAN across logistic regression, Bayesian neural-network classification, and Bayesian matrix factorisation tasks. SGMCMC can achieve similar accuracy or predictive performance with substantially lower computational cost, while control variates reduce convergence time and burn-in.

Experimental design: The experiments evaluate SGMCMC on logistic regression, MNIST image classification with a Bayesian neural network, and movie recommendation using Bayesian probabilistic matrix factorisation.STAN provides a benchmark for the posterior-based comparisons.
Logistic regression model: Control variate algorithms require fewer iterations for convergence because optimisation initializes sampling near the posterior mode, effectively replacing burn-in.SG-HMC and SG-NHT converge faster than SGLD but incur higher computational cost from leapfrog steps.
Logistic regression model: In logistic regression, all SGMCMC algorithms are 10 to 100 times faster than STAN when d = 100, with only slightly worse log-loss and kernel Stein discrepancy.STAN has the lowest diagnostic values, but the SGMCMC results are similar when the samplers are well-tuned.
Bayesian neural network: In the neural-network experiment, control variate methods show some indication of improved predictive performance, while SG-HMC and SG-NHT perform best among standard methods.The comparison uses log-loss on held-out test data and balances computational cost across algorithms.
Bayesian probabilistic matrix factorisation: In matrix factorisation, SG-HMC achieves lower test RMSE than SGLD, and both control variate versions improve results immediately by starting at the posterior mode.The control variate methods remove the burn-in phase, producing lower RMSE overall.
Bayesian probabilistic matrix factorisation: All tested SGMCMC algorithms outperform the STAN baseline on RMSE, although the comparison is affected by STAN using fewer iterations and potentially not converging.STAN took 3.5 days, whereas SGLD, SGLD-CV, SG-HMC and SG-HMC-CV took 3.1, 3.5, 16.4 and 16.8 hours, respectively.

7 Discussion

The discussion concludes that SGMCMC offers theoretically supported, scalable inference with accuracy comparable to traditional MCMC. It identifies broader theoretical coverage, robust tuning, and more user-friendly software as priorities for future development.

Discussion: SGMCMC uses data subsampling to reduce MCMC computational cost while providing parameter inference comparable in accuracy to traditional MCMC.The paper reviews the growing SGMCMC literature, algorithms, and theoretical results.
Future development: Current theoretical error bounds mostly assume log-concave target distributions, limiting established guarantees for broader models such as multimodal posteriors.Relaxing log-concavity is identified as a direction for obtaining bounds for a broader class of models.
Future development: SGMCMC efficacy depends on step-size tuning, while standard acceptance-rate rules do not apply and Stein-discrepancy tuning can be computationally expensive.The paper calls for robust automated tuning methods comparable to adaptive HMC.
Future development: More mature SGMCMC software with automatic differentiation and adaptive step-size tuning could lower the entry barrier for non-experts.The discussion contrasts this need with the user-friendly software ecosystem supporting traditional MCMC.