Subquadratic Overparameterization for Shallow Neural Networks
Year: 2022

Abstract:
Overparameterization refers to the important phenomenon where the width of a neural network is chosen such that learning algorithms can provably attain zero loss in nonconvex training. The existing theory establishes such global convergence using various initialization strategies, training modifications, and width scalings. In particular, the state-of-the-art results require the width to scale quadratically with the number of training data under standard initialization strategies used in practice for best generalization performance. In contrast, the most recent results obtain linear scaling either with requiring initializations that lead to the “lazy-training”, or training only a single layer. In this work, we provide an analytical framework that allows us to adopt standard initialization strategies, possibly avoid lazy training, and train all layers simultaneously in basic shallow neural networks while attaining a desirable subquadratic scaling on the network width. We achieve the desiderata via Polyak-Łojasiewicz condition, smoothness, and standard assumptions on data, and use tools from random matrix theory.

heading: 1 Introduction
text: 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] cannot 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.

heading: 1.1 Summary of contributions
text: • 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.

heading: 1.2 Further related work
text: 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.

heading: 2 Problem, definitions, and assumptions
text: 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.

heading: 3 Gradient flow
text: In this section, we consider gradient flow, which can be viewed as the limit of gradient descent for infinitesimally small learning rates. Inspired by the analysis of gradient flow, we provide an upper bound on the length of the trajectory traversed by gradient descent iterates and then find a sufficient condition in terms of initialization to establish its convergence to a global minimum. We focus on gradient descent in Section 4.
Let t ≥ 0 and consider the gradient flow, which is initialized at w0 ∈ Rd and traverses the curve γ : R+ → Rd, given by
γ̇(t) = dγ(t)
dt = −∇h(γ(t)) (8)
where γ(0) = w0.
We now calculate the length of the curve γ. Suppose that∇Φ∗(w0) is (µΦ, νΦ)-near-isometry. Using Lemma 1, in the following lemma, we control the length inside of ball (w0, ρΦ). See Appendix B for the proof. Lemma 2. Let t ≥ 0 and let `(t) denote the length of the curve γ in (8), restricted to the interval [0, t]. Let tΦ ∈ (0,∞] be the smallest value such that γ(tΦ) /∈ ball(w0, ρΦ). Suppose ∇Φ∗(w0) is (µΦ, νΦ)-near-isometry. Then, for all t ≤ tΦ, we have
`(t) = O
( νΦ √ h(w0)
µ2Φ √ αf
) .
Lemma 2 implies that if the objective value at initialization, h(w0), is sufficiently small, then we can localize gradient flows to a region around w0. Combining with Lemma 1, we show that the limit point of gradient flow is a global minimum. This theorem is formally stated below. Theorem 1 (Gradient flow). Let w0 ∈ Rd. Suppose that Φ and f satisfy Assumption 1 and∇Φ∗(w0) is (µΦ, νΦ)-near-isometry. If w0 satisfies
h(w0) = O
( αfµ 6 Φ
β2Φν 2 Φ
) , (9)
then the gradient flow γ in (8) converges to a global minimum.
Proof of Theorem 1. Proper initialization in (9) ensures `(tΦ) < ρΦ, which implies that ‖γ(tΦ)−w0‖ = ‖γ(tΦ)− γ(0)‖ < ρΦ. (10)
Therefore, γ(t) ∈ ball(w0, ρΦ) for all t ≥ 0, and the length of γ is upper bounded by ρΦ using Lemma 2. Hence, the gradient flow γ converges, i.e., the limit point w ∈ Rd exists and satisfies
‖w −w0‖ ≤ ρΦ. (11) Combining (5) and (11), we have
µΦ 2 ≤ σmin(∇Φ∗(w)) ≤ σmax(∇Φ∗(w)) ≤ 3νΦ 2 .
In particular, we note that∇Φ∗(w) is nonsingular. So we have∇f(z) = 0. Since f satisfies the PL condition in (7), z is a global minimizer of f , and w is a global minimizer of h in (1).

heading: 4 Gradient descent
text: We now view gradient descent as the discretization of gradient flow, and show that a similar argument as in Section 3 holds for gradient descent.
Let η > 0 denote the learning rate and let i ≥ 0. The gradient descent update rule is given by
wi+1 = wi − η∇h(wi). (12)
To study gradient descent, in addition to the previous assumptions on Φ and f for the case of gradient flow described in Theorem 1, we also assume that f is smooth, i.e., there exists βf ≥ 0 such that, for all z, z′ ∈ Rd̃, we have
f(z)− f(z′) ≤ 〈z− z′,∇f(z′)〉+ βf 2 ‖z− z′‖2.
Smoothness of f allows safe discretization of gradient flow without deviating too much from its trajectory. The following result is the analogue of Theorem 1 for gradient descent; see Appendix C for the proof. Theorem 2 (Gradient descent). Let w0 ∈ Rd. Suppose that Φ and f satisfy Assumption 1, f is βf - smooth, and ∇Φ∗(w0) is (µΦ, νΦ)-near-isometry. Suppose that gradient descent is executed with sufficiently small learning rate
η = O
( 1
βΦ‖∇f(Φ(w0)‖+ βfµ2Φ + βfν2Φ
) , (13)
and w0 satisfies (9).
Then the sequence of iterates {wi}i≥0 converges to a global minimum of h exponentially fast. In addition, the rate of convergence is given by
h(wi) ≤ (1− Cηαfµ2Φ)ih(w0) (14)
where C is a universal constant.
To prove Theorem 2, we first compute the length of the trajectory traversed by gradient descent iterates. We then use the smoothness of f and follow the descent inequality to lower bound f(zi)− f(zi+1). Finally, we compute the local Lipschitz constant of f . Remark 1. The idea of initializing a nonconvex problem close to a global minimum has a long history in nonconvex optimization, particularly in matrix factorization; see [6] and references therein. The observation that the length of the learning trajectory is short in the overparameterization regime has a precedent in [12, 37]. From an algorithmic perspective, the idea of linearizing Φ when minimizing h = f ◦ Φ is studied in nonlinear regression and the Gauss-Newton method [35].
In order to apply Theorem 2, the key step is to verify that h(w0) satisfies (9). In Section 5, we focus on the special case of shallow neural networks and improve the state of the art.

heading: 5 Shallow neural networks
text: In this section, we consider the problem of training shallow neural networks with gradient descent. Our strategy is to cast this problem as a special case of problem (1) and then apply Theorem 2 to establish global convergence. We start with the formal problem statement.

heading: 5.1 Setup, assumptions, and initialization
text: Consider a shallow neural network with d0 inputs, one hidden layer that consists of d1 hidden nodes, and d2 outputs. This shallow network is specified by the map
Rd0 7→ Rd2
x 7→ V · φ(Wx), (15)
where W ∈ Rd1×d0 , V ∈ Rd2×d1 , and φ : R → R is an activation function, which is applied entry-wise. Let xi ∈ Rd0 and yi ∈ Rd2 denote the i-th training data and label, respectively, for i ∈ [n]. By concatenating the training data and their labels, we form the matrices X ∈ Rd0×n and Y ∈ Rd2×n. Let denote Θ = (W,V ) ∈ Rd1×d0 × Rd2×d1 and Z = Φ(Θ) = V · φ(WX) ∈ Rd2×n. The fitting problem can be cast as (1) where
h(Θ) = f(Φ(Θ)) = ‖V φ(WX)− Y ‖2. (16)
Remark 2. We assume that the activation function φ : R→ R is twice-differentiable. Despite this assumption excludes the popular ReLU, it is still possible to apply our results to smooth approximations of ReLU such as the softplus or Gaussian error Linear Units (GeLU) [18, 34]. We note that softplus [13] or GeLU [10] often achieve similar or superior performance compared to the ReLU [8, 16, 23, 22, 41]. Definition 4 (Hermite norm [36]). Let φ : R → R. The Hermite norm of φ is given by ‖φ‖H =√∑∞
i=0 c 2 i where ci denotes the i-th Hermite coefficients of φ given by:
ci = 〈φ, qi〉H = 1√ 2π
∫ φ(x)qi(x) exp ( −x 2
2
) dx
and qi : R→ R is the i-th Hermite polynomial for i ≥ 0.
In this section, we assume that φ, f , and data satisfy the following properties: Assumption 2 (Assumptions for shallow neural networks).
• φ is twice-differentiable, φ(0) = 0, supa |φ̇(a)| = φ̇max < ∞, supa |φ̈(a)| = φ̈max < ∞, and ‖φ‖H <∞. The loss function f is quadratic (16).
• ‖xi‖ = 1, ‖Y ‖ ≤ 1, and σmax(Vk) = O ( φ̇max φ̈max ) for i ∈ [n] and k ≥ 0.
The assumption on φ hold for GeLU, sigmoid, and tanh. The assumption φ(0) = 0 is to simplify the derivations and we suspect that it can be removed at the expense of more complicated expressions. The bounded Hermite norm is a mild assumption, which is used to obtain an upper bound on σmax(φ(W
0X)) in terms of the Hermite coefficients of φ. See Appendix E.1 for details. The assumption on the data is fairly mild and standard in the overparameterization literature as we can always normalize the data [29, 20]. Similar boundedness assumptions to the last assumption are commonly used in nonconvex optimization to guarantee convergence [24]. Moreover, such a bound naturally holds by applying a projection step to the gradient descent update rule, which we plan to adopt as a future work.
Initialization. We first consider the initialization scheme: W0 ∼N (0, ω21), V0 ∼N ( 0, ω22 ) . (17)
In Section 6, we study the implications of our initialization and show how to possibly avoid lazy training by varying (ω1, ω2).

heading: 5.2 Main results for shallow neural networks
text: For shallow networks as described above, we verify in Appendix D that the key conditions in Lemma 1 hold with high probability. Combining with Theorem 2, we establish the global convergence guarantees. The proof in Appendix E uses standard tools from random matrix theory to control the random variables involved with initialization. We first estimate variables µΦ, νΦ defined in Definition 1 and βΦ in (4) for the neural network described in Section 5.1. Lemma 3 (Estimation of µΦ, νΦ, βΦ). Suppose that a shallow neural network, which is constructed in Section 5.1, satisfies Assumption 2. Then we have
µΦ = σmin(φ(W0X)),
νΦ = φ̇maxσmax(X)σmax(V0) + σmax(φ(W0X)), βΦ = √ 2σmax(X) ( φ̇max + φ̈maxχmax ) (18) where χmax = supV σmax(V ).
Remark 3. The terms σmin(φ(W0X)) and σmax(φ(W0X)) in (18) play a critical role in our analysis. In [12, 39], strictly positiveness of the eigenvalues of Gram matrix is the primary tool to show the convergence. Oymak and Soltanolkotabi [38] also followed a similar argument using the neural network covariance matrix. The underlying intuition seems similar to Lemma 3. However, the resulting bounds are different since gradient descent updates (W,V ) simultaneously in our problem setup, which is more realistic.
By combining Lemma 3 and the results on global convergence of gradient descent in Section 4, we establish global convergence for shallow neural network. Theorem 3 (Shallow network with gradient descent). Consider the shallow network described in Section 5.1 that satisfies Assumption 2 and τ r1 |φ(a)| ≤ |φ(τa)| ≤ τ r2 |φ(a)| for all a, 0 < τ < 1, and some constants r1, r2.3 Suppose that Θ0 is randomly initialized as in (17) with ω1 and ω2, which satisfy
ω1ω2 . 1√ d0d1 , (19)
and suppose that the hidden layer width d1 satisfies
d1 = Ω̃ ( ξ(Cδ, t, φ, {ci}i≥0) σmax(X) 2 √ n
σmin(X∗t)3
) (20)
where Cδ is a set of constants, ξ is a term independent to d0, n, t is a constant such that n ' dt0, and X∗t ∈ Rdt0×n is derived from Khatri-Rao product with its a-th column defined as vec(xa⊗· · ·⊗xa) ∈ Rdt0 . Then gradient descent converges to a global minimum exponentially fast with probability at least 1− ψ(φ, ξ, d0, d1, d2, X).4See Appendix E.6 for the exact expressions of ξ and ψ. Remark 4. Theorem 3 shows that, with sufficient degree of overparameterization, gradient descent finds a global minimum, except with an arbitrary small probability. Note that we need two conditions for Theorem 3 to hold, both of which are related to the overparameterization of the network. The condition (19) is for the concentration of random matrices, to make ψ arbitrary small, and (20) is for the locality of gradient descent.

heading: 5.3 Order analysis
text: We first decompose the random matrix φ(X>W>0 )φ(W0X) into independent random matrices. We then apply concentration inequalities to establish an upper bound on σmax(φ(W0X)) and a lower bound on σmin(φ(W0X)) through the Hermite decomposition of φ(W0X) and note that with high probability, √
c2t t! d1σmin(X
∗t) . σmin(φ(W0X)) . σmax(φ(W0X)) . √ c20dn.
We also find an upper bound on h(Θ0) at initialization. Substituting νΦ, µΦ, βΦ into (9), we obtain the sufficient condition in (20). We note that ξ(Cδ, t, φ, {ci}i≥0) can be viewed as a constant w.r.t. d0, d1, and n. For t = 1, it requires n ' d0, which is not a common setting in practice. For t ≥ 2, we suppose that n ' dt0, which is the case in practice and estimate σmax(X) ' √ n d0 and
σmin(X ∗t) ' √ n dt0 ' 1 along the lines of [38, Section 2.1]. Substituting σmax(X) and σmin(X∗t) into (20), we have
d1 & n
3 2
d0 . (21)
Therefore, the overall overparameterization degree becomes d0d1 ' Ω̃(n 3 2 ), which is sufficient for gradient descent to find a global minimum at a linear rate except with an arbitrary small probability. We note that an optimal linear scaling for the width d1 ' Õ(n) is sufficient when the number of input features is sufficiently large d0 ' Ω̃( √ n), which improves upon the results of [38] by a factor of Õ(n 1 2 ). Furthermore, unlike [38], we adopt standard initialization strategies in Theorem 3.
3The last assumption holds for popular activation functions such as sigmoid, tanh, and ELU, and can be relaxed if ω1 = 1 in (17).
4ψ can be arbitrary small.

heading: 6 Lazy training and experimental evaluation
text: Following the theoretically motivated initialization in Theorem 3, we set ω1ω2 ' 1√d0d1 . This gives rise to a broad family of initialization schemes as one varies the ratio ω2/ω1. Interestingly, we note that popular initialization schemes such as LeCun [26] and He initialization [17] belong to this family. The purpose of this section is to empirically investigate the impact the choice of this ratio has on generalization of shallow networks.
To this end, we will look at the generalization error of varying initializations in the more practical setting of stochastic gradient descent (SGD). Specifically, we fix the product of the weight initialization ω1ω2 and then proceed by varying
ω2. To ensure that perfect generalization is possible, we adopt the teacher-student setup, where, for the teacher network, we train a two-layer fully connected neural network, on MNIST [25] until SGD reaches zero training error. The student networks are trained for 300 epochs to ensure convergence. The results are shown in Figure 1. We use mean-square loss and a smooth activation function (GeLU [18]) for the student network to match the problem setup as closely as possible.
In Figure 1, we observe that while SGD achieves zero training error for every ω2, as suggested by Theorem 3 applicable in the full batch setting, the generalization ability increases as the ratio ω2/ω1 grows. It is also interesting to observe that the popular He initialization scheme corresponds to a rather balanced ratio that lies at the boundary of the well-performing region. In our experiments, we used He initialization to fix the value ω1ω2. This tendency suggests that a wide family of initialization schemes could generalize well as long as the ratio ω2/ω1 is not too small.
Comment on lazy training. It is important to address the so called lazy regime when generalization is of concern. Let
h̃(Θ) := h(Θ0) + 〈∇h(Θ0),Θ−Θ0〉
be the linearized function of h around Θ0 and let Θi and Θ̃i denote the iterates of gradient descent at time i. The lazy training regime refers to the case where the training trajectory stays close to this linearization, i.e. ‖h(Θi)− h̃(Θ̃i)‖ ' 0 for all i [7]. Such a linearization occurs in infinitely wide neural networks [7], which have been shown to generalize well in some settings [2, 27]. However, in our case of subquadratic (finite) width, the lazy regime might lead to poor generalization. To gain insight on when we cannot avoid it with certainty, let us make a simple rewriting of our network assuming φ is homogeneous:
Φ(Θ) = αV φ(WX)
with V0 ∼ N (0, 1) and W0 ∼ N (0, 1) where the standard deviations are pulled out as a scaling factor α = ω1ω2 ' 1/ √ d0d1. This seems to fit into an example in [7, Appendix A.2] suggesting lazy training as d1 →∞. However, their results require an odd activation function and infinite width, while our activation function is required not to be odd (see the proof in Appendix E) and our results are under subquadratic (finite) width. Instead, to study lazy training, we explicitly compute an upper bound on ‖h(Θi)− h̃(Θ̃i)‖ in Appendix F following [7, Theorem 2.3]. It turns out that the upper bound becomes∞when ω1 ω2, and it becomes zero when ω1 ω2. Our analysis suggests that shallow neural networks can avoid lazy training provided that ω2/ω1 →∞. This analysis is corroborated by the empirical results showing that the generalization capability improves as ω2 grows in Figure 1. On the other hand, if ω2/ω1 → 0, then lazy training is bound to happen asymptotically. For details, see Appendix F. Finally, we note that we have not theoretically claimed that our initialization is guaranteed to be non-lazy, since doing so would require establishing a lower bound on ‖h(Θi)− h̃(Θ̃i)‖, which is an interesting problem for future work. Instead, our discussion above only provides a necessary condition for non-lazy training, and a sufficient condition for lazy training.

heading: 7 Conclusions and future work
text: In this paper, we prove the linear convergence of first-order methods on subquadratically overparameterized two-layer neural networks with smooth activation functions. Our theoretical analysis is compatible with standard initialization strategies, which can potentially avoid lazy training. We train both layers simultaneously and achieve a desirable subquadratic scaling on the width of the network. In particular, we note that a linear scaling for the width d1 ' Õ(n) is sufficient when the number of input features is sufficiently large d0 ' Ω̃( √ n). We use tools from random matrix theory under standard assumptions on data and leverage on the assumption that the loss satisfies Polyak-Łojasiewicz condition. We carefully find an explicit upper bound and lower bound on singular values of the outputs of the first layer at initialization with high probability under general initialization.
It is natural to ask whether we can attain similar degree of overparameterization with nonsmooth activation functions such as ReLU. We plan to adapt our analysis for such extensions as a future work. While our analysis provides a necessary condition for avoiding lazy training, it is interesting to develop sufficient conditions in the future. In particular, developing lower bounds on ‖h(Θi)−h̃(Θ̃i)‖ will be a key to fully characterize lazy training.
Finally, as a theoretical work, we do not anticipate any potential negative societal impacts of our paper. However, the long-term impacts of our work may depend on how machine learning algorithms are used in society.
Acknowledgments and Disclosure of Funding
The authors would like to thank Fabian Latorre, Fanghui Liu, and Paul Rolland for helpful discussions.
This project has received funding from the European Research Council (ERC) under the European Union’s Horizon 2020 research and innovation programme (grant agreement n° 725594 - time-data). This project was sponsored by the Department of the Navy, Office of Naval Research (ONR) under a grant number N62909-17-1-2111. This work was supported by Hasler Foundation Program: Cyber Human Systems (project number 16066). Research was sponsored by the Army Research Office and was accomplished under Grant Number W911NF-19-1-0404.

Review summaries:
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.

This work presents an analysis of the amount of hidden units needed for convergence of a 1 hidden layer neural network to a global minimum under gradient descent (and gradient flow).  In particular, this work establishes a bound of $\tilde{\Omega}(n^{\frac{3}{2}})$ parameters for linear convergence.  Importantly, this work analyzes convergence in a 1 hidden layer network in the case where both weight matrices are trained.  The results also hold for a range of initialization schemes including the Kaiming Normal Initialization.