PaperlayerView exact revision on arXiv

Source-linked AI summary

A Complete Recipe for Stochastic Gradient MCMC

Yi-An Ma, Tianqi Chen, Emily B. Fox

arXiv:1506.04696v2math.STstat.MEstat.ML

TL;DR

Continuous-dynamic MCMC samplers can be difficult to scale and to design with stochastic gradients while preserving the target distribution. This paper gives a complete two-matrix recipe for constructing such samplers, proves that every continuous Markov sampler with the target stationary distribution fits the framework, and uses it to propose SGRHMC, which combines Riemann HMC benefits with stochastic-gradient scalability.

  • Problem

    Stochastic-gradient MCMC samplers require difficult convergence analyses and sampler-specific modifications to account for gradient noise while retaining the target distribution.

  • Method

    The paper parameterizes continuous Markov dynamics with a positive semidefinite diffusion matrix D(z) and a skew-symmetric curl matrix Q(z), then constructs full-data and minibatch variants.

  • Results

    The framework is complete, covering every continuous Markov process with the target stationary distribution, and it produces the new SGRHMC sampler evaluated on synthetic data and streaming Wikipedia analysis.

  • Takeaways & Limitations

    SGRHMC inherits the benefits of Riemann HMC while providing the scalability of stochastic-gradient methods.

Abstract

from arXiv · show

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 non-trivial. 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.

1 Introduction

The paper introduces a complete recipe for continuous-dynamic MCMC that addresses scalability and the difficulty of preserving the target posterior under stochastic-gradient noise. It represents samplers through two matrices and uses the framework to construct SGRHMC.

  • Motivation: MCMC methods can explore distant posterior states efficiently through continuous dynamics, but their computational burden grows substantially with large datasets.HMC and Riemann manifold HMC improve exploration while requiring full-data gradient information.
  • Motivation: Stochastic-gradient MCMC scales continuous-dynamic samplers by replacing full-data gradients with minibatch-based estimates.SGLD established this approach by combining stochastic optimization with Langevin dynamics and appropriately annealed noise.
  • Research gap: Designing continuous dynamics that preserve the target posterior is challenging because stochastic-gradient noise requires sampler-specific physical and geometrical reasoning.The difficulty concerns proving invariance and modifying dynamics without losing the desired stationary distribution.
  • Framework: The framework is complete: any continuous Markov process with the desired invariant distribution can be represented through suitable D(z) and Q(z).This unifies HMC, Riemann manifold HMC, SGLD, SGHMC, their variants, and future continuous samplers.
  • Framework: The proposed recipe parameterizes continuous Markov samplers with a positive semidefinite diffusion matrix D(z) and a skew-symmetric curl matrix Q(z).The dynamics are written directly in terms of the target stationary distribution and these matrices.
  • Applications: The recipe yields practical full-data and minibatch algorithms and enables the construction of stochastic gradient Riemann Hamiltonian Monte Carlo.Experiments use synthetic data and a streaming Wikipedia latent Dirichlet allocation analysis.

2 A Complete Stochastic Gradient MCMC Framework

The framework constructs continuous Markov dynamics with a specified posterior stationary distribution and proves that it covers every continuous Markov process with that distribution. It also provides stochastic-gradient implementations while exposing a finite-step-size bias–speed tradeoff.

  • Target and dynamics: The target is the posterior p(θ|S), with auxiliary variables optionally augmenting the state and later marginalized.The generic sampled state z may be θ itself or an augmented state space.
  • Target and dynamics: Continuous dynamics define Markov transition kernels whose distributional evolution can be analyzed with Fokker–Planck or Liouville equations.The dynamics map the state at time t to the state at time t + h.
  • Motivation: Metropolis-Hastings correction can be costly or infeasible for large-data stochastic-gradient samplers, motivating dynamics with the correct stationary distribution directly.Reverse-path probabilities may be infeasible, and poorly matched dynamics can produce high rejection rates.
  • Specified stationary distribution: The recipe constructs an SDE from a positive semidefinite diffusion matrix D(z) and a skew-symmetric curl matrix Q(z), with f(z) defined from the target distribution.Q(z) represents deterministic traversing effects, while D(z) controls Wiener-process diffusion.
  • Specified stationary distribution: Any admissible D(z) and Q(z) produce dynamics with the desired stationary distribution p_s(z) ∝ exp(−H(z)); uniqueness follows from positive definiteness or ergodicity.For deterministic dynamics with D(z)=0, invariance is analyzed through the Liouville equation, while ergodicity requires an added jump process.
  • Completeness: The framework is complete: under the stated existence and integrability conditions, every continuous Markov process with the target stationary distribution has a corresponding skew-symmetric Q(z).Theorem 2 constructs Q(z) from a general process, establishing equivalence between the framework and all such continuous samplers.
  • Practical algorithm: Stochastic-gradient methods replace computationally intensive full-data potentials with minibatch estimators whose gradients are noisy but unbiased.The resulting gradient noise can alter the stationary distribution and therefore requires analysis or compensation.
  • Practical algorithm: As the step size ϵ_t approaches zero, the stochastic-gradient dynamics approach the correct invariant dynamics without requiring Metropolis-Hastings correction.A finite small step size is faster but produces a biased sampler, yielding a bias–speed tradeoff.

