[Paper Review] Stochastic Gradient Hamiltonian Monte Carlo

Tianqi Chen, Emily B. Fox+1 more|arXiv (Cornell University)|Feb 17, 2014

Markov Chains and Monte Carlo Methods23 references352 citations

TL;DR

This paper proposes Stochastic Gradient Hamiltonian Monte Carlo (SGHMC), a scalable Bayesian inference method that combines Hamiltonian Monte Carlo with stochastic gradients for large-scale and online data. By introducing a friction term in second-order Langevin dynamics, SGHMC maintains the correct target distribution as its invariant measure despite noisy gradients, enabling efficient, high-acceptance sampling without full-data gradient computation.

ABSTRACT

Hamiltonian Monte Carlo (HMC) sampling methods provide a mechanism for defining distant proposals with high acceptance probabilities in a Metropolis-Hastings framework, enabling more efficient exploration of the state space than standard random-walk proposals. The popularity of such methods has grown significantly in recent years. However, a limitation of HMC methods is the required gradient computation for simulation of the Hamiltonian dynamical system-such computation is infeasible in problems involving a large sample size or streaming data. Instead, we must rely on a noisy gradient estimate computed from a subset of the data. In this paper, we explore the properties of such a stochastic gradient HMC approach. Surprisingly, the natural implementation of the stochastic approximation can be arbitrarily bad. To address this problem we introduce a variant that uses second-order Langevin dynamics with a friction term that counteracts the effects of the noisy gradient, maintaining the desired target distribution as the invariant distribution. Results on simulated data validate our theory. We also provide an application of our methods to a classification task using neural networks and to online Bayesian matrix factorization.

Motivation & Objective

• Address the computational infeasibility of full-gradient Hamiltonian Monte Carlo (HMC) in large-scale or streaming data settings.
• Investigate the failure of naive stochastic gradient HMC due to injected noise disrupting the target distribution.
• Develop a modified HMC framework that preserves the desired posterior as the invariant distribution under stochastic gradients.
• Enable efficient, high-acceptance MCMC sampling in big data and online Bayesian inference scenarios.
• Demonstrate practical effectiveness on Bayesian neural networks and online matrix factorization tasks.

Proposed method

• Propose a stochastic gradient HMC variant that replaces full-data gradients with noisy minibatch gradients.
• Introduce a friction term in second-order Langevin dynamics to counteract the effects of stochastic gradient noise.
• Show that the resulting continuous-time dynamics preserve the target posterior as the invariant distribution.
• Use a small, fixed step size in discretized dynamics to avoid the need for Metropolis-Hastings correction.
• Leverage the central limit theorem to model gradient noise as Gaussian, enabling theoretical analysis.
• Validate the method through theoretical analysis and empirical evaluation on synthetic and real-world data.

Experimental results

Research questions

• RQ1 Why does naive stochastic gradient HMC fail to preserve the correct target distribution?
• RQ2 Can a friction term in Langevin dynamics restore the desired invariant distribution under stochastic gradients?
• RQ3 How does SGHMC compare to SGLD and standard HMC in terms of convergence speed and accuracy on large-scale problems?
• RQ4 Can SGHMC be effectively applied to online Bayesian inference tasks like matrix factorization?
• RQ5 What is the trade-off between step size, computational cost, and sampling accuracy in SGHMC?

Key findings

• Naive stochastic gradient HMC fails because the injected noise disrupts the Hamiltonian dynamics, leading to incorrect invariant distributions.
• The proposed friction term in second-order Langevin dynamics successfully counteracts gradient noise, preserving the target posterior as the invariant distribution.
• SGHMC achieves faster convergence to low test error than SGLD and SGD with momentum on Bayesian neural networks for MNIST classification.
• In online Bayesian matrix factorization on the Movielens dataset, SGHMC achieved a predictive RMSE of 0.8411 ± 0.0011, outperforming SGD and SGD with momentum.
• SGHMC demonstrated comparable runtime to SGLD while achieving better or equal performance, confirming its efficiency and scalability.
• Empirical results show that even with a small fixed step size, SGHMC maintains good sampling quality without requiring Metropolis-Hastings correction.