Convergence of Stochastic Gradient Langevin Dynamics in the Lazy Training Regime

TL;DR

This work analyzes stochastic gradient Langevin dynamics (SGLD) in the lazy training regime for overparameterized neural networks, modeling SGD with a scaled Itô SDE that includes state-dependent noise. Under Hessian regularity and a positive neural tangent kernel at initialization, it proves exponential convergence of the expected optimality gap to the empirical minimizer while the parameters remain near initialization, and provides high-probability bounds on the time to leave this regime. It further derives finite-time and finite-width bounds and demonstrates that increasing the output scaling α helps sustain the lazy regime with high probability. Numerical experiments on shallow and deep networks corroborate the theory, showing exponential decay of training loss with larger α and stability of the NTK. Collectively, the results deliver non-asymptotic guarantees for stochastic dynamics in the kernel-like lazy regime, bridging prior deterministic lazy-training analyses with stochastic optimization.

Abstract

Continuous-time models provide important insights into the training dynamics of optimization algorithms in deep learning. In this work, we establish a non-asymptotic convergence analysis of stochastic gradient Langevin dynamics (SGLD), which is an Itô stochastic differential equation (SDE) approximation of stochastic gradient descent in continuous time, in the lazy training regime. We show that, under regularity conditions on the Hessian of the loss function, SGLD with multiplicative and state-dependent noise (i) yields a non-degenerate kernel throughout the training process with high probability, and (ii) achieves exponential convergence to the empirical risk minimizer in expectation, and we establish finite-time and finite-width bounds on the optimality gap. We corroborate our theoretical findings with numerical examples in the regression setting.

Figures (2)

• Figure 1: Test-error evolution during training using Euler-Maruyama scheme with stepsize of $\delta t=0.01$ for $\alpha\in\{1/8,8,32,256\}$. We additionally include curve $f(t)$ with the decay rate $\exp(-\lambda^2t)$ (dotted line) as a reference, where $\lambda=0.01122$ is the minimum eigenvalue of the NTK at initialization.
• Figure 2: Distance from initialization (left) and value of the minimal eigenvalue of the NTK (right) in time for different values $\alpha\in\{1/8,8,32,256\}$.

Theorems & Definitions (39)

• Remark 1
• Lemma 1: Regularity of $R$
• proof : Proof of Lemma \ref{['lemma:Rproperties']}
• Remark 2
• Theorem 1: Convergence of stochastic gradient Langevin dynamics
• proof : Proof of Theorem \ref{['th:ybound']}
• Remark 3
• Corollary 1: Risk convergence for shallow neural networks
• proof : Proof of Corollary \ref{['corollary:snn']}
• Corollary 2: Risk convergence for deep neural networks