Predicting the Outputs of Finite Networks Trained with Noisy Gradients

1 INTRODUCTION . Deep neural networks ( DNNs ) have been rapidly advancing the state-of-the-art in machine learning , yet a complete analytic theory remains elusive . Recently , several exact results were obtained in the highly over-parameterized regime ( N →∞ where N denotes the width or number of channels for fully connected networks ( FCNs ) and convolutional neural networks ( CNNs ) , respectively ) ( Daniely et al. , 2016 ) . This facilitated the derivation of an exact correspondence with Gaussian Processes ( GPs ) known as the Neural Tangent Kernel ( NTK ) ( Jacot et al. , 2018 ) . The latter holds when highly over-parameterized DNNs are trained by gradient flow , namely with vanishing learning rate and involving no stochasticity . The NTK result has provided the first example of a DNN to GP correspondence valid after end-to-end DNN training . This theoretical breakthrough allows one to think of DNNs as inference problems with underlying GPs ( Rasmussen & Williams , 2005 ) . For instance , it provides a quantitative description of the generalization properties ( Cohen et al. , 2019 ; Rahaman et al. , 2018 ) and training dynamics ( Jacot et al. , 2018 ; Basri et al. , 2019 ) of DNNs . Roughly speaking , highly over-parameterized DNNs generalize well because they have a strong implicit bias to simple functions , and train well because low-error solutions in weight space can be reached by making a small change to the random values of the weights at initialization . Despite its novelty and importance , the NTK correspondence suffers from a few shortcomings : ( a ) Its deterministic training is qualitatively different from the stochastic one used in practice , which may lead to poorer performance when combined with a small learning rate ( Keskar et al. , 2016 ) . ( b ) It under-performs , often by a large margin , convolutional neural networks ( CNNs ) trained with SGD ( Arora et al. , 2019 ) . ( c ) Deriving explicit finite width corrections ( FWCs ) is challenging , as it requires solving a set of coupled ODEs ( Dyer & Gur-Ari , 2020 ; Huang & Yau , 2019 ) . Thus , there is a need for an extended theory for end-to-end trained deep networks which is valid for finite width DNNs . Our contribution is three-fold . First , we prove a correspondence between a DNN trained with noisy gradients and a Stochastic Process ( SP ) which at N → ∞ tends to the Neural Network Gaussian Process ( NNGP ) ( Lee et al. , 2018 ; Matthews et al. , 2018 ) . In these works , the NNGP kernel is determined by the distribution of the DNN weights at initialization which are i.i.d . random variables , whereas in our correspondence the weights are sampled across the stochastic training dynamics , drifting far away from their initial values . We call ours the NNSP correspondence , and show that it holds when the training dynamics in output space exhibit ergodicity . Second , we predict the outputs of trained finite-width DNNs , significantly improving upon the corresponding GP predictions . This is done by deriving leading FWCs which are found to scale with width as 1/N . The accuracy at which we can predict the empirical DNNs ’ outputs serves as a strong verification for our aforementioned ergodicity assumption . In the regime where the GP RMSE error scales as 1/n , we find that the leading FWC are a decaying function of n , and thus overall negligible . In the small n regime we find that the FWC is small and grows with n. We thus conclude that finite-width corrections are important for intermediate values of n ( Fig . 1 ) . Third , we propose an explanation for why finite CNNs trained on image classification tasks can outperform their infinite-width counterparts , as observed by Novak et al . ( 2018 ) . The key difference is that in finite CNNs weight sharing is beneficial . Our theory , which accounts for the finite width , quantifies this difference ( §4.2 ) . Overall , the NNSP correspondence provides a rich analytical and numerical framework for exploring the theory of deep learning , unique in its ability to incorporate finite over-parameterization , stochasticity , and depth . We note that there are several factors that make finite SGD-trained DNNs used in practice different from their GP counterparts , e.g . large learning rates , early stopping etc . ( Lee et al. , 2020 ) . Importantly , our framework quantifies the contribution of finite-width effects to this difference , distilling it from the contribution of these other factors . 1.1 RELATED WORK . The idea of leveraging the dynamics of the gradient descent algorithm for approximating Bayesian inference has been considered in various works ( Welling & Teh , 2011 ; Mandt et al. , 2017 ; Teh et al. , 2016 ; Maddox et al. , 2019 ; Ye et al. , 2017 ) . However , to the best of our knowledge , a correspondence with a concrete SP or a non-parametric model was not established nor was a comparison made of the DNN ’ s outputs with analytical predictions . Finite width corrections were studied recently in the context of the NTK correspondence by several authors . Hanin & Nica ( 2019 ) study the NTK of finite DNNs , but where the depth scales together with width , whereas we keep the depth fixed . Dyer & Gur-Ari ( 2020 ) obtained a finite N correction to the linear integral equation governing the evolution of the predictions on the training set . Our work differs in several aspects : ( a ) We describe a different correspondence under different a training protocol with qualitatively different behavior . ( b ) We derive relatively simple formulae for the outputs which become entirely explicit at large n. ( c ) We account for all sources of finite N corrections whereas finite N NTK randomness remained an empirical source of corrections not accounted for by Dyer & Gur-Ari ( 2020 ) . ( d ) Our formalism differs considerably : its statistical mechanical nature enables one to import various standard tools for treating randomness , ergodicity breaking , and taking into account non-perturbative effects . ( e ) We have no smoothness limitation on our activation functions and provide FWCs on a generic data point and not just on the training set . Another recent paper ( Yaida , 2020 ) studied Bayesian inference with weakly non-Gaussian priors induced by finite-N DNNs . Unlike here , there was no attempt to establish a correspondence with trained DNNs . The formulation presented here has the conceptual advantage of representing a distribution over function space for arbitrary training and test data , rather than over specific draws of data sets . This is useful for studying the large n behavior of learning curves , where analytical insights into generalization can be gained ( Cohen et al. , 2019 ) . A somewhat related line of work studied the mean field regime of shallow NNs ( Mei et al. , 2018 ; Chen et al. , 2020 ; Tzen & Raginsky , 2020 ) . We point out the main differences from our work : ( a ) The NN output is scaled differently with width . ( b ) In the mean field regime one is interested in the dynamics ( finite t ) of the distribution over the NN parameters in the form of a PDE of the Fokker-Planck type . In contrast , in our framework we are interested in the distribution over function space at equilibrium , i.e . for t→∞ . ( c ) It seems that the mean field analysis is tailored for two-layer fully-connected NNs and is hard to generalize to deeper nets or to CNNs . In contrast , our formalism generalizes to deeper fully-connected NNs and to CNNs as well , as we showed in section 4.2 . 2 THE NNSP CORRESPONDENCE . In this section we show that finite-width DNNs , trained in a specific manner , correspond to Bayesian inference using a non-parametric model which tends to the NNGP as N → ∞ . We first give a short review of Langevin dynamics in weight space as described by Neal et al . ( 2011 ) , Welling & Teh ( 2011 ) , which we use to generate samples from the posterior over weights . We then shift our perspective and consider the corresponding distribution over functions induced by the DNN , which characterizes the non-parametric model . Recap of Langevin-type dynamics - Consider a DNN trained with full-batch gradient descent while injecting white Gaussian noise and including a weight decay term , so that the discrete time dynamics of the weights read ∆wt : = wt+1 − wt = − ( γwt +∇wL ( zw ) ) dt+ √ 2Tdtξt ( 1 ) where wt is the vector of all network weights at time step t , γ is the strength of the weight decay , L ( zw ) is the loss as a function of the output zw , T is the temperature ( the magnitude of noise ) , dt is the learning rate and ξt ∼ N ( 0 , I ) . As dt → 0 these discrete-time dynamics converge to the continuous-time Langevin equation given by ẇ ( t ) = −∇w ( γ 2 ||w ( t ) || 2 + L ( zw ) ) + √ 2Tξ ( t ) with 〈ξi ( t ) ξj ( t′ ) 〉 = δijδ ( t− t′ ) , so that as t → ∞ the weights will be sampled from the equilibrium distribution in weight space , given by ( Risken & Frank , 1996 ) P ( w ) ∝ exp ( − 1 T ( γ 2 ||w||2 + L ( zw ) ) ) = exp ( − ( 1 2σ2w ||w||2 + 1 2σ2 L ( zw ) ) ) ( 2 ) The above equality holds since the equilibrium distribution of the Langevin dynamics is also the posterior distribution of a Bayesian neural network ( BNN ) with an i.i.d . Gaussian prior on the weights w ∼ N ( 0 , σ2wI ) . Thus we can map the hyper-parameters of the training to those of the BNN : σ2w = T/γ and σ 2 = T/2 . Notice that a sensible scaling for the weight variance at layer ` is σ2w , ` ∼ O ( 1/N ` −1 ) , thus the weight decay needs to scale as γ ` ∼ O ( N ` −1 ) . A transition from weight space to function space - We aim to move from a distribution over weight space Eq . 2 to a one over function space . Namely , we consider the distribution of zw ( x ) implied by the above P ( w ) where for concreteness we consider a DNN with a single scalar output zw ( x ) ∈ R on a regression task with data { ( xα , yα ) } nα=1 ⊂ Rd × R. Denoting by P [ f ] the induced measure on function space we formally write P [ f ] = ∫ dwδ [ f − zw ] P ( w ) ∝ e− 1 2σ2 L [ f ] ∫ dwe − 1 2σ2w ||w||2 δ [ f − zw ] ( 3 ) where ∫ dw denotes an integral over all weights and we denote by δ [ f − zw ] a delta-function in function-space . As common in path-integrals or field-theory formalism ( Schulman , 2012 ) , such a delta function is understood as a limit procedure where one chooses a suitable basis for function space , trims it to a finite subset , treats δ [ f − zw ] as a product of regular delta-functions , and at the end of the computation takes the size of the subset to infinity . To proceed we decompose the posterior over functions Eq . 3 as P [ f ] ∝ e− 1 2σ2 L [ f ] P0 [ f ] where the prior over functions is P0 [ f ] ∝ ∫ dwe − 1 2σ2w ||w||2 δ [ f − zw ] . The integration over weights now obtains a clear meaning : it yields the distribution over functions induced by a DNN with i.i.d . random weights chosen according to the prior P0 ( w ) ∝ e − 1 2σ2w ||w||2 . Thus , we can relate any correlation function in function space and weight space , for instance ( Df is the integration measure over function space ) ∫ DfP0 [ f ] f ( x ) f ( x′ ) = ∫ Df ∫ dwP0 ( w ) δ [ f−zw ] f ( x ) f ( x′ ) = ∫ dwP0 ( w ) zw ( x ) zw ( x ′ ) ( 4 ) As noted by Cho & Saul ( 2009 ) , for highly over-parameterized DNNs the r.h.s . of 4 equals the kernel of the NNGP associated with this DNN , K ( x , x′ ) . Moreover P0 [ f ] tends to a Gaussian and can be written as P0 [ f ] ∝ exp ( −1 2 ∫ dµ ( x ) dµ ( x′ ) f ( x ) K−1 ( x , x′ ) f ( x′ ) ) +O ( 1/N ) ( 5 ) where µ ( x ) is the measure of the input space , and the O ( 1/N ) scaling of the finite-N correction will be explained in §3 . If we now plug 5 in 3 , take the loss to be the total square error1 L [ f ] =∑n α=1 ( yα − f ( xα ) ) 2 , and take N →∞ we have that the posterior P [ f ] is that of a GP . Assuming ergodicity , one finds that training-time averaged output of the DNN is given by the posterior mean of a GP , with measurement noise2 equal to σ2 = T/2 and a kernel given by the NNGP of that DNN . We refer to the above expressions for P0 [ f ] and P [ f ] describing the distribution of outputs of a DNN trained according to our protocol – the NNSP correspondence . Unlike the NTK correspondence , the kernel which appears here is different and no additional initialization dependent terms appear ( as should be the case since we assumed ergodicity ) . Furthermore , given knowledge of P0 [ f ] at finite N , one can predict the DNN ’ s outputs at finite N . Henceforth , we refer to P0 [ f ] as the prior distribution , as it is the prior distribution of a DNN with random weights drawn from P0 ( w ) . Evidence supporting ergodicity - Our derivation relies on the ergodicity of the dynamics . Ergodicity is in general hard to prove rigorously in non-convex settings , and thus we must revert to heuristics . The most robust evidence of ergodicity in function space is the high level of accuracy of our analytical expressions w.r.t . to our numerical results . This is a self-consistency argument : we assume ergodicity in order to derive our analytical results and then indeed find that they agree very well with the experiment , thus validating our original assumption . Another indicator of ergodicity is a small auto-correlation time ( ACT ) of the dynamics . Although short ACT does not logically imply ergodicity ( in fact , the converse is true : exponentially long ACT implies non-ergodic dynamics ) . However , the empirical ACT gives a lower bound on the true correlation time of the dynamics . In our framework , it is sufficient that the dynamics of the outputs zw be ergodic , even if the dynamics of the weights converge much slower to an equilibrium distribution . Indeed , we have found that the ACTs of the outputs are considerably smaller than those of the weights ( see Fig . 2b ) . Full ergodicity may be too strong of a condition and we don ’ t really need it for our purposes , since we are mainly interested in collecting statistics that will allow us to accurately compute the posterior mean of the distribution in function space . Thus , a weaker condition which is sufficient here is ergodicity in the mean ( see App . F ) , and we believe our self-consistent argument above demonstrates that it holds . In a related manner , optimizing the train loss can be seen as an attempt to find a solution to n constraints using far more variables ( roughly M · N2 where M is the number of layers ) . From a different angle , in a statistical mechanical description of satisfiability problems , one typically expects ergodic behavior when the ratio of the number of variables to the number of constraints becomes much larger than one ( Gardner & Derrida , 1988 ) .

Summary: This paper shows a correspondence between deep neural networks (DNN) trained with noisy gradients and NNGP. It provides a general analytical form for the finite width correction (FWC) for NNSP expanding around NNGP. Finally, it argues that this FWC can be used to explain why finite width CNNs can improve the performance relative to their GP counterparts on image classification tasks.