Finite Depth and Width Corrections to the Neural Tangent Kernel

1 INTRODUCTION . Modern neural networks are typically overparameterized : they have many more parameters than the size of the datasets on which they are trained . That some setting of parameters in such networks can interpolate the data is therefore not surprising . But it is a priori unexpected that not only can such interpolating parameter values can be found by stochastic gradient descent ( SGD ) on the highly non-convex empirical risk but also that the resulting network function generalizes to unseen data . In an overparameterized neural network N ( x ) the individual parameters can be difficult to interpret , and one way to understand training is to rewrite the SGD updates ∆θp = − λ ∂L ∂θp , p = 1 , . . . , P of trainable parameters θ = { θp } Pp=1 with a loss L and learning rate λ as kernel gradient descent updates for the values N ( x ) of the function computed by the network : ∆N ( x ) = − λ 〈KN ( x , · ) , ∇L ( · ) 〉 = − λ |B| |B|∑ j=1 KN ( x , xj ) ∂L ∂N ( xj , yj ) . ( 1 ) Here B = { ( x1 , y1 ) , . . . , ( x|B| , y|B| ) } is the current batch , the inner product is the empirical ` 2 inner product over B , and KN is the neural tangent kernel ( NTK ) : KN ( x , x ′ ) = P∑ p=1 ∂N ∂θp ( x ) ∂N ∂θp ( x′ ) . Relation ( 1 ) is valid to first order in λ . It translates between two ways of thinking about the difficulty of neural network optimization : ( i ) The parameter space view where the loss L , a complicated function of θ ∈ R # parameters , is minimized using gradient descent with respect to a simple ( Euclidean ) metric ; ( ii ) The function space view where the loss L , which is a simple function of the network mapping x 7→ N ( x ) , is minimized over the manifold MN of all functions representable by the architecture of N using gradient descent with respect to a potentially complicated Riemannian metric KN on MN . A remarkable observation of Jacot et al . ( 2018 ) is thatKN simplifies dramatically when the network depth d is fixed and its width n tends to infinity . In this setting , by the universal approximation theorem ( Cybenko , 1989 ; Hornik et al. , 1989 ) , the manifold MN fills out any ( reasonable ) ambient linear space of functions . The results in Jacot et al . ( 2018 ) then show that the kernelKN in this limit is frozen throughout training to the infinite width limit of its average E [ KN ] at initialization , which depends on the depth and non-linearity of N but not on the dataset . This mapping between parameter space SGD and kernel gradient descent for a fixed kernel can be viewed as two separate statements . First , at initialization , the distribution of KN converges in the infinite width limit to the delta function on the infinite width limit of its mean E [ KN ] . Second , the infinite width limit of SGD dynamics in function space is kernel gradient descent for this limiting mean kernel for any fixed number of SGD iterations . As long as the loss L is well-behaved with respect to the network outputs N ( x ) and E [ KN ] is non-degenerate in the subspace of function space given by values on inputs from the dataset , SGD for infinitely wide networks will converge with probability 1 to a minimum of the loss . Further , kernel method-based theorems show that even in this infinitely overparameterized regime neural networks will have non-vacuous guarantees on generalization ( Wei et al. , 2018 ) . However , as ( Wei et al. , 2018 ) shows , the regularized neural networks at finite width can have better sample complexity the corresponding infinite width kernel method . But replacing neural network training by gradient descent for a fixed kernel in function space is also not completely satisfactory for several reasons . First , it suggests that no feature learning occurs during training for infinitely wide networks in the sense that the kernel E [ KN ] ( and hence its associated feature map ) is data-independent . In fact , empirically , networks with finite but large width trained with initially large learning rates often outperform NTK predictions at infinite width ( Arora et al. , 2019 ) . One interpretation is that , at finite width , KN evolves through training , learning datadependent features not captured by the infinite width limit of its mean at initialization . In part for such reasons , it is important to study both empirically and theoretically finite width corrections to KN . Another interpretation is that the specific NTK scaling of weights at initialization ( Chizat & Bach , 2018b ; a ; Mei et al. , 2019 ; 2018 ; Rotskoff & Vanden-Eijnden , 2018a ; b ) and the implicit small learning rate limit ( Li et al. , 2019 ) obscure important aspects of SGD dynamics . Second , even in the infinite width limit , although KN is deterministic , it has no simple analytical formula for deep networks , since it is defined via a layer by layer recursion . In particular , the exact dependence , even in the infinite width limit , of KN on network depth is not well understood . Moreover , the joint statistical effects of depth and width on KN in finite size networks remain unclear , and the purpose of this article is to shed light on the simultaneous effects of depth and width on KN for finite but large widths n and any depth d. Our results apply to fully connected ReLU networks at initialization for which our main contributions are : 1 . In contrast to the regime in which the depth d is fixed but the width n is large , KN is not approximately deterministic at initialization so long as d/n is bounded away from 0 . Specifically , for a fixed input x the normalized on-diagonal second moment ofKN satisfies E [ KN ( x , x ) 2 ] E [ KN ( x , x ) ] 2 ' exp ( 5d/n ) ( 1 +O ( d/n2 ) ) . Thus , when d/n is bounded away from 0 , even when both n , d are large , the standard deviation of KN ( x , x ) is at least as large as its mean , showing that its distribution at initialization is not close to a delta function . See Theorem 1 . 2 . Moreover , when L is the square loss , the average of the SGD update ∆KN ( x , x ) to KN ( x , x ) from a batch of size one containing x satisfies E [ ∆KN ( x , x ) ] E [ KN ( x , x ) ] ' d 2 nn0 exp ( 5d/n ) ( 1 +O ( d/n2 ) ) , where n0 is the input dimension . Therefore , if d2/nn0 > 0 , the NTK will have the potential to evolve in a data-dependent way . Moreover , if n0 is comparable to n and d/n > 0 then it is possible that this evolution will have a well-defined expansion in d/n . See Theorem 2 . In both statements above , ' means is bounded above and below by universal constants . We emphasize that our results hold at finite d , n and the implicit constants in both ' and in the error terms O ( d/n2 ) are independent of d , n.Moreover , our precise results , stated in §2 below , hold for networks with variable layer widths . We have denoted network width by n only for the sake of exposition . The appropriate generalization of d/n to networks with varying layer widths is the parameter β : = d∑ i=1 1 nj , which in light of the estimates in ( 1 ) and ( 2 ) plays the role of an inverse temperature . 1.1 PRIOR WORK . A number of articles ( Bietti & Mairal , 2019 ; Dyer & Gur-Ari , 2019 ; Lee et al. , 2019 ; Yang , 2019 ) have followed up on the original NTK work Jacot et al . ( 2018 ) . Related in spirit to our results is the work Dyer & Gur-Ari ( 2019 ) , which uses Feynman diagrams to study finite width corrections to general correlations functions ( and in particular the NTK ) . The most complete results obtained by Dyer & Gur-Ari ( 2019 ) are for deep linear networks but a number of estimates hold general non-linear networks as well . The results there , like in essentially all previous work , fix the depth d and let the layer widths n tend to infinity . In contrast , our results ( as well as those of Hanin ( 2018 ) ; Hanin & Nica ( 2018 ) ; Hanin & Rolnick ( 2018 ) ) , do not treat d as a constant , suggesting that the 1/n expansions ( e.g . in Dyer & Gur-Ari ( 2019 ) ) can be promoted to d/n expansions . Also , the sum-over-path approach to studying correlation functions in randomly initialized ReLU nets was previously taken up for the forward pass by Hanin & Rolnick ( 2018 ) and for the backward pass by Hanin ( 2018 ) and Hanin & Nica ( 2018 ) . We also point the reader to Theorems 3.1 and 3.2 in Arora et al . ( 2019 ) , which provide quantitative rates of convergence for both the neural tangent kernel and the resulting full optimization trajectory of neural networks at large but finite width ( and fixed depth ) . 2 FORMAL STATEMENT OF RESULTS . Consider a ReLU network N with input dimension n0 , hidden layer widths n1 , . . . , nd−1 , and output dimension nd = 1 . We will assume that the output layer of N is linear and initialize the biases in N to zero . Therefore , for any input x ∈ Rn0 , the network N computes N ( x ) = x ( d ) given by x ( 0 ) = x , y ( i ) : = Ŵ ( i ) x ( i−1 ) , x ( i ) : = ReLU ( y ( i ) ) , i = 1 , . . . , d , ( 2 ) where for i = 1 , . . . , d− 1 Ŵ ( d ) : = ( 1/ni−1 ) −1/2W ( i ) , Ŵ ( i ) : = ( 2/ni−1 ) −1/2W ( i ) , W ( i ) α , β ∼ µ i.i.d. , ( 3 ) and µ is a fixed probability measure on R that we assume has a density with respect to Lebesgue measure and satisfies : µ is symmetric around 0 , Var [ µ ] = 1 , ∫ ∞ −∞ x4dµ ( x ) = µ4 < ∞ . ( 4 ) The three assumptions in ( 4 ) hold for virtually all standard network initialization schemes with the exception of orthogonal weight initialization . But believe our results extend hold also for this case but not do take up this issue . The on-diagonal NTK is KN ( x , x ) : = d∑ j=1 nj−1∑ α=1 nj∑ β=1 ( ∂N ∂W ( j ) α , β ( x ) ) 2 + d∑ j=1 nj∑ β=1 ( ∂N ∂b ( j ) β ( x ) ) 2 , ( 5 ) and we emphasize that although we have initialized the biases to zero , they are not removed from the list of trainable parameters . Our first result is the following : Theorem 1 ( Mean and Variance of NKT on Diagonal at Init ) . We have E [ KN ( x , x ) ] = d ( 1 2 + ‖x‖22 n0 ) . Moreover , we have that E [ KN ( x , x ) 2 ] is bounded above and below by universal constants times exp ( 5β ) d2 ‖x‖42n20 + d ‖x‖ 2 2 n0 d∑ j=1 e −5 ∑j i=1 1 ni + d∑ i , j=1 i≤j e −5 ∑j i=1 1 ni , β = d∑ i=1 1 ni times a multiplicative error ( 1 +O ( ∑d i=1 1 n2i ) ) . In particular , if all the hidden layer widths are equal ( i.e . ni = n , for i = 1 , . . . , d− 1 ) , we have E [ KN ( x , x ) 2 ] E [ KN ( x , x ) ] 2 ' exp ( 5β ) ( 1 +O ( β/n ) ) , β = d/n , where f ' g means f is bounded above and below by universal constants times g. This result shows that in the deep and wide double scaling limit ni , d→∞ , 0 < lim ni , d→∞ d∑ i=1 1 ni < ∞ , the NTK does not converge to a constant in probability . This is contrast to the wide and shallow regime where ni →∞ and d < ∞ is fixed . Our next result shows that when L is the square loss KN ( x , x ) is not frozen during training . To state it , fix an input x ∈ Rn0 to N and define ∆KN ( x , x ) to be the update from one step of SGD with a batch of size 1 containing x ( and learning rate λ ) . Theorem 2 ( Mean of Time Derivative of NTK on Diagonal at Init ) . We have that E [ λ−1∆KN ( x , x ) ] is bounded above and below by universal constants times‖x‖42n20 d∑ i1 , i2=1 ii < i2 i2−1∑ ` =i1 e −5/n ` −6 ∑ ` i=i1 1 ni n ` + ‖x‖22 n0 d∑ ii , i2=1 i1 < i2 e −5 ∑i1 i=1 1 ni i2−1∑ ` =i1 e −6 ∑ ` −1 i=i1+1 1 ni n ` exp ( 5β ) times a multiplicative error of size ( 1 +O ( ∑d i=1 1 n2i ) ) , where β = ∑d i=1 1/ni , as in Theorem 1 . In particular , if all the hidden layer widths are equal ( i.e . ni = n , for i = 1 , . . . , d− 1 ) , we find E [ ∆KN ( x , x ) ] E [ KN ( x , x ) ] ' dβ n0 exp ( 5β ) ( 1 +O ( β/n ) ) , β = d/n . Observe that when d is fixed and ni = n → ∞ , the pre-factor in front of exp ( 5β ) scales like 1/n . This is in keeping with the results from Dyer & Gur-Ari ( 2019 ) and Jacot et al . ( 2018 ) . Moreover , it shows that if d , n , n0 grow in any way so that dβ/n0 = d2/nn0 → 0 , the update ∆KN ( x , x ) to KN ( x , x ) from the batch { x } at initialization will have mean 0 . It is unclear whether this will be true also for larger batches and when the arguments of KN are not equal . In contrast , if ni ' n and β = d/n is bounded away from 0 , ∞ , and the n0 is proportional to d , the average update E [ ∆KN ( x , x ) ] has the same order of magnitude as E [ KN ( x ) ] .

Summary: The paper investigates a novel infinite width limit taking depth to infinite at the same time. This is beyond conventional theoretical studies for infinite width networks where depth is kept finite when the width is taken to be infinite. The main object that paper studies is the neural tangent kernel which is of great interest to the theoretical deep learning community as it describes gradient descent dynamics in a tractable way.