3 Applying the Theory to Construct Samplers

The framework represents continuous MCMC samplers through choices of diffusion and curl matrices, unifying existing methods and guiding the construction of SGRHMC. It also identifies why naïve state-adaptive stochastic-gradient dynamics fail and provides corrected constructions.

  • Casting Previous MCMC Algorithms within the Proposed Framework: The framework casts existing continuous samplers through specific choices of D(z), Q(z), and H(z), avoiding sampler-specific proofs for stochastic-gradient methods.It can also identify invalid samplers and guide corrections.
  • Casting Previous MCMC Algorithms within the Proposed Framework: Naïvely replacing the full gradient with a stochastic gradient in HMC does not yield the correct stationary distribution.The framework pairs injected noise with a D(z)∇H(z) term, motivating a friction correction.
  • Casting Previous MCMC Algorithms within the Proposed Framework: SGLD uses constant diffusion with Q(z) = 0, while SGRLD uses an adaptive diffusion matrix based on the Fisher information metric.Both methods emphasize diffusion design rather than Hamiltonian curl dynamics.
  • Casting Previous MCMC Algorithms within the Proposed Framework: HMC uses D(z) = 0 and focuses on Q(z), whereas SGHMC combines nonzero diffusion and curl matrices, and SGNHT makes Q(z) state dependent.These matrix choices summarize how the methods relate within the framework.
  • Stochastic Gradient Riemann Hamiltonian Monte Carlo: The proposed SGRHMC modifies SGHMC using the Fisher information metric, while generalized SGRHMC permits any positive definite G(θ).The generalized form is useful when the Fisher metric is difficult to compute.

4 Experiments

Experiments validate the recipe and assess gSGRHMC on simulated distributions and streaming Wikipedia topic inference. The results show correct invariant behavior, efficient exploration, and benefits from Riemannian and Hamiltonian components at similar runtimes.

  • Experimental Setup: The experiments evaluate recipe correctness and gSGRHMC effectiveness on simulated distributions and an online LDA analysis of Wikipedia entries.The Wikipedia experiment uses minibatches because the dataset is large and continually growing.
  • Synthetic Experiments: Naïve state-adaptive SGHMC does not converge to the target distribution, whereas recipe-derived gSGRHMC has the correct invariant distribution and efficiently explores the simulated distributions.The comparison includes SGLD and SGHMC.
  • Online Latent Dirichlet Allocation: In streaming Wikipedia LDA, Riemannian stochastic-gradient samplers show important benefits for distributions whose mass is concentrated over small regions.The analysis uses an expanded mean parameterization to address the simplex boundary issue.
  • Online Latent Dirichlet Allocation: The HMC term appears beneficial for both Riemannian and non-Riemannian samplers, while average runtimes remain similar because gradient evaluation is the main bottleneck.Results are reported from three random runs.

5 Conclusion

