An Analysis of Constant Step Size SGD in the Non-convex Regime: Asymptotic Normality and Bias
This work addresses uncertainty estimation for practitioners using SGD in non-convex settings, which is an incremental advancement in statistical machine learning.
The authors tackled the problem of quantifying uncertainty for stochastic gradient descent (SGD) in non-convex learning by establishing asymptotic normality for constant step size SGD and characterizing bias relative to critical points, enabling confidence interval construction.
Structured non-convex learning problems, for which critical points have favorable statistical properties, arise frequently in statistical machine learning. Algorithmic convergence and statistical estimation rates are well-understood for such problems. However, quantifying the uncertainty associated with the underlying training algorithm is not well-studied in the non-convex setting. In order to address this shortcoming, in this work, we establish an asymptotic normality result for the constant step size stochastic gradient descent (SGD) algorithm--a widely used algorithm in practice. Specifically, based on the relationship between SGD and Markov Chains [DDB19], we show that the average of SGD iterates is asymptotically normally distributed around the expected value of their unique invariant distribution, as long as the non-convex and non-smooth objective function satisfies a dissipativity property. We also characterize the bias between this expected value and the critical points of the objective function under various local regularity conditions. Together, the above two results could be leveraged to construct confidence intervals for non-convex problems that are trained using the SGD algorithm.