Subquadratic Overparameterization for Shallow Neural Networks

1 Introduction . Training a neural network involves solving a nonconvex optimization problem , which , in theory , might trap first-order methods such as gradient descent to fall in bad local minima or saddle points . However , empirical evidence suggests that first-order methods with random initialization can consistently find a global minimum , even with randomized labels [ 43 ] . Demystifying this observation is of central interest to deep learning . Recently , a line of research [ 45 , 4 , 11 , 29 , 39 , 12 , 38 ] suggests that such an empirical success can possibly be explained by the overparameterization of neural networks , whose number of parameters exceeds the number of training data n. In particular , gradient descent converges linearly fast to a global optimum in a number of problems with models that have wide hidden layers [ 45 , 12 , 39 ] . Despite of these remarkable results , the natural key question “ How much should we overparameterize a neural network ? ” remains open even for the toy example of two-layer neural networks . On one hand , it is widely accepted that , for two-layer neural networks , the number of parameters should grow linearly with n ( e.g. , [ 21 , 38 ] ) . On the other hand , theoretical results either require much more parameters , or they are established under restrictive settings . Specifically , ∗Equal contributions . †This work was done while Armin Eftekhari was at EPFL . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . • Kawaguchi and Huang [ 21 ] has proven the ideal Ω̃ ( n ) scaling for deep neural networks . However , they apply gradient descent only to the last layer , which is not the case in practical scenarios . • A similar issue exists in [ 39 , 38 ] , where the authors have shown that Ω̃ ( n2 ) parameters suffice for two-layer neural networks , but only the first layers are trained . Furthermore , even with infinite width , Oymak and Soltanolkotabi [ 38 ] can not guarantee zero training error with probability approaching to one . The goal of this paper is to close the gap between theory and practice , without resorting to unrealistic assumptions such as those discussed above . We sharpen the results of Oymak and Soltanolkotabi [ 38 ] by proving that , with proper random initialization of each layer , training error approaches to zero with high probability , exponentially fast in the width of the network . In addition , we show that only Ω̃ ( n 3 2 ) parameters suffice such that gradient descent converges to a global minimum with linear rate , which improves upon the state-of-the-art by a factor of Õ ( n 1 2 ) . We summarize the bounds on the number of parameters in terms of n in Table 1 . While our analysis on gradient descent focuses on training error , it has been observed that overparameterization can lead to poor generalization . In particular , [ 7 , 42 , 15 ] have observed the phenomenon of lazy training . Chizat et al . [ 7 ] has explained lazy training as a model behaves similar to its linearization around the initialization . It is known that an overparameterized neural network is likely to be trapped in the lazy regime since the parameters will hardly vary over the course of training with gradient descent [ 12 , 29 , 45 ] . The same phenomenon has been observed for infinitely wide neural networks [ 19 ] . In this paper , we provide theoretical guidance to possibly avoid lazy training through proper initialization . Experimental results confirm that lazy training might be avoided with our theoretically inspired initialization so that the issues reported in [ 7 ] do not apply . 1.1 Summary of contributions . • We first focus on a general minimization problem assuming that the loss function satisfies PolyakŁojasiewicz ( PL ) condition . We find sufficient conditions in terms of initialization for the convergence of gradient flow and gradient descent to a global minimum . • We then focus on the special problem of training a two-layer neural network with quadratic loss and smooth activation , and show that Ω̃ ( n 3 2 ) parameters are sufficient for gradient descent to converge to a global minimum with linear rate and probability approaching to one . We achieve linear scaling for the width when the number of input features is in Ω̃ ( √ n ) . • We theoretically guide how to initialize the parameters of a neural network in the overparameterized regime of interest while possibly avoiding lazy training . 1.2 Further related work . In terms of techniques , our paper is closely related to [ 37 , 38 ] . Similar to our Theorem 3 , Oymak and Soltanolkotabi [ 38 , Theorem 2.1 ] showed that gradient descent converges with linear rate when the Jacobian of the nonlinear mapping has smooth deviations , and the number of parameters grows quadratically with n. However , Oymak and Soltanolkotabi [ 38 ] assumed that gradient descent updates only the first layer . In this paper , we consider the case where gradient descent updates both layers simultaneously , and show that it suffices to have Ω̃ ( n 3 2 ) parameters with a linear rate of convergence . ReLU is an important instance of activation functions that does not satisfy the smoothness assumption . A line of research aims to relax this assumption by instead assuming the data is separable . For shallow neural networks , Du et al . [ 12 ] proved that gradient descent finds a global minimum if the width of the network scales Ω̃ ( n6 ) assuming that no two data points are parallel . In a similar setting , Song and Yang [ 39 ] established convergence to a global minimum with the sufficient width of Ω̃ ( n2 ) . As a result , in the absence of the smoothness assumption , these papers require substantially more number of parameters to guarantee convergence to a global minimum . The theoretical bounds for deep neural networks are even worse . For instance , Allen-Zhu et al . [ 1 ] required the total number of parameters of Ω ( n24L12 ) where L is the number of layers . Zou and Gu [ 44 ] improved the scaling to Ω ( n8L12 ) . In our setting , i.e. , L = 2 , these bounds become vacuous in most interesting regimes . Further , in [ 21 ] , the authors showed that Ω̃ ( n ) parameters is enough to achieve global convergence under the assumption that gradient descent updates only the last layer , which essentially reduces the problem to a simple least-squares regression . Recently , Ji and Telgarsky [ 20 ] , Chen et al . [ 5 ] showed that a polylogarithmic width suffices to achieve convergence for shallow and deep neural networks in an ergodic sense . We note that this is a weaker notion of convergence compared to the one we consider . Li et al . [ 28 ] showed that gradient descent along with early stopping are robust to label noise on a constant fraction of labels in an overparameterized network . However , only the first layer is optimized in [ 28 ] . For possibly overparameterized and linear networks , Eftekhari [ 14 ] showed that gradient flow can successfully avoid lazy training assuming that the network has a layer with a single neuron . We note that our analysis does not require those restrictions . Under an assumption similar to PL condition , Zou et al . [ 45 ] studied the problem of binary classification for a deep network with ReLU activation , which is a different problem compared to ours . In [ 40 ] , the authors proved that gradient descent with overparameterization achieves zero-approximation when the underlying function that generates the labels has low-rank approximation . Their scaling requires perfect information about the target function , which is not the case in our paper . Under a variant of Xavier initialization , Daniely [ 9 ] found near optimal scaling for a binary classification problem trained by stochastic gradient descent . We note that the setting considered in our paper is more challenging than binary classification . Our results establish a new state-of-the-art on the required number of parameters in a nonrestrictive setting when both layers are trained at the same time . Recently , Nguyen and Mondelli [ 34 ] obtained subquadratic scaling for a deep neural network with pyramidal structure under an initialization that leads to lazy training . Our results do not have such restrictions . Mean-field analysis was used to approximate a target distribution of parameters of a neural network by the empirical distributions [ 33 , 32 ] . However , these results do not provide useful bounds on the scaling in terms of n , which is our focus in this paper . Liu et al . [ 31 ] established global convergence when the function to minimize satisfies a variant of PL condition ( local PL condition ) assuming the map is Lipschitz continuous , which is not the case in our paper . Liu et al . [ 30 ] characterized the constancy of the neural tangent kernel via scaling properties of the norm of the Hessian matrix of the network . In this work , we focus on obtaining a sufficient number of parameters for gradient descent to converge to a global minimum with linear rate . Notation . We use ‖ · ‖ to represent the Euclidean norm of a vector and Frobenius norm of a matrix . We use ∇ to denote the Jacobian of a vector-valued and gradient of a scalar-valued function and ∇Φ ( a ) { b } to represent the directional derivative of Φ along b . We use and ⊗ to denote the Hadamard ( entry-wise ) product and Kronecker product , respectively . For A ∈ Rm×n and t ∈ Z+ , we denote A∗t ∈ Rmt×n with its a-th column defined as vec ( xa ⊗ · · · ⊗ xa ) ∈ Rm t . We use lower-case bold font to denote vectors . Sets and scalars are represented by calligraphic and standard fonts , respectively . We use [ n ] to denote { 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 Problem , definitions , and assumptions . In this section , we set up a general compositional optimization problem . Then we focus on the special case of shallow neural networks in Section 5 . Let w ∈ Rd denote a parameter vector where d denotes the number of parameters . In a neural network , w consists of weights and biases of all layers . We consider the minimization problem min w∈Rd h ( w ) ( 1 ) where h : Rd → R+ is the composition of a loss function f : Rd̃ → R+ and a nonlinear and nonconvex function Φ : Rd → Rd̃ : h ( w ) = f ( Φ ( w ) ) = f ( z ) ( 2 ) where z = Φ ( w ) . Before providing the details , let us highlight the simple idea behind the argument ( see also [ 37 ] ) . Let w0 and w denote the initial point and limit point when the gradient descent algorithm is run with some learning rate , respectively . The precise formulation of gradient descent is provided in Section 4 . Let ∇Φ∗ ( w ) : Rd̃ → Rd denote the adjoint operator of ∇Φ ( w ) . Since w is a first-order stationary point of h , we have 0 = ∇h ( w ) = ∇Φ∗ ( w ) { ∇f ( z ) } where z = Φ ( w ) . Suppose that∇Φ∗ ( w ) is a nonsingular operator.Then∇f ( z ) = 0 . If z is a global minimizer of f , then w is a global minimizer of h. To prove global convergence , it suffices to show that∇Φ∗ is nonsingular within a neighborhood of the initialization w0 , and that points reached by gradient descent remain within this neighborhood . We will prove that both statements hold with high probability for shallow neural networks . We first define two notions that are useful to state a key lemma for our main results : Definition 1 ( Near-isometry ) . A linear mapping T : Rd1 → Rd2 is ( µ , ν ) -near-isometry if there exist 0 < µ ≤ ν such that µ ≤ σmin ( T ) ≤ σmax ( T ) ≤ ν . ( 3 ) Definition 2 ( Smoothness ) . Let βψ > 0 . A function ψ : Rd1 → Rd2 is βψ-smooth , if for all u , v ∈ Rd1 , we have σmax ( ∇ψ ( u ) −∇ψ ( v ) ) ≤ βψ‖u− v‖ . ( 4 ) The following lemma shows that a smooth function , which is near-isometry at initialization , remains near-isometry for all nearby points of the initialization . Lemma 1 . Suppose that Φ is βΦ-smooth and ∇Φ∗ ( w0 ) is ( µΦ , νΦ ) -near-isometry . Then , for all w ∈ ball ( w0 , ρΦ ) , we have µΦ 2 ≤ σmin ( ∇Φ∗ ( w ) ) ≤ σmax ( ∇Φ∗ ( w ) ) ≤ 3νΦ 2 ( 5 ) where ρΦ = µΦ 2βΦ . ( 6 ) Intuitively , if ∇Φ∗ ( w0 ) is a ( µΦ , νΦ ) -near-isometry , then one would expect ∇Φ∗ to remain nearisometry for all nearby points . Definition 3 ( PL condition [ 3 ] ) . A function ψ : Rd1 → R satisfies the PL condition if there exists αψ > 0 such that , for all u ∈ Rd1 , we have ψ ( u ) ≤ ‖∇ψ ( u ) ‖ 2 2αψ . ( 7 ) We note that strongly convex functions satisfy a minor variation of the PL condition in ( 7 ) . In our analysis , we will assume that Φ and f satisfy the following properties : Assumption 1 ( Basic assumptions for Φ , f ) . • Φ is twice-differentiable and βΦ-smooth . • f is twice-differentiable , satisfies the PL condition with αf , and min f ( z ) = 0 . Despite f satisfies the PL condition , the nonconvex Φ can render h nonconvex , and hence difficult to minimize in theory . However , we show that fast convergence of gradient descent to a global minimum can be established with appropriate initialization . The intuition behind these assumptions is that to achieve nonsingularity of ∇Φ∗ , we approximate ∇Φ∗ ( w0 ) at initialization and bound∇Φ∗ ( w0 ) −∇Φ∗ ( wi ) at iteration i using the fact that ‖w0−wi‖ is sufficiently small by the overparameterization . In the special case of shallow neural networks , we expect a similar argument applies even when the activation function is ReLU . Adapting our analysis for such extensions is an interesting area of future work .

Summary: In this paper, the author focus on the theoretical framework of minimizing training risk when the loss satisfy PL conditions. They prove the linear convergence in subquadratically over-parameterized shallow networks with smooth activation functions. They conduct experiments on different initialization to support the theory. The theory also provides insights toward avoiding lazy training.