Self-Consistent Dynamical Field Theory of Kernel Evolution in Wide Neural Networks
Anonymous Author(s)
AfﬁliationAddress email
Abstract
We analyze feature learning in inﬁnite-width neural networks trained with gradientﬂow through a self-consistent dynamical ﬁeld theory.
We construct a collectionof deterministic dynamical order parameters which are inner-product kernels forhidden unit activations and gradients in each layer at pairs of time points, providinga reduced description of network activity through training.
These kernel orderparameters collectively deﬁne the hidden layer activation distribution, the evolutionof the neural tangent kernel, and consequently output predictions.
We show that theﬁeld theory derivation recovers the recursive stochastic process of inﬁnite-widthfeature learning networks obtained from Yang & Hu with Tensor Programs [1].
For deep linear networks, these kernels satisfy a set of algebraic matrix equations.
For nonlinear networks, we provide an alternating sampling procedure to selfconsistently solve for the kernel order parameters.
We provide comparisons of theself-consistent solution to various approximation schemes including the static NTKapproximation, gradient independence assumption, and leading order perturbationtheory, showing that each of these approximations can break down in regimeswhere general self-consistent solutions still provide an accurate description.
Lastly,we provide experiments in more realistic settings which demonstrate that the lossand kernel dynamics of CNNs at ﬁxed feature learning strength is preserved acrossdifferent widths on a CIFAR classiﬁcation task.
Introduction
Deep learning has emerged as a successful paradigm for solving challenging machine learning andcomputational problems across a variety of domains [2, 3].
However, theoretical understandingof the training and generalization of modern deep learning methods lags behind current practice.
Ideally, a theory of deep learning would be analytically tractable, efﬁciently computable, capableof predicting network performance and internal features that the network learns, and interpretablethrough a reduced description involving desirably initialization-independent quantities.
Several recent theoretical advances have fruitfully considered the idealization of wide neural networks ,where the number of hidden units in each layer is taken to be large.
Under certain parameterization,
Bayesian neural networks and gradient descent trained networks converge to gaussian processes(NNGPs)
[4–6] and neural tangent kernel (NTK) machines [7–9] in their respective inﬁnite-widthlimits.
These limits provide both analytic tractability as well as detailed training and generalizationanalysis [10–17].
However, in this limit, with these parameterizations, data representations are ﬁxedand do not adapt to data, termed the lazy regime of NN training, to contrast it from the rich regimewhere NNs signiﬁcantly alter their internal features while ﬁtting the data [18, 19].
The fact that therepresentation of data is ﬁxed renders these kernel-based theories incapable of explaining featurelearning, an ingredient which is crucial to the success of deep learning in practice [20, 21].
Thus,alternative theories capable of modeling feature learning dynamics are needed
Recently developed alternative parameterizations such as the mean ﬁeld [22] and the µP [1] parameterizations allow feature learning in inﬁnite-width NNs trained with gradient descent.
Using the
Tensor Programs framework, Yang & Hu identiﬁed a stochastic process that describes the evolutionof preactivation features in inﬁnite-width µP NNs [1].
In this work, we study an equivalent parameterization to µP with self-consistent dynamical mean ﬁeld theory (DMFT) and recover the stochasticprocess description of inﬁnite NNs using this alternative technique.
In the same large width scaling,we include a scalar parameter  0 that allows smooth interpolation between lazy and rich behavior[18].
We provide a new computational procedure to sample this stochastic process and demonstrateits predictive power for wide NNs.
Our novel contributions in this paper are the following:
We develop a path integral formulation of gradient ﬂow dynamics in inﬁnite-width networks in thefeature learning regime.
Our parameterization  0 includes a scalar parameter allows interpolationbetween rich and lazy regimes and comparison to perturbative methods.2.
Using a stationary action argument, we identify a set of saddle point equations that the kernelssatisfy at inﬁnite-width, relating the stochastic processes that deﬁne hidden activation evolution tothe kernels and vice versa.
We show that our saddle point equations recover at
= 1 , from analternative method, the same stochastic process obtained previously with Tensor Programs [1].3.
We develop a polynomial-time numerical procedure to solve the saddle point equations for deepnetworks.
In numerical experiments, we demonstrate that solutions to these self-consistencyequations are predictive of network training at a variety of feature learning strengths, widthsand depths.
We provide comparisons of our theory to various approximate methods, such asperturbation theory.
1.1 Related Works
A natural extension to the lazy NTK/NNGP limit that allows the study of feature learning is to calculateﬁnite width corrections to the inﬁnite-width limit.
Finite width corrections to Bayesian inference inwide networks have been obtained with various perturbative [23–29] and self-consistent techniques[30–33].
In the gradient descent based setting, leading order corrections to the NTK dynamics havebeen analyzed to study ﬁnite width effects [34–36, 27].
These methods give approximate correctionswhich are accurate provided the strength of feature learning is small.
In very rich feature learningregimes, however, the leading order corrections can give incorrect predictions [37, 38].
Another approach to study feature learning is to alter NN parameterization in gradient-based learningto allow signiﬁcant feature evolution even at inﬁnite-width, the mean ﬁeld limit [22, 39].
Works onmean ﬁeld NNs have yielded formal loss convergence results [40, 41] and shown equivalences ofgradient ﬂow dynamics to a partial differential equation (PDE) [42–44].
Our results are most closely related to a set of recent works which studied inﬁnite-width NNs trainedwith gradient descent (GD) using the Tensor Programs (TP) framework [1].
We show that our discretetime ﬁeld theory at unit feature learning strength  0 = 1 recovers the stochastic process which wasderived from TP.
The stochastic process derived from TP has provided insights into practical issues in
NN training such as hyper-parameter search [45].
Computing the exact inﬁnite-width limit of GD hasexponential time requirements [1], which we show can be circumvented with an alternating samplingprocedure.
A projected variant of GD training has provided an inﬁnite-width theory that could bescaled to realistic datasets like CIFAR-10 [46].
Inspired by Chizat and Bach’s work on mechanismsof lazy and rich training [18], our theory interpolates between lazy and rich behavior in the meanﬁeld limit for varying  0 and allows comparison of DMFT to perturbative analysis near small  0 .
Further, our derivation of a DMFT action allows the possibility of pursuing ﬁnite width effects.
Our theory is inspired by self-consistent dynamical mean ﬁeld theory (DMFT) from statistical physics[47–53].
This framework has been utilized in the theory of random recurrent networks [54–59],tensor PCA [60, 61], phase retrieval [62], and high-dimensional linear classiﬁers [63–66], but has yetto be developed for deep feature learning.
By developing a self-consistent DMFT of deep NNs, wegain insight into how features evolve in the rich regime of network training, while retaining manypleasant analytic properties of the inﬁnite-width limit.
Problem Setup and Deﬁnitions
Our theory applies to inﬁnite-width networks, both fully-connected and convolutional.
For notationalease we will relegate convolutional results to later sections.
For input x µ 2 R D , we deﬁne the hiddenpre-activation vectors h ` 2 R N for layers ` 2 { 1 , ..., L } as
[Equation]
where ✓ =
Vec { W 0 , ..., w L } are the trainable parameters of the network and  is a twice differentiable activation function.
Inspired by previous works on the mechanisms of lazy gradient basedtraining, the parameter  will control the laziness or richness of the training dynamics [18, 19, 1, 42].
Each of the trainable parameters are initialized as Gaussian random variables with unit variance W ` ij ⇠
N (0 , 1) .
They evolve under gradient ﬂow ddt ✓ =
2 r ✓ L .
The choice of learning rate  2 causesddt L| t =0 to be independent of  .
To characterize the evolution of weights, we introduce backpropagation variables g ` µ =
[Equation]
p N
[Equation]
h
[Equation]
= ˙
[Equation]
z ` µ , where z ` µ = 1 p
[Equation]
g
[Equation]
is the pre-gradient signal.
The relevant dynamical objects to characterize feature learning are feature and gradient kernels foreach hidden layer ` 2 { 1 , ..., L } , deﬁned as
[Equation]
From the kernels {  ` , G ` } L ` =1 , we can compute the Neural Tangent Kernel K NTKµ ↵ ( t, s ) =
r ✓ f µ ( t ) ·
[Equation]
✓
[Equation]
↵ (
[Equation]
=
P
L ` =
G ` +1 µ ↵ ( t, s )  ` µ ↵ ( t, s ) , [7] and the dynamics of the network function f µ
[Equation]
[Equation]
where we deﬁne base cases G L +1 µ ↵ ( t, s )
= 1 ,  0 µ ↵ ( t, s ) =
K xµ ↵ = 1 D x µ · x ↵ .
We note that the aboveformula holds for any data point µ which may or may not be in the set of P training examples.
Theabove expressions demonstrate that knowledge of the temporal trajectory of the NTK on the t = sdiagonal gives the temporal trajectory of the network predictions f µ ( t ) .
Following prior works on inﬁnite-width networks [22, 1, 40, 19], we study the mean ﬁeld limit
[Equation]
As we demonstrate in the Appendix D and N, this is the only N -scaling which allows feature learningas N ! 1 .
The  0 = 0 limit recovers the static NTK limit [7].
We discuss other scalings andparameterizations in Appendix N, relating our work to the µP -parameterization and TP analysis of[1], showing they have identical feature dynamics in the inﬁnite-width limit.
We also analyze theeffect of different hidden layer widths and initialization variances in the Appendix D.8.
We focus onequal widths and NTK parameterization (as in eq.
(1)) in the main text to reduce complexity.
Self-consistent DMFT
Next, we derive our self-consistent DMFT in a limit where t, P = O N (1) .
Our goal is to builda description of training dynamics purely based on representations, and independent of weights.
Studying feature learning at inﬁnite-width enjoys several analytical properties:
• The kernel order parameters  ` , G ` concentrate over random initializations but are dynamical,allowing ﬂexible adaptation of features to the task structure.
• In each layer ` , each neuron’s preactivation h ` i and pregradient
z
` i become i.i
d. draws from adistribution characterized by a set of order parameters {  ` , G ` , A ` , B ` } .
• The kernels are deﬁned as self-consistent averages (denoted by hi ) over this distribution of neuronsin each layer  ` µ ↵ ( t, s )
=
⌦  ( h ` µ ( t ))
( h `↵ ( s ))
↵ and G ` µ ↵ ( t, s )
=
⌦ g ` µ ( t ) g `↵ ( s ) ↵ .
The next section derives these facts from a path-integral formulation of gradient ﬂow dynamics.
3.1 Path Integral Construction
Gradient ﬂow after a random initialization of weights deﬁnes a high dimensional stochastic processover initalizations for variables { h , g } .
Therefore, we will utilize DMFT formalism to obtain a reduced description of network activity during training.
For a simpliﬁed derivation of the DMFT for thetwo-layer
( L = 1 ) case, see D.2.
Generally, we separate the contribution on each forward/backwardpass between the initial condition and gradient updates to weight matrix W ` , deﬁning new stochasticvariables  ` , ⇠ ` 2 R N as
[Equation]
We let Z represent the moment generating functional (MGF) for these stochastic ﬁelds
[Equation]
which requires, by construction the normalization condition Z [ { 0 , 0 } ]
= 1 .
We enforce our deﬁnitionof  , ⇠ using an integral representation of the delta-function.
Thus for each sample µ 2 [ P ] and eachtime t 2 R + , we multiply Z by
[Equation]
for  and the respective expression for ⇠ .
After making such substitutions, we perform integrationover initial Gaussian weight matrices to arrive at an integral expression for Z , which we derive in theappendix D.4.
We show that Z can be described by set of order-parameters {  ` , ˆ  ` , G ` , ˆ G ` , A ` , B ` }
Z [ { j
[Equation] [Equation]
[Equation]
where S is the DMFT action and Z is a single-site MGF, which deﬁnes the distribution of ﬁelds{  ` , ⇠ ` } over the neural population in each layer.
The kernels A and B are related to the correlationsbetween feedforward and feedback signals in the network.
We provide a detailed formula for Z inthe Appendix D.4 and show that it factorizes over different layers Z =
Q
[Equation]
Z ` .
3.2 Deriving the DMFT Equations from the Path Integral Saddle Point
As N ! 1 , the moment-generating function Z is exponentially dominated by the saddle point of S .
The equations that deﬁne this saddle point also deﬁne our DMFT.
We thus identify the kernels thatrender S locally stationary (  S = 0 ).
The most important equations are those which deﬁne {  ` , G ` }
[Equation]
where hi denotes an average over the stochastic process induced by Z
[Equation]
which is deﬁned below
[Equation]
where we deﬁne base cases  0 µ ↵ ( t, s )
=
K xµ ↵ and G L +1 µ ↵ ( t, s )
= 1 ,
A
= B L = 0 .
Wesee that the ﬁelds { h ` , z ` } , which represent the single site preactivations and pre-gradients,are implicit functionals of the mean-zero Gaussian processes { u ` , r ` } which have covariances⌦ u ` µ ( t ) u `↵ ( s ) ↵ =  `  1 µ ↵ ( t, s ) and ⌦ r
[Equation]
µ ( t ) r `↵ ( s ) ↵ = G ` +1 µ ↵ ( t, s ) .
The other saddle point equations give
[Equation] [Equation]
↵
[Equation]
)  which arise due to coupling betweenthe feedforward and feedback signals.
We note that, in the lazy limit  0 ! 0 , the ﬁelds approach
Gaussian processes h ` !
u ` , z ` !
r ` .
Lastly, the ﬁnal saddle point equations  S   ` = 0 ,  S  G ` = 0imply
that ˆ  ` = ˆ G ` = 0 .
The full set of equations that deﬁne the DMFT are given in D.7.
This theory is easily extended to more general architectures such as networks with varying widths bylayer (App. D.8), trainable bias parameter (App. H), multiple (but O N (1) ) output channels (App.
I), convolutional architectures (App. G), networks trained with weight decay (App. J), Langevinsampling (App. K) and momentum (App. L), discrete time training (App. M).
In Appendix N, wediscuss parameterizations which give equivalent feature and predictor dynamics and show our derivedstochastic process is equivalent to the µP scheme of Yang & Hu [1].
Solving the Self-Consistent DMFT
The saddle point equations obtained from the ﬁeld theory discussed in the previous section must besolved self-consistently.
By this we mean that, given knowledge of the kernels, we can characterize thedistribution of { h ` , z ` } , and given the distribution of { h ` , z ` } , we can compute the kernels [67, 64].
In the Appendix B, we provide Algorithm 1, a numerical procedure based on this idea to efﬁcientlysolve for the kernels with an alternating Monte-Carlo strategy.
The output of the algorithm are thedynamical kernels  ` µ ↵ ( t, s ) , G ` µ ↵ ( t, s ) , A ` µ ↵ ( t, s ) , B ` µ ↵ ( t, s ) , from which any network observablecan be computed as we discuss in Appendix D.
We provide an example of the solution to thesaddle point equations compared to training a ﬁnite NN in Figure 1.
We plot  ` , G ` at the end oftraining and the sample-trace of these kernels through time.
Additionally, we compare the kernels ofﬁnite width N network to the DMFT predicted kernels using a cosine-similarity alignment metric
[Equation]
Additional examples are in Appendix Figures 5 and Figure 6.
4.1 Deep Linear Networks: Closed Form Self-Consistent Equations
Deep linear networks (  ( h ) = h ) are of theoretical interest since they are simpler to analyze thannonlinear networks but preserve non-trivial training dynamics and feature learning [68–72, 25, 32, 23].
In a deep linear network, we can simplify our saddle point equations to algebraic formulas that closein terms of the kernels H ` µ ↵ ( t, s ) =
⌦ h ` µ ( t ) h `↵ ( s ) ↵ , G ` ( t, s )
=
⌦ g ` ( t ) g ` ( s ) ↵ [1].
This is a signiﬁcantsimpliﬁcation since it allows solution of the saddle point equations without a sampling procedure.
To describe the result, we ﬁrst introduce a vectorization notation
h `
= Vec { h ` µ ( t ) } µ 2 [ P ] ,t 2 R + .
Likewise we convert kernels H ` = Mat { H ` µ ↵ ( t, s ) } µ, ↵ 2 [ P ] ,t,s 2 R + into matrices.
The inner productunder this vectorization is deﬁned as a · b =
R 1 0 dt
P
Pµ =1 a µ ( t ) b µ ( t ) .
In a practical computational implementation, the theory would be evaluated on a grid of T time points with discrete timegradient descent, so these kernels H ` 2 R PT ⇥ PT would indeed be matrices of the appropriatesize.
The ﬁelds h ` , g ` are linear functionals of independent Gaussian processes u ` , r ` , giving( I
[Equation]
) g
[Equation]
` +
[Equation]
D ` u ` .
The matrices C ` and D `are causal integral operators which depend on { A `  1 , H `  1 } and { B ` , G ` +1 } respectively whichwe deﬁne in Appendix F.
The saddle point equations which deﬁne the kernels are
[Equation]
Examples of the predictions obtained by solving these systems of equations are provided in Figure 2.
We see that these DMFT equations describe kernel evolution for networks of a variety of depths andthat the change in each layer’s kernel increases with the depth of the network.
Unlike many prior results [68–71], our DMFT does not require any restrictions on the structureof the input data but hold for any K x , y .
However, for whitened data K x
= I we show in
[Figure]
Around N ⇠ 500 DMFT begins to show near perfect agreement with the NN.
(g)-(i)
The same plots but for the gradient kernel G ` .
Whereas ﬁnite width effects for  ` are larger at later layers ` since variance accumulates on the forward pass, ﬂuctuations in G ` are large in early layers
Appendix F.1.1, F.2 that our DMFT learning curves interpolate between NTK dynamics and thesigmoidal trajectories of prior works [68, 69] as  0 is increased.
For example, in the two layer( L = 1 ) linear network with K x =
I , the dynamics of the error norm  ( t ) =
||  ( t ) || takes theform @@
t  ( t ) =
p
+  20 ( y   ( t )) 2  ( t ) where
y =
||
y || .
These dynamics give the linearconvergence rate of the NTK if  0 !
but approaches logistic dynamics of [69] as  0 ! 1 .
Further, 199 H ( t ) =
⌦ h 1 ( t ) h 1 ( t ) > ↵ 2 R P ⇥ P only grows in the yy > direction with H y ( t ) = 1 y 2 y
> H ( t ) y
=p 1 +  20 ( y   ( t )) 2 .
At the end of training H ( t ) !
I + 1 y 2 [ p 1 +  20 y 2  1] yy > , recovering therank one spike which was recently obtained in the small initialization limit [73].
Approximation Schemes
We now compare our exact DMFT with approximations of prior works, providing an explanation ofwhen these approximations give accurate predictions and when they break down.
[Figure]
5.1 Gradient Independence Ansatz
We can study the accuracy of the ansatz A ` = B ` = 0 , which is equivalent to treating the weightmatrices W ` (0) and W ` (0) > which appear in forward and backward passes respectively as independent Gaussian matrices.
This assumption was utilized in prior works on signal propagation in deepnetworks in the lazy regime [74–78].
A consequence of this approximation is the Gaussianity andstatistical independence of  ` and ⇠ ` (conditional on {  ` , G ` } ) in each layer as we show in Appendix
O. This ansatz works very well near  0 ⇡
(the static kernel regime) since d h d r , d z d u ⇠ O (  0 ) oraround initialization t ⇡ 0 but begins to fail at larger values of  0 , t (Figure 3).
5.2 Perturbation theory in  0 at inﬁnite-width
In the  0 !
limit, we recover static kernels, giving linear dynamics identical to the NTK limit[7].
Corrections to this lazy limit can be extracted at small but ﬁnite  0 .
This is conceptually similarto recent works which consider perturbation series for the NTK in powers of 1 /N [35, 27, 28](though not identical, see Appendix P.7 for ﬁnite N effects).
We expand all observables q (  0 )
ina power series in  0 , giving q (  0 ) =
q (0) +  0 q (1) +  20 q (2) + ... and compute corrections up to
O (  20 ) .
We show that the O (  0 ) and O (  30 ) corrections to kernels vanish, giving leading orderexpansions of the form  =
+  20  2 + O (  40 ) and G = G 0 +  20 G 2 + O (  40 ) (see Appendix
P.2).
Further, we show that the NTK has relative change at leading order which scales linearlywith depth |  K NTK | / | K NTK, 0 |
⇠ O  0 ,L ( L  20 ) =
O N,  ,L (  2 LN ) , which is consistent with ﬁnite 223 width effective ﬁeld theory at  =
O N (1) [26–28] (Appendix P.6).
Further, at the leading ordercorrection, all temporal dependencies are controlled by P ( P + 1) functions v ↵ ( t ) =
R t 0 ds  0 ↵ ( s )and v ↵ ( t ) =
[Equation]
0  ( s 0 ) , which is consistent with those derived for ﬁnite width NNsusing a truncation of the Neural Tangent Hierarchy [34, 35, 27].
To lighten notation, we focus ourmain text comparison of our non-perturbative DMFT to perturbation theory in the deep linear case.
Full perturbation theory is in Appendix P.2.
[Figure]
[Figure]
Using the timescales derived in the previous section we ﬁnd that the leading order correction to thekernels in inﬁnite-width deep linear network have the form
[Equation]
We see that the relative change in the NTK | K NTK  K NTK (0) | / | K NTK (0) | ⇠ O (  20 L ) =
O (  2 L/N ) , so that large depth L networks exhibit more signiﬁcant kernel evolution, which agreeswith other perturbative studies [35, 27, 25] as well as the non-perturbative results in Figure 2.
Howeverat large  0 and large L , this theory begins to break down as we show in Figure 3.
Feature Learning Dynamics is Preserved at Fixed  0
Our DMFT suggests that for networks sufﬁciently wide for their kernels to concentrate, the dynamicsof loss and kernels should be invariant under the rescaling N !
RN,  !
/ p R , which keeps  0ﬁxed.
To evaluate how well this idea holds in a realistic deep learning problem, we trained CNNsof varying channel counts N on two-class CIFAR classiﬁcation [79].
We tracked the dynamics ofthe loss and the last layer  L kernel.
The results are provided in Figure 4.
We see that dynamics arelargely independent of rescaling as predicted.
Further, as expected, larger  0 leads to larger changesin kernel norm and faster alignment to the target function y , as was also found in [80].
Consequently,
[Figure]
More details about the experiment can be found in Appendix C.2.
Discussion
We provided a unifying DMFT derivation of feature dynamics in inﬁnite networks trained withgradient based optimization.
Our theory interpolates between lazy inﬁnite-width behavior of a static
NTK in  0 !
and rich feature learning.
At  0 = 1 , our DMFT construction agrees with thestochastic process derived previously with the Tensor Programs framework [1].
Our saddle pointequations give self-consistency conditions which relate the stochastic ﬁelds to the kernels.
Theseequations are exactly solveable in deep linear networks and can be efﬁciently solved with a numericalmethod in the nonlinear case.
Comparisons with other approximation schemes show that DMFT canbe accurate at a much wider range of  0 .
We believe our framework could be a useful perspective forfuture theoretical analyses of feature learning and generalization in wide networks.
Though our DMFT is quite general in regards to the data and architecture, the technique is not entirelyrigorous and relies on heuristic physics techniques.
Our theory holds in the T, P =
O N (1) andmay break down otherwise; other asymptotic regimes (such as P/N, T/ log( N ) =
O N (1) , etc) mayexhibit phenomena relevant to deep learning practice [32, 81].
The computational requirementsof our method, while smaller than the exponential time complexity for exact solution [1], are stillsigniﬁcant for large PT .
In Table 1, we compare the time taken for various theories to compute thefeature kernels throughout T steps of gradient descent.
For a width N network, computation of eachforward pass on all P data points takes O ( PN 2 ) computations.
The static NTK requires computationof O ( P 2 ) entries in the kernel which do not need to be recomputed.
However, the DMFT requiresmatrix multiplications on PT ⇥ PT matrices giving a O ( P 3 T 3 ) time scaling.
Future work couldaim to improve the computational overhead of the algorithm, by considering data averaged theories[64] or one pass SGD [1].
Alternative projected versions of gradient descent have also enabled muchbetter computational scaling in evaluation of the theoretical predictions [46], allowing evaluation onfull CIFAR-10.
Requirements
[Table]