Source-linked AI summary
The Zig-Zag Process and Super-Efficient Sampling for Bayesian Analysis of Big Data
Joris Bierkens, Paul Fearnhead, Gareth Roberts
TL;DR
Big-data Bayesian inference challenges traditional reversible MCMC because repeated full-data evaluations are costly, while sub-sampling methods may lose exactness. The paper develops a continuous-time, non-reversible Zig-Zag process with exact sub-sampling and control variates, obtaining approximately independent samples with O(k+n) cost for k samples rather than O(kn).
Problem
Traditional MCMC is costly for big data, while sub-sampling methods can fail to retain the true posterior as their stationary distribution.
Method
The paper introduces a multi-dimensional, continuous-time Zig-Zag process and an exact sub-sampling variant using unbiased gradient estimates, with control variates for improved efficiency.
Results
For k approximately independent samples, ZZ-CV has computational cost O(k+n), compared with O(kn) for traditional MCMC.
Takeaways & Limitations
After an O(n) preprocessing step, ZZ-CV can produce approximately independent samples at O(1) cost each, including in highly non-Gaussian examples where SGLD and Consensus are biased.
Takeaways & Limitations
Efficient implementation still requires simulating time-inhomogeneous Poisson processes and finding reasonable centering points before MCMC begins.
Abstract
from arXiv · showhide
Standard MCMC methods can scale poorly to big data settings due to the need to evaluate the likelihood at each iteration. There have been a number of approximate MCMC algorithms that use sub-sampling ideas to reduce this computational burden, but with the drawback that these algorithms no longer target the true posterior distribution. We introduce a new family of Monte Carlo methods based upon a multi-dimensional version of the Zig-Zag process of (Bierkens, Roberts, 2017), a continuous time piecewise deterministic Markov process. While traditional MCMC methods are reversible by construction (a property which is known to inhibit rapid convergence) the Zig-Zag process offers a flexible non-reversible alternative which we observe to often have favourable convergence properties. We show how the Zig-Zag process can be simulated without discretisation error, and give conditions for the process to be ergodic. Most importantly, we introduce a sub-sampling version of the Zig-Zag process that is an example of an {\em exact approximate scheme}, i.e. the resulting approximate process still has the posterior as its stationary distribution. Furthermore, if we use a control-variate idea to reduce the variance of our unbiased estimator, then the Zig-Zag process can be super-efficient: after an initial pre-processing step, essentially independent samples from the posterior distribution are obtained at a computational cost which does not depend on the size of the data.
1. Introduction.
The paper develops multi-dimensional Zig-Zag sampling as a non-reversible alternative to conventional MCMC, targeting Bayesian posteriors while addressing the computational burden of massive datasets. Its sub-sampling and control-variate variants preserve exactness and can achieve super-efficient scaling.
- Traditional MCMC relies on reversible chains, although non-reversible chains may offer substantial convergence advantages.The paper identifies reversibility as a limitation and motivates methods that create momentum to improve mixing.
- Big-data Bayesian MCMC is costly because Metropolis-Hastings requires processing the full dataset for likelihood evaluation at every iteration.Sub-sampling alternatives reduce computation but generally introduce approximation error in the resulting posterior samples.
- The paper introduces multi-dimensional Zig-Zag sampling and variants whose dynamics depend on the gradient of the target log density.For Bayesian models, this gradient is a sum that can be estimated unbiasedly using data sub-sampling.
- Zig-Zag with Sub-Sampling retains the exact invariant distribution, while control variates produce the ZZ-CV variant with super-efficient scaling for large datasets.The paper contrasts these properties with methods such as Firefly and discusses related PDMP approaches including Bouncy Particle Sampler and SCALE.
- Zig-Zag is a continuous-time, non-reversible piecewise deterministic Markov process with piecewise linear trajectories and single-component velocity switches.State-dependent switching rates are constrained so that the target density is stationary, and Poisson thinning enables simulation without discretisation error.
2. The Zig-Zag process.
The Zig-Zag process is a continuous-time, non-reversible sampler with piecewise linear trajectories and velocity-component switches governed by state-dependent rates. Under stated rate conditions it has a target-linked invariant distribution, supports posterior sampling, and admits ergodicity results.
- The Zig-Zag process: At each skeleton event, the earliest component-specific switching time determines the next event time and which velocity component flips.Between events, the position follows ξ + θt; the resulting path is assembled from successive skeleton points.
- The Zig-Zag process: The process can be constructed as a Markov process with finitely many switches on every bounded time interval under continuous switching rates.Continuity makes the rates bounded on compact sets, while trajectories travel finite distances over finite intervals.
- Invariant distribution: With rates satisfying the stated balance condition, Zig-Zag has an invariant distribution whose position marginal is proportional to exp(−Ψ) and whose velocity is uniform.The position and velocity are independent under this invariant distribution.
- Invariant distribution: For Bayesian targets, the canonical rates produce a process whose invariant distribution has the posterior density as its marginal, while symmetric nonnegative rate additions preserve that invariant distribution.The added functions γ_i must take the same value before and after flipping velocity component i.
- Ergodicity: Under the one-dimensional assumptions, Zig-Zag is f-exponentially ergodic with total-variation distance bounded by κf(ξ, θ)ρ^t.The bound holds for all initial states and t ≥ 0, for κ > 0 and 0 < ρ < 1.
- Ergodicity: Under product-form and componentwise mixing conditions, the multidimensional process is mixing; positive switching rates also imply at most one invariant measure.The canonical process is mixing in the stated product setting, unlike BPS for an i.i.d. Gaussian without momentum refreshments.
3. Implementation.
Zig-Zag simulation uses Poisson thinning with computable upper bounds on switching rates, enabling exact event simulation without discretization. Its computational efficiency depends on the magnitude of these bounds and can exploit Hessian or global derivative bounds.
- Exact event simulation: Poisson thinning simulates switching events exactly by proposing points from an upper-rate process and retaining each with probability m(τ)/M(τ).The retained points form a non-homogeneous Poisson process with the desired rate m(t).
- Exact event simulation: For each velocity component, the algorithm draws a proposed event time, selects the earliest retained time, and switches the corresponding component.The strong Markov property permits the resulting skeleton point to become the next starting point.
- Exact event simulation: When the integrated upper bound has an explicit inverse, event times are generated by CDF inversion from independent uniform random variables.Piecewise affine and constant bounds yield piecewise linear or quadratic integrated bounds with explicit inverses.
- Computational efficiency: Each simulated event requires evaluating one switching-rate component, so computational bounds determine how far the process moves before another gradient-component evaluation.The paper therefore tracks how these bounds scale with the number of observations.
- Computational efficiency: A global derivative bound supξ |∂iΨ(ξ)| ≤ c_i permits constant upper bounds M_i(t)=c_i for all t ≥ 0.The implementation also derives affine bounds from a positive-definite Hessian domination condition HΨ(ξ) ⪯ Q.
- Computational efficiency: Using Hessian-based bounds may require computing Qθ whenever θ changes, adding O(d) computation.This is an implementation cost associated with the dominated-Hessian construction.
4. Big data Bayesian inference by means of error-free sub-sampling.
The section develops exact sub-sampling methods for Zig-Zag sampling in big-data Bayesian models, preserving the target posterior while reducing per-event computation. Control variates further reduce computational bounds, enabling super-efficient scaling after linear preprocessing under suitable regularity conditions.
- Motivation: Big-data posteriors can decompose into many data contributions, making full-gradient evaluation costly while allowing individual derivative terms to be evaluated O(n) more cheaply.The construction exploits representations in which the negative log posterior is built from many terms and each component derivative is cheaper than the full derivative.
- Zig-Zag with Sub-Sampling: ZZ-SS evaluates one randomly selected data contribution at each simulated time instead of summing all n contributions.This reduces algorithmic complexity by a factor O(n), but increases the computational bound used for exact simulation.
- Zig-Zag with Sub-Sampling: The sub-sampling procedure uses computational upper bounds, draws candidate event times, samples J uniformly, and accepts a switch using the corresponding bounded rate.The resulting algorithm is formally identified as ZZ-SS or ZZ-CV depending on the choice of data contribution estimator.
- Exactness: Under the stated construction, Algorithm 3 generates a Zig-Zag skeleton whose process has invariant distribution µ.The theorem establishes exactness through the specified switching rates rather than discretized approximation.
- Conditions and trade-offs: Larger computational bounds increase switching rates and diffusivity, which can worsen mixing properties while also increasing simulation complexity.This is the principal cost of loose bounds in the exact sub-sampling scheme.
- Conditions and trade-offs: The control-variate construction requires globally and uniformly Lipschitz derivatives, with the choice of norm creating a trade-off between the constants and distance from the reference point.For i.i.d. Gaussian components, the paper states that p = ∞ is optimal, while p = 2 is reasonable when the situation is less clear.
- Zig-Zag with Control Variates: Control variates center each data contribution at a reference point to reduce variability across observations and thereby reduce the computational bound.The reference point is chosen near the posterior mode; finding it and computing the full derivative there each require O(n) preprocessing, after which ZZ-CV can be super-efficient.
5. Scaling analysis.
The scaling analysis shows that standard Zig-Zag requires O(n) computation per independent sample, while control variates reduce ZZ-CV's cost to O(1) under the stated assumptions. The ZZ-CV result relies on O(1) switching behavior after rescaling and retains the true posterior as its target.
- Scaling of Zig-Zag Sampling (ZZ): After rescaling, the canonical Zig-Zag process has unit coordinate speed and switching rates whose dependence on n vanishes asymptotically.The limiting switching rates correspond to a Gaussian distribution with covariance matrix (i(θ0))−1.
- Scaling of Zig-Zag Sampling (ZZ): O(1) proposed switches yield an essentially independent sample when the computational bound is O(n1/2).Each proposed switching time costs O(n) because the full data must be processed to compute the true switching rate.
- Scaling of Zig-Zag Sampling (ZZ): O(n) is the computational complexity of ZZ per independent sample when the computational bound is O(n1/2).This matches the cost of standard Monte Carlo methods that use O(1) iterations with O(n) work per iteration.
- Scaling of Zig-Zag Sampling (ZZ): O(n3/2) is the total computational bound for ZZ with global bounds, because each log-density term costs O(1).More generally, a computational bound O(nα) with α > 1/2 requires O(nα−1/2) proposed switches.
- Scaling of Zig-Zag with Control Variates (ZZ-CV): O(1) proposed switches suffice for ZZ-CV to obtain an independent sample because its computational bounds are O(n1/2).The control-variate construction makes the rescaled switching rate O(1), including when the reference point is not exactly the MLE.
- Scaling of Zig-Zag with Control Variates (ZZ-CV): O(1) is the computational complexity of ZZ-CV per independent sample, providing a factor n increase in efficiency.The method leverages posterior asymptotic normality while continuing to sample from the true posterior; it may be slower when the posterior is far from Gaussian.
6. Examples and experiments.
The experiments compare Zig-Zag variants with established samplers on Gaussian and logistic-regression examples. ZZ-CV achieves favorable scaling in the identifiable settings, while Zig-Zag methods avoid the bias seen in SGLD but can lose super-efficiency under weak posterior contraction and an unbounded Hessian.
- Using Zig-Zag output: The Zig-Zag skeleton can support either sampled trajectory output or continuous-time integration for estimating posterior quantities.Continuous-time integration avoids choosing an arbitrary number of extracted samples, and piecewise linearity can make moment integrals exact.
- Using Zig-Zag output: Switching points must not be treated as posterior samples because they are biased toward the target distribution's tails.Posterior quantiles can instead be estimated from trajectory samples in the usual MCMC manner.
- Gaussian mean experiment: SGLD's second-moment MSE stops decreasing beyond a fixed value, whereas Zig-Zag variants show no such bias.The Zig-Zag result agrees with the theoretical consistency of ergodic averages over Zig-Zag trajectories.
- Gaussian mean experiment: ZZ-CV and ZZ-soCV become significantly more efficient than basic ZZ as the number of observations increases.Using a poor reference point in ZZ-soCV has only a small effect on efficiency.
- Logistic regression: A slope of approximately 0.95 is observed for ZZ-CV's ESS per epoch in both logistic-regression dimensions.The other algorithms have roughly horizontal slopes, corresponding to linear scaling with data size.
- Logistic regression: ZZ-CV is the only tested algorithm whose ESS per CPU second is approximately constant as the data size increases.The comparison includes MALA, Zig-Zag variants with global or Lipschitz bounds, and ZZ-SS.
- Non-identifiable logistic regression: In the non-identifiable example, SGLD and Consensus are strongly biased, whereas ZZ and ZZ-CV target the correct distribution.ZZ-CV loses much of its efficiency because posterior contraction is absent and the Hessian is unbounded, so it is not super-efficient there.
7. Discussion.
The discussion highlights both the promise of Zig-Zag sampling for large datasets and practical implementation challenges, especially Poisson-process simulation and selecting centering points for ZZ-CV.
- Implementation challenges: The main implementation challenges are simulating the relevant time-inhomogeneous Poisson process and finding reasonable centering points before MCMC begins.The efficiency of Poisson simulation depends on the tightness of the rate bounds, while centering-point selection lacks a prescriptive recipe.
- Implementation support: The authors released an R/Rcpp logistic-regression package and code reproducing the Section 6 experiments in response to near-term practical challenges.The release is intended to support implementation while these challenges are addressed.
- Approximate simulation: Approximate shortcuts can speed Zig-Zag simulation but introduce a loss of exactness, leaving the resulting errors to be understood.Approximate Poisson-rate bounds are cited as one route to more efficient simulation for a wider class of models.
- Scope of super-efficiency: Zig-Zag remains exact across target structures, while ZZ-CV can lose super-efficiency without posterior contraction and with an unbounded Hessian.In a non-identifiable logistic regression example, ZZ and ZZ-CV target the correct distribution, but ZZ-CV loses much of its efficiency.
- Computational efficiency: O(k + n) versus O(kn) computational efficiency is reported for obtaining k approximately independent samples with ZZ-CV in big-data settings.A full gradient evaluation costs O(n) initially, after which approximately independent samples cost O(1) each; the comparison is confirmed experimentally.
- Extensions: Multiple reference points and algorithm switching are proposed as natural extensions for multimodal posteriors or parameters lacking posterior concentration.The control-variate construction described is presented as one, possibly simplest, implementation of the broader idea.
Supplement: Supplement to “The Zig-Zag Process and Super-Efficient Sampling for Bayesian Analysis of Big Data”
The supplement covers the mathematical foundations of Zig-Zag, SGLD scaling, and experimental details, including computational-bound calculations.
- The supplement contains mathematics of the Zig-Zag process, SGLD scaling, and experimental details including how to obtain computational bounds.