PaperlayerView exact revision on arXiv

Source-linked AI summary

Stochastic Gradient Hamiltonian Monte Carlo

Tianqi Chen, Emily B. Fox, Carlos Guestrin

arXiv:1402.4102v2stat.MEcs.LGstat.ML

TL;DR

HMC's full-gradient requirement limits Bayesian sampling for massive or streaming datasets, motivating stochastic-gradient alternatives. The paper shows that naive stochastic-gradient HMC can fail, then introduces friction-based second-order Langevin dynamics; simulations and applications validate the proposed approach. In practice, estimating the friction-related quantity inaccurately and using finite step sizes can introduce bias or reduce efficiency.

  • Problem

    Full-gradient computation makes HMC impractical for massive or streaming datasets, motivating scalable sampling with noisy minibatch gradients.

  • Method

    The paper introduces stochastic-gradient HMC using second-order Langevin dynamics with friction to counteract noisy gradients while preserving the target distribution.

  • Results

    Simulated experiments and applications to Bayesian neural-network classification and online Bayesian matrix factorization demonstrate the proposed algorithm's effectiveness.

  • Takeaways & Limitations

    Friction is crucial in this stochastic-gradient setting because it counterbalances noisy gradients and supports high-quality distant sampling steps.

  • Takeaways & Limitations

    Inaccurate estimation of the friction-related quantity and finite nonzero step sizes can introduce sampling bias or reduce efficiency.

Abstract

from arXiv · show

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.

1. Introduction

HMC efficiently explores distant states but requires full-data gradients, limiting its use for massive or streaming datasets. The paper analyzes stochastic-gradient HMC, identifies failure in the naive approach, and introduces friction-based correction validated in experiments.

  • Motivation: HMC enables distant proposals with high acceptance probabilities, supporting rapid exploration of the state space.Its dynamics preserve the target distribution, while discretization is corrected with Metropolis-Hastings.
  • Motivation: Full-gradient computation makes standard HMC infeasible for datasets with millions to billions of observations or streaming data.The paper motivates minibatch-based noisy gradients as a way to scale Bayesian sampling.
  • Problem: The naive stochastic-gradient replacement destroys the desired target distribution's stationarity and can require costly full-data Metropolis-Hastings correction.Long simulations before correction can produce large Hamiltonian deviations and low acceptance rates.
  • Contribution: The proposed stochastic-gradient HMC adds friction through second-order Langevin dynamics to counteract noisy-gradient effects.The method is designed to preserve the desired target distribution as invariant.
  • Evaluation: Simulated experiments, Bayesian neural-network classification, and online Bayesian matrix factorization demonstrate the proposed algorithm's effectiveness.The experiments compare exact HMC, naive stochastic-gradient HMC, friction-based SGHMC, and SGLD.

2. Hamiltonian Monte Carlo

HMC augments target variables with momentum and simulates reversible, energy-preserving Hamiltonian dynamics to generate distant proposals. In practice, leapfrog discretization requires Metropolis-Hastings correction, while tuning extensions improve flexibility.

  • Hamiltonian formulation: HMC samples θ by augmenting it with momentum r and defining a joint distribution whose θ-marginal is the target posterior.The mass matrix M parameterizes the kinetic-energy term and can precondition sampling.
  • Hamiltonian formulation: The Hamiltonian combines potential energy U(θ) with kinetic energy 1/2 r^T M^-1r.The total energy represents the physical system's position and momentum contributions.
  • Dynamics: Hamiltonian dynamics update position using M^-1r and momentum using -∇U(θ), generating proposals across the state space.The continuous dynamics are specified by Eq. (4).
  • Dynamics: Reversibility and energy preservation make the continuous dynamics leave the target invariant and yield accepted proposals.Distant proposals remain possible because the dynamics can move far from the initial state.
  • Implementation: Leapfrog discretization introduces energy errors, so HMC applies a Metropolis-Hastings correction, typically retaining high acceptance rates.The algorithm resamples momentum, simulates discretized dynamics for m steps, and then accepts or rejects the proposal.
  • Extensions: Automatic tuning methods adjust step size and simulation length, while Riemann manifold HMC adapts the mass matrix using curvature information.These extensions make HMC more flexible across target distributions.

