Robust Stochastic Gradient Posterior Sampling with Lattice Based Discretisation

Stochastic-gradient MCMC methods enable scalable Bayesian posterior sampling but often suffer from sensitivity to minibatch size and gradient noise. To address this, we propose Stochastic Gradient Lattice Random Walk (SGLRW), an extension of the Lattice Random Walk discretisation. Unlike conventional Stochastic Gradient Langevin Dynamics (SGLD), SGLRW introduces stochastic noise only through the off-diagonal elements of the update covariance; this yields greater robustness to minibatch size while retaining asymptotic correctness. Furthermore, as a comparison we analyse a natural analogue of SGLD utilising gradient clipping. Experimental validation on Bayesian regression and classification demonstrates that SGLRW remains stable in regimes where SGLD fails, including in the presence of heavy-tailed gradient noise, and matches or improves predictive performance.

AI systems that make medical, scientific, or financial predictions need to honestly quantify how uncertain they are. A standard family of methods, called Bayesian sampling, does this by exploring the full range of plausible answers rather than committing to a single one. But the workhorse algorithm for this task, SGLD, breaks down when modern AI is trained on small chunks of data at a time: random noise from the small samples destabilises the algorithm, sometimes catastrophically. Our paper introduces a new sampler, SGLRW. It replaces the wide, Gaussian random steps used by SGLD with bounded, lattice-like steps. Imagine switching from a random walk that can take arbitrarily large steps to a walk on a grid where each step has a fixed size. This bounded structure keeps the sampler stable in noisy regimes where SGLD fails, while still guaranteeing it explores the right distribution in the long run. SGLRW lets practitioners use Bayesian methods at scales where they were previously impractical. It is also a natural fit for emerging energy-efficient AI hardware, contributing to more sustainable machine learning.