Utilising the CLT Structure in Stochastic Gradient based Sampling : Improved Analysis and Faster Algorithms

TL;DR

This paper leverages the CLT structure in stochastic gradient sampling algorithms to improve convergence analysis, establish new theoretical guarantees, and propose a novel covariance correction method for faster, more reliable sampling.

Contribution

It provides the first stable convergence rates for SGLD without warm start, under milder conditions, and introduces a covariance correction algorithm to enhance stochastic sampling methods.

Findings

Proves stable convergence rate for SGLD in KL divergence without warm start.

Establishes guarantees for SGLD under weaker smoothness and inequality conditions.

Improves theoretical guarantees for RBM, reducing dependence on horizon.

Abstract

We consider stochastic approximations of sampling algorithms, such as Stochastic Gradient Langevin Dynamics (SGLD) and the Random Batch Method (RBM) for Interacting Particle Dynamcs (IPD). We observe that the noise introduced by the stochastic approximation is nearly Gaussian due to the Central Limit Theorem (CLT) while the driving Brownian motion is exactly Gaussian. We harness this structure to absorb the stochastic approximation error inside the diffusion process, and obtain improved convergence guarantees for these algorithms. For SGLD, we prove the first stable convergence rate in KL divergence without requiring uniform warm start, assuming the target density satisfies a Log-Sobolev Inequality. Our result implies superior first-order oracle complexity compared to prior works, under significantly milder assumptions. We also prove the first guarantees for SGLD under even weaker conditions such as H\"{o}lder smoothness and Poincare Inequality, thus bridging the gap between the state-of-the-art guarantees for LMC and SGLD. Our analysis motivates a new algorithm called covariance correction, which corrects for the additional noise introduced by the stochastic approximation by rescaling the strength of the diffusion. Finally, we apply our techniques to analyze RBM, and significantly improve upon the guarantees in prior works (such as removing exponential dependence on horizon), under minimal assumptions.