Linear Convergence of SGD on Overparametrized Shallow Neural Networks

1 INTRODUCTION . Our understanding of the optimization landscape of supervised learning with neural networks has vastly improved in recent years . This is in part due to the observation that overparameterization is key to overcome the pitfalls of first-order methods in general non-convex problems ( Soltanolkotabi et al. , 2019 ) . Under this assumption , a line of research has established convergence of first-order methods such as gradient descent ( GD ) to global optimality , ( Allen-Zhu et al. , 2019 ; Kawaguchi & Huang , 2019 ; Du et al. , 2019 ; Du & Lee , 2018 ; Zou & Gu , 2019 ; Brutzkus & Globerson , 2017 ; Song & Yang , 2019 ; Oymak & Soltanolkotabi , 2020 ) , a phenomenon that has been confirmed in practice . Empirically , as long as the width of a network scales linearly with the size of the training data ( mild overparameterization ) , stochastic gradient descent ( SGD ) enjoys fast convergence to global optimality ( Livni et al. , 2014 ; Safran & Shamir , 2018 ; Oymak & Soltanolkotabi , 2020 ; Kawaguchi & Huang , 2019 ) . Can we explain such behavior theoretically ? Sadly , the available characterizations require a larger degree of overparameterization , or imposes additional assumptions , which do not hold for the algorithms that are used in practice . For example , if GD is applied exclusively to the last layer , Kawaguchi & Huang ( 2019 ) show that an ideal linear scaling of the width is sufficient to guarantee convergence . Song & Yang ( 2019 ) prove quadratic scaling when GD is applied only to the first layer . For two-layer neural networks , when both layers are trained with GD simultaneously , state-of-the-art results show that subquadratic ( not linear ) scaling is enough to converge to global optimality ( Anonymous ) . Despite being close to the ideal linear rate of overparameterization , due to computational constraints , GD is rarely used in modern applications involving huge datasets . Hence , closing the gap between theory and practice requires studying scalable first-order algorithms such as SGD . Our work focuses on mini-batch SGD , which is one of the most common algorithms for training deep models . We study convergence of SGD when it is applied to train both layers of a neural network , which is initialized with standard initialization schemes . Our contributions : • We show that under proper initialization and choice of learning rate , the iterates of SGD converge to a global minimum with high probability and exponentially fast for a general non-convex problem assuming that the loss function satisfies a growth condition . • For the special case a two-layer neural network , we show that a subquadratic scaling on the width is sufficient under standard initialization and training both layers simultaneously , if the minibatch size is sufficiently large and it also grows with the number of training examples . For constant batch size , we show that quadratic overparametrization is sufficient . Our results interpolate between subquadratic and quadratic scalings depending on the batch size . Related work . The majority of the existing literature on overparameterization focuses on GD ( Du et al. , 2019 ; Du & Lee , 2018 ; Allen-Zhu et al. , 2019 ; Zou & Gu , 2019 ) . Allen-Zhu et al . ( 2019 ) provided theoretical bounds for deep networks trained with SGD . However , their results require an overparameterization degree that is too large , compared to what can be achieved for GD . In contrast , we study SGD and how the batch size affects the required degree of overparameterization . Chen et al . ( 2021 ) establish generalization guarantees and sufficient network width when SGD trains deep ReLU networks for binary classification , which is a different setting compared to our paper . We study the case where SGD updates all the parameters of a shallow neural network . In contrast , a number of existing literature assume that only the parameters corresponding to some layers are updated throughout training ( Oymak & Soltanolkotabi , 2020 ; Kawaguchi & Huang , 2019 ; Song & Yang , 2019 ) . When SGD is applied only to the first layer , Oymak & Soltanolkotabi ( 2020 ) showed that quadratic scaling is sufficient for convergence with linear rate . Despite being an interesting theoretical setup , such algorithmic choice rarely happens in practice . There are also differences regarding the choice of activation function . While ReLU can be considered as the default activation function when studying deep neural networks , its non-smoothness may be the reason why results for ReLU networks require substantially more number of parameters or additional assumptions on the data ( like separability ) to guarantee convergence to a global minimum . Moreover , backpropagation on ReLU networks does not correctly calculate the gradient at all points of differentiability ( Kakade & Lee , 2018 ; Bolte & Pauwels , 2021 ) , which raises major technical issues . In contrast , we assume a smooth activation , similar to Anonymous , which avoids such issues and achieves lower overparameterization degrees . The authors of ( Anonymous ) established subquadratic scaling when GD trains a shallow neural network . In this paper , we focus on SGD , which results in substantial technical challenges . Compared to the results in ( Anonymous ) , controlling the length of the trajectory is more involved in this paper , which requires a new analysis technique that bounds the length of the trajectory with high probability . We consider the effect of mini-batch SGD , which shows an interpolation between subquadratic and quadratic scaling . We also improve the estimates at initialization and show that more relaxed assumptions are sufficient to establish sufficient overparameterization degree . We summarize such recent results in the overparametrization literature in Table 1 . Lazy Training . Proving fast convergence to global optimality is not a complete answer . It has been shown that despite fast convergence , it is possible that an algorithm tends towards a solution with poor performance on test data , if the training falls in the so-called Lazy Training regime ( Chizat et al. , 2019 ) . Thus , any useful algorithmic framework for learning neural networks should avoid this regime , usually through careful initialization schemes . For example , despite requiring only linear overparametrization for GD , the initialization studied by Nguyen & Mondelli ( 2020 ) leads to the lazy regime . This is the reason why we omit such result from Table 1 . In this work , we show global convergence and achieve subquadratic scaling under standard initialization schemes , which empirically perform well on test data . Polylogarithmic width is enough to obtain convergence for neural networks of arbitrary depth , according to Ji & Telgarsky ( 2020 ) ; Chen et al . ( 2021 ) . However , in those work , convergence is understood in an ergodic sense . This is a weaker notion than strict convergence with high probability , which is the one we consider , and which better matches practical applications . Given perfect knowledge about the underlying function that generates the labels and under the assumption that such target function has low-rank approximation , Su & Yang ( 2019 ) showed that GD achieves zero-approximation . This is different from the problem considered in our paper . For a binary classification problem , Daniely ( 2020 ) showed that near linear network size is sufficient for SGD to memorize random examples under a variant of Xavier initialization , which is a different setting compared to our paper . For a deep neural network with pyramidal structure and smooth activations , Nguyen & Mondelli ( 2020 ) showed that subquadratic scaling is sufficient for global convergence of GD under a restrictive initialization scheme . In this paper , we establish global convergence for SGD under standard initialization . A recent line of work uses mean-field analysis to approximate the target distribution of the weights in a neural network via their empirical distribution ( Mei et al. , 2019 ; Lu et al. , 2020 ) . Nevertheless , such results do not provide useful overparametrization degree bounds in terms of the number of samples . In contrast , our work does not require such approximations and we focus on deriving explicit sufficient overparametrization rates for global optimality of SGD . Notation . We use ‖ · ‖ to denote the Euclidean norm of a vector and Frobenius norm of a matrix . We use∇ to represent the Jacobian of a vector-valued and gradient of a scalar-valued function . We use and ⊗ to represent the entry-wise Hadamard product and Kronecker product , respectively . We use lower-case bold font to denote vectors . We use calligraphic and standard fonts to represent sets and scalars , respectively . We use σmin ( T ) and σmax ( T ) to denote the smallest and largest singular values of a linear map T . We use [ n ] to represent { 1 , · · · , n } for an integer n. We use Õ and Ω̃ to hide logarithmic factors and use . to ignore terms up to constant and logarithmic factors . 2 A GENERAL GLOBAL CONVERGENCE RESULT FOR SGD . In this section , we consider a general non-convex minimization problem and show that for a certain choice of learning rate and careful initialization , the iterates of SGD converge to a global minimum with high probability and exponentially fast . In Section 3 , we extend our consideration to the training of a shallow neural network and find the hidden layer size , which is sufficient for SGD to converge to a global minimum i.e. , its overparameterization degree . Definition 1 ( Smoothness ) . Let βψ > 0 . A function ψ : Rd1 → Rd2 is βψ-smooth , if for all u , v ∈ Rd1 , we have σmax ( ∇ψ ( u ) −∇ψ ( v ) ) ≤ βψ‖u− v‖ . ( 1 ) Definition 2 ( PL condition ( Bolte et al. , 2017 ) ) . A function ψ : Rd1 → R satisfies the PL condition if there exists αψ > 0 such that , for all u ∈ Rd1 , we have ψ ( u ) ≤ ‖∇ψ ( u ) ‖ 2 2αψ . ( 2 ) We are now ready to state our finite-sum compositional optimization problem : min w∈Rd h ( w ) : = f ( Φ ( w ) ) = 1m m∑ j=1 fj ( Φ ( w ) ) , Φ : Rd → Rd̃ , f , fj : Rd̃ → R+ ( 3 ) where m denotes the number of training examples . Assumption 1 . The functions introduced in Eq . ( 3 ) satisfy the following properties : ( i ) Φ is twicedifferentiable and βΦ-smooth ( Definition 1 ) , ( ii ) f is twice-differentiable and βf -smooth and ( iii ) f satisfies the PL condition with some αf > 0 ( Definition 2 ) . We study the iterates of the stochastic gradient descent ( SGD ) algorithm when applied to the objective function h in Eq . ( 3 ) . For i ≥ 0 , let Ii denote a random minibatch at iteration i drawn uniformly at random , independent of all previous draws . Let b ∈ [ m ] denote the minibatch size , i.e. , |Ii| = b for all i . The SGD iterates are defined by a random variable w0 , referred to as the initialization , and the update rule : wi+1 = wi − λ1 b ∑ j∈Ii ∇hj ( wi ) , ( 4 ) where λ > 0 is the learning rate and hj ( w ) : = fj ( Φ ( w ) ) ∀j , w . An important feature of the SGD iterates is that 1b ∑ j∈Ii ∇hj ( Φ ( wi ) ) in Eq . ( 4 ) is an unbiased estimator of the gradient ∇h ( wi ) given wi , i.e. , E [ 1 b ∑ j∈Ii ∇hj ( Φ ( wi ) ) |wi ] = ∇h ( wi ) . Never- theless , this is not enough for SGD to converge to the first-order optimality . In addition , we will assume that f in Eq . ( 3 ) satisfies the growth condition ( Schmidt & Roux , 2013 ; Vaswani et al. , 2019 ; Cevher & Vu , 2019 ) : Definition 3 ( Growth condition ) . A function ψ : Rd → R with a finite-sum structure satisfies the growth condition with minibatch size b if there exists ηψ > 0 such that , for all u ∈ Rd , we have E [ ∥∥∥1 c ∑ j∈I ∇ψj ( u ) ∥∥∥2 ] ≤ ηψ‖∇ψ ( u ) ‖2 , ( 5 ) where the expectation is over the random choice of set I . Assumption 2 . In Eq . ( 3 ) , f satisfies the growth condition ( Definition 3 ) for some ηf > 0 . We are now ready to state the main result of this section . Theorem 1 . Let Assumptions 1 and 2 hold and let ζ > 1 . Suppose that at initialization , 0 < µΦ ≤ σmin ( ∇Φ∗ ( w0 ) ) ≤ σmax ( ∇Φ∗ ( w0 ) ) ≤ νΦ , h ( w0 ) = O ( αfµ 6 Φ ζβ2Φηfν 2 Φ ) . ( 6 ) Then , for a sufficiently small learning rate λ . min ( µ2Φ ηf ( βΦν2Φ‖∇f ( Φ ( w0 ) ) ‖+ βfν4Φ + βfµΦν3Φ ) , µΦ ζ √ ηfνΦ ( βΦ‖∇f ( Φ ( w0 ) ) ‖+ βfνΦµΦ ) ) ( 7 ) the iterates of SGD { wi } i≥0 ( 4 ) converge to a global minimizer of h ( 3 ) with the optimal value of zero , exponentially fast and with probability at least 1− 1/ζ . The rate of convergence is given by E [ h ( wi ) ] ≤ ( 1− Cλαfµ2Φ ) i · h ( w0 ) for a universal constant C. Remark 1 . The second item in Eq . ( 6 ) suggests initializing close to a global minimum of the nonconvex optimization problem . This feature has precedence in the related literature , e.g. , in matrix factorization ( Chi et al. , 2019 ) . The proof of Theorem 1 is deferred to Appendix B . However , in the remaining of this section we provide a sketch of the main arguments that lead to the result . The first condition in Eq . ( 6 ) is central to our arguments , and we will refer to it as the near-isometry property . Definition 4 ( Near-isometry ) . A linear map T : Rd1 → Rd2 is ( µ , ν ) -near-isometry if there exist 0 < µ ≤ ν such that µ ≤ σmin ( T ) ≤ σmax ( T ) ≤ ν . ( 8 ) Let w denote the limit point when the SGD algorithm is run with some learning rate and let∇Φ∗ ( w ) denote the adjoint operator of∇Φ ( w ) . Convergence of SGD is ensured with high probability due to the strong growth condition ( Definition 3 ) along with proper learning rate and initialization . We note that w is a first-order stationary point of h. Hence we have : 0 = ∇h ( w ) = ∇Φ∗ ( w ) ∇f ( Φ ( w ) ) ( 9 ) Note that if ∇Φ∗ ( w ) is nonsingular , it would follow that ∇f ( Φ ( w ) ) = 0 . The PL condition ( Definition 2 ) would then imply that Φ ( w ) is a global minimizer of f and hence , a global minimizer of h. With this fact in mind , our proof can be summarized in three steps : first , a careful choice of initialization will ensure that ∇Φ∗ is nonsingular for all elements within a certain distance of w0 . Second , we show that under small enough learning rate , the iterates of SGD remain close to the initialization w0 , with high probability regardless of the number of iterations . The third and final step will use the non-singularity of∇Φ∗ at convergence and Eq . ( 9 ) to conclude global optimality . This is akin to the arguments in ( Anonymous ) , however , in our case the stochasticity in the SGD updates poses a challenge for controlling the distance to initialization . We use concentration bounds on the length of the path and show that the SGD trajectory remains in the region where ∇Φ∗ is non-singular , with high probability . A crucial result for the first step of our proof has already been established in by ( Anonymous ) . It shows that a smooth function that is near-isometry at initialization will preserve such property for all points within a certain distance . Lemma 1 ( Anonymous ) . Let Φ be βΦ-smooth and ∇Φ∗ ( w0 ) be a ( µΦ , νΦ ) -near-isometry . Then for all w such that ‖w −w0‖ ≤ µΦ 2βΦ , µΦ 2 ≤ σmin ( ∇Φ∗ ( w ) ) ≤ σmax ( ∇Φ∗ ( w ) ) ≤ 3νΦ 2 ( 10 ) The second step in the proof of Theorem 1 is to compute the expected length of the SGD trajectory which is spent inside the ball defined in Eq . ( 10 ) . We find an upper bound on this expected length depending on the initialization and learning rate , but independent of the number of iterations . Hence , under some proper initialization and learning rate , we can control the expected length of the trajectory for which Lemma 1 holds . In particular , we have Proposition 1 ( Expected length of trajectory ) . Let Assumptions 1 and 2 hold and let ζ > 1 . Let the random variable I denote the first iteration of SGD ( Eq . ( 4 ) ) such that wI /∈ B : = ball ( w0 , ρΦ ) : = { w : ‖w −w0‖ ≤ ρΦ } ( 11 ) or I =∞ if the trajectory does not leave B . Suppose that w0 satisfies Eq . ( 6 ) and SGD is executed with sufficiently small learning rate , which satisfies Eq . ( 7 ) . An upper bound on the expected length of the SGD trajectory is given by E [ ` ( I ) ] ≤ µΦ 2ζβΦ = ρΦ ζ . ( 12 ) We provide the sketch of the proof ( see Appendix A for the complete proof ) . We first find an upper bound on the expected length of the trajectory in terms of the norm of gradients of f . With a proper learning rate , we find an upper bound on the norm of the gradient in terms of the expected decent of f in two consecutive iterates , which are inside the ball . We also ensure that the learning rate is sufficiently small such that E [ ‖wI − wI−1‖ ] is bounded . Finally , under proper initialization , we obtain an upper bound on the expected length of the trajectory for the iterates inside the ball , i.e. , E [ ∑I−1 i=0 ‖wi −wi−1‖ ] . Remark 2 . A similar phenomenon that shows bounded length of the trajectory has been observed in various settings mainly for gradient descent ( Du et al. , 2019 ; Oymak & Soltanolkotabi , 2019 ; Anonymous ) . In this paper , we focus on a compositional non-convex problem trained with SGD , which is more challenging to analyze . Using the upper bound ( 12 ) on the expected length of the trajectory spent inside B = ball ( w0 , ρΦ ) , we can bound the probability that the SGD iterates leaves the ball B . Indeed , in order for the process to leave B starting from w0 , it is required that the length of the trajectory spent inside B satisfies l ( I ) ≥ ρΦ . Hence , using bound ( 12 ) on E [ ` ( I ) ] together with a concentration bound ( Markov inequality in our case ) , we can upper bound the probability of SGD iterates leaving B . Finally , under the event that the SGD iterates remain in B , an upper bound on E [ ` ( I ) ] implies the convergence of the iterates . Remark 3 . With a more involved analysis on the concentration properties of the random variable ` ( I ) , it may be possible to greatly improve the dependence of the initialization and step size on ζ . Indeed , the current analysis assumes the worst-case scenario , where the SGD iterates either remain at the initialization , or directly leave the ball B in a straight line ( this scenario indeed maximizes the probability that the process leaves the ball , given a bound on E [ ` ( I ) ] ) . Although ` ( I ) is obtained as a sum of random variables , the difficulty of obtaining better concentration bounds for ` ( I ) comes from the high level of dependence between all the variables involved . A better analysis would thus need to better understand how the trajectories behave inside the ball B , e.g. , by bounding the variance of ` ( I ) . In the following section , we specify our result to the special case of shallow neural networks . We will show that , in the case of quadratic loss , the strong growth condition naturally holds , with a constant depending on the batch size . Moreover , using Gaussian initialization for the neural network parameters , we can control the initial smoothness and near-isometry parameters involved in Theorem 1 with high probability .

Summary: This paper proved that a two-layer neural network with smooth activation and proper initialization can converge linearly to a global minima of training loss using mini-batch SGD when the width is larger than $\Omega(m^2/\sqrt{b})$ where $m$ is the number of training data and $b$ is batch size. As the batch size increases, this provides an interpolation between the quadratic result for SGD and the sub-quadratic result for full-batch GD. To prove this result, the authors first provide a general convergence result for SGD on a particular class of functions and then apply this framework to 2-layer neural networks to derive the width requirement.