Explicit loss asymptotics in the gradient descent training of neural networks

1 Introduction . A major challenge in the research of neural networks is the quantitative theoretical description of their optimization by gradient descent . At present , many aspects of network training seem to be understood rather well on a qualitative level , or admit convincing heuristic explanations , but we seem to lack tools for making reasonably accurate quantitative predictions , even for relatively simple models and data . In this sense , the theory of neural networks compares unfavorably to physics , which is also an application-driven field but with an apparently much more successful penetration of theoretical methods . The main difficulty here is probably the complex structure of the data and models , which are hard to describe in terms of convenient and simple mathematical abstractions . In recent years , a significant progress in the theoretical analysis of gradient descent of neural networks has been associated with the limit of large networks , which can be studied using various methods from partial differential equations [ 26 , 31 ] , kernel methods [ 20 , 21 ] , spin glass theory [ 12 ] , random matrix theory [ 29 ] , dynamical systems [ 30 ] , and other mathematical fields . In the present work , we consider a setting of large networks and large , smoothly distributed data sets that allows us to obtain explicit leading terms in the long-term evolution of the loss under gradient descent . We are inspired by the spectral theory of singular integral operators [ 5 ] , which we apply to the linearized evolution of the network . While this linearized evolution has been widely studied recently , most related research seems to focus on theoretical convergence guarantees and upper bounds for the loss values [ 3 , 27 ] , or on a highly symmetric problems admitting explicit solution [ 35 , 38 ] . In contrast , we focus on explicit loss evolution formulas , which we find as power laws L ( t ) ∼ Ct−ξ . ( 1 ) We argue that the exponents ξ here exhibit some form of universality , in that they are essentially determined by the input dimension d and by the smoothness classes of the activation function and the target function . In particular , we find that in the case of ReLU networks approximating an indicator 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . function of some region in the d-dimensional space ( a classification problem target ) , the natural value of the exponent is ξ = 1d+1 . On the other hand , in the case of target functions generated by a randomly initialized wide ReLU network , the exponent is ξ = 3d+1 . Our approach also allows us to obtain explicit expressions for the coefficient C in these cases . The power law ( 1 ) is established using similar power laws ( but with different exponents and coefficients ) for the eigenvalues of the evolution operator and for the coefficients in the expansion of the target function over corresponding eigenvectors . These power laws are indeed confirmed by our experiments ( see Figure 1 ) . Our main scenario is approximation by shallow ReLU network in the NTK regime , but we also briefly consider several modifications of this scenario , namely the activation functions ( x+ ) q with q > 0 , approximation by a deep network in the NTK regime , and approximation in the mean field regime . 2 Related work . The approximation of linearized network evolution and its applications were studied in many works , see in particular [ 21 , 10 , 23 , 22 , 19 ] . The role of the eigenvalues and eigenvectors of the NTK in the linearized network was emphasized in [ 3 ] , where the GD dynamic of the finite network was linked to the dynamic of its infinite width counterpart , determined by spectral properties of the corresponding NTK . In subsequent works , the NTK spectrum was central for description of network training [ 27 ] and generalization [ 7 , 4 ] . These papers use the assumption of power law NTK spectrum , but justify it empirically or for highly symmetric problems . Because of the importance of the NTK spectrum , a number of works focused on its description in different settings . We first mention directions and settings that are different from ours . The case of very deep networks was studied in [ 37 , 18 , 17 , 2 ] . This analysis is relying on convergence of the NTK to simple fixed points in the limit of infinite depth . [ 17 ] also studied the choice of activation function , in particular its smoothness . Another line of research [ 1 , 13 ] uses techniques from Random Matrix Theory to analyze the setting where the dataset size M goes to infinity together with data dimension d and layer widths nl . In this work we consider the setting of fixed data dimension and effectively infinite network width and dataset size . In this case [ 35 ] showed that the network evolution can be described by a deterministic integral operator that is easier to analyze than a large but finite matrix . Also , in that paper and in the papers [ 38 , 8 , 27 ] , the integral operator was explicitly diagonalized in the special case of uniform distribution on a sphere . A power law upper bound for the training loss was obtained in [ 27 ] , but with an exponent ξ smaller than ours . The paper [ 4 ] describes a different power law , relating the test loss at the end of training to the dataset size and the network width . This result also relies on a power law assumption for the NTK spectrum . Another related line of research is the case of univariate functions . The gradient descent evolution of univariate shallow networks has been analytically studied in the papers [ 39 , 36 ] . 3 Asymptotic evolution of the loss function . We consider a linearized training of a neural network by gradient descent . Such linear approximations arise naturally in various “ lazy training ” scenarios [ 10 ] . Consider the standard quadratic loss function L ( W ) = 12 ∫ Rd |f̃ ( W , x ) −f ( x ) | 2µ ( x ) dx , where x is the d-dimensional input , f is the approximated function , f̃ is the network , W are the network weights , and µ is the data distribution on which the network is trained . Gradient descent can be written as the differential equation ddtW = −∇WL ( W ) . We assume that the weight vector W is close to a global minimum W∗ where L ( W∗ ) = 0 and f̃ ( W∗ , x ) = f ( x ) µ-a.e. , and that the evolution equation can be linearized at W = W∗ . It is convenient to write this linearized equation in terms of the difference δf ( x ) = f̃ ( W , x ) − f ( x ) between the current output and the target . The corresponding linear equation is ddtδf = −Aδf , where A can be written as the integral operator Aδf ( x ) = ∫ Rd Θ ( x , x ′ ) µ ( x′ ) δf ( x′ ) dx′ with the NTK ( neural tangent kernel ) Θ ( x , x′ ) = ∇Wf̃ ( W∗ , x ) T∇Wf̃ ( W∗ , x′ ) ( 2 ) The evolution operator A is a symmetric non-negative definite operator with respect to the scalar product 〈f1 , f2〉µ = ∫ f1 ( x ) f2 ( x ) µ ( x ) dx . By multiplying functions f by µ1/2 , the operator A can be brought to the form Ã = µ1/2Aµ−1/2 with a symmetric kernel , Ãδf ( x ) = ∫ Rd µ1/2 ( x ) Θ ( x , x′ ) µ1/2 ( x′ ) δf ( x′ ) dx′ . ( 3 ) In this form Ã is symmetric w.r.t . the usual scalar product 〈f1 , f2〉 = ∫ f1 ( x ) f2 ( x ) dx . Observe that the loss at time t can be written as L ( t ) = 12‖gt‖ 2 , where gt = e−tÃg , the norm ‖ · ‖ corresponds to the scalar product 〈· , ·〉 , and the function g is given by g ( x ) = µ1/2 ( x ) ( f̃ ( w ( t = 0 ) , x ) − f ( x ) ) . ( 4 ) We can now describe the evolution of the loss by diagonalizing the operator Ã . We will be interested in the scenario where µ is a smooth function that is compactly supported or falls off at infinity sufficiently fast . ( In particular , in the context of a finite training set this means that this set is large enough to be legitimately approximated by µ . ) In this case , for typical kernels Θ , the operator Ã will have a discrete spectrum with eigenvalues converging to 0 . Let λn denote the eigenvalues of Ã sorted in decreasing order , and let cn be the respective coefficients in the expansion of the initial error g ( given by Eq . ( 4 ) ) over the normalized orthogonal eigenvectors . Then , the loss evolves by L ( t ) = 1 2 ∞∑ n=0 e−2λnt|cn|2 . ( 5 ) To compute the asymptotic of L ( t ) at large times t , we need to know the distribution of the eigenvalues λn and the coefficients cn at large n. The key assumption of our work ( verified later for certain scenarios ) is that these distributions have power law forms . Specifically , regarding the eigenvalues λn we assume that λn ∼ Λn−ν ( 6 ) with some coefficient Λ and exponent ν . Regarding the coefficients cn , we assume that they also have a power law distribution on a large length scale in n , but possibly deviate from this law locally ( e.g. , due to oscillations ) . For this reason , we describe their large n behavior by the partial sums sn ≡ ∑ k≥n |ck|2 . ( 7 ) We then assume that , for some coefficient K and exponent κ , sn ∼ Kn−κ . ( 8 ) Under assumption of the power laws ( 6 ) and ( 8 ) , it is easy to check ( see SM , Section B ) that the loss also has a power law asymptotic ( 1 ) with the constant C and exponent ξ expressible through the constants Λ , K and exponents ν , κ : L ( t ) ∼ K2 Γ ( κ ν + 1 ) ( 2Λt ) −κν , ( 9 ) where Γ ( z ) is the Gamma function . In Figure 1 we illustrate this approach to the long-term loss evolution with several examples of target functions having different smoothness and dimension and , as a result , exhibiting different exponents . In the remainder of the paper we show that the power laws ( 6 ) and ( 8 ) , and hence the large-t asymptotic ( 9 ) of the loss , are indeed valid for some natural network training scenarios . The asymptotic ( 6 ) of the eigenvalues is primarily determined by the singularities of the kernel Θ . These singularities can be explicitly described for shallow neural network with piecewise smooth activations such as ReLU . Then , the power law ( 6 ) can be derived from general results on integral operators with singular kernels . In particular , in the case of ReLU we find that ν = 1 + 1d . The asymptotic ( 8 ) of the coefficients is more subtle , as it depends significantly on the class of the initial error function g ( which in turn depends , by Eq . ( 4 ) , on the target function f and the initial approximation f̃ ) . We derive this power law for one natural class of discontinuous functions g , and for functions g that are realizations of a Gaussian process of a particular “ roughness ” . Moreover , for these two classes , we also find an explicit form of the coefficient C appearing in loss asymptotic ( 1 ) . This requires us , however , to modify the above derivation of the loss asymptotic ( 9 ) by what can be called “ integrated localization ” . Roughly speaking , in the large-t limit we can think of the eigenvectors of Ã as infinitesimally localized in Rd . We then apply the above derivation of Eq . ( 9 ) not to the full set of eigenvectors , but separately to each infinitesimal sub-domain , and then integrate the results . See details in Sections 5.2 and D ( SM ) . The subsequent exposition is structured as follows . In Section 4 we provide the background on the NTK kernel and spectral properties of singular integral operators . Then , in Section 5 , we derive the loss asymptotic ( 1 ) for our main setting – the NTK training with a shallow ReLU network . After that , in Section 6 we consider various modifications of this setting : other activation functions ( Section 6.1 ) , deep networks ( Section 6.2 ) , and training in the mean field regime ( Section 6.3 ) .

Summary: The paper studies the scaling law relation between $t$ (time) and the population loss in the NTK regime. The authors show that the scaling law is determined by two indices, (1) the index of the power law decay of the spectrum of the NTK (2) the index of the power decay of the residual of the loss on the eigenfunctions of the NTK.  The main technical contribution is the computation of the two indices.