Robust Stochastic Gradient Posterior Sampling with Lattice Based Discretisation

TL;DR

This paper introduces SGLRW, a robust stochastic gradient MCMC method that improves stability against minibatch size variations and gradient noise, outperforming traditional SGLD in challenging scenarios.

Contribution

The paper proposes SGLRW, a novel lattice-based stochastic gradient MCMC method that enhances robustness to gradient noise and minibatch size effects, with theoretical and empirical validation.

Findings

SGLRW is more stable than SGLD with heavy-tailed gradient noise.

SGLRW maintains asymptotic correctness despite robustness modifications.

Experimental results show improved predictive performance in Bayesian tasks.

Abstract

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 discretization. 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 comparison we analyze a natural analogue of SGLD utilizing 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.