Stochastic Gradient Leimkuhler-Matthews Integrator with Second Order Discretization Error

Abstract

Stochastic gradient Langevin dynamics (SGLD) has a discretization error of first order in the step size $h$, even with the full gradient. We study the stochastic gradient Leimkuhler-Matthews (SGLM) integrator, which has the same cost as SGLD, one stochastic gradient per step, and reuses each Gaussian vector in consecutive steps. With the full gradient, SGLM is the Leimkuhler-Matthews (LM) integrator, whose discretization error for the invariant distribution is of second order. For a potential strongly convex outside a ball, we prove the bound $\mathcal O(\frac{1}{Nh} + Σ^2 h^2 + h^4)$ on the mean square error (MSE) of time averages of SGLM over $N$ steps, where $Σ$ bounds the variance of the gradient noise. The proof uses a modified class that absorbs the non-Markovian part of SGLM. The discretization error enters this bound as $h^4$ and the MSE bound of SGLD as $h^2$, while a larger batch reduces only the term $Σ^2 h^2$ in both bounds. In Bayesian logistic regression on MNIST, SGLM has a smaller estimated MSE than SGLD. The coefficient of $h$ in the numerical bias of SGLM is given explicitly by the covariance of the gradient noise and the continuous Poisson equation. The coefficient for SGLD contains in addition the Hessian of the potential, which comes from the discretization error. Through Riesz representors, the coefficients of the numerical bias determine the coefficients of the Wasserstein-1 errors of SGLM and SGLD.