3. Stochastic Gradient HMC

Stochastic gradients disrupt HMC’s target invariance, while a friction-corrected second-order Langevin formulation restores it and connects to momentum-based methods.

  • Stochastic-gradient dynamics: Minibatch stochastic gradients replace full-data gradient computation, introducing noise whose covariance may depend on parameters and sample size.The Gaussian noise approximation becomes more accurate as minibatch size increases, but smaller minibatches provide the computational gains sought for large datasets.
  • Naive stochastic-gradient HMC: Naively replacing the exact gradient injects diffusion into momentum updates and destroys invariance of the target distribution before discretization error is considered.The resulting noisy dynamics can increase entropy and move toward distributions far from the target.
  • Naive stochastic-gradient HMC: Naive stochastic-gradient HMC requires costly Metropolis-Hastings corrections or long simulations with low acceptance probabilities, undermining large-data efficiency.Noisy trajectories can diverge substantially from noise-free trajectories, creating a computation-versus-efficiency tradeoff.
  • SGHMC with friction: Adding a friction term to second-order Langevin dynamics counteracts noisy-gradient effects and makes π(θ, r) the unique stationary distribution.The friction term reduces Hamiltonian energy and limits the influence of gradient noise.
  • SGHMC with friction: The proposed SGHMC dynamics preserve the desired target distribution with noise present, provided the estimated noise covariance matches B.This invariance is the central correction to naive stochastic-gradient HMC.
  • Connections and practical limits: With large friction, SGHMC reduces to first-order Langevin dynamics, while its momentum formulation connects naturally to SGD with momentum.The connection to SGLD uses M^-1 as a preconditioning matrix.
  • Connections and practical limits: Finite step sizes and inaccurate estimation of B introduce sampling bias, creating an efficiency–accuracy tradeoff in practice.The paper considers small nonzero step sizes to gain efficiency despite some bias.

4. Experiments

Experiments show that friction stabilizes stochastic-gradient Hamiltonian dynamics, while SGHMC efficiently explores correlated targets and performs well on neural-network and matrix-factorization tasks.

  • Simulated Scenarios: Naïve stochastic-gradient HMC diverges significantly from the target unless corrected by a costly Metropolis-Hastings step.Exact HMC and SGHMC maintain the target as the step size approaches zero.
  • Simulated Scenarios: Friction corrects the divergence of noisy Hamiltonian trajectories and supports the correct stationary distribution.Periodic momentum resampling can also correct trajectory divergence, but does not ensure the correct target distribution for naïve stochastic-gradient HMC.
  • Simulated Scenarios: SGHMC achieves lower covariance-estimation error at very low autocorrelation times than SGLD on a positively correlated bivariate Gaussian.SGLD exhibits random-walk behavior that makes tail exploration challenging, whereas SGHMC moves along distribution contours.
  • Bayesian Neural Networks for Classification: SGHMC converges to low MNIST test error much faster than SGLD, while sampling-based methods achieve better final results than optimization-based methods.Backpropagation dominates runtime in this setting, making the compared methods have similar computational cost.
  • Online Bayesian Probabilistic Matrix Factorization for Movie Recommendations: Both SGHMC and SGLD outperform optimization-based methods on online Bayesian probabilistic matrix factorization, with similar prediction results and comparable per-iteration time.The experiment uses 5-fold cross validation on Movielens with minibatches of 4,000 ratings.

5. Conclusion

The conclusion presents SGHMC as an efficient large-scale sampling method that combines stochastic gradients with Hamiltonian dynamics while correcting the instability of the naïve approach.

  • 5. Conclusion: SGHMC uses stochastic gradient estimates to avoid costly full-gradient computation while generating efficient distant steps for large-scale or online sampling.The method targets the challenge of moving between modes in MCMC-based inference.
  • 5. Conclusion: The naïve stochastic-gradient modification can diverge and behave poorly, whereas a friction term counteracts noisy-gradient effects and preserves the desired invariant distribution.The conclusion reports validation in simulated and real-data experiments.
  • 5. Conclusion: Empirical results demonstrate the practical value of the friction modification across simulated experiments and real-data applications.The paper identifies combining adaptive HMC techniques with SGHMC as a possible next step.

