The asymptotic spectrum of the Hessian of DNN throughout training
Year: 2020

Abstract:
The dynamics of DNNs during gradient descent is described by the so-called Neural Tangent Kernel (NTK). In this article, we show that the NTK allows one to gain precise insight into the Hessian of the cost of DNNs. When the NTK is fixed during training, we obtain a full characterization of the asymptotics of the spectrum of the Hessian, at initialization and during training. In the so-called mean-field limit, where the NTK is not fixed during training, we describe the first two moments of the Hessian at initialization.

The dynamics of DNNs during gradient descent is described by the so-called Neural Tangent Kernel (NTK). In this article, we show that the NTK allows one to gain precise insight into the Hessian of the cost of DNNs. When the NTK is fixed during training, we obtain a full characterization of the asymptotics of the spectrum of the Hessian, at initialization and during training. In the so-called mean-field limit, where the NTK is not fixed during training, we describe the first two moments of the Hessian at initialization.
1 Introduction
The advent of deep learning has sparked a lot of interest in the loss surface of deep neural networks (DNN), and in particular its Hessian. However to our knowledge, there is still no theoretical description of the spectrum of the Hessian. Nevertheless a number of phenomena have been observed numerically.
The loss surface of neural networks has been compared to the energy landscape of different physical models (Choromanska et al., 2015; Geiger et al., 2018; Mei et al., 2018). It appears that the loss surface of DNNs may change significantly depending on the width of the network (the number of neurons in the hidden layer), motivating the distinction between the underand over-parametrized regimes (Baity-Jesi et al., 2018; Geiger et al., 2018; 2019).
The non-convexity of the loss function implies the existence of a very large number of saddle points, which could slow down training. In particular, in (Pascanu et al., 2014; Dauphin et al., 2014), a relation between the rank of saddle points (the number of negative eigenvalues of the Hessian) and their loss has been observed.
For overparametrized DNNs, a possibly more important phenomenon is the large number of flat directions (Baity-Jesi et al., 2018). The existence of these flat minima is conjectured to be related to the generalization of DNNs and may depend on the training procedure (Hochreiter & Schmidhuber, 1997; Chaudhari et al., 2016; Wu et al., 2017).
In (Jacot et al., 2018) it has been shown, using a functional approach, that in the infinitewidth limit, DNNs behave like kernel methods with respect to the so-called Neural Tangent Kernel, which is determined by the architecture of the network. This leads to convergence guarantees for DNNs (Jacot et al., 2018; Du et al., 2019; Allen-Zhu et al., 2018; Huang & Yau, 2019) and strengthens the connections between neural networks and kernel methods (Neal, 1996; Cho & Saul, 2009; Lee et al., 2018).
Our approach also allows one to probe the so-called mean-field/active limit (studied in (Rotskoff & Vanden-Eijnden, 2018; Chizat & Bach, 2018a; Mei et al., 2018) for shallow networks), where the NTK varies during training.
This raises the question: can we use these new results to gain insight into the behavior of the Hessian of the loss of DNNs, at least in the small region explored by the parameters during training?
1.1 Contributions
Following ideas introduced in (Jacot et al., 2018), we consider the training of L+ 1-layered DNNs in a functional setting. For a functional cost C, the Hessian of the loss RP 3 θ 7→ C ( F (L) (θ) ) is the sum of two P × P matrices I and S. We show the following results for large P and for a fixed number of datapoints N :
• The first matrix I is positive semi-definite and its eigenvalues are given by the (weighted) kernel PCA of the dataset with respect to the NTK. The dominating eigenvalues are the principal components of the data followed by a high number of small eigenvalues. The “flat directions” are spanned by the small eigenvalues and the null-space (of dimension at least P − N when there is a single output). Because the NTK is asymptotically constant (Jacot et al., 2018), these results apply at initialization, during training and at convergence.
• The second matrix S can be viewed as residual contribution to H, since it vanishes as the network converges to a global minimum. We compute the limit of the first moment Tr (S) and characterize its evolution during training, of the second moment Tr ( S2 )
which stays constant during training, and show that the higher moments vanish.
• Regarding the sum H = I+S, we show that the matrices I and S are asymptotically orthogonal to each other at initialization and during training. In particular, the moments of the matrices I and S add up: tr(Hk) ≈ tr(Ik) + tr(Sk).
These results give, for any depth and a fairly general non-linearity, a complete description of the spectrum of the Hessian in terms of the NTK at initialization and throughout training. Our theoretical results are consistent with a number of observations about the Hessian (Hochreiter & Schmidhuber, 1997; Pascanu et al., 2014; Dauphin et al., 2014; Chaudhari et al., 2016; Wu et al., 2017; Pennington & Bahri, 2017; Geiger et al., 2018), and sheds a new light on them.
1.2 Related works
The Hessian of the loss has been studied through the decomposition I + S in a number of previous works (Sagun et al., 2017; Pennington & Bahri, 2017; Geiger et al., 2018).
For least-squares and cross-entropy costs, the first matrix I is equal to the Fisher matrix (Wagenaar, 1998; Pascanu & Bengio, 2013), whose moments have been described for shallow networks in (Pennington & Worah, 2018). For deep networks, the first two moments and the operator norm of the Fisher matrix for a least squares loss were computed at initialization in (Karakida et al., 2018) conditionally on a certain independence assumption; our method does not require such assumptions. Note that their approach implicitly uses the NTK.
The second matrix S has been studied in (Pennington & Bahri, 2017; Geiger et al., 2018) for shallow networks, conditionally on a number of assumptions. Note that in the setting of (Pennington & Bahri, 2017), the matrices I and S are assumed to be freely independent, which allows them to study the spectrum of the Hessian; in our setting, we show that the two matrices I and S are asymptotically orthogonal to each other.
2 Setup
We consider deep fully connected artificial neural networks (DNNs) using the setup and NTK parametrization of (Jacot et al., 2018), taking an arbitrary nonlinearity σ ∈ C4b (R) (i.e. σ : R→ R that is 4 times continuously differentiable function with all four derivatives bounded). The layers are numbered from 0 (input) to L (output), each containing n` neurons
for ` = 0, . . . , L. The P = ∑L−1 `=0 (n` + 1)n`+1 parameters consist of the weight matrices W (`) ∈ Rn`+1×n` and bias vectors b(`) ∈ Rn`+1 for ` = 0, . . . , L − 1. We aggregate the parameters into the vector θ ∈ RP .
The activations and pre-activations of the layers are defined recursively for an input x ∈ Rn0 , setting α(0)(x; θ) = x :
α̃(`+1)(x; θ) = 1 √ n` W (`)α(`)(x; θ) + βb(`),
α(`+1)(x; θ) = σ ( α̃(`+1)(x; θ) ) .
The parameter β is added to tune the influence of the bias on training1. All parameters are initialized as iid N (0, 1) Gaussians. We will in particular study the network function, which maps inputs x to the activation of the output layer (before the last non-linearity):
fθ(x) = α̃ (L)(x; θ).
In this paper, we will study the limit of various objects as n1, . . . , nL →∞ sequentially, i.e. we first take n1 →∞, then n2 →∞, etc. This greatly simplifies the proofs, but they could in principle be extended to the simultaneous limit, i.e. when n1 = ... = nL−1 → ∞. All our numerical experiments are done with ‘rectangular’ networks (with n1 = ... = nL−1) and match closely the predictions for the sequential limit.
In the limit we study in this paper, the NTK is asymptotically fixed, as in (Jacot et al., 2018; Allen-Zhu et al., 2018; Du et al., 2019; Arora et al., 2019; Huang & Yau, 2019). By rescaling the outputs of DNNs as the width increases, one can reach another limit where the NTK is not fixed (Chizat & Bach, 2018a;b; Rotskoff & Vanden-Eijnden, 2018; Mei et al., 2019). Some of our results can be extended to this setting, but only at initialization (see Section 3.3). The behavior during training becomes however much more complex.
2.1 Functional viewpoint
The network function lives in a function space fθ ∈ F := [Rn0 → RnL ] and we call the function F (L) : RP → F that maps the parameters θ to the network function fθ the realization function. We study the differential behavior of F (L):
• The derivative DF (L) ∈ RP ⊗ F is a function-valued vector of dimension P . The p-th entry DpF (L) = ∂θpfθ ∈ F represents how modifying the parameter θp modifies the function fθ in the space F . • The Hessian HF (L) ∈ RP ⊗ RP ⊗F is a function-valued P × P matrix.
The network is trained with respect to the cost functional:
C(f) = 1 N N∑ i=1 ci (f(xi)) ,
for strictly convex ci, summing over a finite dataset x1, . . . , xN ∈ Rn0 of size N . The parameters are then trained with gradient descent on the composition C ◦F (L), which defines the usual loss surface of neural networks.
In this setting, we define the finite realization function Y (L) : RP → RNnL mapping parameters θ to be the restriction of the network function fθ to the training set yik = fθ,k(xi). The Jacobian DY (L) is hence an NnL × P matrix and its Hessian HY (L) is a P × P ×NnL tensor. Defining the restricted cost C(y) = 1N ∑ i ci(yi), we have C ◦ F (L) = C ◦ Y (L).
For our analysis, we require that the gradient norm ‖DC‖ does not explode during training. The following condition is sufficient: Definition 1. A loss C : RNnL → R has bounded gradients over sublevel sets (BGOSS) if the norm of the gradient is bounded over all sets Ua = { Y ∈ RNnL : C(Y ) ≤ a } .
For example, the Mean Square Error (MSE) C(Y ) = 12N ‖Y ∗ − Y ‖2 for the labels Y ∗ ∈ RNnL has BGOSS because ‖∇C(Y )‖2 = 1N ‖Y ∗ − Y ‖2 = 2C(Y ). For the binary and softmax cross-entropy the gradient is uniformly bounded, see Proposition 2 in Appendix A. 1In our experiments, we take β = 0.1.
2.2 Neural Tangent Kernel
The behavior during training of the network function fθ in the function space F is described by a (multi-dimensional) kernel, the Neural Tangent Kernel (NTK)
Θ (L) k,k′(x, x ′) = P∑ p=1 ∂θpfθ,k(x)∂θpfθ,k′(x ′).
During training, the function fθ follows the so-called kernel gradient descent with respect to the NTK, which is defined as
∂tfθ(t)(x) = −∇Θ(L)C|fθ(t)(x) := − 1
N N∑ i=1 Θ(L)(x, xi)∇ci(fθ(t)(xi)).
In the infinite-width limit (letting n1 →∞, . . . , nL−1 →∞ sequentially) and for losses with BGOSS, the NTK converges to a deterministic limit Θ(L) → Θ(L)∞ ⊗ IdnL , which is constant during training, uniformly on finite time intervals [0, T ] (Jacot et al., 2018). For the MSE loss, the uniform convergence of the NTK was proven for T =∞ in (Arora et al., 2019).
The limiting NTK Θ (L) ∞ : Rn0 × Rn0 → R is constructed as follows:
1. For f, g : R → R and a kernel K : Rn0 × Rn0 → R, define the kernel Lf,gK : Rn0 × Rn0 → R by
Lf,gK (x0, x1) = E(a0,a1) [f(a0)g(a1)] , for (a0, a1) a centered Gaussian vector with covariance matrix (K(xi, xj))i,j=0,1. For f = g, we denote by LfK the kernel L f,f K .
2. We define the kernels Σ (`) ∞ for each layer of the network, starting with Σ (1) ∞ (x0, x1) =
1/n0(xT0 x1) + β 2 and then recursively by Σ (`+1) ∞ = Lσ
Σ (`) ∞ + β2, for ` = 1, . . . , L − 1, where σ is the network non-linearity.
3. The limiting NTK Θ (L) ∞ is defined in terms of the kernels Σ (`) ∞ and the kernels
Σ̇ (`) ∞ = Lσ̇
Σ (`−1) ∞
:
Θ(L)∞ = L∑ `=1 Σ(`)∞ Σ̇ (`+1) ∞ . . . Σ̇ (L) ∞ .
The NTK leads to convergence guarantees for DNNs in the infinite-width limit, and connect their generalization to that of kernel methods (Jacot et al., 2018; Arora et al., 2019).
2.3 Gram Matrices
For a finite dataset x1, . . . , xN ∈ Rn0 and a fixed depth L ≥ 1, we denote by Θ̃ ∈ RNnL×NnL the Gram matrix of x1, . . . , xN with respect to the limiting NTK, defined by
Θ̃ik,jm = Θ (L) ∞ (xi, xj) δkm.
It is block diagonal because different outputs k 6= m are asymptotically uncorrelated.
Similarly, for any (scalar) kernel K(L) (such as the limiting kernels Σ(L)∞ ,Λ(L)∞ ,Υ(L)∞ ,Φ(L)∞ ,Ξ(L)∞ introduced later), we denote the Gram matrix of the datapoints by K̃.
3 Main Theorems
3.1 Hessian as I + S
Using the above setup, the Hessian H of the loss C ◦ F (L) is the sum of two terms, with the entry Hp,p′ given by
Hp,p′ = HC|fθ (∂θpF, ∂θp′F ) +DC|fθ (∂θp,θp′F ).
For a finite dataset, the Hessian matrix H ( C ◦ Y (L) ) is equal to the sum of two matrices
I = ( DY (L) )T HCDY (L) and S = ∇C · HY (L)
where DY (L) is a NnL×P matrix, HC is a NnL×NnL matrix and HY (L) is a P ×P ×NnL tensor to which we apply a scalar product (denoted by ·) in its last dimension with the NnL vector ∇C to obtain a P × P matrix. Our main contribution is the following theorem, which describes the limiting moments Tr ( Hk ) in terms of the moments of I and S:
Theorem 1. For any loss C with BGOSS and σ ∈ C4b (R), in the sequential limit n1 → ∞, . . . , nL−1 →∞, we have for all k ≥ 1
Tr ( H (t) k ) ≈ Tr ( I (t) k ) + Tr ( S (t) k ) .
The limits of Tr ( I (t) k ) and Tr ( S (t) k ) can be expressed in terms of the NTK Θ (L) ∞ , the
kernels Υ (L) ∞ ,Ξ (L) ∞ and the non-symmetric kernels Φ (L) ∞ , Λ (L) ∞ defined in Appendix C:
• The moments Tr ( I (t) k ) converge to the following limits (with the convention that
ik+1 = i1): Tr ( I (t) k ) → Tr (( HC(Y (t))Θ̃ ) k ) = 1
Nk N∑ i1,...,ik=1 k∏ m=1 c′′im(fθ(t)(xim))Θ (L) ∞ (xim , xim+1).
• The first moment Tr (S (t)) converges to the limit:
Tr (S (t)) = (G(t)) T ∇C(Y (t)).
At initialization (G(0), Y (0)) form a Gaussian pair of NnL-vectors, independent for differing output indices k = 1, ..., nL and with covariance E [Gik(0)Gi′k′(0)] = δkk′Ξ (L) ∞ (xi, xi′) and E [Gik(0)Yi′k′(0)] = δkk′Φ(L)∞ (xi, xi′) for the limiting kernel Ξ (L) ∞ (x, y) and non-symmetric kernel Φ (L) ∞ (x, y). During training, both vectors follow the differential equations
∂tG(t) = −Λ̃∇C(Y (t)) ∂tY (t) = −Θ̃∇C(Y (t)).
• The second moment Tr ( S (t) 2 ) converges to the following limit defined in terms of
the Gram matrix Υ̃: Tr ( S2 ) → (∇C(Y (t)))T Υ̃∇C(Y (t))
• The higher moments Tr ( S (t) k ) for k ≥ 3 vanish.
Proof. The moments of I and S can be studied separately because the moments of their sum is asymptotically equal to the sum of their moments by Proposition 5 below. The limiting moments of I and S are respectively described by Propositions 1 and 4 below.
In the case of a MSE loss C(Y ) = 12N ‖Y − Y ∗‖2, the first and second derivatives take simple forms ∇C(Y ) = 1N (Y − Y ∗) and HC(Y ) = 1N IdNnL and the differential equations can be solved to obtain more explicit formulae:
Corollary 1. For the MSE loss C and σ ∈ C4b (R), in the limit n1, ..., nL−1 →∞, we have uniformly over [0, T ]
Tr ( H(t)k ) → 1 Nk Tr ( Θ̃k ) + Tr ( S(t)k ) where
Tr (S(t))→− 1 N
(Y ∗ − Y (0))T ( IdNnL + e −tΘ̃ ) Θ̃−1Λ̃T e−tΘ̃(Y ∗ − Y (0))
+ 1
N G(0)T e−tΘ̃(Y ∗ − Y (0))
Tr ( S(t)2 ) → 1 N2 (Y ∗ − Y (0))T e−tΘ̃Υ̃e−tΘ̃(Y ∗ − Y (0))
Tr ( S(t)k ) →0 when k > 2.
In expectation we have: E [Tr (S(t))]→− 1 N Tr (( IdNnL + e −tΘ̃ ) Θ̃−1Λ̃T e−tΘ̃ ( Σ̃ + Y ∗Y ∗T )) + 1 N Tr ( e−tΘ̃Φ̃T ) E [ Tr ( S(t)2
)] → 1 N2 Tr ( e−tΘ̃Υ̃e−tΘ̃ ( Σ̃ + Y ∗Y ∗T )) .
Proof. The moments of I are constant because HC = 1N IdNnL is constant. For the moments of S, we first solve the differential equation for Y (t):
Y (t) = Y ∗ − e−tΘ̃(Y ∗ − Y (0)).
Noting Y (t)− Y (0) = −Θ̃ ∫ t
0 ∇C(s)ds, we have
G(t) = G(0)− Λ̃ ∫ t
0
∇C(s)ds
= G(0) + Λ̃Θ̃−1(Y (t)− Y (0)) = G(0) + Λ̃Θ̃−1 ( IdNnL + e −tΘ̃ ) (Y ∗ − Y (0))
The expectation of the first moment of S then follows.
3.2 Mutual Orthogonality of I and S
A first key ingredient to prove Theorem 1 is the asymptotic mutual orthogonality of the matrices I and S
Proposition (Proposition 5 in Appendix D). For any loss C with BGOSS and σ ∈ C4b (R), we have uniformly over [0, T ]
lim nL−1→∞ · · · lim n1→∞ ‖IS‖F = 0.
As a consequence limnL−1→∞ · · · limn1→∞Tr ( [I + S] k ) − [ Tr ( Ik ) + Tr ( Sk )] = 0.
Remark 1. If two matrices A and B are mutualy orthogonal (i.e. AB = 0) the range of A is contained in the nullspace of B and vice versa. The non-zero eigenvalues of the sum A+B are therefore given by the union of the non-zero eigenvalues of A and B. Furthermore the
moments of A and B add up: Tr ( [A+B] k ) = Tr ( Ak ) + Tr ( Bk ) . Proposition 5 shows that
this is what happens asymptotically for I and S.
Note that both matrices I and S have large nullspaces: indeed assuming a constant width w = n1 = ... = nL−1, we have Rank(I) ≤ NnL and Rank(S) ≤ 2(L − 1)wNnL (see Appendix C), while the number of parameters P scales as w2 (when L > 2).
Figure 2 illustrates the mutual orthogonality of I and S. All numerical experiments are done for rectangular networks (when the width of the hidden layers are equal) and agree well with our predictions obtained in the sequential limit.
3.3 Mean-field Limit For a rectangular network with width w, if the output of the network is divided by √ w and the learning rate is multiplied by w (to keep similar dynamics at initialization), the training dynamics changes and the NTK varies during training when w goes to infinity. The new parametrization of the output changes the scaling of the two matrices:
H [ C ( 1√ w Y (L) )] = 1 w ( DY (L) )T HCDY (L) + 1√ w ∇C · HY (L) = 1 w I + 1√ w S.
The scaling of the learning rate essentially multiplies the whole Hessian by w. In this setting, the matrix I is left unchanged while the matrix S is multiplied by √ w (the k-th moment of S is hence multiplied by wk/2). In particular, the two moments of the Hessian are dominated by the moments of S, and the higher moments of S (and the operator norm of S) should not vanish. This suggests that the active regime may be characterised by the fact that ‖S‖F ‖I‖F . Under the conjecture that Theorem 1 holds for the infinite-width limit of rectangular networks, the asymptotic of the two first moments of H is given by:
1/ √ wTr (H)→ N (0,∇CT Ξ̃∇C) 1/wTr ( H2 ) → ∇CT Υ̃∇C,
where for the MSE loss we have ∇C = −Y ∗.
3.4 The matrix S
The matrix S = ∇C · HY (L) is best understood as a perturbation to I, which vanishes as the network converges because ∇C → 0. To calculate its moments, we note that
Tr ( ∇C · HY (L) ) = ( P∑ p=1 ∂2θ2pY )T ∇C = GT∇C,
where the vector G = ∑P k=1 ∂
2 θ2p Y ∈ RNnL is the evaluation of the function gθ(x) =∑P
k=1 ∂ 2 θ2p fθ(x) on the training set.
For the second moment we have
Tr (( ∇C · HY (L) )2) = ∇CT  P∑ p,p′=1 ∂2θpθp′Y ( ∂2θpθp′Y )T∇C = ∇CT Υ̃∇C for Υ̃ the Gram matrix of the kernel Υ(L)(x, y) = ∑P p,p′=1 ∂ 2 θpθp′ fθ(x) ( ∂2θpθp′ fθ(y) )T .
The following proposition desribes the limit of the function gθ and the kernel Υ (L) and the vanishing of the higher moments: Proposition (Proposition 4 in Appendix C). For any loss C with BGOSS and σ ∈ C4b (R), the first two moments of S take the form
Tr (S(t)) = G(t)T∇C(t) Tr ( S(t)2 ) = ∇C(t)T Υ̃(t)∇C(t)
- At initialization, gθ and fθ converge to a (centered) Gaussian pair with covariances
E[gθ,k(x)gθ,k′(x′)] = δkk′Ξ(L)∞ (x, x′) E[gθ,k(x)fθ,k′(x′)] = δkk′Φ(L)∞ (x, x′) E[fθ,k(x)fθ,k′(x′)] = δkk′Σ(L)∞ (x, x′)
and during training gθ evolves according to
∂tgθ,k(x) = N∑ i=1 Λ(L)∞ (x, xi)∂ikC(Y (t))·
- Uniformly over any interval [0, T ], the kernel Υ(L) has a deterministic and fixed limit limnL−1→∞ · · · limn1→∞Υ (L) kk′ (x, x ′) = δkk′Υ (L) ∞ (x, x′) with limiting kernel:
Υ(L)∞ (x, x ′) = L−1∑ `=1 ( Θ(`)∞ (x, x ′)2Σ̈(`)∞ (x, x ′) + 2Θ(`)∞ (x, x ′)Σ̇(`)∞ (x, x ′) ) Σ̇(`+1)∞ (x, x ′) · · · Σ̇(L−1)∞ (x, x′).
- The higher moment k > 2 vanish: limnL−1→∞ · · · limn1→∞ Tr ( Sk ) = 0.
This result has a number of consequences for infinitely wide networks: 1. At initialization, the matrix S has a finite Frobenius norm ‖S‖2F = Tr ( S2 ) =
∇CT Υ̃∇C, because Υ converges to a fixed limit. As the network converges, the derivative of the cost goes to zero ∇C(t)→ 0 and so does the Frobenius norm of S.
2. In contrast the operator norm of S vanishes already at initialization (because for all even k, we have ‖S‖op ≤ k √
Tr (Sk) → 0). At initialization, the vanishing of S in operator norm but not in Frobenius norm can be explained by the matrix S having a growing number of eigenvalues of shrinking intensity as the width grows.
3. When it comes to the first moment of S, Proposition 4 shows that the spectrum of S is in general not symmetric. For the MSE loss the expectation of the first moment at initialization is
E [Tr(S)] = E [ (Y − Y ∗)TG ] = E [ Y TG ] − (Y ∗)T E [G] = Tr ( Φ̃ ) − 0
which may be positive or negative depending on the choice of nonlinearity: with a smooth ReLU, it is positive, while for the arc-tangent or the normalized smooth ReLU, it can be negative (see Figure 1). This is in contrast to the result obtained in (Pennington & Bahri, 2017; Geiger et al., 2018) for the shallow ReLU networks, taking the second derivative of the ReLU to be zero. Under this assumption the spectrum of S is symmetric: if the eigenvalues are ordered from lowest to highest, λi = −λP−i and Tr(S) = 0.
These observations suggest that S has little influence on the shape of the surface, especially towards the end of training, the matrix I however has an interesting structure.
3.5 The matrix I
At a global minimizer θ∗, the spectrum of I describes how the loss behaves around θ∗. Along the eigenvectors of the biggest eigenvalues of I, the loss increases rapidely, while small eigenvalues correspond to flat directions. Numerically, it has been observed that the matrix I features a few dominating eigenvalues and a bulk of small eigenvalues (Sagun et al., 2016; 2017; Gur-Ari et al., 2018; Papyan, 2019). This leads to a narrow valley structure of the loss around a minimum: the biggest eigenvalues are the ‘cliffs’ of the valley, i.e. the directions along which the loss grows fastest, while the small eigenvalues form the ‘flat directions’or the bottom of the valley.
Note that the rank of I is bounded by NnL and in the overparametrized regime, when NnL < P , the matrix I will have a large nullspace, these are directions along which the value of the function on the training set does not change. Note that in the overparametrized regime, global minima are not isolated: they lie in a manifold of dimension at least P −NnL and the nullspace of I is tangent to this solution manifold.
The matrix I is closely related to the NTK Gram matrix: Θ̃ = DY (L) ( DY (L) )T and I = ( DY (L) )T HCDY (L).
As a result, the limiting spectrum of the matrix I can be directly obtained from the NTK2
Proposition 1. For any loss C with BGOSS and σ ∈ C4b (R), uniformly over any interval [0, T ], the moments Tr ( Ik )
converge to the following limit (with the convention that ik+1 = i1):
lim nL−1→∞ · · · lim n1→∞
Tr ( Ik ) = Tr (( HC(Yt)Θ̃ ) k ) = 1
Nk N∑ i1,...,ik=1 k∏ m=1 c′′im(fθ(t)(xim))Θ (L) ∞ (xim , xim+1)
2This result was already obtained in (Karakida et al., 2018), but without identifying the NTK explicitely and only at initialization.
Proof. It follows from Tr ( Ik ) = Tr ((( DY (L) )T HCDY (L))k) = Tr((HCΘ̃)k) and the asymptotic of the NTK (Jacot et al., 2018).
3.5.1 Mean-Square Error
When the loss is the MSE, HC is equal to 1N IdNnL . As a result, Θ̃ and I have the same non-zero eigenvalues up to a scaling of 1/N. Because the NTK is assymptotically fixed, the spectrum of I is also fixed in the limit.
The eigenvectors of the NTK Gram matrix are the kernel principal components of the data. The biggest principal components are the directions in function space which are most favorised by the NTK. This gives a functional interpretation of the narrow valley structure in DNNs: the cliffs of the valley are the biggest principal components, while the flat directions are the smallest components.
Remark 2. As the depth L of the network increases, one can observe two regimes (Poole et al., 2016; Jacot et al., 2019): Order/Freeze where the NTK converges to a constant and Chaos where the NTK converges to a Kronecker delta. In the Order/Freeze the NnL ×NnL Gram matrix approaches a block diagonal matrix with nL constant blocks, and as a result nL eigenvalues of I dominate the other ones, corresponding to constant directions along each outputs (this is in line with the observations of (Papyan, 2019)). This leads to a narrow valley for the loss and slows down training. In contrast, in the Chaos regime, the NTK Gram matrix approaches a scaled identity matrix, and the spectrum of I should hence concentrate around a positive value, hence speeding up training. Figure 3 illustrates this phenomenon: with the smooth ReLU we observe a narrow valley, while with the normalized smooth ReLU (which lies in the Chaos according to (Jacot et al., 2019)) the narrowness of the loss is reduced. A similar phenomenon may explain why normalization helps smoothing the loss surface and speed up training (Santurkar et al., 2018; Ghorbani et al., 2019).
3.5.2 Cross-Entropy Loss
For a binary cross-entropy loss with labels Y ∗ ∈ {−1,+1}N
C(Y ) = 1
N N∑ i=1 log ( 1 + e−Y ∗ i Yi ) ,
HC is a diagonal matrix whose entries depend on Y (but not on Y ∗):
HiiC(Y ) = 1
N
1
1 + e−Yi + eYi .
The eigenvectors of I then correspond to the weighted kernel principal component of the data. The positive weights 1
1+e−Yi+eYi approach 1/3 as Yi goes to 0, i.e. when it is close to
the decision boundary from one class to the other, and as Yi → ±∞ the weight go to zero. The weights evolve in time through Yi, the spectrum of I is therefore not asymptotically fixed as in the MSE case, but the functional interpretation of the spectrum in terms of the kernel principal components remains.
4 Conclusion
We have given an explicit formula for the limiting moments of the Hessian of DNNs throughout training. We have used the common decomposition of the Hessian in two terms I and S and have shown that the two terms are asymptotically mutually orthogonal, such that they can be studied separately.
The matrix S vanishes in Frobenius norm as the network converges and has vanishing operator norm throughout training. The matrix I is arguably the most important as it describes the narrow valley structure of the loss around a global minimum. The eigendecomposition of I is related to the (weighted) kernel principal components of the data w.r.t. the NTK.
Acknowledgements
Clément Hongler acknowledges support from the ERC SG CONSTAMIS grant, the NCCR SwissMAP grant, the NSF DMS-1106588 grant, the Minerva Foundation, the Blavatnik Family Foundation, and the Latsis foundation.
References
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A Convergence Theory for Deep Learning via Over-Parameterization. CoRR, abs/1811.03962, 2018. URL http://arxiv.org/abs/ 1811.03962.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Ruslan Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. arXiv preprint arXiv:1904.11955, 2019.
Marco Baity-Jesi, Levent Sagun, Mario Geiger, Stefano Spigler, Gerard Ben Arous, Chiara Cammarota, Yann LeCun, Matthieu Wyart, and Giulio Biroli. Comparing Dynamics: Deep Neural Networks versus Glassy Systems. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80, pp. 314–323. PMLR, 10–15 Jul 2018. URL http://proceedings.mlr.press/v80/baityjesi18a.html.
Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina. Entropy-sgd: Biasing gradient descent into wide valleys. arXiv preprint arXiv:1611.01838, 2016.
Lénäıc Chizat and Francis Bach. On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport. In Advances in Neural Information Processing Systems 31, pp. 3040–3050. Curran Associates, Inc., 2018a. URL http://papers.nips.cc/paper/7567-on-the-global-convergence-of-gradientdescent-for-over-parameterized-models-using-optimal-transport.pdf.
Lenaic Chizat and Francis Bach. A note on lazy training in supervised differentiable programming. arXiv preprint arXiv:1812.07956, 2018b.
Youngmin Cho and Lawrence K. Saul. Kernel Methods for Deep Learning. In Advances in Neural Information Processing Systems 22, pp. 342–350. Curran Associates, Inc., 2009. URL http://papers.nips.cc/paper/3628-kernel-methods-for-deep-learning.pdf.
Anna Choromanska, Mikael Henaff, Michael Mathieu, Gérard Ben Arous, and Yann LeCun. The Loss Surfaces of Multilayer Networks. Journal of Machine Learning Research, 38: 192–204, nov 2015. URL https://arxiv.org/pdf/1412.0233.pdf.
Yann N. Dauphin, Razvan Pascanu, Caglar Gulcehre, Kyunghyun Cho, Surya Ganguli, and Yoshua Bengio. Identifying and Attacking the Saddle Point Problem in High-dimensional Non-convex Optimization. In Proceedings of the 27th International Conference on Neural Information Processing Systems - Volume 2, NIPS’14, pp. 2933–2941, Cambridge, MA, USA, 2014. MIT Press.
Simon S. Du, Xiyu Zhai, Barnabás Póczos, and Aarti Singh. Gradient Descent Provably Optimizes Over-parameterized Neural Networks. 2019.
Mario Geiger, Stefano Spigler, Stéphane d’Ascoli, Levent Sagun, Marco Baity-Jesi, Giulio Biroli, and Matthieu Wyart. The jamming transition as a paradigm to understand the loss landscape of deep neural networks. arXiv preprint arXiv:1809.09349, 2018.
Mario Geiger, Arthur Jacot, Stefano Spigler, Franck Gabriel, Levent Sagun, Stéphane d’Ascoli, Giulio Biroli, Clément Hongler, and Matthieu Wyart. Scaling description of generalization with number of parameters in deep learning . abs/1901.01608, 2019. URL http://arxiv.org/abs/1901.01608.
Behrooz Ghorbani, Shankar Krishnan, and Ying Xiao. An investigation into neural net optimization via hessian eigenvalue density. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 2232–2241, Long Beach, California, USA, 09–15 Jun 2019. PMLR. URL http://proceedings.mlr.press/v97/ghorbani19b.html.
Guy Gur-Ari, Daniel A. Roberts, and Ethan Dyer. Gradient descent happens in a tiny subspace. CoRR, abs/1812.04754, 2018. URL http://arxiv.org/abs/1812.04754.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural Computation, 9(1):1–42, 1997.
Jiaoyang Huang and Horng-Tzer Yau. Dynamics of deep neural networks and neural tangent hierarchy. arXiv preprint arXiv:1909.08156, 2019.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural Tangent Kernel: Convergence and Generalization in Neural Networks. In Advances in Neural Information Processing Systems 31, pp. 8580–8589. Curran Associates, Inc., 2018. URL http://papers.nips.cc/paper/8076-neural-tangent-kernelconvergence-and-generalization-in-neural-networks.pdf.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Freeze and chaos for dnns: an NTK view of batch normalization, checkerboard and boundary effects. CoRR, abs/1907.05715, 2019. URL http://arxiv.org/abs/1907.05715.
Ryo Karakida, Shotaro Akaho, and Shun-Ichi Amari. Universal Statistics of Fisher Information in Deep Neural Networks: Mean Field Approach. jun 2018. URL http: //arxiv.org/abs/1806.01316.
Jae Hoon Lee, Yasaman Bahri, Roman Novak, Samuel S. Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep Neural Networks as Gaussian Processes. ICLR, 2018.
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115(33): E7665–E7671, 2018.
Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit. arXiv preprint arXiv:1902.06015, 2019.
Radford M. Neal. Bayesian Learning for Neural Networks. Springer-Verlag New York, Inc., Secaucus, NJ, USA, 1996. ISBN 0387947248.
Vardan Papyan. Measurements of three-level hierarchical structure in the outliers in the spectrum of deepnet hessians. CoRR, abs/1901.08244, 2019. URL http://arxiv.org/ abs/1901.08244.
Razvan Pascanu and Yoshua Bengio. Revisiting Natural Gradient for Deep Networks. jan 2013. URL http://arxiv.org/abs/1301.3584.
Razvan Pascanu, Yann N Dauphin, Surya Ganguli, and Yoshua Bengio. On the saddle point problem for non-convex optimization. arXiv preprint, 2014. URL https://arxiv.org/ pdf/1405.4604.pdf.
Jeffrey Pennington and Yasaman Bahri. Geometry of Neural Network Loss Surfaces via Random Matrix Theory. In Proceedings of the 34th International Conference on Machine Learning, volume 70, pp. 2798–2806. PMLR, 06–11 Aug 2017. URL http://proceedings. mlr.press/v70/pennington17a.html.
Jeffrey Pennington and Pratik Worah. The Spectrum of the Fisher Information Matrix of a Single-Hidden-Layer Neural Network. In Advances in Neural Information Processing Systems 31, pp. 5415–5424. Curran Associates, Inc., 2018. URL http://papers.nips.cc/paper/7786-the-spectrum-of-the-fisherinformation-matrix-of-a-single-hidden-layer-neural-network.pdf.
Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Exponential expressivity in deep neural networks through transient chaos. In D. D. Lee, M. Sugiyama, U. V. Luxburg, I. Guyon, and R. Garnett (eds.), Advances in Neural Information Processing Systems 29, pp. 3360–3368. Curran Associates, Inc., 2016. URL http://papers.nips.cc/paper/6322-exponential-expressivityin-deep-neural-networks-through-transient-chaos.pdf.
Grant Rotskoff and Eric Vanden-Eijnden. Parameters as interacting particles: long time convergence and asymptotic error scaling of neural networks. In Advances in Neural Information Processing Systems 31, pp. 7146–7155. Curran Associates, Inc., 2018. URL http://papers.nips.cc/paper/7945-parameters-as-interactingparticles-long-time-convergence-and-asymptotic-error-scaling-of-neuralnetworks.pdf.
Levent Sagun, Léon Bottou, and Yann LeCun. Singularity of the hessian in deep learning. CoRR, abs/1611.07476, 2016. URL http://arxiv.org/abs/1611.07476.
Levent Sagun, Utku Evci, V. Ugur Güney, Yann Dauphin, and Léon Bottou. Empirical Analysis of the Hessian of Over-Parametrized Neural Networks. CoRR, abs/1706.04454, 2017.
Shibani Santurkar, Dimitris Tsipras, Andrew Ilyas, and Aleksander Madry. How does batch normalization help optimization? In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems 31, pp. 2483–2493. Curran Associates, Inc., 2018. URL http://papers.nips. cc/paper/7515-how-does-batch-normalization-help-optimization.pdf.
Daniel Wagenaar. Information geometry of neural networks. 1998. ISSN 0302-9743.
Lei Wu, Zhanxing Zhu, and Weinan E. Towards Understanding Generalization of Deep Learning: Perspective of Loss Landscapes. CoRR, abs/1706.10239, 2017. URL http: //arxiv.org/abs/1706.10239.
A Proofs
For the proofs of the theorems and propositions presented in the main text, we reformulate the setup of (Jacot et al., 2018). For a fixed training set x1, ..., xN , we consider a (possibly random) time-varying training direction D(t) ∈ RNnL which describes how each of the outputs must be modified. In the case of gradient descent on a cost C(Y ), the training direction is D(t) = ∇C(Y (t)). The parameters are updated according to the differential equation
∂tθ(t) = (∂θY (t)) T D(t). Under the condition that ∫ T
0 ‖D(t)‖2 dt is stochastically bounded as the width of the network
goes to infinity, the NTK Θ(L) converges to its fixed limit uniformly over [0, T ].
The reason we consider a general training direction (and not only a gradient of a loss) is that we can split a network in two at a layer ` and the training of the smaller network will be according to the training direction D (`) i (t) given by
D (`) i (t) = diag ( σ̇ ( α(`)(xi) ))( 1 √ n` W (`) )T ...diag ( σ̇ ( α(L−1)(xi) ))( 1 √ nL−1 W (L−1) )T Di(t)
because the derivatives σ̇ are bounded and by Lemma 1 of the Appendix of (Jacot et al., 2018), this training direction satisfies the constraints even though it is not the gradient of a loss. As a consequence, as n1 →∞, ..., n`−1 →∞ the NTK of the smaller network Θ(`) also converges to its limit uniformly over [0, T ]. As we let n` →∞ the pre-activations α̃(`)i and weights W
(`) ij move at a rate of 1/ √ n`. We will use this rate of change to prove that other
types of kernels are constant during training.
When a network is trained with gradient descent on a loss C with BGOSS, the integral∫ T 0 ‖D(t)‖2 dt is stochastically bounded. Because the loss is decreasing during training, the outputs Y (t) lie in the sublevel set UC(Y (0)) for all times t. The norm of the gradient is hence bounded for all times t. Because the distribution of Y (0) converges to a multivariate Gaussian, b(C(Y (0))) is stochastically bounded as the width grows, where b(a) is a bound
on the norm of the gradient on Ua. We then have the bound ∫ T 0 ‖D(t)‖2 dt ≤ Tb(C(Y (0))) which is itself stochastically bounded.
For the binary and softmax cross-entropy losses the gradient is uniformly bounded:
Proposition 2. For the binary cross-entropy loss C and any Y ∈ RN , ‖∇C(Y )‖2 ≤ 1√ N . For the softmax cross-entropy loss C on c ∈ N classes and any Y ∈ RNc, ‖∇C(Y )‖2 ≤ √ 2c√ N .
Proof. The binary cross-entropy loss with labels Y ∗ ∈ {0, 1}N is
C(Y ) = − 1 N N∑ i=1 log eYiY ∗ i 1 + eYi = 1 N N∑ i=1 log ( 1 + eYi ) − YiY ∗i
and the gradient at an input i is
∂iC(Y ) = 1
N eYi − Y ∗i (1 + eYi) 1 + eYi
which is bounded in absolute value by 1N for both Y ∗ i = 0, 1 such that ‖∇C(Y )‖2 ≤ 1√ N .
The softmax cross-entropy loss over c classes with labels Y ∗ ∈ {1, . . . , c}N is defined by
C(Y ) = − 1 N N∑ i=1 log e YiY ∗ i∑c k=1 e Yik = 1 N N∑ i=1 log
( c∑
k=1
eYik ) − YiY ∗i .
The gradient is at an input i and output class m is
∂imC(Y ) = 1
N ( eYim∑c k=1 e Yik − δY ∗i m ) which is bounded in absolute value by 2N such that ‖∇C(Y )‖2 ≤ √ 2c√ N .
B Preliminaries
To study the moments of the matrix S, we first have to show that two tensors vanish as n1, ..., nL−1 →∞:
Ω (L) k0,k1,k2 (x0, x1, x2) = (∇fθ,k0(x0)) T Hfθ,k1(x1)∇fθ,k2(x2)
Γ (L) k0,k1,k2,k3 (x0, x1, x2, x4) = (∇fθ,k0(x0)) T Hfθ,k1(x1)Hfθ,k2(x2)∇fθ,k3(x3).
We study these tensors recursively, for this, we need a recursive definition for the first derivatives ∂θpfθ,k(x) and second derivatives ∂ 2 θpθp′ fθ,k(x). The value of these derivatives depend on the layer ` the parameters θp and θp′ belong to, and on whether they are connection weights W (`) mk or biases b (`) k . The derivatives with respect to the parameters of the last layer are
∂ W
(L−1) mk
fθ,k′(x) = 1
√ nL−1
α(L−1)m (x)δkk′
∂ b (L−1) k
fθ,k′(x) = β 2δkk′
for parameters θp which belong to the lower layers the derivatives can be defined recursively by
∂θpfθ,k(x) = 1
√ nL−1 nL−1∑ m=1 ∂θp α̃ (L−1) m (x)σ̇ ( α̃(L−1)m (x) ) W (L−1) mk .
For the second derivatives, we first note that if either of the parameters θp or θp′ are bias of the last layer, or if they are both connection weights of the last layer, then ∂2θpθp′ fθ,k(x) = 0. Two cases are left: when one parameter is a connection weight of the last layer and the others belong to the lower layers, and when both belong to the lower layers. Both cases can be defined recursively in terms of the first and second derivatives of α̃ (L−1) m :
∂2 θpW (L) mk
fθ,k′(x) = 1
√ nL−1
∂θp α̃ (L−1) m (x)σ̇ ( α̃(L−1)m (x) ) δkk′
∂2θpθp′ fθ,k′(x) = 1
√ nL−1 nL−1∑ m=1 ∂2θpθp′ α̃ (L−1) m (x)σ̇ ( α̃(L−1)m (x) ) W (L−1) mk
+ 1
√ nL−1 nL−1∑ m=1 ∂θp α̃ (L−1) m (x)∂θp′ α̃ (L−1) m (x)σ̈ ( α̃(L−1)m (x) ) W (L−1) mk .
Using these recursive definitions, the tensors Ω(L+1) and Γ(L+1) are given in terms of Θ(L),Ω(L) and Γ(L), in the same manner that the NTK Θ(L+1) is defined recursively in terms of Θ(L) in (Jacot et al., 2018).
Lemma 1. For any loss C with BGOSS and σ ∈ C4b (R), we have uniformly over [0, T ]
lim nL−1→∞ · · · lim n1→∞ Ω (L) k0,k1,k2 (x0, x1, x2) = 0
Proof. The proof is done by induction. When L = 1 the second derivatives ∂2θpθp′ fθ,k(x) = 0 and Ω (L) k0,k1,k2 (x0, x1, x2) = 0.
For the induction step, we write Ω (`+1) k0,k1,k2 (x0, x1, x2) recursively as
n −3/2 ` ∑ m0,m1,m2 Θ(`)m0,m1(x0, x1)Θ (`) m1,m2(x1, x2)σ̇(α̃ (`) m0(x0))σ̈(α̃ (`) m1(x1))σ̇(α̃ (`) m2(x2))W (`) m0k0 W (`) m1k1 W (`) m2k2
+ n −3/2 ` ∑ m0,m1,m2 Ω(`)m0,m1,m2(x0, x1, x2)σ̇(α̃ (`) m0(x0))σ̇(α̃ (`) m1(x1))σ̇(α̃ (`) m2(x2))W (`) m0k0 W (`) m1k1 W (`) m2k2 + n −3/2 `
∑ m0,m1 Θ(`)m0,m1(x0, x1)σ̇(α̃ (`) m0(x0))σ̇(α̃ (`) m1(x1))σ(α̃ (`) m1(x2))W (`) m0k0 δk1k2
+ n −3/2 ` ∑ m1,m2 Θ(`)m1,m2(x1, x2)σ(α̃ (`) m1(x0))σ̇(α̃ (`) m1(x1))σ̇(α̃ (`) m2(x2))δk0k1W (`) m2k2 .
As n1, ..., n`−1 →∞ and for any times t < T , the NTK Θ(`) converges to its limit while Ω(`) vanishes. The second summand hence vanishes and the others converge to
n −3/2 ` ∑ m Θ(`)∞ (x0, x1)Θ (`) ∞ (x1, x2)σ̇(α̃ (`) m (x0))σ̈(α̃ (`) m (x1))σ̇(α̃ (`) m (x2))W (`) mk0 W (`) mk1 W (`) mk2
+ n −3/2 ` ∑ m Θ(`)∞ (x0, x1)σ̇(α̃ (`) m (x0))σ̇(α̃ (`) m (x1))σ(α̃ (`) m (x2))W (`) mk0 δk1k2 + n −3/2 `
∑ m Θ(`)∞ (x1, x2)σ(α̃ (`) m (x0))σ̇(α̃ (`) m (x1))σ̇(α̃ (`) m (x2))δk0k1W (`) mk2 .
At initialization, all terms vanish as n` →∞ because all summands are independent with zero mean and finite variance: in the n1 →∞, . . . , n`−1 →∞ limit, the α̃(`)m (x) are independent for different m, see (Jacot et al., 2018). During training, the weights W (`) and preactivations α̃(`) move at a rate of 1/√n` (see the proof of convergence of the NTK in (Jacot et al., 2018)). Since σ̇ is Lipschitz, we obtain that the motion during training of each of the sums is of order n −3/2+1/2 ` = n −1 ` . As a result, uniformly over times t ∈ [0, T ], all the sums vanish.
Similarily, we have
Lemma 2. For any loss C with BGOSS and σ ∈ C4b (R), we have uniformly over [0, T ]
lim nL−1→∞ · · · lim n1→∞ Γ (L) k0,k1,k2,k3 (x0, x1, x2, x3) = 0
Proof. The proof is done by induction. When L = 1 the hessian HF (1) = 0, such that Γ
(L) k0,k1,k2,k3 (x0, x1, x2, x3) = 0.
For the induction step, Γ(`+1) can be defined recursively:
Γ (L+1) k0,k1,k2,k3 (x0, x1, x2, x3) = n−2L ∑
m0,m1,m2,m3
Γ(L)m0,m1,m2,m3(x0, x1, x2, x3)σ̇(α (L) m0 (x0))σ̇(α (L) m1 (x1))σ̇(α (L) m2 (x2))σ̇(α (L) m3 (x3))
W (L) m0k0 W (L) m1k1 W (L) m2k2 W (L) m3k3 +n−2L ∑
m0,m1,m2,m3
Θ(L)m0,m1(x0, x1)Ω (L) m1,m2,m3(x1, x2, x3)σ̇(α (L) m0 (x0))σ̈(α (L) m1 (x1))
σ̇(α(L)m2 (x2))σ̇(α (L) m3 (x3))W (L) m0k0 W (L) m1k1 W (L) m2k2 W (L) m3k3 +n−2L ∑
m0,m1,m2,m3
Ω(L)m0,m1,m2(x0, x1, x2)Θ (L) m2,m3(x2, x3)σ̇(α (L) m0 (x0))σ̇(α (L) m1 (x1))
σ̈(α(L)m2 (x2))σ̇(α (L) m3 (x3))W (L) m0k0 W (L) m1k1 W (L) m2k2 W (L) m3k3 +n−2L ∑
m0,m1,m2,m3
Θ(L)m0,m1(x0, x1)Θ (L) m1,m2(x1, x2)Θ (L) m2,m3(x2, x3)σ̇(α (L) m0 (x0))σ̈(α (L) m1 (x1))
σ̈(α(L)m2 (x2))σ̇(α (L) m3 (x3))W (L) m0k0 W (L) m1k1 W (L) m2k2 W (L) m3k3 +n−2L ∑
m1,m2,m3
Ω(L)m1,m2,m3(x1, x2, x3)σ(α (L) m1 (x0))σ̇(α (L) m1 (x1))σ̇(α (L) m2 (x2))σ̇(α (L) m3 (x3))
δk0k1W (L) m2k2 W (L) m3k3 +n−2L ∑
m1,m2,m3
Θ(L)m1,m2(x1, x2)Θ (L) m2,m3(x2, x3)σ(α (L) m1 (x0))σ̇(α (L) m1 (x1))σ̈(α (L) m2 (x2))σ̇(α (L) m3 (x3))
δk0k1W (L) m2k2 W (L) m3k3 +n−2L ∑
m0,m1,m2
Ω(L)m0,m1,m2(x0, x1, x2)σ̇(α (L) m0 (x0))σ̇(α (L) m1 (x1))σ̇(α (L) m2 (x2))σ(α (L) m2 (x3))
W (L) m0k0 W (L) m1k1 δk2k3 +n−2L ∑
m0,m1,m2
Θ(L)m0,m1(x0, x1)Θ (L) m1,m2(x1, x2)σ̇(α (L) m0 (x0))σ̈(α (L) m1 (x1))σ̇(α (L) m2 (x2))σ(α (L) m2 (x3))
W (L) m0k0 W (L) m1k1 δk2k3 + n−2L ∑ m1,m2 Θ(L)m1,m2(x1, x2)σ(α (L) m1 (x0))σ̇(α (L) m1 (x1))σ̇(α (L) m2 (x2))σ(α (L) m2 (x3))δk0k1δk2k3
+n−2L ∑
m0,m1,m3
Θ(L)m0,m1(x0, x1)Θ (L) m1,m3(x2, x3)σ̇(α (L) m0 (x0))σ̇(α (L) m1 (x1))σ̇(α (L) m1 (x2))σ̇(α (L) m3 (x3))
W (L) m0k0 δk1k2W (L) m3k3
As n1, ..., n`−1 →∞ and for any times t < T , the NTK Θ(`) converges to its limit while Ω(`) and Γ(`) vanishes. Γ (L+1) k0,k1,k2,k3 (x0, x1, x2, x3) therefore converges to:
+n−2L ∑ m Θ(L)∞ (x0, x1)Θ (L) ∞ (x1, x2)Θ (L) ∞ (x2, x3)σ̇(α (L) m (x0))σ̈(α (L) m (x1))σ̈(α (L) m (x2))σ̇(α (L) m (x3))
W (L) mk0 W (L) mk1 W (L) mk2 W (L) mk3
+n−2L ∑ m Θ(L)∞ (x1, x2)Θ (L) ∞ (x2, x3)σ(α (L) m (x0))σ̇(α (L) m (x1))σ̈(α (L) m (x2))σ̇(α (L) m (x3))
δk0k1W (L) mk2 W (L) mk3 +n−2L ∑ m Θ(L)∞ (x0, x1)Θ (L) ∞ (x1, x2)σ̇(α (L) m (x0))σ̈(α (L) m (x1))σ̇(α (L) m (x2))σ(α (L) m (x3))
W (L) mk0 W (L) mk1 δk2k3 +n−2L ∑ m Θ(L)∞ (x1, x2)σ(α (L) m (x0))σ̇(α (L) m (x1))σ̇(α (L) m (x2))σ(α (L) m (x3))δk0k1δk2k3
+n−2L ∑ m Θ(L)∞ (x0, x1)Θ (L) ∞ (x2, x3)σ̇(α (L) m (x0))σ̇(α (L) m (x1))σ̇(α (L) m (x2))σ̇(α (L) m (x3))
W (L) mk0 δk1k2W (L) mk3
For the convergence during training, we proceed similarily to the proof of Lemma 1. At initialization, all terms vanish as n` → ∞ because all summands are independent (after taking the n1, . . . , nL−1 → ∞ limit) with zero mean and finite variance. During training, the weights W (`) and preactivations α̃(`) move at a rate of 1/√n` which leads to a change of order n −2+1/2 ` = n −1.5 ` , which vanishes for all times t too.
C The Matrix S
We now have the theoretical tools to describe the moments of the matrix S. We first give a bound for the rank of S:
Proposition 3. Rank(S) ≤ 2(n1 + ...+ nL−1)NnL
Proof. We first observe that S is given by a sum of NnL matrices:
Spp′ = N∑ i=1 nL∑ k=1 ∂ikC∂ 2 θpθpfθ,k(xi).
It is therefore sufficiant to show that the rank of each matricesHfθ,k(x) = ( ∂2θpθp′ fθ,k(xi) ) p,p′ is bounded by 2(n1 + ...+ nL).
The derivatives ∂θpfθ,k(x) have different definition depending on whether the parameter θp is a connection weight W (`) ij or a bias b (`) j :
∂ W (`) ij fθ,k(x) = 1 √ n` α (`) i (x; θ)∂α̃(`+1)j (x;θ) fθ,k(x)
∂ b (`) j fθ,k(x) = β∂α̃(`+1)j (x;θ) fθ,k(x) These formulas only depend on θ through the values ( α
(`) i (x; θ) ) `,i and ( ∂ α̃ (`) i (x;θ) fθ,k(x) ) `,i
for ` = 1, ..., L−1 (note that both α(0)i (x) = xi and ∂α̃(L)i (x;θ)fθ,k(x) = δik do not depend on θ). Together there are 2(n1 + ...+nL−1) of them. As a consequence, the map θ 7→ ( ∂θpfθ,k(xi) ) p can be written as a composition
θ ∈ RP 7→ ( α
(`) i (x; θ), ∂α̃(`)i (x;θ) fθ,k(x) ) `,i ∈ R2(n1+...+nL−1) 7→ ( ∂θpfθ,k(xi) ) p ∈ RP
and the matrix Hfθ,k(x) is equal to the Jacobian of this map. By the chain rule, Hfθ,k(x) is the matrix multiplication of the Jacobians of the two submaps, whose rank are bounded by 2(n1 + ... + nL−1), hence bounding the rank of Hfθ,k(x). And because S is a sum of NnL matrices of rank smaller than 2(n1 + ... + nL−1), the rank of S is bounded by 2(n1 + ...+ nL−1)NnL.
C.1 Moments
Let us now prove Proposition 4:
Proposition 4. For any loss C with BGOSS and σ ∈ C4b (R), the first two moments of S take the form
Tr (S(t)) = G(t)T∇C(t) Tr ( S(t)2 ) = ∇C(t)T Υ̃(t)∇C(t)
- At initialization, gθ and fθ converge to a (centered) Gaussian pair with covariances
E[gθ,k(x)gθ,k′(x′)] = δkk′Ξ(L)∞ (x, x′) E[gθ,k(x)fθ,k′(x′)] = δkk′Φ(L)∞ (x, x′) E[fθ,k(x)fθ,k′(x′)] = δkk′Σ(L)∞ (x, x′)
and during training gθ evolves according to
∂tgθ,k(x) = N∑ i=1 Λ(L)∞ (x, xi)∂ikC(Y (t))·
- Uniformly over any interval [0, T ] where ∫ T
0 ‖∇C(t)‖2 dt is stochastically bounded, the kernel
Υ(L) has a deterministic and fixed limit limnL−1→∞ · · · limn1→∞Υ (L) kk′ (x, x ′) = δkk′Υ (L) ∞ (x, x′) with limiting kernel:
Υ(L)∞ (x, x ′) = L−1∑ `=1 ( Θ(`)∞ (x, x ′)2Σ̈(`)(x, x′) + 2Θ(`)∞ (x, x ′)Σ̇(`)(x, x′) ) Σ̇(`+1)(x, x′) · · · Σ̇(L−1)(x, x′).
- The higher moment k > 2 vanish: limnL−1→∞ · · · limn1→∞Tr ( Sk ) = 0.
Proof. The first moment of S takes the form Tr (S) = ∑ p (∇C)T Hp,pY = (∇C)T G
where G is the restriction to the training set of the function gθ(x) = ∑ p ∂ 2 θpθp fθ(x). This process is random at initialization and varies during training. Lemma 3 below shows that, in the infinite width limit, it is a Gaussian process at initialization which then evolves according to a simple differential equation, hence describing the evolution of the first moment during training.
The second moment of S takes the form:
Tr(S2) = P∑ p1,p2=1 N∑ i1,i2=1 ∂2θp1 ,θp2 fθ,k1(x1)∂ 2 θp2 ,θp1 fθ,k2(x2)c ′ i1(xi1)c ′ i2(xi2)
= (∇C)T Υ̃∇C
where Υ (L) k1,k2 (x1, x2) = ∑P p1,p2=1 ∂2θp1 ,θp2 fθ,k1(x1)∂ 2 θp2 ,θp1 fθ,k2(x2) is a multidimensional kernel and Υ̃ is its Gram matrix. Lemma 4 below shows that in the infinite-width limit, Υ (L) k1,k2 (x1, x2) converges to a deterministic and time-independent limit Υ (L) ∞ (x1, x2)δk1k2 .
To show that Tr(Sk) → 0 for all k > 2, it suffices to show that ∥∥S2∥∥
F → 0 as ∣∣Tr(Sk)∣∣ <∥∥S2∥∥ F ‖S‖k−2F and we know that ‖S‖F → (∂Y C)
T Υ̃∂Y C is finite. We have that∥∥S2∥∥
F = N∑ i0,i1,i2,i3=1 nL∑ k0,k1,k2,k3=1 Ψ (L) k0,k1,k2,k3 (xi0 , xi1 , xi2 , xi3)∂fθ,k0 (xi0 )C∂fθ,k1 (xi1 )C
∂fθ,k2 (xi2 )C∂fθ,k3 (xi3 )C
= Ψ̃ · (∂Y C)⊗4
for Ψ̃ the NnL ×NnL ×NnL ×NnL finite version of
Ψ (L) k0,k1,k2,k3 (xi0 , xi1 , xi2 , xi3) = P∑ p0,p1,p2,p3=1 ∂2θp0 ,θp1 fθ,k0(x0)∂ 2 θp1 ,θp2 fθ,k1(x1)
∂2θp2 ,θp3 fθ,k2(x2)∂ 2 θp3 ,θp0 fθ,k3(x3).
which vanishes in the infinite width limit by Lemma 5 below.
Lemma 3. For any loss C with BGOSS and σ ∈ C4b (R), at initialization gθ and fθ converge to a (centered) Gaussian pair with covariances
E[gθ,k(x)gθ,k′(x′)] = δkk′Ξ(L)∞ (x, x′) E[gθ,k(x)fθ,k′(x′)] = δkk′Φ(L)∞ (x, x′) E[fθ,k(x)fθ,k′(x′)] = δkk′Σ(L)∞ (x, x′)
and during training gθ evolves according to
∂tgθ(x) = N∑ i=1 Λ(L)∞ (x, xi)Di(t)
Proof. When L = 1, gθ(x) is 0 for any x and θ.
For the inductive step, the trace g (L+1) θ,k (x) is defined recursively as
1 √ nL nL∑ m=1 g (L) θ,m(x)σ̇ ( α̃(L)m (x) ) W (L) mk + Tr ( ∇fθ,m(x) (∇fθ,m(x))T ) σ̈ ( α̃(L)m (x) ) W (L) mk
First note that Tr ( ∇fθ,m(x) (∇fθ,m(x))T ) = Θ (L) mm(x, x). Now let n1, ...nL−1 →∞, by the
induction hypothesis, the pairs (g (L) θ,m, α̃ (L) m ) converge to iid Gaussian pairs of processes with covariance Φ (L) ∞ at initialization.
At initialization, conditioned on the values of g (L) m , α̃ (L) m the pairs (g (L+1) k , fθ) follow a centered Gaussian distribution with (conditioned) covariance
E[g(L+1)θ,k (x)g (L+1) θ,k′ (x ′)|g(L)θ,m, α̃ (L) m ] =
δkk′
nL nL∑ m=1 ( g (L) θ,m(x)σ̇ ( α̃(L)m (x) ) + Θ(L)∞ (x, x)σ̈ ( α̃(L)m (x) )) ( g
(L) θ,m(x ′)σ̇ ( α̃(L)m (x ′) ) + Θ(L)∞ (x ′, x′)σ̈ ( α̃(L)m (x ′) ))
E[g(L+1)θ,k (x)fθ,k′(x ′)|g(L)θ,m, α̃ (L) m ] =
δkk′
nL nL∑ m=1 ( g (L) θ,m(x)σ̇ ( α̃(L)m (x) ) + Θ(L)∞ (x, x)σ̈ ( α̃(L)m (x) )) σ ( α̃(L)m (x ′) )
E[fθ,k(x)fθ,k′(x′)|g(L)θ,m, α̃ (L) m ] =
δkk′
nL nL∑ m=1 σ ( α̃(L)m (x) ) σ ( α̃(L)m (x ′) ) + β2.
As nL → ∞, by the law of large number, these (random) covariances converge to their expectations which are deterministic, hence the pairs (g
(L+1) k , fθk) have asymptotically the
same Gaussian distribution independent of g (L) m , α̃ (L) m : E [ g
(L) θ,k (x)g (L) θ,k′(x ′) ] → δkk′Ξ(L)∞ (x, x′)
E [ g
(L) θ,k (x)f (L) θ,k′(x ′) ] → δkk′Φ(L)∞ (x, x)
E [ f
(L) θ,k (x)f (L) θ,k′(x ′) ] → δkk′Σ(L)∞ (x, x)
with Ξ (1) ∞ (x, x′) = Φ (1) ∞ (x, x′) = 0 and
Ξ(L+1)∞ (x, x ′) = E [gg′σ̇(α)σ̇(α′)]
+ Θ(L)∞ (x ′, x′)E [gσ̇(α)σ̈(α′)] + Θ(L)∞ (x, x)E [g′σ̇(α′)σ̈(α)] + Θ(L)∞ (x, x)Θ (L) ∞ (x ′, x′)E [σ̈(α′)σ̈(α)]
= Ξ(L)∞ (x, x ′)Σ̇(L)∞ (x, x
′) + (
Φ(L)∞ (x, x ′)Φ(L)∞ (x ′, x) + Φ(L)∞ (x, x)Φ (L) ∞ (x
′, x′) )
Σ̈(L)∞ (x, x ′)
+ Φ(L)∞ (x, x ′)Φ(L)∞ (x ′, x′)E [σ̇(α)...σ (α′)] + Φ(L)∞ (x, x)Φ(L)∞ (x′, x)E [ ... σ (α)σ̇(α′)]
+ Θ(L)∞ (x ′, x′) ( Φ(L)∞ (x, x)Σ̈ (L) ∞ (x, x ′) + Φ(L)∞ (x, x ′)E [σ̇(α)...σ (α′)] ) + Θ(L)∞ (x, x) ( Φ(L)∞ (x ′, x′)Σ̈(L)∞ (x, x ′) + Φ(L)∞ (x ′, x)E [...σ (α)σ̇(α′)] )
+ Θ(L)∞ (x, x)Θ (L) ∞ (x ′, x′)Σ̈(L)∞ (x, x ′)
and
Φ(L+1)∞ (x, x ′) = E [gσ̇(α)σ(α′)] + Θ(L)∞ (x, x)E [σ̈(α)σ(α′)]
= Φ(L)∞ (x, x ′)Σ̇(L+1)(x, x′) + ( Φ(L)∞ (x, x) + Θ (L) ∞ (x, x) ) E [σ̈(α)σ(α′)]
where (g, g′, α, α′) is a Gaussian quadruple of covariance Ξ (L) ∞ (x, x) Ξ (L) ∞ (x, x′) Φ (L) ∞ (x, x) Φ (L) ∞ (x, x′) Ξ (L) ∞ (x, x′) Ξ (L) ∞ (x′, x′) Φ (L) ∞ (x′, x) Φ (L) ∞ (x′, x′)
Φ (L) ∞ (x, x) Φ (L) ∞ (x′, x) Σ (L) ∞ (x, x) Σ (L) ∞ (x, x′) Φ (L) ∞ (x, x′) Φ (L) ∞ (x′, x′) Σ (L) ∞ (x, x′) Σ (L) ∞ (x′, x′)
 .
During training, the parameters follow the gradient ∂tθ(t) = (∂θY (t)) T D(t). By the induction hypothesis, the traces g (L) θ,m then evolve according to the differential equation
∂tg (L) θ,m(x) = 1 √ nL N∑ i=1 nL∑ m=1 Λ (L) mm′(x, xi)σ̇(α̃ (L) m′ (x)) ( W (L) m′ )T Di(t)
and in the limit as n1, ..., nL−1 → ∞, the kernel Λ(L)mm′(x, xi) converges to a deterministic and fixed limit δmm′Λ (L) ∞ (x, xi). Note that as nL grows, the g (L) θ,m(x) move at a rate of 1/√nL just like the pre-activations α̃ (L) m . Even though they move less and less, together they affect the trace g (L+1) θ,k which follows the differential equation
∂tg (L+1) θ,k (x) = N∑ i=1 nL∑ k′=1 Λ (L+1) kk′ (x, xi)Dik′(t)
where
Λ (L+1) kk′ (x, x
′) = 1
nL ∑ m,m′ Λ (L) mm′(x, x ′)σ̇ ( α̃(L)m (x) ) σ̇ ( α̃ (L) m′ (x ′) ) W (L) mkW (L) m′k′
+ 1
nL ∑ m,m′ g (L) θ,m(x)Θ (L) mm′(x, x ′)σ̈ ( α̃(L)m (x) ) σ̇ ( α̃ (L) m′ (x ′) ) W (L) mkW (L) m′k′
+ 1
nL ∑ m g (L) θ,m(x)σ̇ ( α̃(L)m (x) ) σ ( α̃(L)m (x ′) ) δkk′
+ 2
nL ∑ m,m′ Ω (L) m′mm(x ′, x, x)σ̈ ( α̃(L)m (x) ) σ̇ ( α̃ (L) m′ (x ′) ) W (L) mkW (L) m′k′
+ 1
nL ∑ m,m′ Θ(L)mm(x, x)Θ (L) mm′(x, x ′) ... σ ( α̃(L)m (x) ) σ̇ ( α̃ (L) m′ (x ′) ) W (L) mkW (L) m′k′
+ 1
nL ∑ m Θ(L)mm(x, x)σ̈ ( α̃(L)m (x) ) σ ( α̃(L)m (x ′) ) δkk′ .
As n1, ..., nL−1 →∞, the kernels Θ(L)mm′(x, x′) and Λ (L) mm′(x, x ′) converge to their limit and Ω (L) m′mm(x ′, x, x) vanishes:
Λ (L) kk′ (x, x ′)→ 1 nL ∑ m Λ(L)∞ (x, x ′)σ̇ ( α̃(L)m (x) ) σ̇ ( α̃(L)m (x ′) ) W (L) mkW (L) mk′
+ 1
nL ∑ m g (L) θ,m(x)Θ (L) ∞ (x, x ′)σ̈ ( α̃(L)m (x) ) σ̇ ( α̃(L)m (x ′) ) W (L) mkW (L) mk′
+ 1
nL ∑ m g (L) θ,m(x)σ̇ ( α̃(L)m (x) ) σ ( α̃(L)m (x ′) ) δkk′
+ 1
nL ∑ m Θ(L)∞ (x, x)Θ (L) ∞ (x, x ′) ... σ ( α̃(L)m (x) ) σ̇ ( α̃(L)m (x ′) ) W (L) mkW (L) mk′
+ 1
nL ∑ m Θ(L)∞ (x, x)σ̈ ( α̃(L)m (x) ) σ ( α̃(L)m (x ′) ) δkk′
By the law of large numbers, as nL →∞, at initialization Λ(L+1)kk′ (x, x′)→ δkk′Λ (L+1) ∞ (x, x′) where
Λ(L+1)∞ (x, x ′) = Λ(L)∞ (x, x ′)Σ̇(L+1)∞ (x, x ′)
+ Θ(L)∞ (x, x ′)E [gσ̈ (α) σ̇ (α′)] + E [gσ̇ (α)σ (α′)] + Θ(L)∞ (x, x)Θ (L) ∞ (x, x
′)E [...σ (α) σ̇ (α′)] + Θ(L)∞ (x, x)E [σ̈ (α)σ (α′)] = Λ(L)∞ (x, x ′)Σ̇(L+1)∞ (x, x ′)
+ Θ(L)∞ (x, x ′) ( Φ(L)∞ (x, x ′)Σ̈(L+1)∞ (x, x ′) + Φ(L)∞ (x, x)E [ ... σ (α) σ̇ (α′)] ) + Φ(L)∞ (x, x ′)Σ̇(L+1)∞ (x, x ′) + Φ(L)∞ (x, x)E [σ̈ (α)σ (α′)]
+ Θ(L)∞ (x, x)Θ (L) ∞ (x, x ′)E [...σ (α) σ̇ (α′)] + Θ(L)∞ (x, x)E [σ̈ (α) σ̇ (α′)]
During training Θ (L) ∞ and Λ (L) ∞ are fixed in the limit n1, .., nL−1 → ∞, and the values g (L) θ,m(x), α̃ (L) m (x) and W (L) mk vary at a rate of 1/√nL which induce a change of the same rate to Λ (L) kk′ (x, x ′), which is therefore asymptotically fixed during training as nL →∞.
The next lemma describes the asymptotic limit of the kernel Υ(L):
Lemma 4. For any loss C with BGOSS and σ ∈ C4b (R), the second moment of the Hessian of the realization function HF (L) converges uniformly over [0, T ] to a fixed limit as n1, ...nL−1 →∞
Υ (L) kk′ (x, x ′)→ δkk′ L−1∑ `=1 ( Θ(`)∞ (x, x ′)2Σ̈(`)∞ (x, x ′) + 2Θ(`)∞ (x, x ′)Σ̇(`)∞ (x, x ′) ) Σ̇(`+1)∞ (x, x ′) · · · Σ̇(L−1)∞ (x, x′).
Proof. The proof is by induction on the depth L. The case L = 1 is trivially true because ∂2θpθp′ fθ,k(x) = 0 for all p, p ′, k, x. For the induction step we observe that
Υ (L) k,k′(x, x ′)
= P∑ p1,p2=1 ∂2θp1 ,θp2 fθ,k(x)∂ 2 θp2 ,θp1 fθ,k′(x ′)
= 1
nL nL∑ m,m′=1 Υ (L) m,m′(x, x ′)σ̇ ( α̃(L)m (x) ) σ̇ ( α̃ (L) m′ (x ′) ) W (L) mkW (L) m′k′
+ 1
nL nL∑ m,m′=1 Ω (L) m′,m,m′(x ′, x, x′)σ̇ ( α̃(L)m (x) ) σ̈ ( α̃ (L) m′ (x ′) ) W (L) mkW (L) m′k′
+ 1
nL nL∑ m,m′=1 Ω (L) m,m′,m(x, x ′, x)σ̈ ( α̃(L)m (x) ) σ̇ ( α̃ (L) m′ (x ′) ) W (L) mkW (L) m′k′
+ 1
nL nL∑ m,m′=1 Θ (L) m,m′(x, x ′)Θ (L) m′,m(x ′, x)σ̈ ( α̃(L)m (x) ) σ̈ ( α̃ (L) m′ (x ′) ) W (L) mkW (L) m′k′
+ 2
nL nL∑ m=1 Θ (L) m,m′(x, x ′)σ̇ ( α̃(L)m (x) ) σ̇ ( α̃ (L) m′ (x ′) ) δkk′
if we now let the width of the lower layers grow to infinity n1, ...nL−1 →∞, the tensor Ω(L) vanishes and Υ (L) m,m′ and the NTK Θ (L) m,m′ converge to limits which are non-zero only when m = m′. As a result, the term above converges to
1
nL nL∑ m=1 Υ(L)∞ (x, x ′)σ̇ ( α̃(L)m (x) ) σ̇ ( α̃(L)m (x ′) ) W (L) mkW (L) mk′
+ 1
nL nL∑ m=1 Θ(L)∞ (x, x ′)2σ̈ ( α̃(L)m (x) ) σ̈ ( α̃(L)m (x ′) ) W (L) mkW (L) mk′
+ 2
nL nL∑ m=1 Θ(L)∞ (x, x ′)σ̇ ( α̃(L)m (x) ) σ̇ ( α̃(L)m (x ′) ) δkk′
At initialization, we can apply the law of large numbers as nL →∞ such that it converges to Υ (L+1) ∞ (x, x′)δkk′ , for the kernel Υ (L+1) ∞ (x, x′) defined recursively by
Υ(L+1)∞ (x, x ′) =Υ(L)∞ (x, x ′)Σ̇(L)∞ (x, x ′) + Θ(L)∞ (x, x ′)2Σ̈(L)∞ (x, x ′) + 2Θ(L)∞ (x, x ′)Σ̇(L)∞ (x, x ′)
and Υ (1) ∞ (x, x′) = 0.
For the convergence during training, we proceed similarily to the proof of Lemma 1: the activations α̃ (L) m (x) and weights W (L) mk move at a rate of 1/√nL and the change to Υ (L+1) kk′ is therefore of order 1/√nL and vanishes as nL → 0.
Finally, the next lemma shows the vanishing of the tensor Ψ (L) k0,k1,k2,k3
to prove that the higher moments of S vanish.
Lemma 5. For any loss C with BGOSS and σ ∈ C4b (R), uniformly over [0, T ]
lim nL−1→∞ · · · lim n1→∞ Ψ (L) k0,k1,k2,k3 (xi0 , xi1 , xi2 , xi3) = 0
Proof. When L = 1 the Hessian is zero and Ψ (1) k0,k1,k2,k3 (xi0 , xi1 , xi2 , xi3) = 0.
For the induction step, we write Ψ (L+1) k0,k1,k2,k3 (xi0 , xi1 , xi2 , xi3) recursively, because
it contains many terms, we change the notation, writing [ x0 x1 m0 m1 ] for
Θ (L) m0,m1(x0, x1), [ x0 x1 x2 m0 m1 m2 ] for Ω (L) m0,m1,m2(x0, x1, x2) and [ x0 x1 x2 x3 m0 m1 m2 m3 ] for Γ (L) m0,m1,m2,m3(x0, x1, x2, x3). The value Ψ (L+1) k0,k1,k2,k3 (xi0 , xi1 , xi2 , xi3) is then equal to
n−2L ∑
m0,m1,m2,m3
Ψ(L)m0,m1,m2,m3(x0, x1, x2, x3)σ̇ ( α̃(L)m0 (x0) ) σ̇ ( α̃(L)m1 (x1) ) σ̇ ( α̃(L)m2 (x2) ) σ̇ ( α̃(L)m3 (x3) ) W (L) m0k0 W (L) m1k1 W (L) m2k2 W (L) m3k3
+n−2L ∑
m0,m1,m2,m3
[ x0 x1 m0 m1 ] [ x1 x2 m1 m2 ] [ x2 x3 m2 m3 ] [ x3 x0 m3 m0 ] σ̈ ( α̃(L)m0 (x0) ) σ̈ ( α̃(L)m1 (x1) ) σ̈ ( α̃(L)m2 (x2) ) σ̈ ( α̃(L)m3 (x3) ) W (L) m0k0 W (L) m1k1 W (L) m2k2 W (L) m3k3
+n−2L ∑
m0,m1,m2,m3
[ x0 x1 x2 m0 m1 m2 ] [ x2 x3 m2 m3 ] [ x3 x0 m3 m0 ] σ̈ ( α̃(L)m0 (x0) ) σ̇ ( α̃(L)m1 (x1) ) σ̈ ( α̃(L)m2 (x2) ) σ̈ ( α̃(L)m3 (x3) ) W (L) m0k0 W (L) m1k1 W (L) m2k2 W (L) m3k3
+n−2L ∑
m0,m1,m2,m3
[ x0 x1 m0 m1 ] [ x1 x2 x3 m1 m2 m3 ] [ x3 x0 m3 m0 ] σ̈ ( α̃(L)m0 (x0) ) σ̈ ( α̃(L)m1 (x1) ) σ̇ ( α̃(L)m2 (x2) ) σ̈ ( α̃(L)m3 (x3) ) W (L) m0k0 W (L) m1k1 W (L) m2k2 W (L) m3k3
+n−2L ∑
m0,m1,m2,m3
[ x0 x1 m0 m1 ] [ x1 x2 m1 m2 ] [ x2 x3 x0 m2 m3 m0 ] σ̈ ( α̃(L)m0 (x0) ) σ̈ ( α̃(L)m1 (x1) ) σ̈ ( α̃(L)m2 (x2) ) σ̇ ( α̃(L)m3 (x3) ) W (L) m0k0 W (L) m1k1 W (L) m2k2 W (L) m3k3
+n−2L ∑
m0,m1,m2,m3
[ x1 x2 m1 m2 ] [ x2 x3 m2 m3 ] [ x3 x0 x1 m3 m0 m1 ] σ̇ ( α̃(L)m0 (x0) ) σ̈ ( α̃(L)m1 (x1) ) σ̈ ( α̃(L)m2 (x2) ) σ̈ ( α̃(L)m3 (x3) ) W (L) m0k0 W (L) m1k1 W (L) m2k2 W (L) m3k3
+n−2L ∑
m0,m1,m2,m3
[ x0 x1 x2 m0 m1 m2 ] [ x2 x3 x0 m2 m3 m0 ] σ̈ ( α̃(L)m0 (x0) ) σ̇ ( α̃(L)m1 (x1) ) σ̈ ( α̃(L)m2 (x2) ) σ̇ ( α̃(L)m3 (x3) ) W (L) m0k0 W (L) m1k1 W (L) m2k2 W (L) m3k3
+n−2L ∑
m0,m1,m2,m3
[ x1 x2 x3 m1 m2 m3 ] [ x3 x0 x1 m3 m0 m1 ] σ̇ ( α̃(L)m0 (x0) ) σ̈ ( α̃(L)m1 (x1) ) σ̇ ( α̃(L)m2 (x2) ) σ̈ ( α̃(L)m3 (x3) ) W (L) m0k0 W (L) m1k1 W (L) m2k2 W (L) m3k3
+n−2L ∑
m0,m1,m2,m3
[ x0 x1 x2 x3 m0 m1 m2 m3 ] [ x3 x0 m3 m0 ] σ̈ ( α̃(L)m0 (x0) ) σ̇ ( α̃(L)m1 (x1) ) σ̇ ( α̃(L)m2 (x2) ) σ̈ ( α̃(L)m3 (x3) ) W (L) m0k0 W (L) m1k1 W (L) m2k2 W (L) m3k3
+n−2L ∑
m0,m1,m2,m3
[ x0 x1 m0 m1 ] [ x1 x2 x3 x0 m1 m2 m3 m0 ] σ̈ ( α̃(L)m0 (x0) ) σ̈ ( α̃(L)m1 (x1) ) σ̇ ( α̃(L)m2 (x2) ) σ̇ ( α̃(L)m3 (x3) ) W (L) m0k0 W (L) m1k1 W (L) m2k2 W (L) m3k3
+n−2L ∑
m0,m1,m2,m3
[ x1 x2 m1 m2 ] [ x2 x3 x0 x1 m2 m3 m0 m1 ] σ̇ ( α̃(L)m0 (x0) ) σ̈ ( α̃(L)m1 (x1) ) σ̈ ( α̃(L)m2 (x2) ) σ̇ ( α̃(L)m3 (x3) ) W (L) m0k0 W (L) m1k1 W (L) m2k2 W (L) m3k3
+n−2L ∑
m0,m1,m2,m3
[ x2 x3 m2 m3 ] [ x3 x0 x1 x2 m3 m0 m1 m2 ] σ̇ ( α̃(L)m0 (x0) ) σ̇ ( α̃(L)m1 (x1) ) σ̈ ( α̃(L)m2 (x2) ) σ̈ ( α̃(L)m3 (x3) ) W (L) m0k0 W (L) m1k1 W (L) m2k2 W (L) m3k3
+n−2L ∑
m,m1,m2
[ x0 x1 m m1 ] [ x1 x2 m1 m2 ] [ x2 x3 m2 m ] σ̇ ( α̃(L)m (x0) ) σ̈ ( α̃(L)m1 (x1) ) σ̈ ( α̃(L)m2 (x2) ) σ̇ ( α̃(L)m (x3) ) W (L) m1k1 W (L) m2k2 δk0k3
+n−2L ∑
m,m2,m3
[ x1 x2 m m2 ] [ x2 x3 m2 m3 ] [ x3 x0 m3 m ] σ̇ ( α̃(L)m (x0) ) σ̇ ( α̃(L)m (x1) ) σ̈ ( α̃(L)m2 (x2) ) σ̈ ( α̃(L)m3 (x3) ) W (L) m2k2 W (L) m3k3 δk0k1
+n−2L ∑
m,m3,m0
[ x0 x1 m0 m ] [ x2 x3 m m3 ] [ x3 x0 m3 m0 ] σ̈ ( α̃(L)m0 (x0) ) σ̇ ( α̃(L)m (x1) ) σ̇ ( α̃(L)m (x2) ) σ̈ ( α̃(L)m3 (x3) ) W (L) m0k0 W (L) m3k3 δk1k2
+n−2L ∑
m,m0,m1
[ x0 x1 m0 m1 ] [ x1 x2 m1 m ] [ x3 x0 m m0 ] σ̈ ( α̃(L)m0 (x0) ) σ̈ ( α̃(L)m1 (x1) ) σ̇ ( α̃(L)m (x2) ) σ̇ ( α̃(L)m (x3) ) W (L) m0k0 W (L) m1k1 δk2k3
+n−2L ∑
m,m1,m2
[ x0 x1 x2 m m1 m2 ] [ x2 x3 m2 m ] σ̇ ( α̃(L)m (x0) ) σ̇ ( α̃(L)m1 (x1) ) σ̈ ( α̃(L)m2 (x2) ) σ̇ ( α̃(L)m (x3) ) W (L) m1k1 W (L) m2k2 δk0k3
+n−2L ∑
m,m2,m3
[ x1 x2 x3 m m2 m3 ] [ x3 x0 m3 m ] σ̇ ( α̃(L)m (x0) ) σ̇ ( α̃(L)m (x1) ) σ̇ ( α̃(L)m2 (x2) ) σ̈ ( α̃(L)m3 (x3) ) W (L) m2k2 W (L) m3k3 δk0k1
+n−2L ∑
m,m3,m0
[ x0 x1 m0 m ] [ x2 x3 x0 m m3 m0 ] σ̈ ( α̃(L)m0 (x0) ) σ̇ ( α̃(L)m (x1) ) σ̇ ( α̃(L)m (x2) ) σ̇ ( α̃(L)m3 (x3) ) W (L) m0k0 W (L) m3k3 δk1k2
+n−2L ∑
m,m0,m1
[ x1 x2 m1 m ] [ x3 x0 x1 m m0 m1 ] σ̇ ( α̃(L)m0 (x0) ) σ̈ ( α̃(L)m1 (x1) ) σ̇ ( α̃(L)m (x2) ) σ̇ ( α̃(L)m (x3) ) W (L) m0k0 W (L) m1k1 δk2k3
+n−2L ∑
m,m1,m2
[ x0 x1 m m1 ] [ x1 x2 x3 m1 m2 m ] σ̇ ( α̃(L)m (x0) ) σ̈ ( α̃(L)m1 (x1) ) σ̇ ( α̃(L)m2 (x2) ) σ̇ ( α̃(L)m (x3) ) W (L) m1k1 W (L) m2k2 δk0k3
+n−2L ∑
m,m2,m3
[ x1 x2 m m2 ] [ x2 x3 x0 m2 m3 m ] σ̇ ( α̃(L)m (x0) ) σ̇ ( α̃(L)m (x1) ) σ̈ ( α̃(L)m2 (x2) ) σ̇ ( α̃(L)m3 (x3) ) W (L) m2k2 W (L) m3k3 δk0k1
+n−2L ∑
m,m3,m0
[ x2 x3 m m3 ] [ x3 x0 x1 m3 m0 m ] σ̇ ( α̃(L)m0 (x0) ) σ̇ ( α̃(L)m (x1) ) σ̇ ( α̃(L)m (x2) ) σ̈ ( α̃(L)m3 (x3) ) W (L) m0k0 W (L) m3k3 δk1k2
+n−2L ∑
m,m0,m1
[ x0 x1 x2 m0 m1 m ] [ x3 x0 m m0 ] σ̈ ( α̃(L)m0 (x0) ) σ̇ ( α̃(L)m1 (x1) ) σ̇ ( α̃(L)m (x2) ) σ̇ ( α̃(L)m (x3) ) W (L) m0k0 W (L) m1k1 δk2k3
+n−2L ∑
m,m1,m2
[ x0 x1 x2 x3 m m1 m2 m ] σ̇ ( α̃(L)m (x0) ) σ̇ ( α̃(L)m1 (x1) ) σ̇ ( α̃(L)m2 (x2) ) σ̇ ( α̃(L)m (x3) ) W
(L) m1k1 W (L) m2k2 δk0k3
+n−2L ∑
m,m2,m3
[ x1 x2 x3 x0 m m2 m3 m ] σ̇ ( α̃(L)m (x0) ) σ̇ ( α̃(L)m (x1) ) σ̇ ( α̃(L)m2 (x2) ) σ̇ ( α̃(L)m3 (x3) ) W
(L) m2k2 W (L) m3k3 δk0k1
+n−2L ∑
m,m3,m0
[ x2 x3 x0 x1 m m3 m0 m ] σ̇ ( α̃(L)m0 (x0) ) σ̇ ( α̃(L)m (x1) ) σ̇ ( α̃(L)m (x2) ) σ̇ ( α̃(L)m3 (x3) ) W
(L) m0k0 W (L) m3k3 δk1k2
+n−2L ∑
m,m0,m1
[ x3 x0 x1 x2 m m0 m1 m ] σ̇ ( α̃(L)m0 (x0) ) σ̇ ( α̃(L)m1 (x1) ) σ̇ ( α̃(L)m (x2) ) σ̇ ( α̃(L)m (x3) ) W
(L) m0k0 W (L) m1k1 δk2k3
+n−2L ∑ m,m′ [ x0 x1 m m′ ] [ x2 x3 m′ m ] σ̇ ( α̃(L)m (x0) ) σ̇ ( α̃ (L) m′ (x1) ) σ̇ ( α̃ (L) m′ (x2) ) σ̇ ( α̃(L)m (x3) ) δk0k1δk2k3
+n−2L ∑ m,m′ [ x1 x2 m m′ ] [ x3 x0 m′ m ] σ̇ ( α̃(L)m (x0) ) σ̇ ( α̃(L)m (x1) ) σ̇ ( α̃ (L) m′ (x2) ) σ̇ ( α̃ (L) m′ (x3) ) δk0k3δk1k2
Even though this is a very large formula one can notice that most terms are “rotation of each other”. Moreover, as n1, ..., nL−1 →∞, all terms containing either an Ψ(L), an Ω(L) or a Γ(L) vanish. For the remaining terms, we may replace the NTKs Θ(L) by their limit and as a result Ψ (L+1) k0,k1,k2,k3 (xi0 , xi1 , xi2 , xi3) converges to
n−2L ∑ m Θ(L)∞ (x0, x1)Θ (L) ∞ (x1, x2)Θ (L) ∞ (x2, x3)Θ (L) ∞ (x3, x0)σ̈ ( α̃(L)m (x0) ) σ̈ ( α̃(L)m (x1) ) σ̈ ( α̃(L)m (x2) ) σ̈ ( α̃(L)m (x3) ) W (L) mk0 W (L) mk1 W (L) mk2 W (L) mk3
+n−2L ∑ m Θ(L)∞ (x0, x1)Θ (L) ∞ (x1, x2)Θ (L) ∞ (x2, x3)σ̇ ( α̃(L)m (x0) ) σ̈ ( α̃(L)m (x1) ) σ̈ ( α̃(L)m (x2) ) σ̇ ( α̃(L)m (x3) ) W (L) mk1 W (L) mk2 δk0k3
+n−2L ∑ m Θ(L)∞ (x1, x2)Θ (L) ∞ (x2, x3)Θ (L) ∞ (x3, x0)σ̇ ( α̃(L)m (x0) ) σ̇ ( α̃(L)m (x1) ) σ̈ ( α̃(L)m (x2) ) σ̈ ( α̃(L)m (x3) ) W (L) mk2 W (L) mk3 δk0k1
+n−2L ∑ m Θ(L)∞ (x0, x1)Θ (L) ∞ (x2, x3)Θ (L) ∞ (x3, x0)σ̈ ( α̃(L)m (x0) ) σ̇ ( α̃(L)m (x1) ) σ̇ ( α̃(L)m (x2) ) σ̈ ( α̃(L)m (x3) ) W (L) mk0 W (L) mk3 δk1k2
+n−2L ∑ m Θ(L)∞ (x0, x1)Θ (L) ∞ (x1, x2)Θ (L) ∞ (x3, x0)σ̈ ( α̃(L)m (x0) ) σ̈ ( α̃(L)m (x1) ) σ̇ ( α̃(L)m (x2) ) σ̇ ( α̃(L)m (x3) ) W (L) mk0 W (L) mk1 δk2k3
+n−2L ∑ m Θ(L)∞ (x0, x1)Θ (L) ∞ (x2, x3)σ̇ ( α̃(L)m (x0) ) σ̇ ( α̃(L)m (x1) ) σ̇ ( α̃(L)m (x2) ) σ̇ ( α̃(L)m (x3) ) δk0k1δk2k3
+n−2L ∑ m Θ(L)∞ (x1, x2)Θ (L) ∞ (x3, x0)σ̇ ( α̃(L)m (x0) ) σ̇ ( α̃(L)m (x1) ) σ̇ ( α̃(L)m (x2) ) σ̇ ( α̃(L)m (x3) ) δk0k3δk1k2
And all these sums vanish as nL →∞ thanks to the prefactor n−2L , proving the vanishing of Ψ
(L+1) k0,k1,k2,k3 (xi0 , xi1 , xi2 , xi3) in the infinite width limit.
During training, the activations α̃ (L) m (x) and weights W (L) mk move at a rate of 1/√nL which induces a change to Ψ(L+1) of order n −3/2 L which vanishes in the infinite width limit.
D Orthogonality of I and S
From Lemma 2 and the vanishing of the tensor Γ(L) as proven in Lemma 2, we can easily prove the orthogonality of I and S of Proposition 5:
Proposition 5. For any loss C with BGOSS and σ ∈ C4b (R), we have uniformly over [0, T ]
lim nL−1→∞ · · · lim n1→∞ ‖IS‖F = 0.
As a consequence limnL−1→∞ · · · limn1→∞Tr ( [I + S] k ) − [ Tr ( Ik ) + Tr ( Sk )] = 0.
Proof. The Frobenius norm of IS is equal to ‖IS‖2F = ∥∥∥DYHC (DY )T (∇C · HY )∥∥∥2
F
= P∑ p1,p2=1  P∑ p=1 N∑ i1,i2=1 nL∑ k1,k2=1 ∂θp1 fθ,k1(xi1)c ′′ k1(xi1)∂θpfθ,k1(xi1)∂ 2 θp,θp3 fθ,k2(x2)(xi2)c ′ k2(xi2) 2
= N∑ i1,i2,i′1,i ′ 2=1 nL∑ k1,k2,k′1,k ′ 2=1 c′′k1(xi1)c ′′ k′1 (xi′1)c ′ k2(xi2)c ′ k′2 (xi′2)Θk1,k′1(xi1 , xi′1)Γk1,k2,k′2,k′1(xi1 , xi2 , xi′2 , xi′1)
and Γ vanishes as n1, ..., nL−1 →∞ by Lemma 2.
The k-th moment of the sum Tr (I + S) k is equal to the sum over all Tr (A1 · · ·Ak) for any word A1 . . . Ak of Ai ∈ {I, S}. The difference Tr ( [I + S] k ) − [ Tr ( Ik ) + Tr ( Sk )] is hence
equal to the sum over all mixed words, i.e. words A1 . . . Ak which contain at least one I and one S. Such words must contain two consecutive terms AmAm+1 one equal to I and the other equal to S. We can then bound the trace by
|Tr (A1 · · ·Ak)| ≤ ‖A1‖F · · · ‖Am−1‖F ‖AmAm+1‖F ‖Am+2‖F · · · ‖Ak‖F which vanishes in the infinite width limit because ‖I‖F and ‖S‖F are bounded and ‖AmAm+1‖F = ‖IS‖F vanishes.

Review summaries:
This paper uses the Neural Tangent Kernel (NTK) to presents an asymptotic analysis of the evolution of Hessian of the loss (w.r.t. model parameters) throughout training. The authors leverage the Neural Tangent Kernel to analyze the evolution of the Hessian of the loss w.r.t the model parameters. Specifically the authors show that as the width of the neural networks tend to infinity, the Hessian can be decomposed into two components: (1) one that reflects the initialization of the model parameters and reduces as training progresses; and (2) one that captures the principal directions of the data.