Accurate Large-sample Uncertainty Quantification using Stochastic Gradient Markov Chain Monte Carlo
This work provides practical tuning guidance for uncertainty quantification in large-scale Bayesian inference, addressing a key bottleneck for practitioners using stochastic gradient MCMC.
The authors propose new discrete-time approximations for stochastic gradient MCMC algorithms that accurately predict stationary covariance and autocorrelation, enabling practical tuning and uncertainty quantification even with large batches or model misspecification. They provide non-asymptotic error bounds and demonstrate improved tuning across various models where existing methods fail.
Tuning algorithms such as stochastic gradient descent (SGD) and stochastic gradient Langevin dynamics (SGLD) for approximate sampling and uncertainty quantification remains challenging, particularly in the practically relevant settings when the batch size is large or the model is misspecified. Existing theory that provides tuning guidance relies on continuous-time limits or strong statistical assumptions, which can become quantitatively inaccurate in these regimes. We address these shortcomings by proposing new discrete-time approximations to SG(L)D with and without momentum, which enables accurate predictions of the stationary covariance, iterate average covariance, and integrated autocorrelation time. Moreover, we prove quantitative, non-asymptotic error bounds showing that these estimates are sufficiently accurate for practical tuning and uncertainty quantification. Numerical experiments demonstrate that our theory yields improved tuning guidance across a range of models and data-generating distributions where existing approaches fail, including when using the $β$-divergence rather than log-loss to obtain statistically robust inferences.