An Analysis of Constant Step Size SGD in the Non-convex Regime: Asymptotic Normality and Bias

1 Introduction . Non-convex learning problems are prevalent in modern statistical machine learning applications such as matrix and tensor decomposition [ 2 , 3 , 4 , 5 , 6 , 7 , 8 ] , deep neural networks [ 9 , 10 , 11 ] , and robust empirical risk minimization [ 12 , 13 , 14 ] . Developing theoretically principled approaches for tackling such non-convex problems depends critically on the interplay between two aspects . From a computational perspective , variants of stochastic gradient descent ( SGD ) converge to firstorder critical points [ 15 , 16 ] or local minimizers [ 17 , 2 , 18 , 19 ] of the objective function . From a statistical perspective , oftentimes these critical points or local minimizers have nice statistical properties [ 20 , 3 , 12 , 21 , 22 , 5 ] ; see also [ 23 ] for a counterexample . For the purpose of uncertainty quantification in such non-convex settings , studying the fluctuations of iterative algorithms used for training becomes extremely important . In this work , we focus on the widely used constant step size SGD , and develop results for quantifying the uncertainty associated with this algorithm for a class of non-convex problems . We consider minimizing a non-smooth and non-convex objective function f : Rd→ R , min θ∈Rd f ( θ ) . ( 1 ) 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . The iterations of SGD with a constant step size η > 0 , initialized at θ ( η ) 0 ≡ θ0 ∈ Rd , are given by θ ( η ) k+1 = θ ( η ) k − η ( ∇f ( θ ( η ) k ) + ξk+1 ( θ ( η ) k ) ) , k ≥ 0 , ( 2 ) where { ξk } k≥1 is a sequence of random functions from Rd to Rd corresponding to the stochasticity in the gradient estimate . Several problems in machine learning and statistics are naturally formulated as the optimization problem in ( 1 ) , where the function f ( θ ) is given by f ( θ ) : = ∫ F ( θ , Z ) dP ( Z ) , ( 3 ) for an unknown distribution over the random variable Z ∈ Rp . The function F ( θ , Z ) is typically the loss function composed with functions from the hypothesis class parametrized by θ ∈ Rd . In online SGD with batch size b , at each iteration k , b independent samples Zj ∼ P ( Z ) are used to estimate the true gradient with 1b ∑b j=1∇F ( θ ( η ) k , Zj ) . The above iterates are indeed a special case of the iterates in ( 2 ) , with the noise sequence { ξk+1 ( θ ( η ) k ) } k≥0 given by ξk+1 ( θ ( η ) k ) : = 1 b ∑b j=1 [ ∇F ( θ ( η ) k , Zj ) −∇f ( θ ( η ) k ) ] . ( 4 ) Although proposed in the 1950s by [ 24 ] , SGD has been the algorithm of choice for training statistical models due to its simplicity , and superior performance in large-scale settings [ 25 , 1 , 26 , 27 ] . However , the fluctuations of this algorithm is well-understood only when the objective function f is strongly convex and smooth , and the step size η satisfies a specific decreasing schedule so that the iterates asymptotically converge to the unique minimizer [ 28 , 29 , 30 ] . On the other hand , it is well-known that the SGD iterates in ( 2 ) can be viewed as a Markov chain which allows them to converge to a random vector rather than a single critical point [ 1 ] . Building on this analogy between SGD and Markov chains , the aforementioned shortcomings can be alleviated by simply relaxing the global smoothness as well as the strong convexity assumptions to the tails of the objective function f , which allows for a flexible non-convex structure around the region of interest . Similar kinds of tail relaxations have been successfully employed in the diffusion theory when the target potential is non-convex [ 31 , 32 , 33 ] , but they are not studied in the context of non-convex optimization when the algorithm is SGD . In this work , we study the fluctuations and the bias of the averaged SGD iterates in ( 2 ) , around the first-order critical points of the minimization problem ( 1 ) . Our contributions can be summarized as follows . • For a non-convex and non-smooth objective function f with tails growing at least quadratically , we establish the uniqueness of the stationary distribution of the constant step size SGD iterates in Proposition 2.1 , and the asymptotic normality of Polyak-Ruppert averaging in Theorem 2.1 . To the best of our knowledge , these are the first uniqueness and normality results for the SGD algorithm when the objective function is non-convex ( even not strongly convex ) and non-smooth . • We further show in Theorems 3.1 and 3.2 that , with additional local smoothness assumptions on the non-convex objective function f , we can establish a control over the bias in terms of the step size . We further characterize the bias when the objective is ( not strongly ) convex in Theorem 3.3 , providing a thorough bias analysis for the constant step size SGD under various settings that are frequently encountered in statistical learning . Our results provide algorithm-dependent guarantees for uncertainty quantification , and they could be leveraged to obtain confidence intervals ( CIs ) for non-convex and non-smooth learning problems . This is contrary to the majority of the existing results in statistics , which only establish normality results for the true stationary point of the non-convex objective function ; see for example [ 12 , 34 ] . While being useful , such results completely ignore the computational hardships associated with non-convex optimization ; hence , their practical implications are limited . On the other hand , in the optimization and learning theory literature , a majority of the existing results establish the rate of convergence of an algorithm to a critical point , and do not quantify the fluctuations associated with that algorithm . Our work bridges these separate lines of thought by providing asymptotic normality results directly for the SGD algorithm used for minimizing non-convex and non-smooth functions . More Related Works . Establishing asymptotic normality results for the SGD algorithm began with the works of [ 35 , 36 , 37 , 38 , 39 ] , with [ 28 ] providing a definitive result for strongly convex objectives . In particular , [ 28 ] and [ 38 ] established that the averaged SGD iterates with an appropriately chosen decreasing step size is asymptotically normal with the variance achieving the Cramer-Rao lower bound for parameter estimation . Recent works , for example [ 40 , 41 , 29 , 42 , 43 , 44 ] , leverage the asymptotic normality analysis of [ 28 ] , which leads to computing CIs for SGD . The benefits of constant step size SGD for faster convergence under overparametrization has also been demonstrated in the works of [ 45 , 46 , 47 , 48 ] . The use of Markov chain theory to study constant step size stochastic approximation algorithms has been considered in several works [ 49 , 50 , 51 , 25 , 52 , 53 ] . Recently , [ 1 , 54 ] investigated the asymptotic variance of the constant step size SGD . We emphasize here that most of the above works assume strongly convex and smooth objective functions . The non-linear autoregressive ( NLAR ) process [ 55 , 56 , 57 ] is a specification of our general framework ( 2 ) with the noise sequence { ξk } k≥1 being a collection of i.i.d mean-zero random vectors with continuous density supported on Rd . However , the methodology for establishing the geometric ergodicity of NLAR [ 55 , 56 , 57 , 58 ] is by no means straightforward to carry over to the optimization setting , and does not generalize immediately to the state-dependent noise setup considered in our paper ( see Assumption 2.3 ) . In contrast , we establish the geometric ergodicity under easily verifiable assumptions on the objective function using tools from Markov chain theory . Moreover , additional steps are needed to go from geometric ergodicity to CLT results ( especially if the chain starts with an arbitrary initial distribution ) , while we directly obtain a CLT result by explicitly leveraging the Markov chain structure . Furthermore , there exists a vast literature on analyzing Langevin diffusion-based sampling algorithms which relies on the much simpler i.i.d . Gaussian noise sequence . We refer the interested reader to [ 59 , 60 , 32 , 61 , 62 , 63 , 64 , 65 , 66 , 67 , 68 , 69 , 70 ] and the references therein , for details . Finally , recent work [ 71 ] establish the trade-off between the stability , statistical accuracy , and computational efficiency for the non-convex optimization algorithms . Notation . For a , b ∈ R , denote by a ∨ b and a ∧ b the maximum and the minimum of a and b , respectively . We use ‖ · ‖ to denote the Euclidean norm in Rd . We denote the largest eigenvalue of the matrix A as λmax ( A ) , and the smallest one as λmin ( A ) . Let ( Ω , F , P ) represent a probability space , and denote by B ( Rd ) , the Borel σ-field of Rd . Let Pk ( Rd ) : = { ν : ∫ Rd ‖θ‖ kν ( dθ ) < ∞ } denote the set of probability measures with finite k-th moments . For a probability distribution π and a function g on X , we define π ( g ) : = ∫ X g ( x ) dπ ( x ) , and L2 ( π ) : = { g : X → R : π ( g 2 ) < ∞ } . 2 Central Limit Theorem for The Constant Step Size SGD . In this section , we establish an asymptotic central limit theorem ( CLT ) for the Polyak-Ruppert averaging of the constant step size SGD iterates given in ( 2 ) when the objective function is potentially non-convex , non-smooth , and has quadratically growing tails . More specifically , we first prove that there exists a unique stationary distribution πη ∈ P2 ( Rd ) for the Markov chain defined by the SGD algorithm when the objective function is dissipative ( see Assumption 2.2 ) with gradient exhibiting at most linear growth ( see Assumption 2.1 ) . Furthermore , under the same conditions , we prove that a CLT holds for the Polyak-Ruppert averaging , and it is independent of the initialization . In what follows , we list and discuss the main assumptions required to establish a CLT for the SGD iterates , and compare them to those existing in the literature . Assumption 2.1 ( Linear growth ) . The gradient of the objective function f has at most linear growth . That is , for some L ≥ 0 , we have ‖∇f ( θ ) ‖ ≤ L ( 1 + ‖θ‖ ) for all θ ∈ Rd . Majority of the results on SGD focus on smooth functions with gradients satisfying ‖∇f ( θ ) − ∇f ( θ′ ) ‖ ≤ ‖θ − θ′‖ for all θ , θ′ ∈ Rd ; see e.g . [ 28 , 1 ] . The above condition allows for non-smooth objectives , and is a significant relaxation of the standard Lipschitz gradient condition . Assumption 2.2 ( Dissipativity ) . The objective function f is ( α , β ) -dissipative . That is , there exists positive constants α , β such that 〈θ , ∇f ( θ ) 〉 ≥ α ‖θ‖2 − β for all θ ∈ Rd . The dissipativity assumption has its origins in the analysis of dynamical systems , and is used widely in the analysis of optimization and learning algorithms [ 72 , 31 , 33 , 73 ] . It also recently have been used in Bayesian analysis [ 74 ] . This assumption could be viewed as a relaxation of strong convexity since it restricts the quadratic growth assumption to the tails of the function f , enforcing no local growth around the first-order critical points . A canonical example for this condition is the sum of a quadratic and any non-convex function with bounded gradient . For example , consider the function x → x2 + 10 sin ( x ) which is clearly non-convex and ( 1 , 25 ) -dissipative . It is worth mentioning that many non-convex problems that arise in statistical learning such as phase retrieval [ 53 ] satisfy Assumption 2.2 . We provide examples in Section 4 . Assumption 2.3 ( Noise sequence ) . Gradient noise sequence { ξk } k≥1 is a collection of i.i.d . random fields satisfying E [ ξ1 ( θ ) ] = 0 and E1/2 [ ‖ξ1 ( θ ) ‖2 ] ≤ Lξ ( 1 + ‖θ‖ ) , for any θ ∈ Rd and a positive constant Lξ . Moreover , for each θ ∈ Rd the distribution of the random variable ξ1 ( θ ) can be decomposed as µ1 , θ +µ2 , θ where µ1 , θ has a density , say pθ , with respect to Lebesgue measure which satisfies infθ∈C pθ ( t ) > 0 for any bounded set C and any t ∈ Rd . Assumption 2.3 as formulated above is stronger than what is needed in the proofs . It can easily be seen that the lower bound on the density pθ is only required to hold for a specific set whose form depends on η and various constants from Assumptions 2.1–2.3 . The form of this set is complicated , and an exact expression is given in the Appendix – see ( 12 ) . We also emphasize that Assumption 2.3 does not specify any explicit parametric form for the distribution of the noise sequence contrary to recent works in non-convex settings where dissipitavity condition has been heavily utilized [ 31 , 73 , 33 ] . We now establish the existence and uniqueness of the stationary distribution of the SGD iterates . Proposition 2.1 ( Ergodicity of SGD ) . Let the Assumptions 2.1-2.3 hold , and the step size satisfy 0 < η < α− √ ( α2− ( 3L2+Lξ ) ) ∨0 3L2+Lξ . ( a ) SGD ( 2 ) admits a unique stationary distribution πη ∈ P2 ( Rd ) , depending on the step size η . ( b ) For a test function φ : Rd → R satisfying |φ ( θ ) | ≤ Lφ ( 1 + ‖θ‖ ) , ∀θ ∈ Rd and some Lφ > 0 , and for any initialization θ ( η ) 0 = θ0 ∈ Rd of the SGD algorithm , there exists ρ ∈ ( 0 , 1 ) and κ ( both depending on η ) such that we have∣∣E [ φ ( θ ( η ) k ) ] − πη ( φ ) ∣∣ ≤ κ ρk ( 1 + ‖θ0‖2 ) , where πη ( φ ) : = ∫ φ ( x ) dπη ( x ) . The uniqueness of the stationary distribution of the constant step size SGD has been established in [ 1 ] for strongly convex and smooth objectives . In Proposition 2.1 , we relax both of these assumptions allowing for non-convex and non-smooth objectives . Our proof relies on V -uniform ergodicity [ 75 ] , which is fundamentally different from the ergodicity analysis in [ 1 ] . Under the dissipativity condition ( quadratic growth of f ) , geometric ergodicity in Proposition 2.1 is not surprising ; yet , it is worth highlighting that the function f as well as the noise sequence require significantly less structure than what was assumed in the literature . The above step size assumption is almost standard and it is required to obtain a uniform bound on the moments of SGD iterates . We highlight that similar to the gradient descent algorithm , the step size depends on a quantity that serves as a surrogate condition number in our setting , namely , L/α . Note that ρ depends on η and will typically be converging to one if η → 0 . Thus convergence in Proposition 2.1 ( b ) can be expected to be slower when η becomes smaller . However , smaller η leads to a better control of the asymptotic bias ( under additional regularity assumptions ) , see Theorems 3.1-3.3 . Both of those statements ( slower convergence for smaller η but smaller bias eventually ) are confirmed in our numerical experiments , see Figure 1 ( d , h ) . Next , we state our first principal contribution , a central limit theorem for the averaged SGD iterates starting from any initial distribution , for a non-convex objective . For a test function φ : Rd → R , we denote the centered partial sums of φ evaluated at the SGD iterates with Sn ( φ ) , i.e. , Sn ( φ ) : = ∑n−1 k=0 [ φ ( θ ( η ) k ) − πη ( φ ) ] , where πη ( φ ) : = ∫ φ ( x ) dπη ( x ) . Theorem 2.1 ( CLT ) . Let the Assumptions 2.1-2.3 hold . For a step size η and a test function φ satisfying the conditions in Proposition 2.1 , we define σ2πη ( φ ) : = limn→∞ 1 nEπη [ S2n ( φ ) ] . Then , n−1/2Sn ( φ ) d−→ N ( 0 , σ2πη ( φ ) ) . The above result characterizes the fluctuations of a test function φ averaged across SGD iterates , even when the objective function is both non-convex and non-smooth . The asymptotic variance in the above CLT can be equivalently stated in another compact form . If we define the centered test function as h ( θ ) = φ ( θ ) − πη ( φ ) , the asymptotic variance can be written as σ2πη ( φ ) = 2πη ( hĥ ) − πη ( h 2 ) , where ĥ = ∑∞ k=0 E [ h ( θ ( η ) k ) ] . Indeed , this is the variance we compute at the end of our proof in Section A . However , the expression in Theorem 2.1 is obtained by simply applying [ 58 , Thm 21.2.6 ] . For the case of strongly convex functions with decreasing step size schedule , it is well-known from the works of [ 28 , 38 ] that the limiting variance of the averaged SGD iterates achieves the Cramer-Rao lower bound for parameter estimation ; see also [ 76 , 30 ] for non-asymptotic rates in various metrics . The question of providing lower bounds for the limiting variance of the critical points in the non-convex setting is extremely subtle , and is often handled on a case-by-case basis . We refer the interested reader to [ 77 , 78 , 12 ] . There are several important implications of the above CLT for constructing CIs in practice . First note that , following the standard construction in inference , one can write the distribution of the sample mean approximately as n−1Sn ( φ ) ≈ N ( 0 , n−1σ2πη ( φ ) ) . Here , one needs to estimate the population quantity , the asymptotic variance σ2πη ( φ ) , for the purpose of obtaining CIs . In Section 5 , we discuss three strategies for estimating this quantity , which could be eventually used for inference in practice . A theoretical analysis of the proposed approaches in Section 5 is beyond the scope of this work .

Summary: This paper studies the asymptotic normality and the bias for non-convex optimization problem with constant stepsize SGD algorithm. Relaxing previous strongly convex assumption, adopting techniques from Markov chains, they show that under some mild conditions, SGD with constant stepsize will quickly converge to a stationary distribution that is uniquely determined by the stepsize $\eta$. Also, they analyze the bias under various conditions, basically showing that the bias is also (positively) dependent of the stepsize $\eta$. Last, they provide thorough experiments in Regularized MLE and Regularized Blake-Zisserman MLE settings, showing that the numerical results match their theoretical prediction.