Accurate Large-sample Uncertainty Quantification using Stochastic Gradient Markov Chain Monte Carlo

Tuning stochastic gradient methods such as stochastic gradient descent (SGD) and stochastic gradient Langevin dynamics (SGLD) for approximate sampling and uncertainty quantification remains challenging, particularly in relevant settings using a large batch size or when 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 a 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 $\beta$-divergence rather than log-loss to obtain statistically robust inferences.

Many machine learning methods are trained using fast randomized algorithms that process only part of the data at each step. These methods are efficient, but they are hard to tune when we want reliable uncertainty estimates, not just accurate predictions. Existing theory often approximates these algorithms as if they change continuously over time, but this can be inaccurate when the algorithm uses large batches of data, larger learning rates, or imperfect statistical models. We take a different approach by studying these algorithms step by step, matching how they are actually implemented. This allows us to better predict their long-run behavior, including how variable their outputs are and how quickly they explore uncertainty. Based on this analysis, we develop practical tuning guidance for stochastic gradient methods used in uncertainty quantification. Our experiments show that the proposed approach gives more accurate uncertainty estimates than existing theory in both synthetic and real-world examples. This makes stochastic training methods more useful and trustworthy in applications where understanding uncertainty matters.