Supplementary Material

The supplementary material introduces the Fokker–Planck framework used to describe how the stochastic dynamics evolve the joint position–momentum distribution.

  • Fokker–Planck Framework: The Fokker–Planck equation describes the time evolution of distributions induced by a stochastic differential equation.The formulation uses drift and diffusion terms for the random variable z.
  • Fokker–Planck Framework: For the paper’s dynamics, z=(θ,r), with position θ and momentum r, and noise added only to r while depending on θ.This structure determines the associated diffusion representation.
  • Fokker–Planck Framework: The derivation uses the compact Fokker–Planck form together with the fact that diffusion derivatives vanish under the paper’s noise structure.Specifically, the relevant diffusion terms depend on θ rather than the momentum coordinates.

B. Proof of Theorem 3.1

The proof analyzes entropy evolution under Hamiltonian dynamics and stochastic-gradient noise to establish how the noise affects the target distribution.

  • Proof of Theorem 3.1: Hamiltonian dynamics are expressed through a drift matrix acting on the gradient of the Hamiltonian.The corresponding Fokker–Planck equation combines drift and diffusion contributions.
  • Proof of Theorem 3.1: The entropy evolution separates into a noise-free Hamiltonian component and a stochastic-gradient noise component.The Hamiltonian component does not change entropy because the mixed partial derivatives cancel.
  • Proof of Theorem 3.1: The stochastic-gradient noise contribution is computed by integration by parts under boundary conditions on the probability density and its gradient.These conditions allow the boundary terms to vanish and complete the theorem’s argument.

C. Proof of Corollary 3.1

Under the stated regularity and full-rank assumptions, the noisy dynamics cannot preserve the target distribution because entropy increases when initialized at that distribution.

  • Regularity assumptions: The assumptions ensure that π and ln π∇π vanish at infinity, satisfying the conditions required for the entropy argument.The Hamiltonian diverges as the parameter and momentum norms grow, and π is defined from exp(−H).
  • Entropy contradiction: Full-rank Fisher information and f′′(p) > 0 imply positive entropy growth at π.The proof concludes that ∂_t h(p_t(θ,r))|_{p_t=π} > 0.
  • Entropy contradiction: Because an invariant distribution would not change over time, the positive entropy derivative contradicts π being invariant.Thus the stochastic-gradient dynamics fail to retain the desired target distribution even before discretization.

D. FPE for Second-Order Langevin Dynamics

The Fokker–Planck equation for second-order Langevin dynamics exposes how diffusion and antisymmetric dynamics combine, and verifies that the desired distribution is stationary. The framework also clarifies that SGHMC has a time-reversal property without conventional detailed balance.

  • Fokker–Planck formulation: The paper introduces second-order Langevin dynamics and uses a compact Fokker–Planck equation to track distribution evolution.The equation is written in terms of the diffusion matrix D, antisymmetric component G, Hamiltonian H, temperature τ, and density p_t.
  • Fokker–Planck formulation: The identity ∇^T[G∇p_t] = 0 reduces the Fokker–Planck equation to a form suitable for checking stationarity.The cancellation follows from equality of mixed partial derivatives.
  • Stationary distribution: Substituting π(θ,r) ∝ e^(−H(θ,r)/τ) gives ∂_tπ(θ,r) = 0, so π is the stationary distribution.The construction therefore preserves the target distribution under the continuous-time dynamics.
  • Generalized dynamics: Allowing G to depend on θ can incorporate local curvature, but a correction term is then required to preserve the stationary distribution.The paper connects this extension to Riemann-geometric methods.
  • Reversibility: SGHMC is not conventionally reversible, yet reversing time together with momentum signs yields a related reverse process.The target remains stationary without relying on detailed balance.

