Asymptotics of Wide Networks from Feynman Diagrams

1 INTRODUCTION . Neural networks achieve remarkable performance on a wide array of machine learning tasks , yet a complete analytic understanding of deep networks remains elusive . One promising approach is to consider the large width limit , in which the number of neurons in one or several layers is taken to be large . In this limit one can use a mean-field approach to better understand the network ’ s properties at initialization ( Neal ( 1996 ) ; Lee et al . ( 2018 ) ) , as well as its training dynamics ( Daniely ( 2017 ) ; Jacot et al . ( 2018 ) ) . Additional related works are cited below . Suppose that f ( x ) is the network function evaluated at an input x . Let us denote the vector of model parameters by θ , whose elements are initially chosen to be i.i.d . Gaussian . In this work we consider a class of functions we call correlation functions , obtained by taking the ensemble averages of f , its products , and its derivatives with respect to the parameters θ , evaluated on arbitrary inputs . Here are a few examples of correlation functions . Eθ [ f ( x1 ) f ( x2 ) ] , ∑ µ Eθ [ ∂f ( x1 ) ∂θµ ∂f ( x2 ) ∂θµ ] , ∑ µ , ν Eθ [ ∂f ( x1 ) ∂θµ ∂f ( x2 ) ∂θν ∂2f ( x3 ) ∂θµ∂θν f ( x4 ) ] . ( 1 ) Correlation functions often show up in the study of wide networks . For example , the first correlation function in ( 1 ) plays a central role in the Gaussian Process picture of wide networks ( Lee et al . ( 2018 ) ) , and has been used to diagnose signal propagation in wide networks ( Pennington et al . ( 2017 ) ) . The second example in ( 1 ) is the ensemble average of the Neural Tangent Kernel ( NTK ) , which controls the evolution of wide networks under gradient flow ( Jacot et al . ( 2018 ) ) , and the third example shows up when computing the time derivative of the NTK with MSE loss . While correlation functions can be computed analytically in some special cases ( Cho & Saul ( 2009 ) ) , they are not analytically tractable in general . In this work , we present a method for bounding the asymptotic behavior of such functions at large width . In many cases of interest , the bound we obtain is tight . Derivation of the method relies on Feynman diagrams ( Feynman ( 1949 ) ) , a technique for calculating multivariate Gaussian integrals , and specifically on the ’ t Hooft expansion ( ’ t Hooft ( 1974 ) ) . However , applying the method is straightforward and does not require any knowledge of Feynman diagrams . Motivation . The method presented here is a general-purpose tool that can be added to a researcher ’ s toolbox . It allows a researcher to quickly identify ways in which network behavior and training ∗Authors listed alphabetically . dynamics simplify at large width , for example by deriving which quantities vanish in this limit . As we demonstrate , using this method one is able to dramatically cut down on the amount of effort required to derive several key results related to wide networks , as well as to go beyond the infinite width limit and investigate ordinary neural networks . Our contribution .. 1 . We present a general method for bounding the asymptotic behavior of correlation functions . The method is an adaptation of Feynman diagrams to the case of wide neural networks . The adaptation involves a novel treatment of derivatives of the network function , an element that is not present in the original theoretical physics formulation . 2 . We apply the method to the study of wide network evolution under gradient descent . We improve on existing results for gradient flow ( Jacot et al . ( 2018 ) ) by deriving tighter bounds , and extending the analysis to the case of stochastic gradient descent ( SGD ) . Going beyond the infinite-width limit , we present a formalisn for deriving finite-width corrections to network evolution , and present explicit formulas for the first order correction . To our knowledge , this is the first time this correction has been calculated . 3 . As additional applications of our method , in Appendix E.2 we show that in the large width limit the SGD updates are linear in the learning rate , and in Appendix E.3 we discuss finite width corrections to the spectrum of the Hessian . Limitations of our approach . The main result of this paper is a conjecture . We test our predictions extensively using numerical experiments , and prove the conjecture in some cases , but we do not have a proof that applies to all the cases we tested , including for deep networks with general non-linearities . Furthermore , our method can only be used to derive asymptotic bounds at large width ; it does not produce the width-independent coefficient , which is often of interest . Related work . For additional works on wide networks , including relating them to Gaussian processes , see de G. Matthews et al . ( 2018 ) ; Novak et al . ( 2019 ) ; Garriga-Alonso et al . ( 2019 ) ; Yang ( 2019 ) ; Pennington & Bahri ( 2017 ) ; Pennington & Worah ( 2018 ) ; Schoenholz et al . ( 2017 ) ; Xiao et al . ( 2018 ) ; Chen et al . ( 2018 ) ; Daniely et al . ( 2016 ) ; Lee et al . ( 2019 ) ; Du et al . ( 2018b ; a ) ; Allen-Zhu et al . ( 2018 ) . For additional works discussing the training dynamics of wide networks see Geiger et al . ( 2019a ) ; Arora et al . ( 2019 ) . For a previous use of diagrams in this context , see Pennington & Worah ( 2017 ) . The Neural Tangent Hierarchy presented in Huang & Yau ( 2019 ) , published during the completion of this version , has significant overlap with the recursive differential equations ( 10 ) presented below . The rest of the paper is organized as follows . In Section 2 we present our main conjecture and supporting evidence . In Section 3 we apply the method to gradient descent evolution of wide networks , and in Section 4 we present details on Feynman diagrams , which is the basic technique used in our proofs . We conclude with a Discussion . Proofs , additional applications , and details can be found in the Appendices . Note : An earlier version of this work appeared in the ICML 2019 workshop , Theoretical Physics for Deep Learning ( Dyer & Gur-Ari ( 2019 ) ) . 2 CORRELATION FUNCTION ASYMPTOTICS . In this section we present our main result : a method for computing asymptotic bounds on correlation functions of wide networks . We present the result as a conjecture , supported by analytic and empirical evidence . 2.1 NOTATION . Let f ( x ) ∈ R be the network output of a deep network with d hidden layers and input x ∈ RDin , defined by f ( x ) = n−1/2V Tσ ( n−1/2W d−1 · · ·σ ( n−1/2W 1σ ( Ux ) ) ) . ( 2 ) Here U ∈ Rn×Din , V ∈ Rn , W 1 , . . . , W d−1 are weight matrices of dimension n , and σ : R→ R is the non-linearity.1 We denote the vector of all model parameters by θ . At initialization , the elements of θ are independent Gaussian variables , with each element θµ ∼ N ( 0 , 1 ) . The corresponding distribution of θ is denoted by P0.2 Let us now define correlation functions , the class of functions that is the focus of this work . These functions involve derivative tensors of the network function . We denote the rank-k derivative tensor by Tµ1 ... µk ( x ; f ) : = ∂ kf ( x ) /∂θµ1 · · · ∂θµk . For k = 0 we define T ( x ; f ) : = f ( x ) , and still refer to this as a derivative tensor for consistency . Definition 1 . A correlation function is the expectation value of a product of derivative tensors , evaluated at arbitrary inputs , where the tensor indices are summed in pairs over all the model parameters . A general correlation function C takes the form C ( x1 , . . . , xm ) : = ∑ µ1 , ... , µkm ∆ ( π ) µ1 ... µkmEθ [ Tµ1 ... µk1 ( x1 ) Tµk1+1 ... µk2 ( x2 ) · · ·Tµkm−1+1 ... µkm ( xm ) ] . ( 3 ) Here , 0 ≤ k1 ≤ · · · ≤ km−1 ≤ km are integers,3 m and km are even , π ∈ Skm is a permutation , and ∆ ( π ) µ1 ... µkm = δµπ ( 1 ) µπ ( 2 ) · · · δµπ ( km−1 ) µπ ( km ) . We use δ to denote the Kronecker delta . If two derivative tensors in a correlation function have matching indices that are summed over , we say that they are contracted . For example , the correlation function ∑ µ Eθ [ ∂f ( x1 ) /∂θµ · ∂f ( x2 ) /∂θµ ] has one pair of contracted tensors . See ( 1 ) for additional examples of correlation functions . 2.2 ASYMPTOTIC BOUNDS ON WIDE NETWORKS . We now present our main conjecture , which allows us to place asymptotic bounds on general correlation functions of wide networks . Conjecture 1 . Let C ( x1 , . . . , xm ) be a correlation function . The cluster graph GC ( V , E ) of C is a graph with vertices V = { v1 , . . . , vm } , where vi = T ( xi ) for all i , and edges E = { ( vi , vj ) | ( T ( xi ) , T ( xj ) ) contracted in C } . Suppose that the cluster graph GC has ne connected components with an even size ( even number of vertices ) , and no components of odd size . Then C ( x1 , . . . , xm ) = O ( nsC ) , where sC = ne + no 2 − m 2 . ( 4 ) We will refer to the connected components of a cluster graph GC as the clusters of C. Table 1 lists examples of bounds derived using the Conjecture for several correlation functions . The intuition behind Conjecture 1 comes from the following result for deep linear networks . Theorem 1 . Conjecture 1 holds for correlation functions of networks with linear activations , including for deep linear networks with biases ( affine transformations ) . Let us discuss the intuition behind this theorem . Computing correlation functions of deep linear networks amounts to evaluating Gaussian integrals with polynomial integrands in θ . One can evaluate such integrals using Isserlis ’ theorem , which tells us how to express moments of multivariate Gaussian variables in terms of their second moments . For example , given centered Gaussian variables z1 , ... , z4 , Ez [ z1z2z3z4 ] = Ez [ z1z2 ] Ez [ z3z4 ] + Ez [ z1z3 ] Ez [ z2z4 ] + Ez [ z1z4 ] Ez [ z2z3 ] . ( 5 ) Therefore , correlation functions of deep linear networks can be expressed in terms of the covariances Eθ [ UiαUjβ ] = δijδαβ , Eθ [ ViVj ] = δij , and Eθ [ W ( l ) ij W ( l ) kl ] = δikδjl . For example , for a deep linear network with 2 hidden layers , we have Eθ [ f ( x1 ) f ( x2 ) ] = 1 n2 Eθ [ V TWUx1V TWUx2 ] = xT1 x2 n2 n∑ i , k δikδik n∑ j , l δjlδjl = x T 1 x2 . ( 6 ) 1We take all layers widths to be equal to n for simplicity , but our results hold in the more general case where all widths scale linearly with n. 2We use µ , ν , . . . to denote θ indices , i , j , . . . to denote individual weight matrix and weight vector indices , and α , β , . . . for input dimension indices . 3When ka = ka−1 , the tensor T ( xa ) has no derivatives . Every correlation function of a deep linear network can be similarly reduced to sums over products of Kronecker delta functions and width-independent functions of the inputs . The asymptotic large width behavior is determined by these sums over delta functions , which are tedious to compute by hand . Feynman diagrams are a graphical tool for computing these sums , allowing us to obtain the asymptotic behavior with minimal effort . This tool , which is described in detail in Section 4 , is used to prove Theorem 1 . For networks with non-linear activations we further show the following Theorem 2 . Conjecture 1 holds for ( 1 ) networks with ReLU activations , where all inputs are set to be equal , and for ( 2 ) networks with one hidden layer and smooth activation . For case ( 1 ) , the idea behind the proof is to put an asymptotic bound on the ReLU network in terms of a corresponding deep linear network . For case ( 2 ) , the basic idea is that each network function contains a single sum over the width , and by keeping track of these sums using Feynman diagrams we are able to bound the asymptotic behavior . We refer the reader to Appendix C for details .

Summary: The paper investigates the asymptotic behavior of correlation functions of wide networks. The main part of the work revolves around Conjecture 1 which assesses how the correlation function scales depending on the number of connected components of even/odd size of the constructed cluster graph. The authors then study training dynamics of wide networks under gradient flow and (stochastic) gradient descent and present empirical evidence to support their claims.