SULA: Robust Langevin Sampling via Drift-Diffusion Saturation

Abstract

The unadjusted Langevin algorithm (ULA) and its stochastic-gradient variant, SGLD, are workhorses of modern Bayesian machine learning, but can be fragile when faced with a target distribution possessing a non-Lipschitz drift. In such settings, the expected squared jump distance can grow rapidly, leading to numerical instability or the need for prohibitively small step-sizes. We propose the saturated unadjusted Langevin algorithm (SULA) that pairs a saturated gradient drift with a gradient-informed state-dependent diffusion, while preserving first-order consistency with the Langevin diffusion. The drift and diffusion terms are chosen to guarantee that the one-step squared update is independent of the gradient. Consequently, SULA is robust to the choice of step-size and preserves estimation and prediction quality. We provide sufficient conditions under which SULA with a step-size $h$ converges weakly to the Langevin diffusion at the standard $O(h)$ rate, the SULA chain is geometrically ergodic, and that the bias of its invariant measure relative to the target is $O(h)$. Further, we develop SG-SULA, a stochastic-gradient version of SULA that rescales the stochastic gradient using a Gaussian CDF correction under Gaussianity of the gradient noise. Numerical experiments on various examples demonstrate SULA's robustness to tuning and SG-SULA's improved predictive performance.

arXiv comment57 pages, 4 figures