F. Convergence Analysis

The convergence analysis relates inaccurate noise estimation to error in the final stationary distribution through χ²-divergence. Smaller estimation error improves accuracy, but SGHMC mixing-rate bounds remain unresolved because the process is irreversible.

  • Error from noise estimation: SGHMC efficiency involves choosing a small nonzero step size while approximating the noise matrix B with ˆB.The discrepancy δS = B − ˆB enters the continuous-time dynamics as an unaccounted error term.
  • Error from noise estimation: In the special case V = C, the stationary distribution can be calculated exactly, and the approximation approaches the true distribution as ϵ becomes small.This establishes a direct small-step-size accuracy relationship in that special case.
  • Convergence bound: The χ²-divergence measures distance from the true distribution, and Theorem F.1 shows that this divergence decreases as δ becomes smaller.The analysis decomposes the divergence change into contributions from the correct process and the noise-estimation error.
  • Limitations: The same theorem can quantify SGLD error, but corresponding SGHMC mixing-rate bounds are unclear because SGHMC is irreversible.The paper leaves these SGHMC bounds for future work.
  • Convergence bound: The contraction argument is formulated for a continuous-time Markov process, although an analogous guarantee can be established for a discrete-time transition kernel.The paper points to prior work for discrete-time details.

G. Setting SGHMC Parameters

SGHMC can be parameterized like SGD with momentum using a learning rate, momentum decay, and minibatch size. Keeping the stochastic-gradient noise term small requires trading learning rate against dataset size and minibatch choices.

  • Momentum parameterization: Rewriting SGHMC yields an SGD-with-momentum form in which η is the learning rate and 1−α is the momentum term.The paper recommends this equivalent parameterization as cleaner.
  • Noise estimation: The stochastic-gradient noise can be modeled with ˆβ, either ignored by setting ˆβ = 0 or estimated using empirical Fisher information.The estimated alternative is ˆβ = η ˆV / 2.
  • Parameter trade-offs: The main tunable parameters are learning rate η, momentum decay α, and minibatch size |D̃|, while β quantifies noise induced by the stochastic gradient.Maintaining a small β keeps the stationary distribution close to the target.
  • Parameter trade-offs: As training size grows, one must use a small learning rate η = O(1/|D|) or increase the minibatch size to control stochastic-gradient error.The paper describes this as a trade-off rather than a free improvement.
  • Practical settings: In practice, minibatches of hundreds, α around 0.01 or 0.1, and η = γ/|D| with γ around 0.1 or 0.01 work well.These settings are also commonly used for SGD with momentum.

H. Experimental Setup

The experiments use Bayesian neural network and probabilistic matrix factorization models, sampling parameters and precision hyperparameters alternately with SGHMC or SGLD and Gibbs updates. Method-specific training settings are selected by validation or cross-validation.

  • Bayesian neural network: The Bayesian neural network models digit labels with two weight matrices, bias terms, Gaussian parameter priors, and gamma priors on precision terms.For MNIST, the input dimension is d = 784.
  • Bayesian neural network: Neural-network sampling alternates minibatch-based SGHMC or SGLD weight updates with Gibbs sampling of gamma-distributed precision variables.Weights use minibatches of 500 instances for 100 steps before hyperparameter updates.
  • Parameter selection: Training parameters are selected empirically: neural-network settings use validation, while matrix-factorization settings use cross-validation.The reported SGHMC settings are η = 0.2 × 10^-5 for the neural network and η = 0.4 × 10^-6 for matrix factorization.
  • Bayesian probabilistic matrix factorization: The Bayesian probabilistic matrix factorization model uses latent user and movie vectors, bias terms, and Bayesian priors on precision variables.Integrating over the precision variables provides Bayesian treatment of the key regularization parameter.
  • Bayesian probabilistic matrix factorization: Matrix-factorization sampling alternates minibatch-based SGHMC or SGLD weight updates with Gibbs sampling of precision variables.Weights use minibatches of 4,000 ratings for 2,000 steps before hyperparameter updates.
Loading 1402.4102v2…