Generalization of Two-layer Neural Networks: An Asymptotic Viewpoint

1 INTRODUCTION . In modern neural networks , the number of parameters can easily exceed the number of training samples , yet in many circumstances , there is little sign of overfitting even in the absence of explicit regularization ( Zhang et al. , 2016 ) . This phenomenon is usually explained by the interplay between the model architecture and the optimization method . Existing works have analyzed the implicit regularization of gradient descent on simple models ( Gunasekar et al. , 2018 ; Ji and Telgarsky , 2018 ) , and provided generalization guarantees ( Arora et al. , 2018 ; Bartlett et al. , 2017 ; Dziugaite and Roy , 2017 ) that align with the empirical observations . Recently , a series of works highlighted the implicit regularization of interpolators in the overparameterized regime ( Belkin et al. , 2018 ; Spigler et al. , 2018 ; Geiger et al. , 2018 ; Advani and Saxe , 2017 ) . Specifically , a second decrease in the population risk is observed when the model is further overparameterized beyond the interpolation limit , i.e . when the model achieves zero training error . This phenomenon is known as double descent , and can be precisely quantified for certain linear models ( Hastie et al. , 2019 ; Mei and Montanari , 2019 ; Belkin et al. , 2019 ; Bartlett et al. , 2019 ; Xu and Hsu , 2019 ) . Among the recent works , Hastie et al . ( 2019 ) and Mei and Montanari ( 2019 ) explicitly derived the population risk of linear regression and random features regression models in high dimensions using tools from random matrix theory . However , there is still a gap between the practical benefit of overparameterization and the recently proved double descent phenomenon , which is typically established under models that exhibits the following structure : the trained model solves a linear inverse problem , and the “ cusp ” in the risk arises from the instability of the inverse at the interpolation threshold . Moreover , given a dataset or fixed n , d , the number of parameters in the linear regression model is also fixed , i.e . the level of overparameterization can not be altered . It is therefore unclear if the trend persists in the optimization of more complex models , for instance in two-layer neural networks where overparameterization can be controlled simply by adding more neurons . In this work , we analyze the generalization properties of two-layer neural networks in the unregularized least squares regression setting and examine the presence/absence of the double descent phenomenon . We consider the proportional asymptotic limit where the number of samples n , input features d , and neurons h tend to infinity at the same rate , under which overparameterization corresponds to increasing the limit of h/n ( network “ width ” ) . This regime is particularly interesting because even though n→∞ , the empirical risk is not equivalent to the population risk . In addition , the joint scaling of n , d , h is parallel to the practical choice of model architectures , where it is common to train a larger network when the number of samples and input features are larger . Following Hastie et al . ( 2019 ) , we assume unit Gaussian input and noisy linear observations , and analytically derive the population risk of the solution of gradient flow on either the first or the second layer parameters when the flow is initialized close to zero . 0.5 1.0 1.5 2.0 2.5 2 = h/n 0.5 1.0 1.5 2.0 2.5 3.0 ris k empirical prediction 1 2 3 0.0 0.5 Our findings can be summarized as follows ( see Figure 1 ) : 1.1 RELATED WORKS . Global Convergence of Two-layer Networks . A plethora of recent works have explored the global convergence of shallow neural networks . Mei et al . ( 2018 ; 2019 ) ; Chizat and Bach ( 2018a ) ; Rotskoff and Vanden-Eijnden ( 2018 ) ; Sirignano and Spiliopoulos ( 2018 ) ; Nitanda and Suzuki ( 2017 ) studied the mean-field limit where the number of neurons h→∞ and the second layer scaled by 1/h , and established correspondence between the main-particle limit of gradient descent and Wasserstein gradient flow to demonsrate global convergence . On the other hand , Jacot et al . ( 2018 ) ; Du et al . ( 2018 ) ; Oymak and Soltanolkotabi ( 2019 ) ; Allen-Zhu et al . ( 2018b ) ; Song and Yang ( 2019 ) considered a different scaling and showed that gradient descent on overparameterized models converges to global minimizer at a linear rate ; key to these results is an observation that optimization via gradient descent is asymptotically equivalent to kernel regression with respect to the neural tangent kernel . Active vs. Lazy Training . Following Chizat and Bach ( 2018b ) , we refer to the two aforementioned scalings as the active and lazy ( kernel ) regime . It has been observed that different regimes lead to contrasting inductive biases . Williams et al . ( 2019 ) ; Woodworth et al . ( 2019 ) ; Li et al . ( 2017 ) showed that for certain two-layer network or overparameterized linear model , the scale of initialization controls the implicit regularization of gradient descent ( from sparse to smooth solution ) . In the student-teacher setup ( Tian , 2017 ; Zhong et al. , 2017 ) , Ghorbani et al . ( 2019b ; a ) showed that kernel models in high dimensions perform no better than low-degree polynomials on the input or fullytrained two-layer network . Additionally , Suzuki ( 2018 ) ; Allen-Zhu and Li ( 2019 ) ; Yehudai and Shamir ( 2019 ) ; Wei et al . ( 2018 ) demonstrated that neural network outperforms linear estimators ( including kernel method ) in learning various target functions . The difference between fixed bases and adaptive bases mirrors the difference in optimizing the first or second layer in our setup . Generalization of Overparameterized Models . It is often observed that overparameterization does not result in overfitting ( Neyshabur et al. , 2014 ) . In the lazy regime , generalization guarantees can be derived from the distance traveled by the parameters ( Neyshabur et al. , 2018 ; Nagarajan and Kolter , 2019 ) , which becomes small if the model is sufficiently overparameterized ( Arora et al. , 2019b ; Li and Liang , 2018 ; Allen-Zhu et al. , 2018a ; Cao and Gu , 2019 ) . Compared to these guarantees that usually require significant overparameterization , our result relies on stronger data assumptions , but consequently we obtain the exact population risk instead of a vacuous upper-bound . Beyond the kernel regime , Advani and Saxe ( 2017 ) ; Goldt et al . ( 2019 ) analyzed the generalization dynamics of overparameterized models in the student-teacher setup . Double Descent . The term double descent refers to the phenomenon that the population risk of an empirical risk minimizer manifests a `` cusp '' at the interpolation threshold , and further overparameterization decreases the risk . First observed in Krogh and Hertz ( 1992 ) , the phenomenon has been recently connected to the benefit of overparameterization ( Belkin et al. , 2018 ; Geiger et al. , 2018 ; Spigler et al. , 2018 ; Advani and Saxe , 2017 ) , and can be precisely characterized for certain simple models ( Hastie et al. , 2019 ; Belkin et al. , 2019 ; Bartlett et al. , 2019 ; Xu and Hsu , 2019 ) . Our work is inspired by Hastie et al . ( 2019 ) which uses random matrix theory to derive the asymptotic risk for linear and random feature models . Concurrent to our work , Mei and Montanari ( 2019 ) analyzed the random features model and derived its population risk for which double descent occurs both in bias and variance . This aligns with our results on optimizing the second layer in Section 4 although we do not derive the bias component explicitly . Compared to Hastie et al . ( 2019 ) ; Mei and Montanari ( 2019 ) , the focus of this work is to highlight the different generalization property of models obtained from optimizing different layers of the network and from different initialization . Random Matrix Theory . High-dimensional models , including kernel models and neural networks , can be analyzed by studying the properties of random matrices . El Karoui et al . ( 2010 ) ; Cheng and Singer ( 2013 ) ; Fan and Montanari ( 2019 ) studied the spectral properties of kernel matrix via decomposing the nonlinearity with Taylor series or Hermite polynomials , which in turn explains the generalization of high-dimensional kernel ridgeless interpolators ( Liang and Rakhlin , 2018 ) . In addition , similar tools have been used to study two-layer neural networks ( Louart et al. , 2018 ; Pennington and Worah , 2017 ) and related quantities such as the Fisher information matrix ( Karakida et al. , 2018 ; Pennington and Worah , 2018 ) . 2 PRELIMINARIES : TWO-LAYER NEURAL NETWORK . Consider the following bias-free two-layer neural network f : Rd → R with h hidden units f ( x ) = h∑ i=1 aiφ ( 〈x , wi〉 ) , ( 1 ) where x ∈ Rd is the input , wi ∈ Rd is the weights corresponding to neuron i , ai ∈ R is the i-th coefficient of the second layer , and φ : R → R is a Lipschitz continuous activation function with bounded Gaussian moments , i.e . E [ φ ( G ) k ] < ∞ , ∀k ∈ Z+ for G ∼ N ( 0 , 1 ) . For concise notation , we write W = [ w1 , ... wh ] ∈ Rd×h for the weight matrix , a = [ a1 , ... ah ] ∈ Rh for the coefficient vector , X = [ x1 , ... xn ] ∈ Rd×n for the data matrix , y ∈ Rn for the corresponding vector of labels , and Φ = φ ( W > X ) ∈ Rh×n for the feature matrix at the first layer . We omit arguments of f when they are clear from the context . We consider a student-teacher setup , in which data is generated by a teacher model F : Rd → R with additive noise , and the student model aims to minimize the squared loss : ( xi , εi ) i.i.d.∼ Px × Pε , yi = F ( xi ) + εi , L ( X ; f ) = 1 2n n∑ i=1 ( yi − f ( xi ) ) 2 , ( 2 ) where E [ xi ] = 0 , Cov ( xi ) = Σ , E [ εi ] = 0 , Var ( εi ) = σ2 . We are interested in the population risk R ( f ) = EPx [ ( F ( x ) − f ( x ) ) 2 ] . Our analysis will be made under the proportional asymptotics : n , d , h→∞ ; d/n→ γ1 , h/n→ γ2 ; γ1 , γ2 ∈ ( 0 , ∞ ) , in which overparameterization corresponds to increasing γ2 . Thus the characteristics of double descent considered in this work are : 1 ) large population risk as γ2 → 1 ; 2 ) decrease in the risk for γ2 > 1 . While the empirical risk can be minimized in various ways , we analyze the solution of gradient flow , in which we update either the first layer W or the second layer a : dW ( t ) = −∇WL ( X ; f ) dt or da ( t ) = −∇aL ( X ; f ) dt , ( 3 ) from small initialization . The rest of the paper is organized as follows . In Section 3 , we start with a simple example of two-layer linear network as warm-up . In Section 4 , we consider optimizing the second layer coefficients ( flow over a ) of a non-linear two-layer neural network under fixed Gaussian first layer , which is a random feature model . Section 5 considers optimizing the first layer weights ( flow over W ) of such network under fixed Rademacher second layer . We defer all proofs and details on experiments to appendix .

Summary: This paper provides exact bounds on the risk when training a two-layer neural network in an asymptotic regime. Namely, the paper considers training under the square-loss objective, a two-layer neural network with $h$ hidden units on inputs of dimension $d$ and training on $n$ samples. The asymptotic regime is considered by making all of $d$, $h$, $n$ go to $\infty$, in a way that the ratio $d/n$ approaches $\gamma_1$ and the ratio $h/n$ approaches $\gamma_2$.