Model, sample, and epoch-wise descents: exact solution of gradient flow in the random feature model

1 Introduction . Deep learning models have vastly increased in terms of number of parameters in the architecture and data sample sizes with recent applications using unprecedented numbers with as much as 175 billions parameters trained over billions of tokens [ 1 ] . Such massive amounts of data and growing training budgets have spurred research seeking empirical power laws to scale model sizes appropriately with available resources [ 2 ] , and nowadays it is common wisdom among practitioners that `` larger models are better '' . This ongoing trend has been challenging classical statistical modeling where it is thought that increasing the number of parameters past an interpolation threshold ( at which the training error vanishes while the test error usually increases ) is doomed to over-fit the data [ 3 ] . We refer to [ 4 ] for a recent extensive discussion on this contradictory state of affairs . Progress towards rationalizing this situation came from a series of papers [ 5 , 6 , 7 , 8 , 9 , 10 , 11 ] evidencing the existence of phases where increasing the number of parameters beyond the interpolation threshold can actually achieve good generalization , and the characteristic U curve of the bias-variance tradeoff is followed by a `` descent '' of the generalization error . This phenomenon has been called the double descent and was analytically corroborated in linear models [ 12 , 13 , 14 , 15 , 16 ] as well as random feature ( RF ) ( or random feature regression ) shallow network models [ 17 , 18 , 19 , 20 ] . Many of these works provide rigorous proofs with precise asymptotic expressions of double descent curves . Further developments have brought forward rich phenomenology , for example , a triple-descent phenomenon [ 21 ] linked to the degree of non-linearity of the activation function . Further empirical evidence [ 22 ] has also shown that a similar effect occurs while training ( ResNet18s on CIFAR10 trained using Adam ) and has been 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . called epoch-wise double descent . Moreover the authors of [ 22 ] extensively test various CIFAR data sets , architectures ( CNNs , ResNets , Transformers ) and optimizers ( SGD , Adam ) and classify their observations into three types of double descents : ( i ) model-wise double descent when the number of network parameters is varied ; ( ii ) sample-wise double descent when the data set size is varied ; and ( iii ) epoch-wise double descent which occurs while training . We wish to note that sample-wise double descent was derived long ago in precursor work on single layer perceptron networks [ 23 , 24 ] . An important theoretical challenge is to unravel all these structures in a unified analytical way and understand how generalization error evolves in time . In this contribution we achieve a detailed analytical analysis of the gradient flow dynamics of the RF model ( or regression ) in the high-dimensional asymptotic limit . The model was initially introduced in [ 25 ] as an approximation of kernel machines ; more recently it has been recognized as an important playground for theoretical analysis of the model-wise double descent phenomenon , using tools from random matrix theory [ 17 , 18 , 26 ] . Following [ 17 ] we view the RF model as a 2-layer neural network with fixed-random-first-layer-weights and dynamical second layer learned weights . The data is given by n training pairs constituted of d-dimensional input vectors and output given by a linear function with additive gaussian noise . The data is fed through N neurons with a non-linear activation function and followed by one linear neuron whose weights we learn by gradient descent over a quadratic loss function . The high-dimensional asymptotic limit is defined as the regime n , d , N → +∞ while the ratios tend to finite values Nd → ψ and n d → φ . As the training loss is convex one expects that the least-squares predictor ( with Moore-Penrose inversion ) gives the long time behavior of gradient descent . This has led to the calculation of highly non-trivial analytical algebraic expressions for training and generalization errors which describe ( model-wise and sample-wise ) double and triple descent curves [ 17 , 21 ] . However , to the best of our knowledge , there is no complete analytical derivation of the whole time evolution of the two errors . We analyze the gradient flow equations in the high-dimensional regime and deduce the whole time evolution of the training and generalization errors . Numerical simulations show that the gradient flow is an excellent approximation of gradient descent in the high-dimensional regime as long as the step size is small enough ( see Fig . 2 ) . Main contributions presented in detail in Sect . 3 comprise : a . Results 3.1 and 3.2 give expressions of the time evolution of the errors in terms of one and two-dimensional integrals over spectral densities whose Stieltjes transforms are given by a closed set of purely algebraic equations . The expressions lend themselves to numerical computation as illustrated in Fig . 1 and more extensively in Sect . 3 and the supplementary material . b . Model and sample-wise double descents develop after some definite time at the interpolation threshold and are preceded by a dip or minimum before the spike develops . This indicates that early stopping is beneficial for some parameter regimes . A similar behavior also occurs for the triple descent . ( See Figs . 3 , 4 and the 3D version Fig . 1 ) . c. We observe two kinds of epoch-wise `` descent '' structures . The first is a double plateau monotonously descending structure at widely different time scales in the largely overparameterized regime ( see Fig . 3 ) . The second is an epoch-wise double descent similar to the one found in [ 22 ] . In fact , as in [ 22 ] , rather than a spike , this double descent appears to be an elongated bump over a wide time scale ( see Fig . 5 and the 3D version Fig . 1 ) . Let us say a few words about the techniques used in this work . We first translate the gradient flow equations for the learned weights of the second layer into a set of integro-differential equations for generating functions , as in [ 27 ] , involving the resolvent of a random matrix ( constructed out of the fixed first layer weights , the data , and the non-linear activation ) . The solution of the integrodifferential equations and the time evolution of the errors can then be expressed in terms of Cauchy complex integral representation which has the advantage to decouple the time dependence and static contributions involving traces of algebraic combinations of standard random matrices ( see [ 28 ] for related methods ) . This is the content of propositions 2.0.1 and 2.0.2 . With a natural concentration hypothesis in the high-dimensional regime , it remains to carry out averages over the static traces involving random matrices . This is resolved using traces of sub-blocks from the inverse of a larger nontrivial block-matrix , a so-called linear pencil . To the best of our knowledge linear pencils have been introduced in the machine learning community only recently in [ 29 ] . This theory is developed in the context of random matrix theory in [ 30 , 31 ] and [ 32 ] using operator valued free-probability . A side-contribution in the SM is also an independent ( but non-rigorous ) derivation of a basic result of this theory ( a fixed point equation ) using a replica symmetric calculation from statistical physics . The non-linearity of the activation function is addressed using the gaussian equivalence principle [ 33 , 34 , 29 ] . Finally , our analysis is not entirely mathematically controlled mainly due to the concentration hypothesis in Sect . 2.3 but comparison with simulations ( see Fig . 2 and SM ) confirms that the analytical results are exact . In the conclusion we briefly discuss possible extensions of the present analysis and open problems among which is the comparison with a dynamical mean-field theory approach . 2 Random feature model . 2.1 Model description . Generative model and neural network : We consider the problem of learning a linear function fd ( x ) = d − 12 βTx with x , β ∈ Rd column vectors . The vector x is interpreted as a random input and β as a random hidden vector ; both with distribution N ( 0 , Id ) , Id the d× d identity matrix . We assume having access to the hidden function through the noisy outputs y = fd ( x ) + with additive gaussian noise ∼ N ( 0 , s2 ) , s ∈ R+ . We suppose that we have n data-points ( xi , yi ) 1≤i≤n . This data can be represented as the n×dmatrixX ∈ Rn×d where xTi is the i-th row ofX , and the column vector vector Y ∈ Rn with i-th entry yi . Therefore , we have the matrix notation Y = d− 1 2Xβ + ξ where ξ ∼ N ( 0 , s2In ) and In the n× n identity matrix . We learn the data with a shallow 2-layer neural network . There are N hidden neurons with weight ( column ) vectors θi ∈ Rd , i = 1 , · · · , N each connected to the d input neurons . Out of these we form the matrix ( of the first layer connecting input and hidden neurons ) Θ ∈ RN×d where θTi is the i-th row of Θ . Its entries are assumed independent and sampled through a standard gaussian distribution N ( 0 , 1 ) ; they are not learned but fixed once for all . The data-points in X are applied linearly to the parameters Θ , and the output Z ∈ Rn×N of the first layer is the pointwise application of an activation function σ : R→ R , Z = σ ( d− 12XΘT ) . We use the notation zTi to express the i-th row of Z . The second layer consists in a weight ( column ) vector at ∈ RN to be learned , indexed by time t ≥ 0 , with components initially sampled at t = 0 i.i.dN ( 0 , r2 ) , r ∈ R+ . The prediction vector is expressed as Ŷt = N− 1 2Zat . We assume that the activation function belongs to L2 ( e− x2 2 dx ) with inner product denoted 〈 , 〉 . It can be expanded on the basis of Hermite polynomials , so σ ∈ Span ( ( Hek ) k≥0 ) , where Hek ( x ) = ( −1 ) ke x 2 2 dk dxk e− x2 2 ( so He0 ( x ) = 1 , He1 ( x ) = x , He2 ( x ) = x 2 − 1 , He3 ( x ) = x3 − 3x , ... ) . Furthermore we take σ centered with 〈σ , He0〉 = 0 , and set µ = 〈σ , He1〉 , ν2 = 〈σ , σ〉 − µ2 . For instance , σ = id has ( µ , ν ) = ( 1 , 0 ) while σ = Relu − 1√ 2π has ( µ , ν ) = ( 12 , 1 2 ( 1 − 2 π ) 1/2 ) ' ( 0.5 , 0.3 ) . Finally , we recall that we are interested in the high dimensional regime where the parameters N , n , d tend to infinity with the ratios Nd → ψ and n d → φ . Training and test errors : For a new input x0 ∈ Rd , we define the predictor ŷt ( x0 ) = 1√ N z ( x0 ) Tat where z ( x0 ) = σ ( 1√dΘx0 ) . We will further define the standard training and test errors with a penalization term λ > 0 and the quadratic loss : Htrain ( a ) = 1 n ∥∥∥Y − Ŷ ∥∥∥2 + λ N ‖a‖2 , Htest ( a ) = Ex0∼N ( 0,1 ) [ ( y ( x0 ) − ŷ ( x0 ) ) 2 ] ( 1 ) Note that because of the λ-penalization term , in this context , the training error can be above the test error in some configurations of parameters . Also , we will slightly abuse this notation throughout the paper by usingHtraint , Htestt to designateHtrain ( at ) , Htest ( at ) . Gradient flow : Minimizing the training error of this shallow-network is equivalent to a standard Tikhonov regularization problem with a design matrix Z for which the optimal weights are given by a∞ = ( ZTZ N + n N λIN ) −1 ZT√ N Y . The errors generated by the predictors with weights a∞ have been analytically calculated in the high-dimensional regime in [ 17 ] and further analyzed in [ 21 ] . Here we study the whole time evolution of the gradient flow and thus introduce an additional time dimension in our model . Of course as t→ +∞ one recovers the errors generated by the predictors with weights a∞ . The output vector at is updated through the ordinary differential equation datdt = −η∇aH train ( at ) with a fixed learning rate parameter η > 0 . As η can be absorbed in the time parameter , from now on we consider without loss of generality that η = n2 . Setting δ = λ n N , we find that the gradient flow for at is a first order linear matrix differential equation , dat dt = − ( ZTZ N + δIN ) at + ZTY√ N . ( 2 ) Recall the initial condition a0 is a vector with i.i.d N ( 0 , r2 ) components .

Summary: This article provides, modulo some important technical assumptions, computationally tractable formulas for the test and train error of a random features model trained by gradient flow on \ell_2-regularized MSE during the entire course of training. The features are those computed by the first layer of a randomly initialized neural network and the regime of interest for this article is a triple scaling limit in which the input dimension, dataset size, and hidden layer width tend to infinity but their ratios remain order 1. Unlike prior work this article provides formulas for the evolution in time of both the training and test error, instead only of only their behaviors are t = \infty. While the results are not completely rigorous (they assume concentration for spectral statistics of certain random matrices), I found them intriguing and am fairly convinced by plots in which the predictions fit empirics quite well. The main weakness for me is a lack of further formulation of takeaways from the theoretical results (see weakness #2). Overall, I thought this was a solid paper.