The paper presents a complete two-matrix recipe for continuous MCMC processes, uses it to unify and correct existing stochastic-gradient samplers, and proposes SGRHMC. Experiments demonstrate SGRHMC’s efficiency and scalability on simulated data and streaming Wikipedia analysis.

  • Conclusion: The framework specifies continuous MCMC dynamics with a positive semidefinite diffusion matrix D(z) and skew-symmetric curl matrix Q(z).It proves both construction for arbitrary matrices and representation of any continuous Markov process with the target stationary distribution.
  • Conclusion: The recipe is used to reinvent previous stochastic-gradient MCMC samplers and propose stochastic-gradient Riemann Hamiltonian Monte Carlo.The conclusion highlights its usefulness when sampler-specific convergence reasoning is challenging.
  • Conclusion: SGRHMC’s efficiency and scalability are demonstrated on simulated data and a streaming Wikipedia analysis.

A Proof of Stationary Distribution

The proof establishes stationary-distribution invariance for the proposed continuous dynamics by analyzing the associated Fokker–Planck equation. It concludes that the target density is invariant under the specified dynamics.

  • Proof Strategy: The proof analyzes the Fokker–Planck equation associated with the dynamics when the drift satisfies the framework’s stated condition.
  • Proof Strategy: The derivation rewrites and decomposes terms, using the antisymmetry of Q to establish the required cancellation.
  • Stationary Distribution: The proof verifies that p(z | S) ∝ e^-H(z) is invariant under the compact dynamics and concludes Theorem 1.

B Proof of Completeness

The proof constructs the skew-symmetric matrix Q(z) by transforming the defining equations into Fourier space, solving linear constraints, and applying an inverse Fourier transform. It establishes that the resulting framework can represent arbitrary continuous Markov processes with the target stationary distribution.

  • The proof reformulates finding Q(z) as solving for Q(z)ps(z), whose right-hand side is divergence-free.
  • Fourier transformation converts the equation and its constraint into a set of linear equations.
  • The Fourier-space matrix Q̂ is constructed as a skew-symmetric projection between the spans of k and F̂, which are orthogonal.
  • The final Q(z) is obtained by multiplying the inverse Fourier transform of Q̂(k) by ps(z)^−1.
  • Under the stated L1 condition, any continuous-time Markov process can be converted into the framework, and the constructed projector has real entries.
  • The framework recasts prior samplers, including HMC, SGHMC, SGLD, SGRLD, and SGNHT, through specific choices of D(z) and Q(z).

E Discussion of Choice of D and Q

The framework treats D(z) and Q(z) as design variables whose choices determine different continuous-dynamic samplers. Adaptive diffusion can trade exploration speed in low-probability regions against accuracy near high-probability modes.

  • Different choices of D(z) and Q(z) define different regions of the sampler design space, with gSGRHMC covering a previously unexplored portion.
  • An adaptive diffusion matrix can increase noise in low-probability regions to facilitate escape from local modes and decrease noise near global modes to improve accuracy.

F.1 Synthetic Experiments

Synthetic experiments compare SGLD, SGHMC, naïve SGRHMC, and gSGRHMC on two one-dimensional target distributions. The naïve implementation is incorrect, whereas gSGRHMC has the correct invariant distribution and efficiently explores the targets.

  • The constant C makes eU(θ) + C positive in most cases, so fluctuation is smaller where probability density is higher.
  • The comparison indicates that the naïve SGRHMC implementation is incorrect.
  • The experiments compare SGLD, SGHMC, naïve SGRHMC, and gSGRHMC on U(θ) = θ^2/2 and U(θ) = θ^4−2θ^2.

F.2 Online Latent Dirichlet Allocation Experiment

The online LDA experiment applies SGRHMC to streaming Wikipedia data using minibatch stochastic gradients. It learns topic-word associations from 10,000 documents while comparing sampler settings and posterior parameterizations.

  • The experiment uses minibatches of 50 documents and K = 50 topics to compute stochastic gradients for online LDA.
  • The analysis uses Gibbs sampling within documents to estimate latent topic-assignment counts needed for the stochastic gradient.
  • Riemannian samplers use γ = 0.01 and α = 0.0001, whereas SGLD and SGHMC use γ = 0.1 and α = 0.01.
  • The SGRHMC stepsize is selected as ϵt = 0.02, while SGLD uses ϵt = 0.01 with D = I.
  • Topic-word distributions are learned from 10,000 documents, approximately 0.3% of Wikipedia articles, with 11 randomly selected topics summarized by their top seven words.
Loading 1506.04696v2…