Explicit loss asymptotics in the gradient descent training of neural networks
Year: 2021

Abstract:
Current theoretical results on optimization trajectories of neural networks trained by gradient descent typically have the form of rigorous but potentially loose bounds on the loss values. In the present work we take a different approach and show that the learning trajectory of a wide network in a lazy training regime can be characterized by an explicit asymptotic at large training times. Specifically, the leading term in the asymptotic expansion of the loss behaves as a power law L(t) ∼ Ct−ξ with exponent ξ expressed only through the data dimension, the smoothness of the activation function, and the class of function being approximated. Our results are based on spectral analysis of the integral operator representing the linearized evolution of a large network trained on the expected loss. Importantly, the techniques we employ do not require a specific form of the data distribution, for example Gaussian, thus making our findings sufficiently universal.

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

heading: 2 Related work
text: 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].

heading: 3 Asymptotic evolution of the loss function
text: 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).

heading: 4 Background

heading: 4.1 Infinitely wide networks
text: Lazy training scenarios discussed in section 3 naturally arise for the networks in the limit of infinite width. However, there are several ways to scale parameters of the network with width, which lead to different operating regimes of infinitely wide networks [15, 16].
NTK regime. The first option we consider is the NTK regime [20], for which the NTK Θ defined in (2) is deterministic and constant during training. This simplification immediately leads to a linear dynamic in the space of network outputs. We consider feed-forward networks parametrized as z1j = ∑d i=1 σww 1 ijxi + σbb 1 j zlj = ∑nl−1 i=1 σw√ nl−1 wlijx l−1 i + σbb l j , l > 1
xlj = φ(z l j)
(10)
Here nl is the width of layer l, xi is a network input and zLj is the network output. We also consider the last layer without bias term bLj and having width nL = 1 (scalar output). Trainable parameters wlij , b l j are initialized as i.i.d. normal Gaussians.
The output of each layer l at initialisation is a Gaussian process with covariance 〈zlj(x)zlj′(x′)〉 = δjj′Σ
(l)(x,x′). By introducing the NTK’s Θ(l)(x,x′) of intermediate layers, one can recursively compute [21] the both NTK and covariance{
Σ(l+1) = σ2w 〈 φ(zl)φ(z′l 〉 + σ2b
Θ(l+1) = Σ(l+1) + σ2wΘ (l) 〈 φ̇(zl)φ̇(z′l) 〉 (11) Here zl = zl(x) and z′l = zl(x′) are draws from the Gaussian process with covariance Σ(l)(x,x′), and 〈. . .〉 is the averaging. It is also convenient to parametrize the covariance at input points x,x′ in some layer l as
Σ(l)(x,x′) =
( r2l rlr ′ l cosϕl
rlr ′ l cosϕl r ′2 l
) (12)
To analyze the spectrum of the evolution operator Ã given by (3) we will intensively use the explicit form of the NTK and covariance. For the ReLU activation φ(z) = (z)+ the averages in (11) can be computed analytically [11]. In the case of shallow network (L = 2) the result is:
Σ(x,x′) = σ2w 2π rr ′( sinϕ+ cosϕ(π − ϕ)) (13) Θ(x,x′) = Σ(x,x′) +
σ2w 2π rr ′ cosϕ(π − ϕ) (14)
Here r, r′, ϕ are parameters from (12) with l = 1 (we dropped index 1). They have a clear interpretation in terms of extended input vectors x̃ = (σwx, σb) ∈ Rd+1. Specifically, r = ‖x̃‖, r′ = ‖x̃′‖ and ϕ is the angle between x̃ and x̃′.
Mean Field regime. This operating regime of infinitely wide networks is naturally defined [26, 31] for shallow networks of the form f(W,x) = 1N ∑N i=1 ciφ(wi · x + bi) = 1 N ∑N i=1 φ̃(w̃i,x). Here w̃i = (ci,wi, bi) denotes the collection of parameters, associated with a single neuron. In the infinite width limit N →∞ the evolution is described by a PDE on parameter density distribution p(w̃):
∂tp = ∇w̃ [ p∇w̃ ∫ K(w̃, w̃′)(p(w̃′)− p∞(w̃′))dw̃′ ] , K(w̃, w̃′) = ∫ φ̃(w̃,x)µ(x)φ̃(w̃′,x)dx
Under mild assumptions [9], solution of the MF equation converges to the global optimum p∞(w̃′). At large times t the MF equation can be linearized [39] around p∞, thus bringing the network to the lazy training regime discussed in section 3. The network NTK in this case is equal to
Θ(x,x′) = ∫ ∇w̃φ̃(w̃,x)p∞(w̃)∇w̃φ̃(w̃,x′)dw̃. (15)

heading: 4.2 Singular integral operators
text: Consider the evolution operator Ã given by Eq. (3). Under our assumptions, the kernel µ1/2(x)Θ(x,x′)µ1/2(x′) of this operator quickly falls off at infinity and is smooth outside the diagonal x = x′, but, as we will see later, has a homogeneous singularity on this diagonal. In this setting, a general theory developed in [5] allows us to obtain the asymptotic distribution (6) of the eigenvalues with explicit constants Λ and ν.
Specifically, suppose that in a neighborhood of the diagonal the kernel Θ has a representation
Θ(x,x′) = θx(x− x′) + . . . , (16)
where θx(·) is a (possibly x-dependent) even (θx(z) = θx(−z)) homogeneous function of degree α:
θx(cz) = |c|αθx(z), (17)
and the dots . . . denote terms of higher smoothness.
Let Nλ denote the number of eigenvalues of Ã greater than λ. Then, it is shown in [5] that for small λ, the leading term of Nλ is given by
Nλ ∼ (∫ γxµ d d+α (x)dx ) λ − dd+α . (18)
Here, γx = (2π)−d|{k ∈ Rd : θ̃x(k) > 1}|, where | · | denotes the Lebesgue measure and θ̃x is a suitably defined Fourier transform of the homogeneous function θx (see SM, Section C).
Formula (18) can be derived as follows. Divide the domain Rd into multiple small subsets Ωm, and think of Ã as an operator matrix corresponding to the decomposition L2(Rd) = ⊕mL2(Ωm). Using the fall off and smoothness of the kernel outside the diagonal, one can show that the leading term of Nλ is determined only by the diagonal elements of this operator matrix. Then, the leading term can be found by considering each restriction Ã|L2(Ωm) separately and summing the respective contributions to Nλ, i.e., Nλ(Ã) ∼ ∑ mNλ(Ã|L2(Ωm)). If we decrease the size of each Ωm by a factor M , then the number of terms in this sum increases Md-fold, but at the same time each Nλ(Ã|L2(Ωm)) decreases also roughly Md-fold, due to rescaling of eigenvalues. In the limit of infinitely small subsets Ωm, the operator Ã|L2(Ωm) can be approximated by the convolution with the homogeneous function θx. The asymptotic form of its eigenvalues can then be written in terms of the Fourier transform θ̃x as given above. The power dd+α in Eq. (18) can be deduced by observing that the volume of the k-space corresponding to θ̃x > λ scales as λ − dd+α .
The formula (18) can be translated into the power law (6) by inverting the relation between λ and n (note that Nλn = n for any n). Specifically, we find that the law (6) holds with
ν = 1 + α
d , Λ =
(∫ γxµ 1/ν(x)dx )ν . (19)
In Section 5.2 we will show that this approach can be extended to yield the loss asymptotic (1).

heading: 5 Asymptotic analysis of wide networks

heading: 5.1 NTK operators and their singularities
text: In this section we demonstrate the asymptotic law (6) of the NTK spectrum as well as outputs covariance spectrum, and find constants ν,Λ . We consider a shallow ReLU network in the NTK regime and with the data distribution µ(x) as described in Section 3. The NTK Θ(x,x′) and outputs covariance Σ(x,x′) of such network is given by (14), (13) with r, r′ and ϕ explicitly depending on input points x,x′:
r(x) = √ σ2w|x|2 + σ2b , ϕ(x,x ′) = arccos (σ2wx·x′+σ2b r(x)r(x′) ) . (20)
To use the spectral theory described in section 4.2 we analyze the smoothness of the kernels Θ and Σ. Firstly, they are smooth (infinitely differentiable) functions of r, r′, ϕ on the whole domain. Now suppose that the bias term is not absent: σb > 0. Then r(x) is a smooth function for all x ∈ Rd. The argument of arccos in (20) is also smooth everywhere, but arccos(z) itself is smooth on (−1, 1) and has divergent derivative at the end points z = 1,−1 corresponding to ϕ = 0, π. We see that condition σb > 0 implies that the case ϕ = π is never realized, while ϕ(x,x′) = 0 at all coinciding inputs. Thus we established that Θ(x,x′) and Σ(x,x′) are smooth everywhere except the diagonal x = x′, where they might have a singularity.
To analyze the behavior at the diagonal, we first expand ϕ in Eq. (20) for small δx = x− x′:
ϕ(x,x′) = σw √ r2(x)−σ2w|x|2 cos2 ψ
r2(x) |δx|+O(|δx| 2) (21)
Here ψ is the angle between δx and x. If σb > 0, then the expression under the square root is always positive, therefore the angle ϕ(x,x′) has a homogeneous singularity of degree 1 on the diagonal. Now we expand expressions (13) and (14) for small ϕ and find that both NTK and covariance indeed have a singularity on the diagonal with leading singular terms.
Σsing(x,x ′) = σ2w 2π r 2(x) 13ϕ 3(x,x′) ∝ |δx|3 (22)
Θsing(x,x ′) = − σ
2 w
2π r 2(x)ϕ(x,x′) ∝ |δx| (23)
101 103 n
10−10
10−8
10−6
10−4
10−2
100 d= 1
M= 100 M= 1000 M= 10000 Singular Asym
101 103 n
10−7
10−5
10−3
10−1
d= 3 M= 100 M= 1000 M= 10000 Singular Asym
101 103 n
10−6
10−4
10−2
100 d= 6
M= 100 M= 1000 M= 10000 Singular Asym
101 103 n
10−4
10−2
100 d= 10
M= 100 M= 1000 M= 10000 Singular Asym
101 103 n
10−10
10−8
10−6
10−4
10−2
100 d= 1
M= 100 M= 1000 M= 10000 Singular Asym
101 103 n
10−7
10−5
10−3
10−1
d= 3 M= 100 M= 1000 M= 10000 Singular Asym
101 103 n
10−6
10−4
10−2
100 d= 6
M= 100 M= 1000 M= 10000 Singular Asym
101 103 n
10−4
10−2
100 d= 10
M= 100 M= 1000 M= 10000 Singular Asym
Figure 2: NTK eigenvalues λn for networks with different input dimension d and for different data set sizes M . Asym shows the theoretical power law (6) with parameters Λ, ν derived in Section 5.1; in particular, ν = 1 + 1d . Singular corresponds to eigenvalues of the singular part of NTK (23). We see that for moderately big n the singular part has the same eigenvalues as the full NTK. Observe that the number n at which the spectrum converges to its asymptotic form increases with dimension d.
Now we see that both NTK Θ and covariance Σ have the form needed for application of spectral theory discussed in section 4.2. For NTK we have singularity degree α = 1, which by Eq. (19) leads to the already announced exponent ν = 1 + 1d (see Figure 1). Interestingly, the singularity degree for the covariance is higher (namely, 3), resulting in a faster fall off of the corresponding eigenvalues. In the sequel (see Section 5.2), this latter degree will appear in the analysis of loss asymptotic for target functions generated by the neural network Gaussian process, and will be denoted by β.
In the case of NTK Θ(x,x′) the singularity is essentially given by (21). The corresponding Fourier transform θ̃x and the function γx can be calculated analytically (see Section C). This leads to an explicit expression for the constant Λ:
Λ = σ3wσ
1 d b
(2π)2 Γ( d+1 2 )Γ( d 2 + 1)
−(1+ 1d ) 〈 µ(x)− 1 d+1 (r(x)) d−1 d+1 〉1+ 1d µ , (24)
where 〈u(x)〉µ denotes the integral ∫ µ(x)u(x)dx. In the experiment, the value 〈u(x)〉µ can be computed by averaging u over a µ-distributed data set. In Figure 2 we compare theoretical and numerical NTK eigenvalue distributions for several dimensions d and data set sizes M .

heading: 5.2 The loss function
text: We extend now the arguments of Section 4.2 to derive the loss asymptotic (1). We consider two classes of functions g representing the initial error (4).
Scenario 1: a discontinuous function g. We assume that g is supported on a bounded subset Ω ⊂ Rd with a smooth boundary ∂Ω so that g has a discontinuity on this boundary but is smooth inside Ω. An obvious example of g is the indicator function of Ω. In this scenario we obtain
L(t) ∼ ∫ ∂Ω |∆g(x)|2(µ(x)θ̃x(n))− 1 d+α dS · 12πΓ( 1 d+α + 1) · (2t) − 1d+α . (25)
Here, x is the point on the surface ∂Ω, n is the unit normal to ∂Ω, and ∆g(x) is the size of the jump of g at x, given by the limit of g(y) as y approaches x from inside Ω. The eigenvalue, coefficient, and loss exponents in this scenario are, respectively,
κ = 1d , ν = 1 + α d , ξ = κ ν = 1 d+α . (26)
Scenario 2: g generated by a Gaussian process. Suppose that g is a realization of a Gaussian process with a covariance matrix Σ(x,x′) = 〈g(x)g(x′)〉 and that Σ has a homogeneous singularity ζx of degree β at the diagonal (in the sense of Eqs. (16),(17)). In this scenario we find
L(t) ∼ ∫ Rd ∫ |n|=1 ζ̃x(n)(µ(x)θ̃x(n)) − βd+α dxdS · 1 2(2π)dβ Γ( βd+α + 1) · (2t) − βd+α . (27)
Here, ζ̃x is the Fourier transform of ζx (defined similarly to the Fourier transform θ̃x). The eigenvalue, coefficient, and loss exponents in this scenario are, respectively,
κ = βd , ν = 1 + α d , ξ = κ ν = β d+α . (28)
In our experiments we model GP by a large network in the NTK regime. The corresponding covariance is analyzed in Section 5.1 and has the singularity degree β = 3.
We provide the full derivations of Eqs. (25) and (27) in Section D of SM, and sketch now the main ideas. Our general strategy is to complement the localized eigenvalue analysis of Section 4.2 by the analysis of expansion coefficients. Note, however, that the simple approximation of Ã by a direct sum of localized operators Ã|L2(Ωm) (as performed in [5] and sketched in Section 4.2) is “too rough” for the study of expansion coefficients (due to a stronger effect of boundary conditions in each Ωm). Accordingly, we replace this approximation by the short-time Fourier transform of g:
F (y,k) = (2π)−d/2 ∫ Rd g(x)ω(x− y)e−ik·xdx,
where ω is a window function such that ∫ ω2 = 1. The coefficient F (y,k) describes the component of g having the wave number k and localized at y. Then at large t, using the stationary phase method,
gt(x) = e −tÃg(x) ∼ (2π)− d2 ∫ Rd ∫ Rd F (y,k)e−tµxθ̃x(k)ω(x− y)eik·xdydk.
The leading contribution to this integral comes from large k. For such k, we can write the coefficients F (y,k) in an asymptotic form, primarily determined by the singularities of g. By integrating out y and k, we then arrive at the desired Eqs. (25) and (27).
The above argument establishes the loss asymptotic (1) while bypassing the computation of the asymptotic (8) of the expansion coefficients aligned with the sorted “global” eigenvalues. This latter asymptotic (including the coefficient K) can be found by a similar computation, see Section D.4.

heading: 6 Extensions
text: The power law asymptotic of eigenvalues obtained in section 5.1 was based on the analysis of diagonal singularity of the NTK in the setting of ReLU activation, shallow depth 2, and the NTK regime. We argue now that our general approach is not restricted to this narrow setting. To show this, we separately consider several modifications of the network from Section 5.1 and an application to MNIST (see Figure 3). In this section we mostly limit ourselves to describing final results, with the derivations provided in SM, Section E.

heading: 6.1 Activations of different smoothness
text: Let’s consider a shallow network in NTK regime with activation function φq(z) = (z) q
+, q > 0 (“a ReLU with the altered smoothness q”). Similarly to the ReLU case, one can show that NTK in current setting has a singularity on the diagonal for all values of q except half-integers q = 12 , 3 2 , . . .. The leading singular term is
Θq(x,x ′) = σ2w 2π rqr′qaqϕ 2q−1 (29) aq = Γ2(q)Γ( 12 − q)√
π2q (30)
Here r, r′, ϕ are the same as in section 5.
The singularity of NTK with degree α = 2q − 1 implies the eigenvalue power law asymptotic (6) with the exponent νq = 1 + 2q−1d . Thus, the singularity degree in the NTK is determined by the singularity degree of the activation function. The coefficient Λq can be explicitly computed as
Λq = σ α+2 w σ α d b q 2(2π)d+q−2 Γ ( d+α 2 ) Γ2(q)(
Γ(d2 + 1) ) d+α d
〈 µ(x)− α d+α r(x) 2d−αd−α d+α 〉 d+α d
µ (31)
In Figure 3a we compare the theoretical and numerical eigenvalue distributions for several values of d and q.
101 103 10−13
10−10
10−7
10−4
10−1
d=1
q=0.5 q=0.75 q=1.0 q=1.5
101 103 10−8
10−6
10−4
10−2
100 d=3
q=0.5 q=0.75 q=1.0 q=1.5
101 103 10−13
10−10
10−7
10−4
10−1
d=1
q=0.5 q=0.75 q=1.0 q=1.5
101 103 10−8
10−6
10−4
10−2
100 d=3
q=0.5 q=0.75 q=1.0 q=1.5
(a) Activations (z) q+.
10−2 100 102 104 t
10−6
10−4
10−2
100
Loss
L=1 L=2 L=3
101 103 n
10−6
10−4
10−2
100 102 Eigenvalues λn L=1 L=2 L=3
10−2 100 102 104 t
10−6
10−4
10−2
100
Loss
L=1 L=2 L=3
101 103 n
10−6
10−4
10−2
100 102 Eigenvalues λn L=1 L=2 L=3
(b) Deep networks
100 102 104 t
10−3
10−2
10−1
Loss exp asym
101 103 n
10−7
10−5
10−3
10−1
101 Eigenvalues
t=0 t=100 t=10000 asym
100 102 104 t
10−3
10−2
10−1
Loss exp asym
101 103 n
10−7
10−5
10−3
10−1
101 Eigenvalues
t=0 t=100 t=10000 asym
(c) Mean field
10−1 101 103 105 t
10−3
10−2
10−1 Loss
exp power fit, ξ= 0.26
10−1 101 103 105 t
10−3
10−2
10−1 Loss
exp power fit, ξ= 0.26
101 103 n
10−6
10−4
10−2
100 Eigenvalues λn
exp power fit, ν= 1.35
10−1 101 103 105 10−3
10−2
10−1 Loss
exp power fit, ξ= 0.26
101 103 n
6
10−4
10−2
100 Eigenvalues λn
exp power fit, ν= 1.35
101 103 n
10−3
10−2
Coefficient sums sn
exp power fit, κ= 0.33
(d) MNIST

heading: 6.2 Deep networks
text: We consider now a network of arbitrary depth L > 2 in the NTK regime and with the ReLU activation function. Similarly to the shallow ReLU case, the angles ϕl (see Eq. (12)) are singular on the diagonal: ϕl ∝ |x− x′|. Using relations (11) one can obtain recursive relations for ϕl, rl and finally for the singular part of NTK Θ(l)sing
Θ (l+1) diag = r 2 l+1 + σ2w 2 Θ (l) diag
Θ (l+1) sing = −
1
2π Θ
(l) diagϕl + σ2w 2 Θ (l) sing
(32)
Here Θ(l+1)diag is the value of NTK on the diagonal. Since Θ (1) sing = 0, we can see from (32) that Θ (L) sing is a weighted sum of −ϕl for l = 1, . . . , L − 1. Thus, the singularity degree of the NTK is α = 1 and the eigenvalue power law (6) holds with ν = 1 + 1d . However, obtaining explicit formula for Λ is harder for this case and we leave it for future work. In Figure 3b we compare the theoretical and numerical eigenvalue distributions for NTK’s of deep networks.

heading: 6.3 MF regime
text: In this section we consider a shallow network in the MF regime and with the ReLU activation function, with the NTK given by (15). The neural tangent kernel (14) of the shallow NTK network can be obtained from (15) with the distribution p∞(c,w, b) taken as a product of Gaussians for each variable. Therefore, neural tangent kernels in MF regime represent a broader class of kernels, containing our basic example (14). It turns out that the diagonal singularity is present for all sufficiently smooth and quickly decaying distributions p∞ of this broader class:
Θsing(x,x ′) = −
√ 1 + |x|2
2
∫ w·x=−b ∣∣w · δx∣∣dp∞,2(w, b) Here δx = x − x′ and p∞,2 is the second moment of the distribution p∞ w.r.t. the variable c, i.e. p∞,2(w, b) = ∫ c2p∞(c,w, b)dc. Thus, the eigenvalue power law exponent is the same as in the NTK regime: ν = 1 + 1d .
In Figure 3c we show the loss dynamic and the eigenvalue distribution at different moments of time for a network in the MF regime. Since the neural tangent kernel significantly changes during the training of the MF network, picking NTK’s at different moments of training provides sufficiently distinct and general kernels of the form (15).

heading: 7 Discussion
text: The main results of the paper are theoretical derivations of power laws for the NTK spectrum and the loss asymptotic. Spectral properties of neural networks is an active area of current research [24, 34, 25, 28]. It has been observed (e.g., [32, 33, 14]) that, close to a global minimum, the Hessian of the loss of neural networks in practical tasks typically has eigenvalues accumulating near 0. This implies, in particular, that the Hessian is ill-conditioned, which has profound implications for network training: convergence of gradient descent in this case is much slower than for well-conditioned optimization problems. A power law spectral distribution is assumed in some recent works (e.g., for the study of generalization in [6]), but we are not aware of any previous general theoretical derivations of such a law. Our work essentially provides such a derivation (since at the global minimum the NTK spectrum coincides with the Hessian spectrum up to the eigenvalue 0). In a sense, we prove theoretically (under some reasonable assumptions) that a slow (power-law) convergence of the gradient descent is inevitable. Moreover, we precisely quantify this convergence.
Our work provides simple explicit formulas for the asymptotic eigenvalue distribution and loss evolution. This allows, in principle, to make some quantitative predictions about optimization not feasible with theoretical methods only providing upper bounds. For example, the loss exponent ξ allows to directly estimate the number of GD steps needed to decrease the loss by one order of magnitude. In principle, the full loss asymptotic (25) can be used to analyze the effect of different data distributions µ(x) on convergence speed of GD.
The key idea of our asymptotic analysis is the extraction of the diagonal singularity from the NTK kernel. This is a new and promising way to relatively easily obtain information about the model spectrum: it only requires simple Taylor expansion and, as we show, is applicable in a wide range of scenarios, extending to deep networks, the mean field regime, and different activation functions. We expect this operation to be important for future studies of network training.
Our loss exponents ξ = 1α+d and ξ = β α+d tend to 0 as the dimension d increases. However, this does not contradict the practical trainability of networks for high-dimensional tasks such as image recognition, since real images occupy only tiny (i.e., effectively low-dimensional) subsets in the ambient high-dimensional space. While our current theory does not predict specific exponents in this setting, we do empirically observe power laws for eigenvalues, coefficients and the loss, and the fitted exponents satisfy our theoretical relation ξ = κν (see Figure 3d). This suggests that our method can be extended to this setting, which will be the topic of a future work.
Funding transparency statement
Research was supported by Russian Science Foundation, grant 21-11-00373.

Review summaries:
The paper identifies a polynomial decay rate in the loss function during gradient descent under the condition that both the eigenvalues of the semigroup operator, and the coefficients in the eigenfunction expansion of an initial loss (g), exhibit power law decay. These assumptions are verified for the semigroup operator in the infinitely wide (NTK) regime, a mean field regime, and also, interestingly, under a homogeneity assumption in the transition kernel. Two scenarios for g are also considered. Throughout, the authors have taken care to be precise with coefficients and compute quantities exactly. 

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.