Analysis of one-hidden-layer neural networks via the resolvent method

1 Introduction . Machine learning has seen many successful achievements in recent years . Applications in face identification , object and speech recognition , translation , email spam filtering , navigation , medical diagnosis , etc . have proved the enormous potential of machine learning for day-to-day live [ 16 , 11 ] . Deep neural networks have turned out to be a particularly powerful machine learning method , and understanding the theoretical underpinning of their success has received tremendous attention in mathematics , physics and computer science . A fully-connected , feed-forward neural network with L hidden layers of dimensions n1 , . . . , nL can be modelled as follows : fθ ( x ) = β ∗f ( W ( L ) f ( W ( L−1 ) f ( . . . f ( W ( 1 ) x ) . . . ) ) ) ∈ Rd , where x ∈ Rn0 denotes the input data vector and f : R → R is a non-linear activation function which is applied entry-wise . We denote the parameters of the network by θ : = ( W ( 1 ) , . . . , W ( L ) , β ) , where W ( l ) ∈ Rnl×nl−1 for 1 ≤ l ≤ L and β ∈ RnL×d are the matrices of the weights . In the classical setting of supervised learning , we are given a training set of ( say ) m samples of input feature vectors xi ∈ Rn0 with associated target vectors zi ∈ Rd . For example , xi may encode the pixels of a photograph of an animal and the target zi may label the species of the 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . animal in the image . Roughly speaking , the goal of supervised learning is to learn the mapping between the feature and the target vectors based on a given training set in order to predict the output of new unlabelled data . Let X = ( x1 . . . xm ) ∈ Rn0×m be the matrix of the data and let Z = ( z1 . . . zm ) ∈ Rd×m be the target matrix . Then , the aim of the network is to find optimal parameters θ such that fθ ( X ) = ( fθ ( x1 ) , . . . , fθ ( xm ) ) ∈ Rd×m approximates the target Z optimally . During the training phase , weights are adjusted in order to minimize the empirical risk R ( θ ) = EL ( fθ ( X ) , Z ) , where L ( · , · ) is a given loss function , usually involving some penalty for large weights θ in order to avoid over-fitting . Stochastic gradient descent ( SGD ) and its variants with back-propagation are the most commonly used algorithms for training multilayer networks by iteratively updating the parameters into the direction of the negative of the gradient of the empirical risk . For a much more complete survey , we refer the reader to [ 11 ] . In the present paper , we will focus on a single-hidden-layer neural network of the form fθ ( X ) = β∗Y with Y = f ( WX ) . This model was first studied by Louart , Liao , and Couillet [ 19 ] for the case where the data matrix X is deterministic and W is a matrix of random weights ( in particular , the weights are given by functions of standard Gaussian random variables ) , and by Pennington and Worah [ 22 ] for the case where X and W are independent random matrices with both centred Gaussian entries . In both papers , the matrix β ∈ Rn1×d is the only parameter to be learned and is chosen as the unique minimizer of the ridge-regularized least squares loss function L ( fθ ( X ) , Z ) = 1 2dm ∥Z − β∗Y ∥2F + γ∥β∥2F , where γ > 0 is the learning rate . The unique minimizing weight matrix β̂ is then equal to β̂ = Y G ( −γ ) Z∗ , where G ( −γ ) = ( 1 m Y ∗Y + γ ) −1 is the resolvent of 1mY ∗Y . As proved in [ 19 , 22 ] , the expected training loss Etrain is related to −γ ∂∂γG ( −γ ) , and thus also to the Stieltjes transform of the limiting spectral measure of 1 mY ∗Y . Here the Stieltjes transformmµ of a probability measure µ on R is defined asmµ ( z ) : = ∫ R ( x−z ) −1 dµ ( x ) for z ∈ C such that ℑz ≥ 0 , and for µn1 being the empirical probability measure of the n1 eigenvalues of 1mY ∗Y is related to the resolvent via mµn1 ( z ) = 1 n1 TrG ( z ) . The performance of one-hiddenlayer neural networks depends on the asymptotic spectral properties of the matrix 1mY ∗Y . Pennington and Worah [ 22 ] investigated the limiting spectral measure of the random matrix M = 1mY Y ∗ and derived the quartic self-consistent equation 1 + zg∞ = θ1 ( f ) g∞ ( 1− ϕ ψ ( 1 + zg∞ ) ) − θ2 ( f ) ψ g∞ ( 1 + zg∞ ) ( 1− ϕ ψ ( 1 + zg∞ ) ) + θ2 ( f ) ( θ1 ( f ) − θ2 ( f ) ) ψ g2∞ ( 1− ϕ ψ ( 1 + zg∞ ) ) 2 , ( 1 ) where g∞ ( z ) : = limn1→∞ g ( z ) and g ( z ) : = 1 n1 TrG ( z ) is the Stieltjes transform , which is approximately satisfied , g∞ ( z ) ≈ g ( z ) , by g ( z ) in case of Gaussian W , X . It is notable that the asymptotic spectrum of Y ∗Y for large dimensions such that n0/m → ϕ ∈ ( 0 , ∞ ) and n0/n1 → ψ ∈ ( 0 , ∞ ) depends on the non-linear function f only through two integral parameters θ1 ( f ) and θ2 ( f ) , where θ1 ( f ) is the Gaussian mean of f2 and θ2 ( f ) is the square of the Gaussian mean of f ′ ( c.f . ( 5 ) ) . Benigni and Péché [ 6 ] then extended this model to random matrices W and X with general i.i.d . centred entries , and obtained the same self-consistent equation ( 1 ) . We mention that ( 1 ) may be reduced , for some special cases , to the quadratic equation that is satisfied by the Stieltjes transform mµMP of the Marchenko-Pastur distribution µMP [ 20 ] . This means that for some activation functions , the non-linear random matrix model has the same limiting spectral distribution as that of sample covariance matrices XX∗ ( in other cases , the equation can simplify to the cubic equation approximately satisfied by product Wishart matrices [ 7 , 10 ] ) . This can be generalised to multilayer networks : [ 22 ] found experimentally that the singular value distribution is preserved through multiple layers by activation functions with θ2 ( f ) = 0 and is given by the Marchenko-Pastur distribution in each layer . This conjecture was proved in [ 6 ] for the general case of bounded activation functions . Moreover , [ 19 ] performed a spectral analysis on the Gram matrix model with general training data and proved that , in the large dimensional regime , the resolvent of Y ∗Y has a similar behaviour as that observed in sample covariance matrix models . This was extended in [ 17 ] by considering Gaussian mixture of data . We also refer to the recent paper [ 18 ] . In the context of multilayer feedforward neural networks , Fan and Wang [ 9 ] analysed the eigenvalue distribution of the Gram matrix model , where the weights are at random and the input vectors are assumed to be approximately pairwise orthogonal . In particular , they showed that the limiting spectral distribution converges to a deterministic limit and , at each intermediate layer , this limit corresponds to the Marchenko-Pastur map of a linear transformation of that of the previous layer . In recent years , there has been some progress in the asymptotic analysis of the eigenvalue distribution of another Gram matrix , the so-called Neural Tangent Kernel ( NTK ) . Consider a multilayer neural network and denote by J = ∇θfθ ( X ) the Jacobian matrix of the network outputs with respect to the weights θ . Then , the NTK is the Gram matrix of J , defined by KNTK = J∗J . It was shown in [ 14 ] that the NTK at random initialization converges , in the infinite-width limit , to a deterministic kernel and it remains constant during the whole training time of the network . Subsequently , [ 23 ] analysed the spectrum of the sample covariance matrix JJ∗ in a single-hidden-layer neural network , and provided an exact asymptotic characterization of the spectral distribution of JJ∗ with random Gaussian weights and data . Recently , [ 9 ] proved that the limiting spectral measure of the NTK converges to a deterministic measure , which may be described by recursive fixed-point equations that extend the Marchenko-Pastur distribution . The present paper is structured as follows . In the first part we consider the non-linear random matrix model studied in [ 6 ] and we compute its asymptotic spectral density . We follow the resolvent method via the cumulant expansion which , together with the moment method , is a standard approach to obtain the asymptotics for the empirical spectral density . In particular , we compute the self-consistent equation that is approximately satisfied by the Stieltjes transform of the limiting spectral distribution . This is a quartic equation and is the same as that found in [ 6 ] . In [ 6 , 22 ] the authors relied on the method of moments : they approximated general non-linear functions by polynomial ones and then computed the asymptotics of high moments ETr ( Y k ) with Y = f ( WX ) to obtain the limiting measure via its moments . Conversely , we approach matters in a more robust and less combinatorial fashion by applying the resolvent method : we consider Y as a random matrix with correlated entries and then we directly derive a self-consistent equation for its resolvent . In particular , we prove that the random matrix Y has cycle correlations , in the sense that the joint cumulant does not vanish when the random variables Yij ’ s are joined by a cycle graph . We find that the variance of Yij is given by the parameter θ1 ( f ) , whereas for k > 1 the cumulants κ ( Yi1i2 , Y ∗ i2i3 , Yi3i4 , . . . , Y ∗ i2ki1 ) are powers of θ2 ( f ) . We note that in the random matrix literature matrices with general decaying correlations have been studied previously , see e.g . [ 3 , 8 , 1 ] . However , the cycle correlations of Y considered in the present paper are much stronger compared to these previous results . The second part of this paper concerns the additive bias case which is a more realistic model for machine learning applications . More precisely , we consider the random feature matrix Y = f ( WX +B ) , where B is a rectangular rank-one Gaussian random matrix , and derive a characterization of the Stieltjes transform of the limiting spectral density . We chose B to be rank-one since for the most commonly used neural network architectures the added bias is chosen equal for each sample . [ 2 ] studied the bias case for deterministic data and i.i.d . Gaussian random weights , and computed the exact training error of a ridge-regularized noisy autoenconder in the high-dimensional regime . Interestingly we find that in the case of additive bias it is impossible to choose an activation function f such that the eigenvalue distribution is preserved throughout multiple layers , unlike in the bias-free case where θ2 ( f ) = 0 yields the Marchenko-Pastur distribution in each layer . Finally , we remark that in the bias-free case our proof via the resolvent method has no significant advantage compared to the moment method , beyond requiring less combinatorics . The main advantage of the resolvent approach is that it allows to include an additive bias without much additional effort .

Summary: The authors study the asymptotic eigenvalue distribution of the post-activation Gram matrix in a one-hidden-layer neural network, in a regime where all matrix dimensions are growing proportionally. For input data with i.i.d. centered entries and a weight matrix with i.i.d. centered weights, the authors provide a new proof (under weaker distributional assumptions) of a result first due to Pennington and Worah that the asymptotic eigenvalue distribution is characterized by a quartic equation in its Stieltjes transform. The authors then provide an extension of this result to a model with i.i.d. Gaussian biases in its hidden layer. A qualitative difference in the setting with biases is that---in contrast to the bias-free setting, where special choices of activation function can allow the eigenvalue distribution to remain the same as that of the input Gram matrix, and hence remain constant across multiple layers---the authors show that the eigenvalue distribution with biases must necessarily evolve over layers (cf. Remark 2.6 and the discussion on page 7).  The proof is based on deriving the fixed-point equation for the expected Stieltjes transform using the cumulant-expansion generalization of Gaussian integration-by-parts, and explicitly analyzing the asymptotic mixed cumulants of the entries of the post-activation matrix. A technical challenge is to translate the more easily-obtained understanding of mixed cumulants of the pre-activation entries to the post-activations, and the authors accomplish this using a Fourier-analytic technique.