Subsampling Error in Stochastic Gradient Langevin Diffusions
Kexin Jin, Chenguang Liu, Jonas Latz
TL;DR
The paper analyzes subsampling error in Langevin-based MCMC by introducing the idealised Stochastic Gradient Langevin Diffusion (SGLDiff), a continuous-time switched diffusion that alternates data subsets at rate 1. It proves exponential ergodicity of SGLDiff and derives a bound in Wasserstein distance between the target posterior $\mu(d\theta)=\frac{1}{Z}e^{-\bar{\Phi}(\theta)}d\theta$ (with $\bar{\Phi}(\theta)=\frac{1}{N}\sum_i \Phi_i(\theta)$) and its invariant measure $\mu^\eta$, with the bound scaling as a fractional power of the mean waiting time $\eta$. The main results establish (i) strong finite-time convergence of $\theta_t$ to the Langevin limit $\zeta_t$ as $\eta\to0$, (ii) exponential ergodicity and explicit contraction rates in Wasserstein distance for the switched diffusion, and (iii) an asymptotic subsampling error bound $\mathcal{W}_{\|\cdot\|}(\mu^\eta,\mu) \le C_{\Phi,d}\eta^{c_\Phi}$ with a dimension-dependent prefactor $C_{\Phi,d}=\mathcal{O}(\sqrt{d})$. The analysis uses reflection coupling to obtain dimension-free contraction and demonstrates that the continuous-time model captures key aspects of SGLD’s bias, suggesting the Euler–Maruyama discretisation is appropriate in this subsampling regime. The authors also outline open directions, including optimization with temperature, momentum, and epoch-based or without-replacement subsampling schemes.
Abstract
The Stochastic Gradient Langevin Dynamics (SGLD) are popularly used to approximate Bayesian posterior distributions in statistical learning procedures with large-scale data. As opposed to many usual Markov chain Monte Carlo (MCMC) algorithms, SGLD is not stationary with respect to the posterior distribution; two sources of error appear: The first error is introduced by an Euler--Maruyama discretisation of a Langevin diffusion process, the second error comes from the data subsampling that enables its use in large-scale data settings. In this work, we consider an idealised version of SGLD to analyse the method's pure subsampling error that we then see as a best-case error for diffusion-based subsampling MCMC methods. Indeed, we introduce and study the Stochastic Gradient Langevin Diffusion (SGLDiff), a continuous-time Markov process that follows the Langevin diffusion corresponding to a data subset and switches this data subset after exponential waiting times. There, we show the exponential ergodicity of SLGDiff and that the Wasserstein distance between the posterior and the limiting distribution of SGLDiff is bounded above by a fractional power of the mean waiting time. We bring our results into context with other analyses of SGLD.